Add Filter/Split/Project planner and add Datafusion usage of this (#10706)

Signed-off-by: Joe Isaacs <joe.isaacs@live.co.uk>

Author: joseph-isaacs

benchmarks/datafusion-bench/src/lib.rs

@@ -98,13 +98,15 @@ pub fn make_object_store(
     }
 }
 
-pub fn format_to_df_format(format: Format) -> anyhow::Result<Arc<dyn FileFormat>> {
+pub fn format_to_df_format(format: Format, scan_v2: bool) -> anyhow::Result<Arc<dyn FileFormat>> {
+    let mut options = vortex_table_options();
+    options.scan_v2 = scan_v2;
     Ok(match format {
         Format::Csv => Arc::new(CsvFormat::default()) as _,
         Format::Parquet => Arc::new(ParquetFormat::new()),
-        Format::OnDiskVortex | Format::VortexCompact | Format::VortexSpatialNative => Arc::new(
-            VortexFormat::new_with_options(SESSION.clone(), vortex_table_options()),
-        ),
+        Format::OnDiskVortex | Format::VortexCompact | Format::VortexSpatialNative => {
+            Arc::new(VortexFormat::new_with_options(SESSION.clone(), options))
+        }
         Format::ArrowIpc | Format::OnDiskDuckDB | Format::Lance | Format::LitData => {
             anyhow::bail!("Format {format} cannot be turned into a DataFusion `FileFormat`")
         }

benchmarks/datafusion-bench/src/main.rs

@@ -120,6 +120,10 @@ struct Args {
     #[arg(long, value_delimiter = ',', value_parser = value_parser!(Format))]
     formats: Vec<Format>,
 
+    /// Use the planner-driven V2 Vortex scan (defaults to true for TPC-H).
+    #[arg(long, action = clap::ArgAction::Set)]
+    scan_v2: Option<bool>,
+
     #[arg(long = "opt", value_delimiter = ',', value_parser = value_parser!(Opt))]
     options: Vec<Opt>,
 }
@@ -132,6 +136,10 @@ async fn main() -> anyhow::Result<()> {
     vortex_bench::configure_allocator();
     let args = Args::parse();
     let opts = Opts::from(args.options);
+    let scan_v2 = args.scan_v2.unwrap_or_else(|| {
+        matches!(args.benchmark, BenchmarkArg::TpcH)
+            && !std::env::var("VORTEX_SCAN_V2").is_ok_and(|v| v == "0")
+    });
 
     set_join_set_tracer(get_static_tracer())?;
     setup_logging_and_tracing(args.verbose, args.tracing)?;
@@ -191,7 +199,7 @@ async fn main() -> anyhow::Result<()> {
                         session.sql(&sql).await?.collect().await?;
                     }
                     datafusion_bench::make_object_store(&session, benchmark.data_url())?;
-                    register_benchmark_tables(&session, benchmark, format).await?;
+                    register_benchmark_tables(&session, benchmark, format, scan_v2).await?;
                     Ok((session, format))
                 }
             },
@@ -255,12 +263,14 @@ async fn register_benchmark_tables<B: Benchmark + ?Sized>(
     session: &SessionContext,
     benchmark: &B,
     format: Format,
+    scan_v2: bool,
 ) -> anyhow::Result<()> {
-    if use_scan_api() && matches!(format, Format::OnDiskVortex | Format::VortexCompact) {
+    if !scan_v2 && use_scan_api() && matches!(format, Format::OnDiskVortex | Format::VortexCompact)
+    {
         register_v2_tables(session, benchmark, format).await
     } else {
         let benchmark_base = benchmark.data_url().join(&format!("{}/", format.name()))?;
-        let file_format = format_to_df_format(format)?;
+        let file_format = format_to_df_format(format, scan_v2)?;
 
         for table in benchmark.table_specs().iter() {
             let pattern = benchmark.pattern(table.name, format);

vortex-datafusion/src/persistent/format.rs

@@ -197,6 +197,13 @@ extensions_options! {
     /// [`ParquetFormatFactory::new_with_options`]: https://docs.rs/datafusion/latest/datafusion/datasource/file_format/parquet/struct.ParquetFormatFactory.html#method.new_with_options
     /// [`JsonFormatFactory::new_with_options`]: https://docs.rs/datafusion/latest/datafusion/datasource/file_format/json/struct.JsonFormatFactory.html#method.new_with_options
     pub struct VortexTableOptions {
+        /// Use the planner-driven V2 scan executor for supported layouts.
+        ///
+        /// Disabled by default. Select V2 before creating a table with
+        /// `SET vortex.scan_v2 = true` (with the `vortex` options extension registered),
+        /// or for one table with `OPTIONS(scan_v2 'true')`. In Rust, set
+        /// `VortexTableOptions { scan_v2: true, ..Default::default() }` on the factory.
+        pub scan_v2: bool, default = false
         /// The number of bytes to read when parsing a file footer.
         ///
         /// Values smaller than `MAX_POSTSCRIPT_SIZE + EOF_SIZE` will be clamped to that minimum
@@ -1033,17 +1040,18 @@ mod tests {
         Ok(())
     }
 
+    #[rstest]
     #[tokio::test]
-    async fn configure_format_source() -> anyhow::Result<()> {
+    async fn configure_format_source(#[values(false, true)] scan_v2: bool) -> anyhow::Result<()> {
         let ctx = TestSessionContext::default();
 
         ctx.session
             .sql(
-                "CREATE EXTERNAL TABLE my_tbl \
+                &format!("CREATE EXTERNAL TABLE my_tbl \
                 (c1 VARCHAR NOT NULL, c2 INT NOT NULL) \
                 STORED AS vortex \
                 LOCATION 'table/' \
-                OPTIONS( footer_initial_read_size_bytes '12345', predicate_pushdown 'false', scan_concurrency '3' );",
+                OPTIONS( footer_initial_read_size_bytes '12345', predicate_pushdown 'false', scan_concurrency '3', scan_v2 '{scan_v2}' );"),
             )
             .await?
             .collect()
@@ -1061,6 +1069,35 @@ mod tests {
         assert_eq!(format.options().footer_initial_read_size_bytes, 12345);
         assert!(!format.options().predicate_pushdown);
         assert_eq!(format.options().scan_concurrency, Some(3));
+        assert_eq!(format.options().scan_v2, scan_v2);
+        Ok(())
+    }
+
+    #[rstest]
+    #[tokio::test]
+    async fn session_scan_version(#[values(false, true)] scan_v2: bool) -> anyhow::Result<()> {
+        let factory = Arc::new(VortexFormatFactory::new());
+        let ctx = TestSessionContext::new_with_factory(Arc::clone(&factory));
+        ctx.session
+            .state_ref()
+            .write()
+            .config_mut()
+            .options_mut()
+            .extensions
+            .insert(VortexTableOptions::default());
+
+        for enabled in [scan_v2, !scan_v2] {
+            ctx.session
+                .sql(&format!("SET vortex.scan_v2 = {enabled}"))
+                .await?
+                .collect()
+                .await?;
+            let format = factory.create(&ctx.session.state(), &Default::default())?;
+            let format = format
+                .downcast_ref::<VortexFormat>()
+                .ok_or_else(|| anyhow::anyhow!("expected VortexFormat"))?;
+            assert_eq!(format.options().scan_v2, enabled);
+        }
         Ok(())
     }
 
@@ -1076,6 +1113,7 @@ mod tests {
     #[test]
     fn format_plumbs_source_options() -> anyhow::Result<()> {
         let opts = VortexTableOptions {
+            scan_v2: true,
             projection_pushdown: true,
             predicate_pushdown: false,
             scan_concurrency: Some(3),
@@ -1095,6 +1133,7 @@ mod tests {
         );
         assert_eq!(source.options().predicate_pushdown, opts.predicate_pushdown);
         assert_eq!(source.options().scan_concurrency, opts.scan_concurrency);
+        assert_eq!(source.options().scan_v2, opts.scan_v2);
         Ok(())
     }