Some changes

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

Author: AdamGS

vortex-datafusion/src/persistent/opener.rs

@@ -42,8 +42,6 @@ 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;
@@ -298,7 +296,8 @@ impl FileOpener for VortexOpener {
                 }
             };
 
-            let mut scan_builder = ScanBuilder::new(session.clone(), layout_reader);
+            let mut scan_builder = ScanBuilder::new(session.clone(), layout_reader)
+                .with_segment_source(vxf.segment_source());
 
             if let Some(extensions) = file.extensions
                 && let Some(vortex_plan) = extensions.downcast_ref::<VortexAccessPlan>()
@@ -363,8 +362,6 @@ impl FileOpener for VortexOpener {
                 .with_projection(scan_projection)
                 .with_some_filter(filter)
                 .with_ordered(has_output_ordering)
-                .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-file/src/file.rs

@@ -103,18 +103,19 @@ impl VortexFile {
                 self.session.clone(),
             ));
         }
-        Ok(Arc::new(LayoutReaderDataSource::new(
-            reader,
-            self.session.clone(),
-        )))
+        Ok(Arc::new(
+            LayoutReaderDataSource::new(reader, self.session.clone())
+                .with_segment_source(self.segment_source()),
+        ))
     }
 
     /// Initiate a scan of the file, returning a builder for configuring the scan.
     pub fn scan(&self) -> VortexResult<ScanBuilder<ArrayRef>> {
         Ok(ScanBuilder::new(
             self.session.clone(),
             self.layout_reader()?,
-        ))
+        )
+        .with_segment_source(self.segment_source()))
     }
 
     /// Returns true if the expression will never match any rows in the file.

vortex-file/src/read/driver.rs

@@ -140,6 +140,7 @@ impl State {
                     tracing::debug!(?req, "ReadRequest dropped before registration");
                     return;
                 }
+                self.metrics.registered_requests.add(1);
                 self.requests_by_offset.insert((req.offset, req.id));
                 self.requests.insert(req.id, req);
             }
@@ -149,20 +150,43 @@ impl State {
                         self.requests_by_offset.remove(&(req.offset, req_id));
                         tracing::debug!(?req, "ReadRequest dropped before poll");
                     } else {
+                        self.metrics.polled_requests.add(1);
                         self.polled_requests.insert(req_id, req);
                     }
                 }
             }
             ReadEvent::Dropped(req_id) => {
                 if let Some(req) = self.requests.remove(&req_id) {
+                    self.metrics.dropped_requests.add(1);
                     self.requests_by_offset.remove(&(req.offset, req_id));
                     tracing::debug!(?req, "ReadRequest dropped before poll");
                 }
                 if let Some(req) = self.polled_requests.remove(&req_id) {
+                    self.metrics.dropped_requests.add(1);
                     self.requests_by_offset.remove(&(req.offset, req_id));
                     tracing::debug!(?req, "ReadRequest dropped after poll");
                 }
             }
+            ReadEvent::BatchBoundary => {
+                // Promote all registered-but-unpolled requests to polled status.
+                // This tells the coalescer that the entire batch is needed now,
+                // allowing it to form optimal coalesced reads.
+                let promoted = self.requests.len();
+                if promoted > 0 {
+                    tracing::debug!(
+                        promoted,
+                        "BatchBoundary: promoting registered requests to polled"
+                    );
+                    self.metrics.polled_requests.add(promoted as u64);
+                    for (req_id, req) in std::mem::take(&mut self.requests) {
+                        if req.callback.is_closed() {
+                            self.requests_by_offset.remove(&(req.offset, req_id));
+                        } else {
+                            self.polled_requests.insert(req_id, req);
+                        }
+                    }
+                }
+            }
         }
     }
 
