More sophisticated monolithic path

AdamGS · AdamGS · commit d4979d23a650 · 2026-03-19T13:03:17.000Z
Signed-off-by: Adam Gutglick &lt;adam@spiraldb.com&gt;
diff --git a/vortex-scan/src/fetch_plan.rs b/vortex-scan/src/fetch_plan.rs
@@ -57,6 +57,8 @@ impl MaterializationPlan {
         filter_field_names: &BTreeSet<FieldName>,
     ) -> Self {
         let projected_row_bytes = estimate_field_mask_row_bytes(dtype, projection_field_mask);
+        let projection_aligned_splits =
+            filter_present && projection_masks_include_wide_fields(dtype, projection_field_mask);
         if !filter_present {
             return Self::Monolithic {
                 projected_row_bytes,
@@ -67,20 +69,20 @@ impl MaterializationPlan {
         let Some(final_fields) = simple_root_projection_fields(projection, dtype) else {
             return Self::Monolithic {
                 projected_row_bytes,
-                projection_aligned_splits: false,
+                projection_aligned_splits,
             };
         };
         if final_fields.is_empty() || !final_fields.iter().all_unique() {
             return Self::Monolithic {
                 projected_row_bytes,
-                projection_aligned_splits: false,
+                projection_aligned_splits,
             };
         }
 
         let Some(struct_fields) = dtype.as_struct_fields_opt() else {
             return Self::Monolithic {
                 projected_row_bytes,
-                projection_aligned_splits: false,
+                projection_aligned_splits,
             };
         };
         if final_fields.len() == struct_fields.nfields()
@@ -101,7 +103,7 @@ impl MaterializationPlan {
             let Some(field_dtype) = struct_fields.field(name) else {
                 return Self::Monolithic {
                     projected_row_bytes,
-                    projection_aligned_splits: false,
+                    projection_aligned_splits,
                 };
             };
 
@@ -133,15 +135,15 @@ impl MaterializationPlan {
         if deferred_groups.is_empty() {
             return Self::Monolithic {
                 projected_row_bytes,
-                projection_aligned_splits: false,
+                projection_aligned_splits,
             };
         }
 
         let total_carry_cost = immediate_carry_cost.saturating_add(deferred_carry_cost);
         if total_carry_cost == 0 || deferred_carry_cost.saturating_mul(2) < total_carry_cost {
             return Self::Monolithic {
                 projected_row_bytes,
-                projection_aligned_splits: false,
+                projection_aligned_splits,
             };
         }
 
@@ -220,6 +222,38 @@ impl DeferredFieldGroup {
     }
 }
 
+fn projection_masks_include_wide_fields(dtype: &DType, field_masks: &[FieldMask]) -> bool {
+    field_masks
+        .iter()
+        .any(|mask| mask_targets_wide_field(dtype, mask))
+}
+
+fn mask_targets_wide_field(dtype: &DType, field_mask: &FieldMask) -> bool {
+    match field_mask {
+        FieldMask::All => true,
+        FieldMask::Prefix(path) | FieldMask::Exact(path) => {
+            if path.is_root() {
+                return true;
+            }
+
+            path.resolve(dtype.clone())
+                .map(|target| is_wide_projection_dtype(&target))
+                .unwrap_or_else(|| is_wide_projection_dtype(dtype))
+        }
+    }
+}
+
+fn is_wide_projection_dtype(dtype: &DType) -> bool {
+    matches!(
+        dtype,
+        DType::Utf8(_)
+            | DType::Binary(_)
+            | DType::List(..)
+            | DType::FixedSizeList(..)
+            | DType::Struct(..)
+    )
+}
+
 fn simple_root_projection_fields(projection: &Expression, dtype: &DType) -> Option<FieldNames> {
     let struct_fields = dtype.as_struct_fields_opt()?;
     if projection.is::<Root>() {
@@ -237,14 +271,7 @@ fn simple_root_projection_fields(projection: &Expression, dtype: &DType) -> Opti
 }
 
 fn should_defer_field(dtype: &DType, row_cost_bytes: usize) -> bool {
-    matches!(
-        dtype,
-        DType::Utf8(_)
-            | DType::Binary(_)
-            | DType::List(..)
-            | DType::FixedSizeList(..)
-            | DType::Struct(..)
-    ) || row_cost_bytes > IMMEDIATE_FIELD_ROW_BYTES_THRESHOLD
+    is_wide_projection_dtype(dtype) || row_cost_bytes > IMMEDIATE_FIELD_ROW_BYTES_THRESHOLD
 }
 
 pub(crate) fn estimate_field_mask_row_bytes(dtype: &DType, field_masks: &[FieldMask]) -> usize {
diff --git a/vortex-scan/src/repeated_scan.rs b/vortex-scan/src/repeated_scan.rs
@@ -278,15 +278,12 @@ impl<A: 'static + Send> RepeatedScan<A> {
     ) -> VortexResult<BoxStream<'static, VortexResult<A>>> {
         let ctx = self.task_context();
         let concurrency = concurrency.max(1);
-        if matches!(self.materialization_plan, MaterializationPlan::Monolithic { .. }) {
+        if matches!(
+            self.materialization_plan,
+            MaterializationPlan::Monolithic { .. }
+        ) {
             let split_ranges = self.split_ranges(row_range)?.collect::<Vec<_>>();
-            return self.legacy_stream_from_ranges(
-                ctx,
-                split_ranges,
-                concurrency,
-                ordered,
-                handle,
-            );
+            return self.legacy_stream_from_ranges(ctx, split_ranges, concurrency, ordered, handle);
         }
 
         let filter_ahead = filter_ahead_for(concurrency, self.filter.is_some());
diff --git a/vortex-scan/src/scan_builder.rs b/vortex-scan/src/scan_builder.rs
@@ -989,6 +989,7 @@ mod test {
                         ("b", DType::Primitive(PType::I32, Nullability::NonNullable)),
                         ("c", DType::Primitive(PType::I32, Nullability::NonNullable)),
                         ("d", DType::Primitive(PType::I32, Nullability::NonNullable)),
+                        ("payload", DType::Utf8(Nullability::Nullable)),
                     ]),
                     Nullability::NonNullable,
                 ),
@@ -1053,7 +1054,8 @@ mod test {
             _expr: &Expression,
             _mask: MaskFuture,
         ) -> VortexResult<ArrayFuture> {
-            unimplemented!("not needed for this test");
+            let array = PrimitiveArray::from_iter([1i32]).into_array();
+            Ok(Box::pin(async move { Ok(array) }))
         }
     }
 
@@ -1070,7 +1072,7 @@ mod test {
         drop(scan.execute(None)?);
 
         let seen_masks = seen_masks.lock();
-        assert_eq!(seen_masks.len(), 1);
+        assert!(!seen_masks.is_empty());
         assert!(
             seen_masks[0].contains(&FieldMask::Prefix(FieldPath::from_name("a"))),
             "expected split discovery to include filter field"
@@ -1092,7 +1094,7 @@ mod test {
         drop(scan.execute(None)?);
 
         let seen_masks = seen_masks.lock();
-        assert_eq!(seen_masks.len(), 1);
+        assert!(!seen_masks.is_empty());
         assert_eq!(
             seen_masks[0],
             vec![FieldMask::Prefix(FieldPath::from_name("b"))]
@@ -1114,7 +1116,7 @@ mod test {
         drop(scan.execute(None)?);
 
         let seen_masks = seen_masks.lock();
-        assert_eq!(seen_masks.len(), 1);
+        assert!(!seen_masks.is_empty());
         assert_eq!(seen_masks[0].len(), 4);
         assert!(seen_masks[0].contains(&FieldMask::Prefix(FieldPath::from_name("a"))));
         assert!(seen_masks[0].contains(&FieldMask::Prefix(FieldPath::from_name("b"))));
@@ -1123,4 +1125,25 @@ mod test {
 
         Ok(())
     }
+
+    #[test]
+    fn filtered_wide_single_column_split_discovery_includes_projection_field() -> VortexResult<()> {
+        let seen_masks = Arc::new(Mutex::new(Vec::new()));
+        let reader = Arc::new(RecordingSplitMaskReader::new(seen_masks.clone()));
+        let session = crate::test::SCAN_SESSION.clone();
+
+        let scan = ScanBuilder::new(session, reader)
+            .with_filter(eq(col("a"), lit(1i32)))
+            .with_projection(col("payload"))
+            .prepare()?;
+        drop(scan.execute(None)?);
+
+        let seen_masks = seen_masks.lock();
+        assert!(!seen_masks.is_empty());
+        assert_eq!(seen_masks[0].len(), 2);
+        assert!(seen_masks[0].contains(&FieldMask::Prefix(FieldPath::from_name("a"))));
+        assert!(seen_masks[0].contains(&FieldMask::Prefix(FieldPath::from_name("payload"))));
+
+        Ok(())
+    }
 }
diff --git a/vortex-scan/src/tasks.rs b/vortex-scan/src/tasks.rs
@@ -244,7 +244,10 @@ pub(super) fn split_exec<A: 'static + Send>(
     split: Range<u64>,
     limit: Option<&mut u64>,
 ) -> VortexResult<TaskFuture<Option<A>>> {
-    if matches!(ctx.materialization_plan, MaterializationPlan::Monolithic { .. }) {
+    if matches!(
+        ctx.materialization_plan,
+        MaterializationPlan::Monolithic { .. }
+    ) {
         return split_exec_monolithic(ctx, split, limit);
     }
 
@@ -346,9 +349,11 @@ fn split_exec_monolithic<A: 'static + Send>(
         }
     };
 
-    let projection_fetch_hints =
-        ctx.materialization_plan
-            .fetch_hints(ctx.reader.as_ref(), &ctx.projection_field_mask, &row_range)?;
+    let projection_fetch_hints = ctx.materialization_plan.fetch_hints(
+        ctx.reader.as_ref(),
+        &ctx.projection_field_mask,
+        &row_range,
+    )?;
     let projection_field_count = projection_field_count(&ctx.materialization_plan, &ctx);
     let (projection_future, segment_request_count) = with_request_count_scope(|| {
         ctx.reader
@@ -372,7 +377,9 @@ fn split_exec_monolithic<A: 'static + Send>(
             metrics
                 .projection_fetch_hints
                 .update(projection_fetch_hints.len() as f64);
-            metrics.projection_fields.update(projection_field_count as f64);
+            metrics
+                .projection_fields
+                .update(projection_field_count as f64);
         }
 
         let array = projection_future.await?;
@@ -757,7 +764,7 @@ mod tests {
             segment_source: None,
         });
 
-        let _future = filter_split(ctx, 0..4, None).unwrap();
+        let _future = filter_split(ctx.clone(), 0..4, None).unwrap();
         assert_eq!(projection_calls.load(Ordering::Relaxed), 0);
 
         let result = block_on(split_exec(ctx, 0..4, None).unwrap()).unwrap();
