[experiment] V2 scan with one segment source per partition

.github/workflows/sql-bench-matrix.yml

@@ -253,9 +253,10 @@ jobs:
           cat comment.md >> "$GITHUB_STEP_SUMMARY"
 
       - name: Comment PR
-        if: inputs.mode == 'pr' && github.event.pull_request.head.repo.fork == false
+        if: inputs.mode == 'pr'
         uses: thollander/actions-comment-pull-request@24bffb9b452ba05a4f3f77933840a6a841d1b32b  # v3
         with:
+          pr-number: 10765  # EXPERIMENT: run by workflow_dispatch, which has no PR context
           file-path: comment.md
           # There is exactly one comment per comment-tag. If a comment with this tag already exists,
           # this action will *update* the comment instead of posting a new comment. The preset is
@@ -264,9 +265,10 @@ jobs:
           comment-tag: bench-pr-comment-${{ matrix.id }}-${{ inputs.matrix_preset }}
 
       - name: Comment PR on failure
-        if: failure() && inputs.mode == 'pr' && github.event.pull_request.head.repo.fork == false
+        if: failure() && inputs.mode == 'pr'
         uses: thollander/actions-comment-pull-request@24bffb9b452ba05a4f3f77933840a6a841d1b32b  # v3
         with:
+          pr-number: 10765  # EXPERIMENT: run by workflow_dispatch, which has no PR context
           message: |
             # 🚨🚨🚨❌❌❌ SQL BENCHMARK FAILED ❌❌❌🚨🚨🚨
 

benchmarks/compress-bench/src/vortex.rs

@@ -3,22 +3,25 @@
 
 use std::io::Cursor;
 use std::path::Path;
-use std::sync::Arc;
 use std::time::Duration;
 use std::time::Instant;
 
 use anyhow::Result;
 use anyhow::bail;
+use arrow_schema::Field;
 use async_trait::async_trait;
 use bytes::Bytes;
 use futures::StreamExt;
 use futures::pin_mut;
 use vortex::array::IntoArray;
+use vortex::array::VortexSessionExecute;
 use vortex::dtype::FieldNames;
 use vortex::expr::root;
 use vortex::expr::select;
 use vortex::file::OpenOptionsSessionExt;
 use vortex::file::WriteOptionsSessionExt;
+use vortex::file::v2::ScanDriver;
+use vortex::io::session::RuntimeSessionExt;
 use vortex::layout::scan::v2::ScanBuilder;
 use vortex_arrow::ArrowSessionExt;
 use vortex_bench::Format;
@@ -96,13 +99,17 @@ async fn vortex_decompress_read(data: Bytes) -> Result<()> {
             .optimize_recursive()?;
         scan = scan.with_projection(projection);
     }
-    let schema = Arc::new(SESSION.arrow().to_arrow_schema(scan.dtype())?);
+    let schema = SESSION.arrow().to_arrow_schema(scan.dtype())?;
+    let field = Field::new_struct("", schema.fields().clone(), false);
 
-    let stream = scan.into_record_batch_stream(schema)?;
+    let stream = ScanDriver::new(SESSION.handle(), 4, true).run(scan);
     pin_mut!(stream);
 
-    while let Some(batch) = stream.next().await {
-        let _batch = batch?;
+    while let Some(chunk) = stream.next().await {
+        let mut ctx = SESSION.create_execution_ctx();
+        let _batch = SESSION
+            .arrow()
+            .execute_arrow(chunk?, Some(&field), &mut ctx)?;
     }
     Ok(())
 }

benchmarks/datafusion-bench/src/lib.rs

@@ -16,8 +16,10 @@ use datafusion::execution::runtime_env::RuntimeEnvBuilder;
 use datafusion::prelude::SessionConfig;
 use datafusion::prelude::SessionContext;
 use datafusion_common::GetExt;
+use object_store::ClientConfigKey;
 use object_store::ObjectStore;
 use object_store::aws::AmazonS3Builder;
+use object_store::aws::AmazonS3ConfigKey;
 use object_store::gcp::GoogleCloudStorageBuilder;
 use object_store::local::LocalFileSystem;
 use url::Url;