@@ -215,6 +239,9 @@ impl State {
         let first_req = self.next_uncoalesced()?;
 
         let mut requests = vec![first_req];
+        let mut payload_bytes = requests[0].length as u64;
+        let mut registered_only_requests = 0usize;
+        let mut polled_requests = 1usize;
         let mut current_start = requests[0].offset;
         let mut current_end = requests[0].offset + requests[0].length as u64;
         let align = *self.coalesced_buffer_alignment as u64;
@@ -269,18 +296,28 @@ impl State {
                     let new_total_size = new_end - aligned_start;
 
                     if new_total_size > window.max_size {
+                        self.metrics.batched_skipped_max_size.add(1);
                         // Skip it but keep it available for future coalescing operations.
                         continue;
                     }
 
                     current_start = new_start;
                     current_end = new_end;
-                    let req = self
-                        .polled_requests
-                        .remove(&req_id)
-                        .or_else(|| self.requests.remove(&req_id))
-                        .vortex_expect("Missing request in requests_by_offset");
+                    let (req, was_polled) = if let Some(req) = self.polled_requests.remove(&req_id)
+                    {
+                        (req, true)
+                    } else if let Some(req) = self.requests.remove(&req_id) {
+                        (req, false)
+                    } else {
+                        unreachable!("Missing request in requests_by_offset");
+                    };
 
+                    payload_bytes = payload_bytes.saturating_add(req.length as u64);
+                    if was_polled {
+                        polled_requests = polled_requests.saturating_add(1);
+                    } else {
+                        registered_only_requests = registered_only_requests.saturating_add(1);
+                    }
                     requests.push(req);
                     if ids_to_remove.insert(req_id) {
                         keys_to_remove.push((req_offset, req_id));
@@ -302,6 +339,18 @@ impl State {
         requests.sort_unstable_by_key(|r| r.offset);
 
         let aligned_start = current_start - (current_start % align);
+        let range_bytes = current_end - aligned_start;
+
+        self.metrics.batched_range_bytes.update(range_bytes as f64);
+        self.metrics
+            .batched_payload_bytes
+            .update(payload_bytes as f64);
+        self.metrics
+            .batched_registered_only_requests
+            .update(registered_only_requests as f64);
+        self.metrics
+            .batched_polled_requests
+            .update(polled_requests as f64);
 
         tracing::debug!(
             "Coalesced {} requests into range {}..{} (len={})",
@@ -808,4 +857,86 @@ mod tests {
         assert_eq!(individual_count, 2, "Expected 2 individual requests");
         assert_eq!(coalesced_operations, 0, "Expected 0 coalesced operations");
     }
+
+    #[tokio::test]
+    async fn test_metrics_record_registered_only_batch_members() {
+        let (req1, _rx1) = create_request(1, 0, 10);
+        let (req2, _rx2) = create_request(2, 50, 10);
+        let (req3, _rx3) = create_request(3, 100, 10);
+
+        let events = vec![
+            ReadEvent::Request(req1),
+            ReadEvent::Request(req2),
+            ReadEvent::Request(req3),
+            ReadEvent::Polled(2),
+        ];
+
+        let event_stream = stream::iter(events);
+        let metrics_registry = DefaultMetricsRegistry::default();
+        let metrics = RequestMetrics::new(&metrics_registry, vec![]);
+        let io_stream = IoRequestStream::new(
+            event_stream,
+            Some(CoalesceConfig {
+                distance: 60,
+                max_size: 1024,
+            }),
+            Alignment::none(),
+            metrics,
+        );
+
+        let outputs: Vec<IoRequest> = io_stream.collect().await;
+        assert_eq!(outputs.len(), 1);
+
+        let snapshot = metrics_registry.snapshot();
+        let mut registered = 0u64;
+        let mut polled = 0u64;
+        let mut coalesced = 0u64;
+        let mut registered_only_count = 0usize;
+        let mut registered_only_total = 0.0;
+        let mut polled_in_batch_count = 0usize;
+        let mut polled_in_batch_total = 0.0;
+
+        for metric in snapshot.iter() {
+            match metric.value() {
+                MetricValue::Counter(counter) => match metric.name().as_ref() {
+                    "io.requests.registered" => registered = counter.value(),
+                    "io.requests.polled" => polled = counter.value(),
+                    "io.requests.coalesced" => coalesced = counter.value(),
+                    _ => {}
+                },
+                MetricValue::Histogram(histogram) => match metric.name().as_ref() {
+                    "io.requests.batched.registered_only_requests" => {
+                        registered_only_count = histogram.count();
+                        registered_only_total = histogram.total();
+                    }
+                    "io.requests.batched.polled_requests" => {
+                        polled_in_batch_count = histogram.count();
+                        polled_in_batch_total = histogram.total();
+                    }
+                    _ => {}
+                },
+                _ => {}
+            }
+        }
+
+        assert_eq!(registered, 3, "Expected 3 registered requests");
+        assert_eq!(polled, 1, "Expected 1 polled request");
+        assert_eq!(coalesced, 1, "Expected 1 coalesced operation");
+        assert_eq!(
+            registered_only_count, 1,
+            "Expected one histogram sample for registered-only requests"
+        );
+        assert_eq!(
+            registered_only_total, 2.0,
+            "Expected two registered-only requests in the coalesced batch"
+        );
+        assert_eq!(
+            polled_in_batch_count, 1,
+            "Expected one histogram sample for polled requests"
+        );
+        assert_eq!(
+            polled_in_batch_total, 1.0,
+            "Expected one polled request in the coalesced batch"
+        );
+    }
 }

vortex-file/src/segments/source.rs

@@ -39,6 +39,10 @@ pub enum ReadEvent {
     Request(ReadRequest),
     Polled(RequestId),
     Dropped(RequestId),
+    /// Signals that a logical batch of requests has been fully registered.
+    /// The driver promotes all registered-but-unpolled requests to polled status,
+    /// allowing the coalescer to form optimal reads over the entire batch.
+    BatchBoundary,
 }
 
 /// A [`SegmentSource`] for file-like IO.
@@ -134,6 +138,10 @@ impl FileSegmentSource {
 }
 
 impl SegmentSource for FileSegmentSource {
+    fn flush(&self) {
+        drop(self.events.unbounded_send(ReadEvent::BatchBoundary));
+    }
+
     fn request(&self, id: SegmentId) -> SegmentFuture {
         // We eagerly register the read request here assuming the behaviour of [`FileRead`], where
         // coalescing becomes effective prior to the future being polled.
@@ -232,23 +240,55 @@ impl Drop for ReadFuture {
 }
 
 pub struct RequestMetrics {
+    pub registered_requests: Counter,
+    pub polled_requests: Counter,
+    pub dropped_requests: Counter,
     pub individual_requests: Counter,
     pub coalesced_requests: Counter,
     pub num_requests_coalesced: Histogram,
+    pub batched_range_bytes: Histogram,
+    pub batched_payload_bytes: Histogram,
+    pub batched_registered_only_requests: Histogram,
+    pub batched_polled_requests: Histogram,
+    pub batched_skipped_max_size: Counter,
 }
 
 impl RequestMetrics {
     pub fn new(metrics_registry: &dyn MetricsRegistry, labels: Vec<Label>) -> Self {
         Self {
+            registered_requests: MetricBuilder::new(metrics_registry)
+                .add_labels(labels.clone())
+                .counter("io.requests.registered"),
+            polled_requests: MetricBuilder::new(metrics_registry)
+                .add_labels(labels.clone())
+                .counter("io.requests.polled"),
+            dropped_requests: MetricBuilder::new(metrics_registry)
+                .add_labels(labels.clone())
+                .counter("io.requests.dropped"),
             individual_requests: MetricBuilder::new(metrics_registry)
                 .add_labels(labels.clone())
                 .counter("io.requests.individual"),
             coalesced_requests: MetricBuilder::new(metrics_registry)
                 .add_labels(labels.clone())
                 .counter("io.requests.coalesced"),
             num_requests_coalesced: MetricBuilder::new(metrics_registry)
-                .add_labels(labels)
+                .add_labels(labels.clone())
                 .histogram("io.requests.coalesced.num_coalesced"),
+            batched_range_bytes: MetricBuilder::new(metrics_registry)
+                .add_labels(labels.clone())
+                .histogram("io.requests.batched.range_bytes"),
+            batched_payload_bytes: MetricBuilder::new(metrics_registry)
+                .add_labels(labels.clone())
+                .histogram("io.requests.batched.payload_bytes"),
+            batched_registered_only_requests: MetricBuilder::new(metrics_registry)
+                .add_labels(labels.clone())
+                .histogram("io.requests.batched.registered_only_requests"),
+            batched_polled_requests: MetricBuilder::new(metrics_registry)
+                .add_labels(labels.clone())
+                .histogram("io.requests.batched.polled_requests"),
+            batched_skipped_max_size: MetricBuilder::new(metrics_registry)
+                .add_labels(labels)
+                .counter("io.requests.batched.skipped_max_size"),
         }
     }
 }

vortex-layout/public-api.lock

@@ -632,6 +632,8 @@ pub fn vortex_layout::layouts::row_idx::RowIdxLayoutReader::name(&self) -> &allo
 
 pub fn vortex_layout::layouts::row_idx::RowIdxLayoutReader::projection_evaluation(&self, row_range: &core::ops::range::Range<u64>, expr: &vortex_array::expr::expression::Expression, mask: vortex_array::mask_future::MaskFuture) -> vortex_error::VortexResult<futures_core::future::BoxFuture<'static, vortex_error::VortexResult<vortex_array::array::ArrayRef>>>
 
+pub fn vortex_layout::layouts::row_idx::RowIdxLayoutReader::projection_fetch_hints(&self, field_mask: alloc::vec::Vec<vortex_array::dtype::field_mask::FieldMask>, row_range: core::ops::range::Range<u64>, row_bytes_hint: usize) -> vortex_error::VortexResult<alloc::vec::Vec<vortex_layout::ProjectionFetchHint>>
+
 pub fn vortex_layout::layouts::row_idx::RowIdxLayoutReader::pruning_evaluation(&self, row_range: &core::ops::range::Range<u64>, expr: &vortex_array::expr::expression::Expression, mask: vortex_mask::Mask) -> vortex_error::VortexResult<vortex_array::mask_future::MaskFuture>
 
 pub fn vortex_layout::layouts::row_idx::RowIdxLayoutReader::register_splits(&self, field_mask: &[vortex_array::dtype::field_mask::FieldMask], row_range: &core::ops::range::Range<u64>, splits: &mut alloc::collections::btree::set::BTreeSet<u64>) -> vortex_error::VortexResult<()>
@@ -1540,6 +1542,28 @@ pub fn vortex_layout::LazyReaderChildren::get(&self, idx: usize) -> vortex_error
 
 pub fn vortex_layout::LazyReaderChildren::new(children: alloc::sync::Arc<dyn vortex_layout::LayoutChildren>, dtypes: alloc::vec::Vec<vortex_array::dtype::DType>, names: alloc::vec::Vec<alloc::sync::Arc<str>>, segment_source: alloc::sync::Arc<dyn vortex_layout::segments::SegmentSource>, session: vortex_session::VortexSession) -> Self
 
+pub struct vortex_layout::ProjectionFetchHint
+
+pub vortex_layout::ProjectionFetchHint::estimated_fetch_bytes: usize
+
+pub vortex_layout::ProjectionFetchHint::row_range: core::ops::range::Range<u64>
+
+impl core::clone::Clone for vortex_layout::ProjectionFetchHint
+
+pub fn vortex_layout::ProjectionFetchHint::clone(&self) -> vortex_layout::ProjectionFetchHint
+
+impl core::cmp::Eq for vortex_layout::ProjectionFetchHint
+
+impl core::cmp::PartialEq for vortex_layout::ProjectionFetchHint
+
+pub fn vortex_layout::ProjectionFetchHint::eq(&self, other: &vortex_layout::ProjectionFetchHint) -> bool
+
+impl core::fmt::Debug for vortex_layout::ProjectionFetchHint
+
+pub fn vortex_layout::ProjectionFetchHint::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
+impl core::marker::StructuralPartialEq for vortex_layout::ProjectionFetchHint
+
 pub trait vortex_layout::ArrayFutureExt
 
 pub fn vortex_layout::ArrayFutureExt::masked(self, mask: vortex_array::mask_future::MaskFuture) -> Self
@@ -1670,6 +1694,8 @@ pub fn vortex_layout::LayoutReader::name(&self) -> &alloc::sync::Arc<str>
 
 pub fn vortex_layout::LayoutReader::projection_evaluation(&self, row_range: &core::ops::range::Range<u64>, expr: &vortex_array::expr::expression::Expression, mask: vortex_array::mask_future::MaskFuture) -> vortex_error::VortexResult<vortex_layout::ArrayFuture>
 
+pub fn vortex_layout::LayoutReader::projection_fetch_hints(&self, field_mask: alloc::vec::Vec<vortex_array::dtype::field_mask::FieldMask>, row_range: core::ops::range::Range<u64>, row_bytes_hint: usize) -> vortex_error::VortexResult<alloc::vec::Vec<vortex_layout::ProjectionFetchHint>>
+
 pub fn vortex_layout::LayoutReader::pruning_evaluation(&self, row_range: &core::ops::range::Range<u64>, expr: &vortex_array::expr::expression::Expression, mask: vortex_mask::Mask) -> vortex_error::VortexResult<vortex_array::mask_future::MaskFuture>
 
 pub fn vortex_layout::LayoutReader::register_splits(&self, field_mask: &[vortex_array::dtype::field_mask::FieldMask], row_range: &core::ops::range::Range<u64>, splits: &mut alloc::collections::btree::set::BTreeSet<u64>) -> vortex_error::VortexResult<()>
@@ -1688,6 +1714,8 @@ pub fn vortex_layout::layouts::row_idx::RowIdxLayoutReader::name(&self) -> &allo
 
 pub fn vortex_layout::layouts::row_idx::RowIdxLayoutReader::projection_evaluation(&self, row_range: &core::ops::range::Range<u64>, expr: &vortex_array::expr::expression::Expression, mask: vortex_array::mask_future::MaskFuture) -> vortex_error::VortexResult<futures_core::future::BoxFuture<'static, vortex_error::VortexResult<vortex_array::array::ArrayRef>>>
 
+pub fn vortex_layout::layouts::row_idx::RowIdxLayoutReader::projection_fetch_hints(&self, field_mask: alloc::vec::Vec<vortex_array::dtype::field_mask::FieldMask>, row_range: core::ops::range::Range<u64>, row_bytes_hint: usize) -> vortex_error::VortexResult<alloc::vec::Vec<vortex_layout::ProjectionFetchHint>>
+
 pub fn vortex_layout::layouts::row_idx::RowIdxLayoutReader::pruning_evaluation(&self, row_range: &core::ops::range::Range<u64>, expr: &vortex_array::expr::expression::Expression, mask: vortex_mask::Mask) -> vortex_error::VortexResult<vortex_array::mask_future::MaskFuture>
 
 pub fn vortex_layout::layouts::row_idx::RowIdxLayoutReader::register_splits(&self, field_mask: &[vortex_array::dtype::field_mask::FieldMask], row_range: &core::ops::range::Range<u64>, splits: &mut alloc::collections::btree::set::BTreeSet<u64>) -> vortex_error::VortexResult<()>