share config

Signed-off-by: Adam Gutglick <adam@spiraldb.com>

Author: AdamGS

vortex-datafusion/src/persistent/opener.rs

@@ -42,6 +42,8 @@ use vortex::layout::LayoutReader;
 use vortex::metrics::Label;
 use vortex::metrics::MetricsRegistry;
 use vortex::scan::ScanBuilder;
+use vortex::scan::api::DEFAULT_TARGET_OUTPUT_BYTES_HINT;
+use vortex::scan::api::DEFAULT_TARGET_OUTPUT_ROWS_HINT;
 use vortex::session::VortexSession;
 use vortex_utils::aliases::dash_map::DashMap;
 use vortex_utils::aliases::dash_map::Entry;
@@ -361,8 +363,8 @@ impl FileOpener for VortexOpener {
                 .with_projection(scan_projection)
                 .with_some_filter(filter)
                 .with_ordered(has_output_ordering)
-                .with_target_output_rows(64 * 1024)
-                .with_target_output_bytes(4 << 20)
+                .with_target_output_rows(DEFAULT_TARGET_OUTPUT_ROWS_HINT)
+                .with_target_output_bytes(DEFAULT_TARGET_OUTPUT_BYTES_HINT)
                 .map(move |chunk| {
                     let mut ctx = session.create_execution_ctx();
                     chunk.execute_record_batch(&stream_schema, &mut ctx)

vortex-scan/src/api.rs

@@ -41,6 +41,18 @@ use crate::Selection;
 /// A sendable stream of partitions.
 pub type PartitionStream = BoxStream<'static, VortexResult<PartitionRef>>;
 
+/// Shared default lower bound for rows accumulated before materializing projected data.
+///
+/// Engines that drive scans via [`ScanRequest`] can use this to align with the tuned Vortex
+/// execution path used by DataFusion.
+pub const DEFAULT_TARGET_OUTPUT_ROWS_HINT: usize = 64 * 1024;
+
+/// Shared default lower bound for projected payload bytes accumulated before materialization.
+///
+/// Engines that drive scans via [`ScanRequest`] can use this to align with the tuned Vortex
+/// execution path used by DataFusion.
+pub const DEFAULT_TARGET_OUTPUT_BYTES_HINT: usize = 4 << 20;
+
 /// Opens a Vortex [`DataSource`] from a URI.
 ///
 /// Configuration can be passed via the URI query parameters, similar to JDBC / ADBC.
@@ -126,6 +138,10 @@ pub struct ScanRequest {
     /// Optional limit on the number of rows returned by scan. Limits are applied after all
     /// filtering and row selection.
     pub limit: Option<u64>,
+    /// Preferred lower bound for rows accumulated before materializing projected data.
+    pub target_output_rows: Option<usize>,
+    /// Preferred lower bound for projected payload bytes accumulated before materialization.
+    pub target_output_bytes: Option<usize>,
 }
 
 impl Default for ScanRequest {
@@ -137,6 +153,8 @@ impl Default for ScanRequest {
             selection: Selection::default(),
             ordered: false,
             limit: None,
+            target_output_rows: None,
+            target_output_bytes: None,
         }
     }
 }

vortex-scan/src/layout.rs

@@ -34,6 +34,8 @@ use vortex_session::VortexSession;
 
 use crate::ScanBuilder;
 use crate::Selection;
+use crate::api::DEFAULT_TARGET_OUTPUT_BYTES_HINT;
+use crate::api::DEFAULT_TARGET_OUTPUT_ROWS_HINT;
 use crate::api::DataSource;
 use crate::api::DataSourceScan;
 use crate::api::DataSourceScanRef;
@@ -166,6 +168,8 @@ impl DataSource for LayoutReaderDataSource {
             limit: scan_request.limit,
             selection: scan_request.selection,
             ordered: scan_request.ordered,
+            target_output_rows: scan_request.target_output_rows,
+            target_output_bytes: scan_request.target_output_bytes,
             metrics_registry: self.metrics_registry.clone(),
             next_row: row_range.start,
             end_row: row_range.end,
@@ -187,6 +191,8 @@ struct LayoutReaderScan {
     limit: Option<u64>,
     ordered: bool,
     selection: Selection,
+    target_output_rows: Option<usize>,
+    target_output_bytes: Option<usize>,
     metrics_registry: Option<Arc<dyn MetricsRegistry>>,
     next_row: u64,
     end_row: u64,
@@ -254,6 +260,8 @@ impl Stream for LayoutReaderScan {
             ordered: this.ordered,
             row_range,
             selection: this.selection.clone(),
+            target_output_rows: this.target_output_rows,
+            target_output_bytes: this.target_output_bytes,
             metrics_registry: this.metrics_registry.clone(),
         }) as PartitionRef;
 
@@ -281,6 +289,8 @@ struct LayoutReaderSplit {
     ordered: bool,
     row_range: Range<u64>,
     selection: Selection,
+    target_output_rows: Option<usize>,
+    target_output_bytes: Option<usize>,
     metrics_registry: Option<Arc<dyn MetricsRegistry>>,
 }
 
@@ -312,8 +322,17 @@ impl Partition for LayoutReaderSplit {
             .with_projection(self.projection)
             .with_some_filter(self.filter)
             .with_some_limit(self.limit)
-            .with_some_metrics_registry(self.metrics_registry)
             .with_ordered(self.ordered);
+        let builder = builder
+            .with_target_output_rows(
+                self.target_output_rows
+                    .unwrap_or(DEFAULT_TARGET_OUTPUT_ROWS_HINT),
+            )
+            .with_target_output_bytes(
+                self.target_output_bytes
+                    .unwrap_or(DEFAULT_TARGET_OUTPUT_BYTES_HINT),
+            );
+        let builder = builder.with_some_metrics_registry(self.metrics_registry);
 
         let dtype = builder.dtype()?;
         // Use into_stream() which creates a LazyScanStream that spawns individual I/O

vortex-scan/src/multi.rs

@@ -48,6 +48,8 @@ use vortex_mask::Mask;
 use vortex_session::VortexSession;
 
 use crate::ScanBuilder;
+use crate::api::DEFAULT_TARGET_OUTPUT_BYTES_HINT;
+use crate::api::DEFAULT_TARGET_OUTPUT_ROWS_HINT;
 use crate::api::DataSource;
 use crate::api::DataSourceScan;
 use crate::api::DataSourceScanRef;
@@ -397,7 +399,17 @@ impl Partition for MultiLayoutPartition {
             .with_projection(request.projection)
             .with_some_filter(request.filter)
             .with_some_limit(request.limit)
-            .with_ordered(request.ordered);
+            .with_ordered(request.ordered)
+            .with_target_output_rows(
+                request
+                    .target_output_rows
+                    .unwrap_or(DEFAULT_TARGET_OUTPUT_ROWS_HINT),
+            )
+            .with_target_output_bytes(
+                request
+                    .target_output_bytes
+                    .unwrap_or(DEFAULT_TARGET_OUTPUT_BYTES_HINT),
+            );
 
         if let Some(row_range) = request.row_range {
             builder = builder.with_row_range(row_range);