@@ -65,11 +67,18 @@ pub fn make_object_store(
     match source.scheme() {
         "s3" => {
             let bucket_name = &source[url::Position::BeforeHost..url::Position::AfterHost];
-            let s3 = Arc::new(
-                AmazonS3Builder::from_env()
-                    .with_bucket_name(bucket_name)
-                    .build()?,
-            );
+            // A scan's burst of reads opens a few hundred connections, and queries reusing
+            // them all from an unbounded idle pool alternate between a fast and a slow mode.
+            // A bounded pool keeps every query in the fast mode; `AWS_POOL_MAX_IDLE_PER_HOST`
+            // overrides it like the other `AWS_` client options.
+            let mut builder = AmazonS3Builder::from_env().with_bucket_name(bucket_name);
+            if std::env::var_os("AWS_POOL_MAX_IDLE_PER_HOST").is_none() {
+                builder = builder.with_config(
+                    AmazonS3ConfigKey::Client(ClientConfigKey::PoolMaxIdlePerHost),
+                    "32",
+                );
+            }
+            let s3 = Arc::new(builder.build()?);
             session.register_object_store(
                 &Url::parse(&format!("s3://{bucket_name}/"))?,
                 Arc::<object_store::aws::AmazonS3>::clone(&s3),
@@ -98,13 +107,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-bench/src/random_access/take.rs

@@ -35,6 +35,8 @@ use vortex::array::stream::ArrayStreamExt;
 use vortex::buffer::Buffer;
 use vortex::file::OpenOptionsSessionExt;
 use vortex::file::VortexFile;
+use vortex::file::v2::ScanDriver;
+use vortex::io::session::RuntimeSessionExt;
 use vortex::layout::scan::v2::ScanBuilder;
 use vortex::scan::strict_sorted_buffer::StrictSortedBuffer;
 
@@ -190,9 +192,9 @@ impl RandomAccessor for VortexRandomAccessor {
 
     async fn take(&self, indices: &[u64]) -> anyhow::Result<RandomAccessorRet> {
         let indices = StrictSortedBuffer::try_new(Buffer::from(indices.to_vec()))?;
-        let stream = ScanBuilder::new(SESSION.clone(), self.file.scan_file())
-            .with_row_indices(indices)
-            .into_stream()?;
+        let scan =
+            ScanBuilder::new(SESSION.clone(), self.file.scan_file()).with_row_indices(indices);
+        let stream = ScanDriver::new(SESSION.handle(), 4, true).run(scan);
         let array = ArrayStreamAdapter::new(self.file.dtype().clone(), stream)
             .read_all()
             .await?;

vortex-datafusion/src/persistent/cache.rs

@@ -5,12 +5,18 @@ use datafusion_common::HashMap;
 use datafusion_execution::cache::cache_manager::FileMetadata;
 use vortex::file::Footer;
 use vortex::file::VortexFile;
+use vortex::layout::scan::v2::ScanCache;
 
 /// Cached Vortex file metadata for use with DataFusion's [`FileMetadataCache`].
 ///
 /// [`FileMetadataCache`]: datafusion_execution::cache::cache_manager::FileMetadataCache
 pub struct CachedVortexMetadata {
     footer: Footer,
+    /// What V2 scans of the file derived from its layout: the lowered plan, each projection's
+    /// and filter's optimized plans, the zones a filter has pruned and the dictionary values
+    /// read. Kept with the footer so later queries of the file plan nothing again. Its size
+    /// is not counted in [`FileMetadata::memory_size`], which is taken when the entry is put.
+    scan_cache: ScanCache,
 }
 
 impl CachedVortexMetadata {
@@ -19,9 +25,21 @@ impl CachedVortexMetadata {
         let footer = vortex_file.footer();
         Self {
             footer: footer.clone(),
+            scan_cache: ScanCache::default(),
         }
     }
 
+    /// Keeps `scan_cache` with the entry, the one the partitions opening the file share.
+    pub fn with_scan_cache(mut self, scan_cache: ScanCache) -> Self {
+        self.scan_cache = scan_cache;
+        self
+    }
+
+    /// The scan cache of V2 scans of the file.
+    pub fn scan_cache(&self) -> &ScanCache {
+        &self.scan_cache
+    }
+
     /// Get the cached footer.
     pub fn footer(&self) -> &Footer {
         &self.footer

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(())
     }
 

vortex-datafusion/src/persistent/metrics.rs

@@ -159,6 +159,7 @@ fn metric_value_to_datafusion(name: &str, metric: &MetricValue) -> Vec<Datafusio
         MetricValue::Timer(timer) => {
             let mut res = Vec::new();
             res.push(df_counter(format!("{name}_count"), timer.count()));
+            res.push(df_timer(format!("{name}_total"), timer.total()));
 
             if !timer.is_empty() {
                 let max = timer.quantile(1.0).vortex_expect("must not be empty");

vortex-datafusion/src/persistent/mod.rs

@@ -28,6 +28,7 @@ mod format;
 pub mod metrics;
 mod opener;
 pub mod reader;
+mod scan_v2;
 mod sink;
 mod source;
 mod stream;