From 6e4dda49928d44cc737c25d7b2ed8d5d928d18ec Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Fri, 31 Jul 2026 17:14:57 +0100 Subject: [PATCH 1/4] refactor: propagate bound expressions through layouts Signed-off-by: Joe Isaacs --- vortex-array/src/expression.rs | 27 +++ vortex-cuda/src/layout.rs | 14 +- vortex-file/src/v2/file_stats_reader.rs | 23 ++- vortex-layout/src/layouts/chunked/reader.rs | 21 +- vortex-layout/src/layouts/dict/reader.rs | 74 ++++--- vortex-layout/src/layouts/flat/reader.rs | 36 ++-- vortex-layout/src/layouts/flat/writer.rs | 24 ++- vortex-layout/src/layouts/list/expr.rs | 91 +++++++-- vortex-layout/src/layouts/list/reader.rs | 111 ++++++---- vortex-layout/src/layouts/partitioned.rs | 6 +- vortex-layout/src/layouts/row_idx/mod.rs | 215 ++++++++++++-------- vortex-layout/src/layouts/struct_/reader.rs | 132 +++++++----- vortex-layout/src/layouts/zoned/pruning.rs | 76 +++---- vortex-layout/src/layouts/zoned/reader.rs | 27 ++- vortex-layout/src/layouts/zoned/zone_map.rs | 5 +- vortex-layout/src/reader.rs | 8 +- vortex-layout/src/scan/filter.rs | 65 +++++- vortex-layout/src/scan/layout.rs | 3 +- vortex-layout/src/scan/multi.rs | 3 +- vortex-layout/src/scan/repeated_scan.rs | 10 +- vortex-layout/src/scan/scan_builder.rs | 25 ++- vortex-layout/src/scan/split_by.rs | 8 +- vortex-layout/src/scan/tasks.rs | 4 +- vortex-test/compat-gen/src/adapter.rs | 3 +- vortex-tui/src/browse/app.rs | 5 +- vortex-tui/src/wasm.rs | 5 +- 26 files changed, 664 insertions(+), 357 deletions(-) diff --git a/vortex-array/src/expression.rs b/vortex-array/src/expression.rs index d6b8739abf6..5e5f8d7ba41 100644 --- a/vortex-array/src/expression.rs +++ b/vortex-array/src/expression.rs @@ -8,12 +8,39 @@ use crate::ArrayRef; use crate::IntoArray; use crate::arrays::ConstantArray; use crate::arrays::ScalarFnArray; +use crate::expr::BoundExpression; +use crate::expr::BoundKind; use crate::expr::Expression; use crate::optimizer::ArrayOptimizer; use crate::scalar_fn::fns::literal::Literal; use crate::scalar_fn::fns::root::Root; impl ArrayRef { + /// Apply a bound expression to this array, producing a new array in constant time. + pub fn apply_bound(self, expr: &BoundExpression) -> VortexResult { + let BoundKind::Scalar { + scalar_fn, + children, + } = expr.kind() + else { + return Ok(self); + }; + + if let Some(scalar) = scalar_fn.as_opt::() { + return Ok(ConstantArray::new(scalar.clone(), self.len()).into_array()); + } + + let children: Vec<_> = children + .iter() + .map(|child| self.clone().apply_bound(child)) + .try_collect()?; + + let array = + ScalarFnArray::try_new_with_len(scalar_fn.clone(), children, self.len())?.into_array(); + + array.optimize() + } + /// Apply the expression to this array, producing a new array in constant time. pub fn apply(self, expr: &Expression) -> VortexResult { // If the expression is a root, return self. diff --git a/vortex-cuda/src/layout.rs b/vortex-cuda/src/layout.rs index 16cef49cf0b..216b3369bd1 100644 --- a/vortex-cuda/src/layout.rs +++ b/vortex-cuda/src/layout.rs @@ -19,7 +19,7 @@ use vortex::array::MaskFuture; use vortex::array::ProstMetadata; use vortex::array::VortexSessionExecute; use vortex::array::arrays::Constant; -use vortex::array::expr::Expression; +use vortex::array::expr::BoundExpression; use vortex::array::expr::stats::Precision; use vortex::array::expr::stats::Stat; use vortex::array::expr::stats::StatsProvider; @@ -268,7 +268,7 @@ impl LayoutReader for CudaFlatReader { fn pruning_evaluation( &self, _row_range: &Range, - _expr: &Expression, + _expr: &BoundExpression, mask: Mask, ) -> VortexResult { Ok(MaskFuture::ready(mask)) @@ -277,7 +277,7 @@ impl LayoutReader for CudaFlatReader { fn filter_evaluation( &self, row_range: &Range, - expr: &Expression, + expr: &BoundExpression, mask: MaskFuture, ) -> VortexResult { let row_range = usize::try_from(row_range.start) @@ -299,13 +299,13 @@ impl LayoutReader for CudaFlatReader { let mask_density = mask.density(); let array_mask = if mask_density < EXPR_EVAL_THRESHOLD { - let array = array.apply(&expr)?; + let array = array.apply_bound(&expr)?; let array = array.filter(mask.clone())?; let mut ctx = session.create_execution_ctx(); let array_mask = array.null_as_false().execute(&mut ctx)?; mask.intersect_by_rank(&array_mask) } else { - let array = array.apply(&expr)?; + let array = array.apply_bound(&expr)?; let mut ctx = session.create_execution_ctx(); let array_mask = array.null_as_false().execute(&mut ctx)?; mask.bitand(&array_mask) @@ -326,7 +326,7 @@ impl LayoutReader for CudaFlatReader { fn projection_evaluation( &self, row_range: &Range, - expr: &Expression, + expr: &BoundExpression, mask: MaskFuture, ) -> VortexResult>> { let row_range = usize::try_from(row_range.start) @@ -351,7 +351,7 @@ impl LayoutReader for CudaFlatReader { array = array.filter(mask)?; } - array = array.apply(&expr)?; + array = array.apply_bound(&expr)?; Ok(array) } diff --git a/vortex-file/src/v2/file_stats_reader.rs b/vortex-file/src/v2/file_stats_reader.rs index 141b6b8ca73..034ceb3b4af 100644 --- a/vortex-file/src/v2/file_stats_reader.rs +++ b/vortex-file/src/v2/file_stats_reader.rs @@ -14,6 +14,7 @@ use vortex_array::MaskFuture; use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; use vortex_array::dtype::StructFields; +use vortex_array::expr::BoundExpression; use vortex_array::expr::Expression; use vortex_error::VortexResult; use vortex_layout::ArrayFuture; @@ -41,7 +42,7 @@ pub struct FileStatsLayoutReader { file_stats: FileStatistics, struct_fields: StructFields, session: VortexSession, - prune_cache: DashMap, + prune_cache: DashMap, } impl FileStatsLayoutReader { @@ -113,7 +114,7 @@ impl LayoutReader for FileStatsLayoutReader { fn pruning_evaluation( &self, row_range: &Range, - expr: &Expression, + expr: &BoundExpression, mask: Mask, ) -> VortexResult { // Check cache first with read-only lock. @@ -125,7 +126,8 @@ impl LayoutReader for FileStatsLayoutReader { } // Evaluate and cache. - let pruned = self.evaluate_file_stats(expr)?; + let expression = expr.unbind(); + let pruned = self.evaluate_file_stats(&expression)?; self.prune_cache.insert(expr.clone(), pruned); if pruned { @@ -138,7 +140,7 @@ impl LayoutReader for FileStatsLayoutReader { fn filter_evaluation( &self, row_range: &Range, - expr: &Expression, + expr: &BoundExpression, mask: MaskFuture, ) -> VortexResult { self.child.filter_evaluation(row_range, expr, mask) @@ -147,7 +149,7 @@ impl LayoutReader for FileStatsLayoutReader { fn projection_evaluation( &self, row_range: &Range, - expr: &Expression, + expr: &BoundExpression, mask: MaskFuture, ) -> VortexResult { self.child.projection_evaluation(row_range, expr, mask) @@ -260,7 +262,7 @@ mod tests { FileStatsLayoutReader::new(child, test_file_stats(0, 100), SESSION.clone()); // col > 200 should be prunable since max is 100. - let expr = gt(get_item("col", root()), lit(200i32)); + let expr = gt(get_item("col", root()), lit(200i32)).bind(reader.dtype())?; let mask = Mask::new_true(5); let result = reader.pruning_evaluation(&(0..5), &expr, mask)?.await?; assert_eq!(result, Mask::new_false(5)); @@ -299,7 +301,7 @@ mod tests { FileStatsLayoutReader::new(child, test_file_stats(0, 100), SESSION.clone()); // col > 50 should NOT be prunable since max is 100 (some rows could match). - let expr = gt(get_item("col", root()), lit(50i32)); + let expr = gt(get_item("col", root()), lit(50i32)).bind(reader.dtype())?; let mask = Mask::new_true(5); let result = reader.pruning_evaluation(&(0..5), &expr, mask)?.await?; // Should delegate to child, which returns the mask unchanged (struct reader doesn't prune). @@ -336,7 +338,8 @@ mod tests { let reader = FileStatsLayoutReader::new(child, test_file_stats(0, 100), SESSION.clone()); - let expr = gt(checked_add(get_item("col", root()), lit(5i32)), lit(102i32)); + let expr = gt(checked_add(get_item("col", root()), lit(5i32)), lit(102i32)) + .bind(reader.dtype())?; let mask = Mask::new_true(2); let result = reader.pruning_evaluation(&(0..2), &expr, mask)?.await?; @@ -391,7 +394,7 @@ mod tests { let reader = FileStatsLayoutReader::new(child, file_stats, SESSION.clone()); // `is_null(deleted_at)` — should NOT panic or error due to dtype mismatch. - let expr = is_null(get_item("deleted_at", root())); + let expr = is_null(get_item("deleted_at", root())).bind(reader.dtype())?; let mask = Mask::new_true(3); let result = reader.pruning_evaluation(&(0..3), &expr, mask)?.await?; // null_count is 1 (non-zero), so is_null is not falsified => not pruned. @@ -435,7 +438,7 @@ mod tests { let reader = FileStatsLayoutReader::new(child, test_file_null_count_stats(5), SESSION.clone()); - let expr = is_not_null(get_item("col", root())); + let expr = is_not_null(get_item("col", root())).bind(reader.dtype())?; let mask = Mask::new_true(5); let result = reader.pruning_evaluation(&(0..5), &expr, mask)?.await?; assert_eq!(result, Mask::new_false(5)); diff --git a/vortex-layout/src/layouts/chunked/reader.rs b/vortex-layout/src/layouts/chunked/reader.rs index 978dad4cbe8..949ae24d01a 100644 --- a/vortex-layout/src/layouts/chunked/reader.rs +++ b/vortex-layout/src/layouts/chunked/reader.rs @@ -18,7 +18,7 @@ use vortex_array::MaskFuture; use vortex_array::arrays::ChunkedArray; use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; -use vortex_array::expr::Expression; +use vortex_array::expr::BoundExpression; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -213,7 +213,7 @@ impl LayoutReader for ChunkedReader { fn pruning_evaluation( &self, row_range: &Range, - expr: &Expression, + expr: &BoundExpression, mask: Mask, ) -> VortexResult { if row_range.is_empty() { @@ -259,7 +259,7 @@ impl LayoutReader for ChunkedReader { fn filter_evaluation( &self, row_range: &Range, - expr: &Expression, + expr: &BoundExpression, mask: MaskFuture, ) -> VortexResult { if row_range.is_empty() { @@ -298,14 +298,11 @@ impl LayoutReader for ChunkedReader { fn projection_evaluation( &self, row_range: &Range, - expr: &Expression, + expr: &BoundExpression, mask: MaskFuture, ) -> VortexResult>> { if row_range.is_empty() { - return Ok(future::ready(Ok( - Canonical::empty(&expr.return_dtype(self.dtype())?).into_array() - )) - .boxed()); + return Ok(future::ready(Ok(Canonical::empty(expr.dtype()).into_array())).boxed()); } let mut chunk_evals = vec![]; @@ -473,12 +470,14 @@ mod test { ) { block_on(|_h| async { let mut ctx = SESSION.create_execution_ctx(); - let result = layout + let reader = layout .new_reader("".into(), segments, &SESSION, &Default::default()) - .unwrap() + .unwrap(); + let expr = root().bind(reader.dtype()).unwrap(); + let result = reader .projection_evaluation( &(0..layout.row_count()), - &root(), + &expr, MaskFuture::new_true(usize::try_from(layout.row_count()).unwrap()), ) .unwrap() diff --git a/vortex-layout/src/layouts/dict/reader.rs b/vortex-layout/src/layouts/dict/reader.rs index 36d4bf04abf..235c61a3307 100644 --- a/vortex-layout/src/layouts/dict/reader.rs +++ b/vortex-layout/src/layouts/dict/reader.rs @@ -19,6 +19,7 @@ use vortex_array::arrays::SharedArray; use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; use vortex_array::dtype::Nullability; +use vortex_array::expr::BoundExpression; use vortex_array::expr::Expression; use vortex_array::expr::direct_annotations; use vortex_array::expr::is_root; @@ -53,7 +54,7 @@ pub struct DictReader { /// Cached dict values array values_array: OnceLock, /// Cache of expression evaluation results on the values array by expression - values_evals: DashMap, + values_evals: DashMap, values: LayoutReaderRef, codes: LayoutReaderRef, @@ -103,12 +104,15 @@ impl DictReader { // We capture the name, so it may be wrong if we re-use the same reader within multiple // different parent readers. But that's rare... let values_len = self.values_len; + let root = root() + .bind(self.values.dtype()) + .vortex_expect("root must bind against the dictionary values dtype"); self.values_array .get_or_init(move || { self.values .projection_evaluation( &(0..values_len as u64), - &root(), + &root, MaskFuture::new_true(values_len), ) .vortex_expect("must construct dict values array evaluation") @@ -125,11 +129,14 @@ impl DictReader { // We capture the name, so it may be wrong if we re-use the same reader within multiple // different parent readers. But that's rare... let values_len = self.values_len; + let root = root() + .bind(self.values.dtype()) + .vortex_expect("root must bind against the dictionary values dtype"); self.values_array.get().cloned().unwrap_or_else(|| { self.values .projection_evaluation( &(0..values_len as u64), - &root(), + &root, MaskFuture::new_true(values_len), ) .vortex_expect("must construct dict values array evaluation") @@ -139,7 +146,7 @@ impl DictReader { }) } - fn values_eval(&self, expr: Expression) -> SharedArrayFuture { + fn values_eval(&self, expr: BoundExpression) -> SharedArrayFuture { // This is unsound since we cannot be sure that all the values are referenced in the query // after applying the filter, so if the expression is fallible this might fail when it // shouldn't. @@ -155,7 +162,7 @@ impl DictReader { .or_insert_with(|| { self.values_array_uncanonical() .map(move |array| { - let array = array?.apply(&expr)?; + let array = array?.apply_bound(&expr)?; Ok(SharedArray::new(array).into_array()) }) .boxed() @@ -228,7 +235,7 @@ impl LayoutReader for DictReader { fn pruning_evaluation( &self, _row_range: &Range, - _expr: &Expression, + _expr: &BoundExpression, mask: Mask, ) -> VortexResult { // NOTE: we can get the values here, convert expression to the codes domain, and push down @@ -241,7 +248,7 @@ impl LayoutReader for DictReader { fn filter_evaluation( &self, row_range: &Range, - expr: &Expression, + expr: &BoundExpression, mask: MaskFuture, ) -> VortexResult { // TODO(joe): fix up expr partitioning with fallibility and strictness annotations @@ -250,11 +257,10 @@ impl LayoutReader for DictReader { // We register interest on the entire codes row_range for now, there // is no straightforward shift into the codes domain we can do to the expression // without reading values. - let codes_eval = self.codes.projection_evaluation( - row_range, - &root(), - MaskFuture::new_true(mask.len()), - )?; + let root = root().bind(self.codes.dtype())?; + let codes_eval = + self.codes + .projection_evaluation(row_range, &root, MaskFuture::new_true(mask.len()))?; let session = self.session.clone(); @@ -273,36 +279,39 @@ impl LayoutReader for DictReader { fn projection_evaluation( &self, row_range: &Range, - expr: &Expression, + expr: &BoundExpression, mask: MaskFuture, ) -> VortexResult>> { // TODO: fix up expr partitioning with fallibility and strictness annotations + let codes_root = root().bind(self.codes.dtype())?; let codes_eval = self .codes - .projection_evaluation(row_range, &root(), mask) + .projection_evaluation(row_range, &codes_root, mask) .map_err(|err| err.with_context("While evaluating projection on codes"))?; - let (expr_outer, expr_inner) = split_expression_for_pushdown(expr.clone(), self.dtype())?; + let (expr_outer, expr_inner) = split_expression_for_pushdown(expr.unbind(), self.dtype())?; let values_eval = if let Some(inner) = expr_inner { // "outer" takes a struct field with PUSHDOWN_ANNOTATION name, so // pack inner with this name as well - let inner = pack([(PUSHDOWN_ANNOTATION, inner)], Nullability::NonNullable); + let inner = pack([(PUSHDOWN_ANNOTATION, inner)], Nullability::NonNullable) + .bind(self.values.dtype())?; // We can't use values_eval as it uses values_array_uncanonical // which in turn gets populated from self.values. If // self.values_array() is called first, it will populate // self.values with uncompressed data. Supply uncached data let values_len = self.values_len; + let values_root = root().bind(self.values.dtype())?; self.values .projection_evaluation( &(0..values_len as u64), - &root(), + &values_root, MaskFuture::new_true(values_len), ) .vortex_expect("must construct dict values array evaluation") .map_err(Arc::new) - .map(move |array| Ok(SharedArray::new(array?.apply(&inner)?).into_array())) + .map(move |array| Ok(SharedArray::new(array?.apply_bound(&inner)?).into_array())) .boxed() .shared() } else { @@ -324,7 +333,8 @@ impl LayoutReader for DictReader { .into_array() .optimize()?; - array.apply(&expr_outer) + let expr_outer = expr_outer.bind(array.dtype())?; + array.apply_bound(&expr_outer) } .boxed()) } @@ -483,9 +493,11 @@ mod tests { Nullability::NonNullable, ); assert!(layout.encoding_id() == LayoutId::new("vortex.dict")); - let actual = layout + let reader = layout .new_reader("".into(), segments, &session, &Default::default()) - .unwrap() + .unwrap(); + let expression = expression.bind(reader.dtype()).unwrap(); + let actual = reader .projection_evaluation( &(0..layout.row_count()), &expression, @@ -569,9 +581,11 @@ mod tests { Nullability::Nullable, )), ); - let mask = layout + let reader = layout .new_reader("".into(), segments, &session, &Default::default()) - .unwrap() + .unwrap(); + let filter = filter.bind(reader.dtype()).unwrap(); + let mask = reader .filter_evaluation(&(0..3), &filter, MaskFuture::new_true(3)) .unwrap() .await @@ -632,9 +646,11 @@ mod tests { let expression = is_not_null(root()); assert_eq!(layout.encoding_id(), LayoutId::new("vortex.dict")); - let actual = layout + let reader = layout .new_reader("".into(), segments, &session, &Default::default()) - .unwrap() + .unwrap(); + let expression = expression.bind(reader.dtype()).unwrap(); + let actual = reader .projection_evaluation( &(0..layout.row_count()), &expression, @@ -687,12 +703,14 @@ mod tests { let mut ctx = session.create_execution_ctx(); let (layout, segments) = write_dict_layout(array, &session).await; assert_eq!(layout.encoding_id(), LayoutId::new("vortex.dict")); - let actual = layout + let reader = layout .new_reader("".into(), segments, &session, &Default::default()) - .unwrap() + .unwrap(); + let expression = byte_length(root()).bind(reader.dtype()).unwrap(); + let actual = reader .projection_evaluation( &(0..layout.row_count()), - &byte_length(root()), + &expression, MaskFuture::new_true(layout.row_count().try_into().unwrap()), ) .unwrap() diff --git a/vortex-layout/src/layouts/flat/reader.rs b/vortex-layout/src/layouts/flat/reader.rs index 12191b35429..aa7609f1659 100644 --- a/vortex-layout/src/layouts/flat/reader.rs +++ b/vortex-layout/src/layouts/flat/reader.rs @@ -13,7 +13,7 @@ use vortex_array::MaskFuture; use vortex_array::VortexSessionExecute; use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; -use vortex_array::expr::Expression; +use vortex_array::expr::BoundExpression; use vortex_array::serde::SerializedArray; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -115,7 +115,7 @@ impl LayoutReader for FlatReader { fn pruning_evaluation( &self, _row_range: &Range, - _expr: &Expression, + _expr: &BoundExpression, mask: Mask, ) -> VortexResult { Ok(MaskFuture::ready(mask)) @@ -124,7 +124,7 @@ impl LayoutReader for FlatReader { fn filter_evaluation( &self, row_range: &Range, - expr: &Expression, + expr: &BoundExpression, mask: MaskFuture, ) -> VortexResult { let row_range = usize::try_from(row_range.start) @@ -153,7 +153,7 @@ impl LayoutReader for FlatReader { // We have the choice to apply the filter or the expression first, we apply the // expression first so that it can try pushing down itself and then the filter // after this. - let array = array.apply(&expr)?; + let array = array.apply_bound(&expr)?; let array = array.filter(mask.clone())?; let mut ctx = session.create_execution_ctx(); let array_mask = array.null_as_false().execute(&mut ctx)?; @@ -161,7 +161,7 @@ impl LayoutReader for FlatReader { mask.intersect_by_rank(&array_mask) } else { // Run over the full array, with a simpler bitand at the end. - let array = array.apply(&expr)?; + let array = array.apply_bound(&expr)?; let mut ctx = session.create_execution_ctx(); let array_mask = array.null_as_false().execute(&mut ctx)?; @@ -183,7 +183,7 @@ impl LayoutReader for FlatReader { fn projection_evaluation( &self, row_range: &Range, - expr: &Expression, + expr: &BoundExpression, mask: MaskFuture, ) -> VortexResult>> { let row_range = usize::try_from(row_range.start) @@ -214,7 +214,7 @@ impl LayoutReader for FlatReader { } // Evaluate the projection expression. - array = array.apply(&expr)?; + array = array.apply_bound(&expr)?; Ok(array) } @@ -278,11 +278,12 @@ mod test { "vortex.flat(i32?, rows=5, segments=[0])" ); - let result = layout - .new_reader("".into(), segments, &session, &Default::default())? + let reader = layout.new_reader("".into(), segments, &session, &Default::default())?; + let expr = root().bind(reader.dtype())?; + let result = reader .projection_evaluation( &(0..layout.row_count()), - &root(), + &expr, MaskFuture::new_true(layout.row_count().try_into()?), )? .await?; @@ -315,10 +316,11 @@ mod test { .await .unwrap(); - let expr = gt(root(), lit(3i32)); - let result = layout + let reader = layout .new_reader("".into(), segments, &session, &Default::default()) - .unwrap() + .unwrap(); + let expr = gt(root(), lit(3i32)).bind(reader.dtype()).unwrap(); + let result = reader .projection_evaluation( &(0..layout.row_count()), &expr, @@ -354,10 +356,12 @@ mod test { .await .unwrap(); - let result = layout + let reader = layout .new_reader("".into(), segments, &session, &Default::default()) - .unwrap() - .projection_evaluation(&(2..4), &root(), MaskFuture::new_true(2)) + .unwrap(); + let expr = root().bind(reader.dtype()).unwrap(); + let result = reader + .projection_evaluation(&(2..4), &expr, MaskFuture::new_true(2)) .unwrap() .await .unwrap(); diff --git a/vortex-layout/src/layouts/flat/writer.rs b/vortex-layout/src/layouts/flat/writer.rs index 83c2101f12c..9761c71f9ae 100644 --- a/vortex-layout/src/layouts/flat/writer.rs +++ b/vortex-layout/src/layouts/flat/writer.rs @@ -239,12 +239,14 @@ mod tests { .await .unwrap(); - let result = layout + let reader = layout .new_reader("".into(), segments, &SESSION, &Default::default()) - .unwrap() + .unwrap(); + let expr = root().bind(reader.dtype()).unwrap(); + let result = reader .projection_evaluation( &(0..layout.row_count()), - &root(), + &expr, MaskFuture::new_true(layout.row_count().try_into().unwrap()), ) .unwrap() @@ -290,12 +292,14 @@ mod tests { .await .unwrap(); - let result = layout + let reader = layout .new_reader("".into(), segments, &SESSION, &Default::default()) - .unwrap() + .unwrap(); + let expr = root().bind(reader.dtype()).unwrap(); + let result = reader .projection_evaluation( &(0..layout.row_count()), - &root(), + &expr, MaskFuture::new_true(layout.row_count().try_into().unwrap()), ) .unwrap() @@ -363,12 +367,14 @@ mod tests { }; // We should be able to read the array we just wrote. - let result: ArrayRef = layout + let reader = layout .new_reader("".into(), segments, &SESSION, &Default::default()) - .unwrap() + .unwrap(); + let expr = root().bind(reader.dtype()).unwrap(); + let result: ArrayRef = reader .projection_evaluation( &(0..layout.row_count()), - &root(), + &expr, MaskFuture::new_true(layout.row_count().try_into().unwrap()), ) .unwrap() diff --git a/vortex-layout/src/layouts/list/expr.rs b/vortex-layout/src/layouts/list/expr.rs index 152b7ea5dce..39e2d1d604e 100644 --- a/vortex-layout/src/layouts/list/expr.rs +++ b/vortex-layout/src/layouts/list/expr.rs @@ -1,13 +1,19 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::expr::BoundExpression; +#[cfg(test)] use vortex_array::expr::Expression; +#[cfg(test)] use vortex_array::expr::is_root; -use vortex_array::expr::not; -use vortex_array::expr::root; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ScalarFnVTableExt; use vortex_array::scalar_fn::fns::is_not_null::IsNotNull; use vortex_array::scalar_fn::fns::is_null::IsNull; use vortex_array::scalar_fn::fns::list_length::ListLength; +use vortex_array::scalar_fn::fns::not::Not; use vortex_error::VortexResult; /// The minimal set of list children an expression needs for evaluation. @@ -27,6 +33,7 @@ pub(super) enum ListChildrenNeeded { } /// The minimal set of list children needed to evaluate `expr`, where `root()` is a field with list dtype. +#[cfg(test)] pub(super) fn get_necessary_list_children(expr: &Expression) -> ListChildrenNeeded { if is_null_root(expr) { return ListChildrenNeeded::Validity; @@ -49,12 +56,48 @@ pub(super) fn get_necessary_list_children(expr: &Expression) -> ListChildrenNeed .unwrap_or(ListChildrenNeeded::Validity) } +/// The minimal set of list children needed to evaluate a bound expression. +pub(super) fn get_necessary_bound_list_children(expr: &BoundExpression) -> ListChildrenNeeded { + if is_bound_null_root(expr) { + return ListChildrenNeeded::Validity; + } + + if is_bound_list_length_root(expr) { + return ListChildrenNeeded::OffsetsAndValidity; + } + + if expr.is_root() { + return ListChildrenNeeded::All; + } + + expr.children() + .iter() + .map(get_necessary_bound_list_children) + .max() + .unwrap_or(ListChildrenNeeded::Validity) +} + +fn is_bound_null_root(expr: &BoundExpression) -> bool { + (expr.as_scalar().is_some_and(|f| f.is::()) + || expr.as_scalar().is_some_and(|f| f.is::())) + && expr.children().len() == 1 + && expr.children()[0].is_root() +} + +fn is_bound_list_length_root(expr: &BoundExpression) -> bool { + expr.as_scalar().is_some_and(|f| f.is::()) + && expr.children().len() == 1 + && expr.children()[0].is_root() +} + +#[cfg(test)] fn is_null_root(expr: &Expression) -> bool { (expr.is::() || expr.is::()) && expr.children().len() == 1 && is_root(expr.child(0)) } +#[cfg(test)] fn is_list_length_root(expr: &Expression) -> bool { expr.is::() && expr.children().len() == 1 && is_root(expr.child(0)) } @@ -62,17 +105,38 @@ fn is_list_length_root(expr: &Expression) -> bool { /// Rewrite a validity-class expression so it can be evaluated against the list's validity bool /// array (`true` == valid row): `is_not_null(root())` becomes `root()` and `is_null(root())` /// becomes `not(root())`. All other nodes are rebuilt with rewritten children. -pub(super) fn rewrite_validity_expr(expr: &Expression) -> VortexResult { - if expr.is::() && expr.children().len() == 1 && is_root(expr.child(0)) { - return Ok(root()); +pub(super) fn rewrite_validity_expr(expr: &BoundExpression) -> VortexResult { + let validity_dtype = DType::Bool(Nullability::NonNullable); + rewrite_validity_expr_with_root(expr, &validity_dtype) +} + +fn rewrite_validity_expr_with_root( + expr: &BoundExpression, + root_dtype: &DType, +) -> VortexResult { + if expr.as_scalar().is_some_and(|f| f.is::()) + && expr.children().len() == 1 + && expr.children()[0].is_root() + { + return Ok(BoundExpression::new_root(root_dtype.clone())); } - if expr.is::() && expr.children().len() == 1 && is_root(expr.child(0)) { - return Ok(not(root())); + if expr.as_scalar().is_some_and(|f| f.is::()) + && expr.children().len() == 1 + && expr.children()[0].is_root() + { + return BoundExpression::try_new( + Not.bind(EmptyOptions), + [BoundExpression::new_root(root_dtype.clone())], + ); + } + if expr.is_root() { + return Ok(BoundExpression::new_root(root_dtype.clone())); } + let children = expr .children() .iter() - .map(rewrite_validity_expr) + .map(|child| rewrite_validity_expr_with_root(child, root_dtype)) .collect::>>()?; expr.clone().with_children(children) } @@ -81,15 +145,18 @@ pub(super) fn rewrite_validity_expr(expr: &Expression) -> VortexResult VortexResult { - if is_list_length_root(expr) { - return Ok(root()); +pub(super) fn rewrite_offsets_expr( + expr: &BoundExpression, + lengths_dtype: &DType, +) -> VortexResult { + if is_bound_list_length_root(expr) || expr.is_root() { + return Ok(BoundExpression::new_root(lengths_dtype.clone())); } let children = expr .children() .iter() - .map(rewrite_offsets_expr) + .map(|child| rewrite_offsets_expr(child, lengths_dtype)) .collect::>>()?; expr.clone().with_children(children) } diff --git a/vortex-layout/src/layouts/list/reader.rs b/vortex-layout/src/layouts/list/reader.rs index 4bea8e74ca0..53227635de6 100644 --- a/vortex-layout/src/layouts/list/reader.rs +++ b/vortex-layout/src/layouts/list/reader.rs @@ -19,7 +19,7 @@ use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; -use vortex_array::expr::Expression; +use vortex_array::expr::BoundExpression; use vortex_array::expr::root; use vortex_array::scalar_fn::fns::operators::Operator; use vortex_array::validity::Validity; @@ -36,7 +36,7 @@ use crate::RowSplits; use crate::SplitRange; use crate::layouts::list::ListLayout; use crate::layouts::list::expr::ListChildrenNeeded; -use crate::layouts::list::expr::get_necessary_list_children; +use crate::layouts::list::expr::get_necessary_bound_list_children; use crate::layouts::list::expr::rewrite_offsets_expr; use crate::layouts::list::expr::rewrite_validity_expr; use crate::segments::SegmentSource; @@ -108,7 +108,7 @@ impl ListReader { fn project_validity( &self, row_range: &Range, - expr: &Expression, + expr: &BoundExpression, mask: MaskFuture, ) -> VortexResult { let validity_reader = self.validity.clone(); @@ -127,16 +127,19 @@ impl ListReader { }; let validity_array = match validity_reader.as_ref() { - Some(v) => Some( - v.projection_evaluation(&row_range, &root(), MaskFuture::ready(mask))? - .await?, - ), + Some(v) => { + let root = root().bind(v.dtype())?; + Some( + v.projection_evaluation(&row_range, &root, MaskFuture::ready(mask))? + .await?, + ) + } None => None, }; let validity = create_validity(validity_array, nullability).to_array(out_len); - validity.apply(&rewritten) + validity.apply_bound(&rewritten) } .boxed()) } @@ -148,7 +151,7 @@ impl ListReader { fn project_all( &self, row_range: &Range, - expr: &Expression, + expr: &BoundExpression, mask: MaskFuture, ) -> VortexResult { let is_full_range = row_range.start == 0 && row_range.end == self.layout.row_count(); @@ -167,7 +170,7 @@ impl ListReader { } /// Fetch the complete `elements`, `offsets`, and `validity` children concurrently. - fn project_all_full(&self, expr: &Expression) -> VortexResult { + fn project_all_full(&self, expr: &BoundExpression) -> VortexResult { let row_count = self.layout.row_count(); let elements_row_count = self.elements.row_count(); let nullability = self.layout.dtype().nullability(); @@ -189,7 +192,7 @@ impl ListReader { ListArray::new_unchecked(elements, offsets, create_validity(validity, nullability)) } .into_array(); - list.apply(&expr) + list.apply_bound(&expr) } .boxed()) } @@ -202,14 +205,13 @@ impl ListReader { fn project_all_bounded( &self, row_range: &Range, - expr: &Expression, + expr: &BoundExpression, mask: Mask, ) -> VortexResult { // Crop to the smallest contiguous row range containing every selected list. let Some(selected_rows) = selected_row_range(&mask) else { - let empty = Canonical::empty(self.layout.dtype()).into_array(); - let expr = expr.clone(); - return Ok(async move { empty.apply(&expr) }.boxed()); + let empty = Canonical::empty(expr.dtype()).into_array(); + return Ok(async move { Ok(empty) }.boxed()); }; let selected_mask = mask.slice(selected_rows.clone()); @@ -246,7 +248,7 @@ impl ListReader { } else { list.filter(selected_mask)? }; - list.apply(&expr) + list.apply_bound(&expr) } .boxed()) } @@ -255,13 +257,14 @@ impl ListReader { fn project_offsets_validity( &self, row_range: &Range, - expr: &Expression, + expr: &BoundExpression, mask: MaskFuture, ) -> VortexResult { let offsets = self.fetch_raw_offsets(row_range)?; let reader = self.clone(); let row_range = row_range.clone(); - let rewritten = rewrite_offsets_expr(expr)?; + let lengths_dtype = DType::Primitive(PType::U64, self.layout.dtype().nullability()); + let rewritten = rewrite_offsets_expr(expr, &lengths_dtype)?; Ok(async move { let mask = mask.await?; @@ -285,7 +288,7 @@ impl ListReader { let validity = validity_fut.await?; let lengths = apply_lengths_validity(lengths, validity, nullability)?; - lengths.apply(&rewritten) + lengths.apply_bound(&rewritten) } .boxed()) } @@ -297,9 +300,10 @@ impl ListReader { fn fetch_raw_offsets(&self, row_range: &Range) -> VortexResult { let offsets_range = row_range.start..(row_range.end + 1); let offsets_count = usize::try_from(offsets_range.end - offsets_range.start)?; + let root = root().bind(self.offsets.dtype())?; self.offsets.projection_evaluation( &offsets_range, - &root(), + &root, MaskFuture::new_true(offsets_count), ) } @@ -309,8 +313,9 @@ impl ListReader { /// No mask or expression is applied. fn fetch_raw_elements(&self, row_range: &Range) -> VortexResult { let row_count = usize::try_from(row_range.end - row_range.start)?; + let root = root().bind(self.elements.dtype())?; self.elements - .projection_evaluation(row_range, &root(), MaskFuture::new_true(row_count)) + .projection_evaluation(row_range, &root, MaskFuture::new_true(row_count)) } } @@ -407,7 +412,7 @@ impl LayoutReader for ListReader { fn pruning_evaluation( &self, _row_range: &Range, - _expr: &Expression, + _expr: &BoundExpression, mask: Mask, ) -> VortexResult { Ok(MaskFuture::ready(mask)) @@ -416,7 +421,7 @@ impl LayoutReader for ListReader { fn filter_evaluation( &self, row_range: &Range, - expr: &Expression, + expr: &BoundExpression, mask: MaskFuture, ) -> VortexResult { let len = mask.len(); @@ -455,11 +460,11 @@ impl LayoutReader for ListReader { fn projection_evaluation( &self, row_range: &Range, - expr: &Expression, + expr: &BoundExpression, mask: MaskFuture, ) -> VortexResult { // Read as little as possible based on which list children the expression needs. - match get_necessary_list_children(expr) { + match get_necessary_bound_list_children(expr) { ListChildrenNeeded::Validity => self.project_validity(row_range, expr, mask), ListChildrenNeeded::OffsetsAndValidity => { self.project_offsets_validity(row_range, expr, mask) @@ -523,7 +528,10 @@ fn fetch_validity( mask: MaskFuture, ) -> VortexResult { let fut = validity - .map(|v| v.projection_evaluation(row_range, &root(), mask)) + .map(|v| { + let root = root().bind(v.dtype())?; + v.projection_evaluation(row_range, &root, mask) + }) .transpose()?; Ok(async move { match fut { @@ -608,6 +616,7 @@ mod tests { use vortex_array::arrays::ListArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; + use vortex_array::expr::Expression; use vortex_array::expr::cast; use vortex_array::expr::gt; use vortex_array::expr::is_not_null; @@ -651,14 +660,16 @@ mod tests { let (segments, layout, session) = write_layout(&flat_list_strategy(), list).await?; let reader = layout.new_reader("".into(), segments, &session, &ctx)?; + let not_null_expr = is_not_null(root()).bind(reader.dtype())?; let not_null = reader - .projection_evaluation(&(0..3), &is_not_null(root()), MaskFuture::new_true(3))? + .projection_evaluation(&(0..3), ¬_null_expr, MaskFuture::new_true(3))? .await?; let mut exec_ctx = session.create_execution_ctx(); assert_arrays_eq!(not_null, BoolArray::from_iter(valid.clone()), &mut exec_ctx); + let is_null_expr = is_null(root()).bind(reader.dtype())?; let is_null_res = reader - .projection_evaluation(&(0..3), &is_null(root()), MaskFuture::new_true(3))? + .projection_evaluation(&(0..3), &is_null_expr, MaskFuture::new_true(3))? .await?; assert_arrays_eq!( is_null_res, @@ -676,8 +687,9 @@ mod tests { let (segments, layout, session) = write_layout(&flat_list_strategy(), list).await?; let reader = layout.new_reader("".into(), segments, &session, &ctx)?; + let expr = list_length(root()).bind(reader.dtype())?; let result = reader - .projection_evaluation(&(0..3), &list_length(root()), MaskFuture::new_true(3))? + .projection_evaluation(&(0..3), &expr, MaskFuture::new_true(3))? .await?; let mut exec_ctx = session.create_execution_ctx(); @@ -692,8 +704,9 @@ mod tests { let (segments, layout, session) = write_layout(&flat_list_strategy(), list).await?; let reader = layout.new_reader("".into(), segments, &session, &ctx)?; + let expr = list_length(root()).bind(reader.dtype())?; let result = reader - .projection_evaluation(&(0..3), &list_length(root()), MaskFuture::new_true(3))? + .projection_evaluation(&(0..3), &expr, MaskFuture::new_true(3))? .await?; let expected = @@ -711,8 +724,9 @@ mod tests { let reader = layout.new_reader("".into(), segments, &session, &ctx)?; let mask = Mask::from_iter([false, true, true]); + let expr = list_length(root()).bind(reader.dtype())?; let result = reader - .projection_evaluation(&(0..3), &list_length(root()), MaskFuture::ready(mask))? + .projection_evaluation(&(0..3), &expr, MaskFuture::ready(mask))? .await?; let expected = PrimitiveArray::from_option_iter::([None, Some(1)]).into_array(); @@ -731,7 +745,8 @@ mod tests { let expr = cast( list_length(root()), DType::Primitive(PType::I64, Nullability::Nullable), - ); + ) + .bind(reader.dtype())?; let result = reader .projection_evaluation(&(0..3), &expr, MaskFuture::new_true(3))? .await?; @@ -750,12 +765,9 @@ mod tests { let (segments, layout, session) = write_layout(&flat_list_strategy(), list).await?; let reader = layout.new_reader("".into(), segments, &session, &ctx)?; + let expr = gt(list_length(root()), lit(1u64)).bind(reader.dtype())?; let result = reader - .filter_evaluation( - &(0..3), - >(list_length(root()), lit(1u64)), - MaskFuture::new_true(3), - )? + .filter_evaluation(&(0..3), &expr, MaskFuture::new_true(3))? .await?; assert_eq!(result, Mask::from_iter([true, false, false])); @@ -778,6 +790,7 @@ mod tests { let (segments, layout, session) = write_layout(&flat_list_strategy(), list).await?; let reader = layout.new_reader("".into(), segments, &session, &ctx)?; + let expr = expr.bind(reader.dtype())?; let result = reader .filter_evaluation(&(0..3), &expr, MaskFuture::new_true(3))? .await?; @@ -794,8 +807,9 @@ mod tests { let reader = layout.new_reader("".into(), segments, &session, &ctx)?; let input_mask = Mask::from_iter([true, true, false]); + let expr = is_not_null(root()).bind(reader.dtype())?; let result = reader - .filter_evaluation(&(0..3), &is_not_null(root()), MaskFuture::ready(input_mask))? + .filter_evaluation(&(0..3), &expr, MaskFuture::ready(input_mask))? .await?; assert_eq!(result, Mask::from_iter([true, false, false])); @@ -810,8 +824,9 @@ mod tests { let reader = layout.new_reader("".into(), segments, &session, &ctx)?; let input_mask = Mask::from_iter([false, false, false, false, true, false]); + let expr = is_not_null(root()).bind(reader.dtype())?; let result = reader - .filter_evaluation(&(0..6), &is_not_null(root()), MaskFuture::ready(input_mask))? + .filter_evaluation(&(0..6), &expr, MaskFuture::ready(input_mask))? .await?; assert_eq!( @@ -928,8 +943,9 @@ mod tests { let (segments, layout, session) = write_layout(&flat_list_strategy(), list.clone()).await?; let reader = layout.new_reader("".into(), segments, &session, &ctx)?; + let expr = root().bind(reader.dtype())?; let result = reader - .projection_evaluation(&row_range, &root(), MaskFuture::new_true(len))? + .projection_evaluation(&row_range, &expr, MaskFuture::new_true(len))? .await?; let expected = @@ -947,8 +963,9 @@ mod tests { let reader = layout.new_reader("".into(), segments, &session, &ctx)?; let mask = Mask::from_iter([true, false, true]); + let expr = root().bind(reader.dtype())?; let result = reader - .projection_evaluation(&(0..3), &root(), MaskFuture::ready(mask.clone()))? + .projection_evaluation(&(0..3), &expr, MaskFuture::ready(mask.clone()))? .await?; let expected = list.filter(mask)?; @@ -992,8 +1009,9 @@ mod tests { let (segments, layout, session) = write_layout(&flat_list_strategy(), list.clone()).await?; let reader = layout.new_reader("".into(), segments, &session, &ctx)?; + let expr = root().bind(reader.dtype())?; let result = reader - .projection_evaluation(&(0..5), &root(), MaskFuture::ready(mask.clone()))? + .projection_evaluation(&(0..5), &expr, MaskFuture::ready(mask.clone()))? .await?; let expected = list.filter(mask)?; @@ -1012,8 +1030,9 @@ mod tests { let (segments, layout, session) = write_layout(&flat_list_strategy(), list.clone()).await?; let reader = layout.new_reader("".into(), segments, &session, &ctx)?; + let expr = root().bind(reader.dtype())?; let result = reader - .projection_evaluation(&(1..4), &root(), MaskFuture::new_true(3))? + .projection_evaluation(&(1..4), &expr, MaskFuture::new_true(3))? .await?; let expected = list.slice(1..4)?; @@ -1133,8 +1152,9 @@ mod tests { let reader = layout.new_reader("".into(), source, &session, &ctx)?; let mask = Mask::from_iter([true, false, false, false, false]); + let expr = root().bind(reader.dtype())?; let result = reader - .projection_evaluation(&(0..5), &root(), MaskFuture::ready(mask.clone()))? + .projection_evaluation(&(0..5), &expr, MaskFuture::ready(mask.clone()))? .await?; let expected = list.filter(mask)?; @@ -1172,8 +1192,9 @@ mod tests { write_layout(&chunked_elements_list_strategy(), list.clone()).await?; let reader = layout.new_reader("".into(), segments, &session, &ctx)?; + let expr = root().bind(reader.dtype())?; let result = reader - .projection_evaluation(&row_range, &root(), MaskFuture::ready(mask.clone()))? + .projection_evaluation(&row_range, &expr, MaskFuture::ready(mask.clone()))? .await?; let sliced = diff --git a/vortex-layout/src/layouts/partitioned.rs b/vortex-layout/src/layouts/partitioned.rs index 76327a7ee5f..09636c4c339 100644 --- a/vortex-layout/src/layouts/partitioned.rs +++ b/vortex-layout/src/layouts/partitioned.rs @@ -89,8 +89,9 @@ impl PartitionedExprEval

for PartitionedExpr

{ .into_array(); let mut ctx = session.create_execution_ctx(); + let root = self.root.bind(root_scope.dtype())?; let root_mask = root_scope - .apply(&self.root)? + .apply_bound(&root)? .null_as_false() .execute(&mut ctx)?; @@ -126,7 +127,8 @@ impl PartitionedExprEval

for PartitionedExpr

{ )? .into_array(); - root_scope.apply(&self.root) + let root = self.root.bind(root_scope.dtype())?; + root_scope.apply_bound(&root) })) } } diff --git a/vortex-layout/src/layouts/row_idx/mod.rs b/vortex-layout/src/layouts/row_idx/mod.rs index 7eaa820ce2a..f4d2da9ec7a 100644 --- a/vortex-layout/src/layouts/row_idx/mod.rs +++ b/vortex-layout/src/layouts/row_idx/mod.rs @@ -24,7 +24,8 @@ use vortex_array::dtype::FieldMask; use vortex_array::dtype::FieldName; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; -use vortex_array::expr::ExactExpr; +use vortex_array::expr::BoundExpression; +use vortex_array::expr::ExactBoundExpr; use vortex_array::expr::Expression; use vortex_array::expr::is_root; use vortex_array::expr::root; @@ -50,7 +51,7 @@ pub struct RowIdxLayoutReader { name: Arc, row_offset: u64, child: Arc, - partition_cache: DashMap>>, + partition_cache: DashMap>>, session: VortexSession, } @@ -65,8 +66,8 @@ impl RowIdxLayoutReader { } } - fn partition_expr(&self, expr: &Expression) -> VortexResult { - let key = ExactExpr(expr.clone()); + fn partition_expr(&self, expr: &BoundExpression) -> VortexResult { + let key = ExactBoundExpr(expr.clone()); // Check cache first with read-only lock. if let Some(entry) = self.partition_cache.get(&key) @@ -75,7 +76,7 @@ impl RowIdxLayoutReader { return Ok(partitioning.clone()); } - let result = self.compute_partitioning(expr)?; + let result = self.compute_partitioning(&expr.unbind())?; self.partition_cache .entry(key) @@ -180,18 +181,24 @@ impl LayoutReader for RowIdxLayoutReader { fn pruning_evaluation( &self, row_range: &Range, - expr: &Expression, + expr: &BoundExpression, mask: Mask, ) -> VortexResult { Ok(match &self.partition_expr(expr)? { - Partitioning::RowIdx(expr) => row_idx_mask_future( - self.row_offset, - row_range, - expr, - MaskFuture::ready(mask), - self.session.clone(), - ), - Partitioning::Child(expr) => self.child.pruning_evaluation(row_range, expr, mask)?, + Partitioning::RowIdx(expr) => { + let expr = expr.bind(&row_idx_dtype())?; + row_idx_mask_future( + self.row_offset, + row_range, + expr, + MaskFuture::ready(mask), + self.session.clone(), + ) + } + Partitioning::Child(expr) => { + let expr = expr.bind(self.child.dtype())?; + self.child.pruning_evaluation(row_range, &expr, mask)? + } Partitioning::Partitioned(..) => MaskFuture::ready(mask), }) } @@ -199,35 +206,50 @@ impl LayoutReader for RowIdxLayoutReader { fn filter_evaluation( &self, row_range: &Range, - expr: &Expression, + expr: &BoundExpression, mask: MaskFuture, ) -> VortexResult { match &self.partition_expr(expr)? { // Since this is run during pruning, we skip re-evaluating the row index expression // during the filter evaluation. Partitioning::RowIdx(_) => Ok(mask), - Partitioning::Child(expr) => self.child.filter_evaluation(row_range, expr, mask), + Partitioning::Child(expr) => { + let expr = expr.bind(self.child.dtype())?; + self.child.filter_evaluation(row_range, &expr, mask) + } Partitioning::Partitioned(p) => Arc::clone(p).into_mask_future( mask, |annotation, expr, mask| match annotation { - Partition::RowIdx => Ok(row_idx_mask_future( - self.row_offset, - row_range, - expr, - mask, - self.session.clone(), - )), - Partition::Child => self.child.filter_evaluation(row_range, expr, mask), + Partition::RowIdx => { + let expr = expr.bind(&row_idx_dtype())?; + Ok(row_idx_mask_future( + self.row_offset, + row_range, + expr, + mask, + self.session.clone(), + )) + } + Partition::Child => { + let expr = expr.bind(self.child.dtype())?; + self.child.filter_evaluation(row_range, &expr, mask) + } }, |annotation, expr, mask| match annotation { - Partition::RowIdx => Ok(row_idx_array_future( - self.row_offset, - row_range, - expr, - mask, - self.session.clone(), - )), - Partition::Child => self.child.projection_evaluation(row_range, expr, mask), + Partition::RowIdx => { + let expr = expr.bind(&row_idx_dtype())?; + Ok(row_idx_array_future( + self.row_offset, + row_range, + expr, + mask, + self.session.clone(), + )) + } + Partition::Child => { + let expr = expr.bind(self.child.dtype())?; + self.child.projection_evaluation(row_range, &expr, mask) + } }, self.session.clone(), ), @@ -237,28 +259,40 @@ impl LayoutReader for RowIdxLayoutReader { fn projection_evaluation( &self, row_range: &Range, - expr: &Expression, + expr: &BoundExpression, mask: MaskFuture, ) -> VortexResult>> { match &self.partition_expr(expr)? { - Partitioning::RowIdx(expr) => Ok(row_idx_array_future( - self.row_offset, - row_range, - expr, - mask, - self.session.clone(), - )), - Partitioning::Child(expr) => self.child.projection_evaluation(row_range, expr, mask), + Partitioning::RowIdx(expr) => { + let expr = expr.bind(&row_idx_dtype())?; + Ok(row_idx_array_future( + self.row_offset, + row_range, + expr, + mask, + self.session.clone(), + )) + } + Partitioning::Child(expr) => { + let expr = expr.bind(self.child.dtype())?; + self.child.projection_evaluation(row_range, &expr, mask) + } Partitioning::Partitioned(p) => { Arc::clone(p).into_array_future(mask, |annotation, expr, mask| match annotation { - Partition::RowIdx => Ok(row_idx_array_future( - self.row_offset, - row_range, - expr, - mask, - self.session.clone(), - )), - Partition::Child => self.child.projection_evaluation(row_range, expr, mask), + Partition::RowIdx => { + let expr = expr.bind(&row_idx_dtype())?; + Ok(row_idx_array_future( + self.row_offset, + row_range, + expr, + mask, + self.session.clone(), + )) + } + Partition::Child => { + let expr = expr.bind(self.child.dtype())?; + self.child.projection_evaluation(row_range, &expr, mask) + } }) } } @@ -269,6 +303,10 @@ impl LayoutReader for RowIdxLayoutReader { } } +fn row_idx_dtype() -> DType { + DType::Primitive(PType::U64, NonNullable) +} + // Returns a SequenceArray representing the row indices for the given row range, fn idx_array(row_offset: u64, row_range: &Range) -> SequenceArray { Sequence::try_new( @@ -285,17 +323,19 @@ fn idx_array(row_offset: u64, row_range: &Range) -> SequenceArray { fn row_idx_mask_future( row_offset: u64, row_range: &Range, - expr: &Expression, + expr: BoundExpression, mask: MaskFuture, session: VortexSession, ) -> MaskFuture { let row_range = row_range.clone(); - let expr = expr.clone(); MaskFuture::new(mask.len(), async move { let array = idx_array(row_offset, &row_range).into_array(); let mut ctx = session.create_execution_ctx(); - let result_mask = array.apply(&expr)?.null_as_false().execute(&mut ctx)?; + let result_mask = array + .apply_bound(&expr)? + .null_as_false() + .execute(&mut ctx)?; Ok(result_mask.bitand(&mask.await?)) }) @@ -304,18 +344,17 @@ fn row_idx_mask_future( fn row_idx_array_future( row_offset: u64, row_range: &Range, - expr: &Expression, + expr: BoundExpression, mask: MaskFuture, session: VortexSession, ) -> ArrayFuture { let row_range = row_range.clone(); - let expr = expr.clone(); async move { let array = idx_array(row_offset, &row_range).into_array(); let filtered = array.filter(mask.await?)?; let mut ctx = session.create_execution_ctx(); let array = filtered.execute::(&mut ctx)?.into_array(); - array.apply(&expr) + array.apply_bound(&expr) } .boxed() } @@ -370,21 +409,23 @@ mod tests { .unwrap(); let expr = eq(root(), lit(3i32)); - let result = RowIdxLayoutReader::new( + let reader = RowIdxLayoutReader::new( 0, layout .new_reader("".into(), segments, &session, &Default::default()) .unwrap(), session.clone(), - ) - .projection_evaluation( - &(0..layout.row_count()), - &expr, - MaskFuture::new_true(layout.row_count().try_into().unwrap()), - ) - .unwrap() - .await - .unwrap(); + ); + let expr = expr.bind(reader.dtype()).unwrap(); + let result = reader + .projection_evaluation( + &(0..layout.row_count()), + &expr, + MaskFuture::new_true(layout.row_count().try_into().unwrap()), + ) + .unwrap() + .await + .unwrap(); assert_arrays_eq!( result, @@ -415,21 +456,23 @@ mod tests { .unwrap(); let expr = gt(row_idx(), lit(3u64)); - let result = RowIdxLayoutReader::new( + let reader = RowIdxLayoutReader::new( 0, layout .new_reader("".into(), segments, &session, &Default::default()) .unwrap(), session.clone(), - ) - .projection_evaluation( - &(0..layout.row_count()), - &expr, - MaskFuture::new_true(layout.row_count().try_into().unwrap()), - ) - .unwrap() - .await - .unwrap(); + ); + let expr = expr.bind(reader.dtype()).unwrap(); + let result = reader + .projection_evaluation( + &(0..layout.row_count()), + &expr, + MaskFuture::new_true(layout.row_count().try_into().unwrap()), + ) + .unwrap() + .await + .unwrap(); assert_arrays_eq!( result, @@ -464,21 +507,23 @@ mod tests { or(gt(row_idx(), lit(3u64)), eq(root(), lit(1i32))), ); - let result = RowIdxLayoutReader::new( + let reader = RowIdxLayoutReader::new( 0, layout .new_reader("".into(), segments, &session, &Default::default()) .unwrap(), session.clone(), - ) - .projection_evaluation( - &(0..layout.row_count()), - &expr, - MaskFuture::new_true(layout.row_count().try_into().unwrap()), - ) - .unwrap() - .await - .unwrap(); + ); + let expr = expr.bind(reader.dtype()).unwrap(); + let result = reader + .projection_evaluation( + &(0..layout.row_count()), + &expr, + MaskFuture::new_true(layout.row_count().try_into().unwrap()), + ) + .unwrap() + .await + .unwrap(); assert_arrays_eq!( result, diff --git a/vortex-layout/src/layouts/struct_/reader.rs b/vortex-layout/src/layouts/struct_/reader.rs index ee939ecd6bd..073cb2bd6ec 100644 --- a/vortex-layout/src/layouts/struct_/reader.rs +++ b/vortex-layout/src/layouts/struct_/reader.rs @@ -19,7 +19,8 @@ use vortex_array::dtype::FieldMask; use vortex_array::dtype::FieldName; use vortex_array::dtype::Nullability; use vortex_array::dtype::StructFields; -use vortex_array::expr::ExactExpr; +use vortex_array::expr::BoundExpression; +use vortex_array::expr::ExactBoundExpr; use vortex_array::expr::Expression; use vortex_array::expr::col; use vortex_array::expr::make_free_field_annotator; @@ -59,7 +60,7 @@ pub struct StructReader { expanded_root_expr: Expression, field_lookup: Option>, - partitioned_expr_cache: DashMap>>, + partitioned_expr_cache: DashMap>>, } impl StructReader { @@ -155,8 +156,8 @@ impl StructReader { } /// Utility for partitioning an expression over the fields of a struct. - fn partition_expr(&self, expr: Expression) -> VortexResult { - let key = ExactExpr(expr.clone()); + fn partition_expr(&self, expr: &BoundExpression) -> VortexResult { + let key = ExactBoundExpr(expr.clone()); // Look up the cell under a shared shard lock; only a miss takes the write lock, and // only for as long as it takes to insert an empty cell. @@ -175,7 +176,7 @@ impl StructReader { if let Some(value) = cell.get() { return Ok(value.clone()); } - let result = self.compute_partitioned_expr(expr)?; + let result = self.compute_partitioned_expr(expr.unbind())?; Ok(cell.get_or_init(|| result).clone()) } @@ -266,17 +267,22 @@ impl LayoutReader for StructReader { fn pruning_evaluation( &self, row_range: &Range, - expr: &Expression, + expr: &BoundExpression, mask: Mask, ) -> VortexResult { // Partition the expression into expressions that can be evaluated over individual fields - match &self.partition_expr(expr.clone())? { - Partitioned::Single(name, partition) => self - .field_reader(name)? - .pruning_evaluation(row_range, partition, mask) - .map_err(|err| { - err.with_context(format!("While evaluating pruning filter partition {name}")) - }), + match &self.partition_expr(expr)? { + Partitioned::Single(name, partition) => { + let reader = self.field_reader(name)?; + let partition = partition.bind(reader.dtype())?; + reader + .pruning_evaluation(row_range, &partition, mask) + .map_err(|err| { + err.with_context(format!( + "While evaluating pruning filter partition {name}" + )) + }) + } Partitioned::Multi(_) => { // TODO(ngates): if all partitions are boolean, we can use a pruning evaluation. Otherwise // there's not much we can do? Maybe... it's complicated... @@ -288,29 +294,36 @@ impl LayoutReader for StructReader { fn filter_evaluation( &self, row_range: &Range, - expr: &Expression, + expr: &BoundExpression, mask: MaskFuture, ) -> VortexResult { // Partition the expression into expressions that can be evaluated over individual fields - match &self.partition_expr(expr.clone())? { - Partitioned::Single(name, partition) => self - .field_reader(name)? - .filter_evaluation(row_range, partition, mask) - .map_err(|err| { - err.with_context(format!("While evaluating filter partition {name}")) - }), + match &self.partition_expr(expr)? { + Partitioned::Single(name, partition) => { + let reader = self.field_reader(name)?; + let partition = partition.bind(reader.dtype())?; + reader + .filter_evaluation(row_range, &partition, mask) + .map_err(|err| { + err.with_context(format!("While evaluating filter partition {name}")) + }) + } Partitioned::Multi(partitioned) => Arc::clone(partitioned).into_mask_future( mask, |name, expr, mask| { - self.field_reader(name)? - .filter_evaluation(row_range, expr, mask) + let reader = self.field_reader(name)?; + let expr = expr.bind(reader.dtype())?; + reader + .filter_evaluation(row_range, &expr, mask) .map_err(|err| { err.with_context(format!("While evaluating filter partition {name}")) }) }, |name, expr, mask| { - self.field_reader(name)? - .projection_evaluation(row_range, expr, mask) + let reader = self.field_reader(name)?; + let expr = expr.bind(reader.dtype())?; + reader + .projection_evaluation(row_range, &expr, mask) .map_err(|err| { err.with_context(format!( "While evaluating projection partition {name}" @@ -325,29 +338,40 @@ impl LayoutReader for StructReader { fn projection_evaluation( &self, row_range: &Range, - expr: &Expression, + expr: &BoundExpression, mask_fut: MaskFuture, ) -> VortexResult { let validity_fut = self .validity()? - .map(|reader| reader.projection_evaluation(row_range, &root(), mask_fut.clone())) + .map(|reader| { + let root = root().bind(reader.dtype())?; + reader.projection_evaluation(row_range, &root, mask_fut.clone()) + }) .transpose()?; // Partition the expression into expressions that can be evaluated over individual fields - let (projected, is_pack_merge) = match &self.partition_expr(expr.clone())? { - Partitioned::Single(name, partition) => ( - self.field_reader(name)? - .projection_evaluation(row_range, partition, mask_fut) - .map_err(|err| { - err.with_context(format!("While evaluating projection partition {name}")) - })?, - partition.is::() || partition.is::(), - ), + let (projected, is_pack_merge) = match &self.partition_expr(expr)? { + Partitioned::Single(name, partition) => { + let reader = self.field_reader(name)?; + let bound_partition = partition.bind(reader.dtype())?; + ( + reader + .projection_evaluation(row_range, &bound_partition, mask_fut) + .map_err(|err| { + err.with_context(format!( + "While evaluating projection partition {name}" + )) + })?, + partition.is::() || partition.is::(), + ) + } Partitioned::Multi(partitioned) => ( Arc::clone(partitioned).into_array_future(mask_fut, |name, expr, mask| { - self.field_reader(name)? - .projection_evaluation(row_range, expr, mask) + let reader = self.field_reader(name)?; + let expr = expr.bind(reader.dtype())?; + reader + .projection_evaluation(row_range, &expr, mask) .map_err(|err| { err.with_context(format!( "While evaluating projection partition {name}" @@ -630,7 +654,9 @@ mod tests { let filt = or( eq(col("a"), lit(7)), or(eq(col("b"), lit(5)), eq(col("a"), lit(3))), - ); + ) + .bind(reader.dtype()) + .unwrap(); let result = block_on(|_| { reader .filter_evaluation(&(0..3), &filt, MaskFuture::new_true(3)) @@ -648,7 +674,9 @@ mod tests { let reader = layout .new_reader("".into(), segments, &SESSION, &Default::default()) .unwrap(); - let expr = gt(get_item("a", root()), get_item("b", root())); + let expr = gt(get_item("a", root()), get_item("b", root())) + .bind(reader.dtype()) + .unwrap(); let result = block_on(|_| { reader .projection_evaluation(&(0..3), &expr, MaskFuture::new_true(3)) @@ -667,7 +695,9 @@ mod tests { let reader = layout .new_reader("".into(), segments, &SESSION, &Default::default()) .unwrap(); - let expr = gt(get_item("a", root()), get_item("b", root())); + let expr = gt(get_item("a", root()), get_item("b", root())) + .bind(reader.dtype()) + .unwrap(); let result = block_on(|_| { reader .projection_evaluation( @@ -694,7 +724,9 @@ mod tests { let expr = pack( [("a", get_item("a", root())), ("b", get_item("b", root()))], Nullability::NonNullable, - ); + ) + .bind(reader.dtype()) + .unwrap(); let result = block_on(|_| { reader .projection_evaluation( @@ -735,7 +767,7 @@ mod tests { let reader = layout .new_reader("".into(), segments, &SESSION, &Default::default()) .unwrap(); - let expr = get_item("a", root()); + let expr = get_item("a", root()).bind(reader.dtype()).unwrap(); let project = reader .projection_evaluation(&(0..3), &expr, MaskFuture::new_true(3)) .unwrap(); @@ -772,8 +804,10 @@ mod tests { let result = block_on(move |handle| { let session = new_session().with_handle(handle); async move { - layout - .new_reader("".into(), segments, &session, &Default::default())? + let reader = + layout.new_reader("".into(), segments, &session, &Default::default())?; + let expr = expr.bind(reader.dtype())?; + reader .projection_evaluation(&(0..3), &expr, MaskFuture::new_true(3))? .await } @@ -832,7 +866,9 @@ mod tests { let reader = layout .new_reader("".into(), segments, &SESSION, &Default::default()) .unwrap(); - let expr = pack(Vec::<(String, Expression)>::new(), Nullability::Nullable); + let expr = pack(Vec::<(String, Expression)>::new(), Nullability::Nullable) + .bind(reader.dtype()) + .unwrap(); let project = reader .projection_evaluation(&(0..5), &expr, MaskFuture::new_true(5)) @@ -849,7 +885,7 @@ mod tests { /// A filter expression whose DType is incompatible with the scanned schema /// (e.g. comparing a u8 column to an i32 literal) must return an error, not panic. #[test] - fn test_struct_filter_dtype_mismatch_returns_error() { + fn test_struct_filter_dtype_mismatch_fails_binding() { let ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); let (ptr, eof) = SequenceId::root().split(); @@ -889,7 +925,7 @@ mod tests { // DType mismatch: "age" is u8 but literal is i32 let filt = eq(col("age"), lit(67i32)); - let result = reader.filter_evaluation(&(0..3), &filt, MaskFuture::new_true(3)); + let result = filt.bind(reader.dtype()); assert!(result.is_err()); let err = result.err().unwrap().to_string(); assert!(err.contains("Cannot compare different DTypes"), "{err}"); diff --git a/vortex-layout/src/layouts/zoned/pruning.rs b/vortex-layout/src/layouts/zoned/pruning.rs index 7764378eb05..51d59b19cd0 100644 --- a/vortex-layout/src/layouts/zoned/pruning.rs +++ b/vortex-layout/src/layouts/zoned/pruning.rs @@ -18,6 +18,7 @@ use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::AggregateFnRef; use vortex_array::arrays::StructArray; use vortex_array::dtype::DType; +use vortex_array::expr::BoundExpression; use vortex_array::expr::Expression; use vortex_array::expr::root; use vortex_array::scalar_fn::fns::dynamic::DynamicExprUpdates; @@ -47,7 +48,7 @@ pub(super) struct PruningState { aggregate_fns: Arc<[AggregateFnRef]>, lazy_children: Arc, session: VortexSession, - pruning_result: LazyLock>>, + pruning_result: LazyLock>>, zone_map: OnceLock, pruning_predicates: LazyLock>>, } @@ -77,45 +78,48 @@ impl PruningState { } } - pub(super) fn pruning_mask_future(&self, expr: Expression) -> Option { + pub(super) fn pruning_mask_future(&self, expr: BoundExpression) -> Option { if let Some(result) = self.pruning_result.get(&expr) { return result.value().clone(); } self.pruning_result .entry(expr.clone()) - .or_insert_with(|| match self.pruning_predicate(expr.clone()) { - None => { - trace!(%expr, "no pruning predicate"); - None - } - Some(predicate) => { - trace!(%expr, ?predicate, "constructed pruning predicate"); - let zone_map = self.zone_map(); - let dynamic_updates = DynamicExprUpdates::new(&expr); - let session = self.session.clone(); - - Some( - async move { - let zone_map = zone_map.await?; - let initial_mask = - zone_map.prune(&predicate, &session).map_err(|err| { - err.with_context(format!( + .or_insert_with(|| { + let expr = expr.unbind(); + match self.pruning_predicate(expr.clone()) { + None => { + trace!(%expr, "no pruning predicate"); + None + } + Some(predicate) => { + trace!(%expr, ?predicate, "constructed pruning predicate"); + let zone_map = self.zone_map(); + let dynamic_updates = DynamicExprUpdates::new(&expr); + let session = self.session.clone(); + + Some( + async move { + let zone_map = zone_map.await?; + let initial_mask = + zone_map.prune(&predicate, &session).map_err(|err| { + err.with_context(format!( "While evaluating pruning predicate {} (derived from {})", predicate, expr )) - })?; - Ok(Arc::new(PruningResult { - zone_map, - predicate, - dynamic_updates, - latest_result: RwLock::new((0, initial_mask)), - session, - })) - } - .boxed() - .shared(), - ) + })?; + Ok(Arc::new(PruningResult { + zone_map, + predicate, + dynamic_updates, + latest_result: RwLock::new((0, initial_mask)), + session, + })) + } + .boxed() + .shared(), + ) + } } }) .clone() @@ -139,13 +143,17 @@ impl PruningState { self.zone_map .get_or_init(move || { let zone_count = self.zone_count; - let zones_eval = self + let zones_reader = self .lazy_children .get(1) - .vortex_expect("failed to get zone child") + .vortex_expect("failed to get zone child"); + let root = root() + .bind(zones_reader.dtype()) + .vortex_expect("root must bind against the zone-map dtype"); + let zones_eval = zones_reader .projection_evaluation( &(0..zone_count as u64), - &root(), + &root, MaskFuture::new_true(zone_count), ) .vortex_expect("Failed construct zone map evaluation"); diff --git a/vortex-layout/src/layouts/zoned/reader.rs b/vortex-layout/src/layouts/zoned/reader.rs index dbcef1f7a68..52a57c1b92b 100644 --- a/vortex-layout/src/layouts/zoned/reader.rs +++ b/vortex-layout/src/layouts/zoned/reader.rs @@ -12,7 +12,7 @@ use vortex_array::ArrayRef; use vortex_array::MaskFuture; use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; -use vortex_array::expr::Expression; +use vortex_array::expr::BoundExpression; use vortex_buffer::BitBufferMut; use vortex_error::VortexError; use vortex_error::VortexResult; @@ -136,7 +136,7 @@ impl LayoutReader for ZonedReader { fn pruning_evaluation( &self, row_range: &Range, - expr: &Expression, + expr: &BoundExpression, mask: Mask, ) -> VortexResult { trace!("Stats pruning evaluation: {} - {}", &self.name, expr); @@ -209,7 +209,7 @@ impl LayoutReader for ZonedReader { fn filter_evaluation( &self, row_range: &Range, - expr: &Expression, + expr: &BoundExpression, mask: MaskFuture, ) -> VortexResult { self.data_child()?.filter_evaluation(row_range, expr, mask) @@ -218,7 +218,7 @@ impl LayoutReader for ZonedReader { fn projection_evaluation( &self, row_range: &Range, - expr: &Expression, + expr: &BoundExpression, mask: MaskFuture, ) -> VortexResult>> { // TODO(ngates): there are some projection expressions that we may also be able to @@ -325,12 +325,14 @@ mod test { block_on(|handle| async { let mut ctx = array_session().create_execution_ctx(); let session = session_with_handle(handle); - let result = layout + let reader = layout .new_reader("".into(), segments, &session, &Default::default()) - .unwrap() + .unwrap(); + let expr = root().bind(reader.dtype()).unwrap(); + let result = reader .projection_evaluation( &(0..layout.row_count()), - &root(), + &expr, MaskFuture::new_true(layout.row_count().try_into().unwrap()), ) .unwrap() @@ -354,7 +356,7 @@ mod test { .unwrap(); // Choose a prune-able expression - let expr = gt(root(), lit(7)); + let expr = gt(root(), lit(7)).bind(reader.dtype()).unwrap(); let result = reader .pruning_evaluation( @@ -420,10 +422,11 @@ mod test { .new_reader("".into(), segments, &session, &Default::default()) .unwrap(); + let expr = is_not_null(root()).bind(reader.dtype()).unwrap(); let result = reader .pruning_evaluation( &(0..row_count), - &is_not_null(root()), + &expr, Mask::new_true(row_count.try_into().unwrap()), ) .unwrap() @@ -483,20 +486,22 @@ mod test { let reader = legacy_layout.new_reader("".into(), segments, &session, &Default::default())?; + let expr = gt(root(), lit(7)).bind(reader.dtype())?; let result = reader .pruning_evaluation( &(0..row_count), - >(root(), lit(7)), + &expr, Mask::new_true(row_count.try_into().unwrap()), )? .await?; assert_eq!(result, Mask::from_iter(expected)); + let root = root().bind(reader.dtype())?; let projected = reader .projection_evaluation( &(0..row_count), - &root(), + &root, MaskFuture::new_true(row_count.try_into().unwrap()), )? .await?; diff --git a/vortex-layout/src/layouts/zoned/zone_map.rs b/vortex-layout/src/layouts/zoned/zone_map.rs index e5f2af494de..9c7ff87b693 100644 --- a/vortex-layout/src/layouts/zoned/zone_map.rs +++ b/vortex-layout/src/layouts/zoned/zone_map.rs @@ -135,7 +135,7 @@ impl ZoneMap { /// `true` means the zone cannot contain matching rows and can be skipped. /// /// If the predicate contains [`row_count`][vortex_array::scalar_fn::internal::row_count] - /// placeholders, they are replaced after [`ArrayRef::apply`] with per-zone + /// placeholders, they are replaced after [`ArrayRef::apply_bound`] with per-zone /// counts derived from `zone_len` and `row_count`. Uniform zones use a /// [`ConstantArray`]; a short final zone uses a run-end encoded array. /// `row_count` is a layout property rather than a stored stats field, and the @@ -146,7 +146,8 @@ impl ZoneMap { let num_zones = self.array.len(); let predicate = self.lower_stats(predicate.clone())?; - let applied = self.array.clone().into_array().apply(&predicate)?; + let array = self.array.clone().into_array(); + let applied = array.apply(&predicate)?; if !contains_row_count(&applied) { return applied.null_as_false().execute(&mut ctx); diff --git a/vortex-layout/src/reader.rs b/vortex-layout/src/reader.rs index 9318aaf3016..b576a993d22 100644 --- a/vortex-layout/src/reader.rs +++ b/vortex-layout/src/reader.rs @@ -14,7 +14,7 @@ use vortex_array::MaskFuture; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; -use vortex_array::expr::Expression; +use vortex_array::expr::BoundExpression; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_mask::Mask; @@ -159,7 +159,7 @@ pub trait LayoutReader: 'static + Send + Sync { fn pruning_evaluation( &self, row_range: &Range, - expr: &Expression, + expr: &BoundExpression, mask: Mask, ) -> VortexResult; @@ -175,7 +175,7 @@ pub trait LayoutReader: 'static + Send + Sync { fn filter_evaluation( &self, row_range: &Range, - expr: &Expression, + expr: &BoundExpression, mask: MaskFuture, ) -> VortexResult; @@ -191,7 +191,7 @@ pub trait LayoutReader: 'static + Send + Sync { fn projection_evaluation( &self, row_range: &Range, - expr: &Expression, + expr: &BoundExpression, mask: MaskFuture, ) -> VortexResult; } diff --git a/vortex-layout/src/scan/filter.rs b/vortex-layout/src/scan/filter.rs index 12a5a66eb77..3b2ab931c0b 100644 --- a/vortex-layout/src/scan/filter.rs +++ b/vortex-layout/src/scan/filter.rs @@ -7,9 +7,10 @@ use bit_vec::BitVec; use itertools::Itertools; use parking_lot::RwLock; use sketches_ddsketch::DDSketch; -use vortex_array::expr::Expression; -use vortex_array::expr::forms::conjuncts; +use vortex_array::expr::BoundExpression; +use vortex_array::scalar_fn::fns::binary::Binary; use vortex_array::scalar_fn::fns::dynamic::DynamicExprUpdates; +use vortex_array::scalar_fn::fns::operators::Operator; use vortex_error::VortexExpect; use vortex_error::vortex_err; use vortex_error::vortex_panic; @@ -22,7 +23,7 @@ const DEFAULT_SELECTIVITY_QUANTILE: f64 = 0.1; /// conjunctions in an attempt to minimize the work done. pub struct FilterExpr { /// The conjuncts involved in the filter expression. - conjuncts: Vec, + conjuncts: Vec, /// A histogram for the selectivity of each conjunct. conjunct_selectivity: Vec>, /// Dynamic expression trackers for each conjunct, incase they contain dynamic expressions. @@ -33,12 +34,34 @@ pub struct FilterExpr { selectivity_quantile: f64, } +fn bound_conjuncts(expr: &BoundExpression) -> Vec { + let mut conjuncts = Vec::new(); + let mut pending = vec![expr]; + + while let Some(expr) = pending.pop() { + if expr + .as_scalar() + .and_then(|scalar_fn| scalar_fn.as_opt::()) + .is_some_and(|operator| *operator == Operator::And) + { + pending.extend(expr.children().iter().rev()); + } else { + conjuncts.push(expr.clone()); + } + } + + conjuncts +} + impl FilterExpr { - pub fn new(expr: Expression) -> Self { - let conjuncts = conjuncts(&expr); + pub fn new(expr: BoundExpression) -> Self { + let conjuncts = bound_conjuncts(&expr); let num_conjuncts = conjuncts.len(); - let dynamic_conjuncts = conjuncts.iter().map(DynamicExprUpdates::new).collect_vec(); + let dynamic_conjuncts = conjuncts + .iter() + .map(|expr| DynamicExprUpdates::new(&expr.unbind())) + .collect_vec(); Self { conjuncts, @@ -55,7 +78,7 @@ impl FilterExpr { /// The conjuncts that make up this filter expression. #[inline] - pub fn conjuncts(&self) -> &[Expression] { + pub fn conjuncts(&self) -> &[BoundExpression] { &self.conjuncts } @@ -132,3 +155,31 @@ impl FilterExpr { ); } } + +#[cfg(test)] +mod tests { + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::expr::and; + use vortex_array::expr::lit; + use vortex_array::expr::not; + use vortex_array::expr::root; + use vortex_error::VortexResult; + + use super::FilterExpr; + + #[test] + fn bound_conjuncts_preserve_order_and_types() -> VortexResult<()> { + let expr = and(root(), and(not(root()), lit(true))); + let bound = expr.bind(&DType::Bool(Nullability::NonNullable))?; + let filter = FilterExpr::new(bound); + + assert!( + filter + .conjuncts() + .iter() + .all(|expr| { expr.dtype() == &DType::Bool(Nullability::NonNullable) }) + ); + Ok(()) + } +} diff --git a/vortex-layout/src/scan/layout.rs b/vortex-layout/src/scan/layout.rs index 0b1f150c67e..f0e6366a8fc 100644 --- a/vortex-layout/src/scan/layout.rs +++ b/vortex-layout/src/scan/layout.rs @@ -140,12 +140,13 @@ impl DataSource for LayoutReaderDataSource { // Check file-level pruning: if the filter can be proven false for the entire row range // using file-level statistics (e.g. via FileStatsLayoutReader), skip the scan entirely. if let Some(filter) = &scan_request.filter { + let filter = filter.bind(self.reader.dtype())?; let mask = Mask::new_true( usize::try_from(row_range.end - row_range.start).unwrap_or(usize::MAX), ); let pruning_result = self .reader - .pruning_evaluation(&row_range, filter, mask)? + .pruning_evaluation(&row_range, &filter, mask)? .now_or_never(); if let Some(Ok(result_mask)) = pruning_result && result_mask.all_false() diff --git a/vortex-layout/src/scan/multi.rs b/vortex-layout/src/scan/multi.rs index 8f28372bbb4..d251ee15617 100644 --- a/vortex-layout/src/scan/multi.rs +++ b/vortex-layout/src/scan/multi.rs @@ -446,7 +446,8 @@ fn reader_partition( if let Some(filter) = &request.filter { let mask_len = usize::try_from(row_range.end - row_range.start).unwrap_or(usize::MAX); let mask = Mask::new_true(mask_len); - if let Ok(pruning_future) = reader.pruning_evaluation(&row_range, filter, mask) + if let Ok(filter) = filter.bind(reader.dtype()) + && let Ok(pruning_future) = reader.pruning_evaluation(&row_range, &filter, mask) && let Some(Ok(result_mask)) = pruning_future.now_or_never() && result_mask.all_false() { diff --git a/vortex-layout/src/scan/repeated_scan.rs b/vortex-layout/src/scan/repeated_scan.rs index 681f33639bc..10e6c0a1b18 100644 --- a/vortex-layout/src/scan/repeated_scan.rs +++ b/vortex-layout/src/scan/repeated_scan.rs @@ -12,7 +12,7 @@ use itertools::Either; use itertools::Itertools; use vortex_array::ArrayRef; use vortex_array::dtype::DType; -use vortex_array::expr::Expression; +use vortex_array::expr::BoundExpression; use vortex_array::iter::ArrayIterator; use vortex_array::iter::ArrayIteratorAdapter; use vortex_array::stream::ArrayStream; @@ -38,8 +38,8 @@ use crate::scan::tasks::split_exec; pub struct RepeatedScan { session: VortexSession, layout_reader: LayoutReaderRef, - projection: Expression, - filter: Option, + projection: BoundExpression, + filter: Option, ordered: bool, /// Optionally read a subset of the rows in the file. row_range: Option>, @@ -92,8 +92,8 @@ impl RepeatedScan { pub fn new( session: VortexSession, layout_reader: LayoutReaderRef, - projection: Expression, - filter: Option, + projection: BoundExpression, + filter: Option, ordered: bool, row_range: Option>, selection: Selection, diff --git a/vortex-layout/src/scan/scan_builder.rs b/vortex-layout/src/scan/scan_builder.rs index 03f8c49649d..c45155cdc62 100644 --- a/vortex-layout/src/scan/scan_builder.rs +++ b/vortex-layout/src/scan/scan_builder.rs @@ -309,6 +309,11 @@ impl ScanBuilder { )?) }; + let projection = projection.bind(layout_reader.dtype())?; + let filter = filter + .map(|expr| expr.bind(layout_reader.dtype())) + .transpose()?; + Ok(RepeatedScan::new( self.session.clone(), layout_reader, @@ -476,7 +481,7 @@ mod test { use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::StructFields; - use vortex_array::expr::Expression; + use vortex_array::expr::BoundExpression; use vortex_array::expr::eq; use vortex_array::expr::get_item; use vortex_array::expr::is_not_null; @@ -597,7 +602,7 @@ mod test { fn pruning_evaluation( &self, _row_range: &Range, - _expr: &Expression, + _expr: &BoundExpression, _mask: Mask, ) -> VortexResult { unimplemented!("not needed for this test"); @@ -606,7 +611,7 @@ mod test { fn filter_evaluation( &self, _row_range: &Range, - _expr: &Expression, + _expr: &BoundExpression, _mask: MaskFuture, ) -> VortexResult { unimplemented!("not needed for this test"); @@ -615,7 +620,7 @@ mod test { fn projection_evaluation( &self, _row_range: &Range, - _expr: &Expression, + _expr: &BoundExpression, _mask: MaskFuture, ) -> VortexResult { Ok(Box::pin(async move { @@ -688,7 +693,7 @@ mod test { fn pruning_evaluation( &self, _row_range: &Range, - _expr: &Expression, + _expr: &BoundExpression, mask: Mask, ) -> VortexResult { Ok(MaskFuture::ready(mask)) @@ -697,7 +702,7 @@ mod test { fn filter_evaluation( &self, _row_range: &Range, - _expr: &Expression, + _expr: &BoundExpression, mask: MaskFuture, ) -> VortexResult { Ok(mask) @@ -706,7 +711,7 @@ mod test { fn projection_evaluation( &self, row_range: &Range, - _expr: &Expression, + _expr: &BoundExpression, _mask: MaskFuture, ) -> VortexResult { let start = usize::try_from(row_range.start) @@ -800,7 +805,7 @@ mod test { fn pruning_evaluation( &self, _row_range: &Range, - _expr: &Expression, + _expr: &BoundExpression, _mask: Mask, ) -> VortexResult { unimplemented!("not needed for this test"); @@ -809,7 +814,7 @@ mod test { fn filter_evaluation( &self, _row_range: &Range, - _expr: &Expression, + _expr: &BoundExpression, _mask: MaskFuture, ) -> VortexResult { unimplemented!("not needed for this test"); @@ -818,7 +823,7 @@ mod test { fn projection_evaluation( &self, _row_range: &Range, - _expr: &Expression, + _expr: &BoundExpression, _mask: MaskFuture, ) -> VortexResult { Ok(Box::pin(async move { diff --git a/vortex-layout/src/scan/split_by.rs b/vortex-layout/src/scan/split_by.rs index be145572524..6106d4f671d 100644 --- a/vortex-layout/src/scan/split_by.rs +++ b/vortex-layout/src/scan/split_by.rs @@ -133,7 +133,7 @@ mod test { use vortex_array::dtype::FieldPath; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; - use vortex_array::expr::Expression; + use vortex_array::expr::BoundExpression; use vortex_buffer::buffer; use vortex_io::runtime::single::block_on; use vortex_mask::Mask; @@ -242,7 +242,7 @@ mod test { fn pruning_evaluation( &self, _: &Range, - _: &Expression, + _: &BoundExpression, _: Mask, ) -> VortexResult { unimplemented!() @@ -251,7 +251,7 @@ mod test { fn filter_evaluation( &self, _: &Range, - _: &Expression, + _: &BoundExpression, _: MaskFuture, ) -> VortexResult { unimplemented!() @@ -260,7 +260,7 @@ mod test { fn projection_evaluation( &self, _: &Range, - _: &Expression, + _: &BoundExpression, _: MaskFuture, ) -> VortexResult>> { unimplemented!() diff --git a/vortex-layout/src/scan/tasks.rs b/vortex-layout/src/scan/tasks.rs index a86546e15ef..218efb64a0d 100644 --- a/vortex-layout/src/scan/tasks.rs +++ b/vortex-layout/src/scan/tasks.rs @@ -11,7 +11,7 @@ use futures::FutureExt; use futures::future::BoxFuture; use vortex_array::ArrayRef; use vortex_array::MaskFuture; -use vortex_array::expr::Expression; +use vortex_array::expr::BoundExpression; use vortex_error::VortexResult; use vortex_mask::Mask; use vortex_scan::row_mask::RowMask; @@ -159,7 +159,7 @@ pub struct TaskContext { /// The layout reader. pub reader: Arc, /// The projection expression to apply to gather the scanned rows. - pub projection: Expression, + pub projection: BoundExpression, /// Function that maps into an A. pub mapper: Arc VortexResult + Send + Sync>, } diff --git a/vortex-test/compat-gen/src/adapter.rs b/vortex-test/compat-gen/src/adapter.rs index 1c8ee9bdc6a..be23417e44b 100644 --- a/vortex-test/compat-gen/src/adapter.rs +++ b/vortex-test/compat-gen/src/adapter.rs @@ -142,8 +142,9 @@ pub fn read_layout_tree(bytes: ByteBuffer) -> VortexResult<()> { )?; let len = usize::try_from(row_count).map_err(|e| vortex_err!("row count overflow: {e}"))?; + let expr = root().bind(reader.dtype())?; reader - .projection_evaluation(&(0..row_count), &root(), MaskFuture::new_true(len))? + .projection_evaluation(&(0..row_count), &expr, MaskFuture::new_true(len))? .await?; } diff --git a/vortex-tui/src/browse/app.rs b/vortex-tui/src/browse/app.rs index 4280e6f81f2..eeea77a62bc 100644 --- a/vortex-tui/src/browse/app.rs +++ b/vortex-tui/src/browse/app.rs @@ -373,10 +373,13 @@ impl AppState { &Default::default(), ) .vortex_expect("Failed to create reader"); + let expr = root() + .bind(reader.dtype()) + .vortex_expect("root must bind against the layout dtype"); let array = reader .projection_evaluation( &(0..row_count), - &root(), + &expr, MaskFuture::new_true( usize::try_from(row_count).vortex_expect("row_count overflowed usize"), ), diff --git a/vortex-tui/src/wasm.rs b/vortex-tui/src/wasm.rs index 5fd7dd9054b..4598a00101b 100644 --- a/vortex-tui/src/wasm.rs +++ b/vortex-tui/src/wasm.rs @@ -85,10 +85,13 @@ async fn load_flat_array( &Default::default(), ) .vortex_expect("Failed to create reader"); + let expr = root() + .bind(reader.dtype()) + .vortex_expect("root must bind against the layout dtype"); reader .projection_evaluation( &(0..row_count), - &root(), + &expr, MaskFuture::new_true( usize::try_from(row_count).vortex_expect("row_count overflowed usize"), ), From c85008ea81ee7eba476eab9f78b7e12206a395ba Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 3 Aug 2026 18:00:43 +0100 Subject: [PATCH 2/4] fix: preserve bound expression identity in layout caches Signed-off-by: Joe Isaacs --- .../expr/analysis/referenced_field_paths.rs | 78 +++---- vortex-array/src/expr/bound_expression.rs | 15 +- .../src/expr/transform/bound_partition.rs | 5 +- vortex-file/src/v2/file_stats_reader.rs | 9 +- vortex-layout/src/layouts/dict/reader.rs | 71 ++++--- vortex-layout/src/layouts/partitioned.rs | 60 +++--- vortex-layout/src/layouts/row_idx/mod.rs | 181 ++++++++-------- vortex-layout/src/layouts/struct_/reader.rs | 197 ++++++++++++++---- vortex-layout/src/scan/scan_builder.rs | 75 ++++--- 9 files changed, 409 insertions(+), 282 deletions(-) diff --git a/vortex-array/src/expr/analysis/referenced_field_paths.rs b/vortex-array/src/expr/analysis/referenced_field_paths.rs index 5239e4a691c..48ef285c20a 100644 --- a/vortex-array/src/expr/analysis/referenced_field_paths.rs +++ b/vortex-array/src/expr/analysis/referenced_field_paths.rs @@ -4,17 +4,15 @@ use vortex_error::VortexResult; use vortex_error::vortex_err; -use crate::dtype::DType; use crate::dtype::Field; use crate::dtype::FieldPath; use crate::dtype::FieldPathSet; -use crate::expr::Expression; +use crate::expr::BoundExpression; use crate::expr::traversal::FoldDownContext; use crate::expr::traversal::FoldUp; use crate::expr::traversal::NodeExt; use crate::expr::traversal::NodeFolderContext; use crate::scalar_fn::fns::get_item::GetItem; -use crate::scalar_fn::fns::root::Root; use crate::scalar_fn::fns::select::Select; /// Returns the rooted field paths referenced by an expression. @@ -24,47 +22,13 @@ use crate::scalar_fn::fns::select::Select; /// expression is represented by [`FieldPath::root`], which conservatively selects all fields. /// Scalar functions other than `GetItem` and `Select` conservatively reference each complete child /// output. -pub fn referenced_field_paths(expr: &Expression, scope: &DType) -> VortexResult { - // Validate the whole expression so plain GetItem paths and Select paths behave consistently. - expr.return_dtype(scope)?; - +pub fn referenced_field_paths(expr: &BoundExpression) -> VortexResult { let mut collector = ReferencedFieldPaths { - scope, field_paths: FieldPathSet::default(), }; expr.clone() .fold_context(&vec![FieldPath::root()], &mut collector)?; - let field_paths = collector.field_paths; - - // The top-level field of every referenced path must be one of the immediately accessed scope - // fields: this analysis only refines *which nested fields* are read, never which top-level - // fields. `FieldPath::root()` stands in for "all fields", so it expands to the whole scope. - #[cfg(debug_assertions)] - if let Some(scope_fields) = scope.as_struct_fields_opt() { - use vortex_utils::aliases::hash_set::HashSet; - - use crate::dtype::FieldName; - use crate::expr::analysis::immediate_access::immediate_scope_access; - - let referenced_heads: HashSet = if field_paths.iter().any(FieldPath::is_root) { - scope_fields.names().iter().cloned().collect() - } else { - field_paths - .iter() - .filter_map(|path| match path.parts().first() { - Some(Field::Name(name)) => Some(name.clone()), - _ => None, - }) - .collect() - }; - debug_assert_eq!( - referenced_heads, - immediate_scope_access(expr, scope_fields), - "referenced field path heads must match the immediately accessed scope fields" - ); - } - - Ok(field_paths) + Ok(collector.field_paths) } /// Threads the set of currently-requested field paths down the expression tree, narrowing it at @@ -78,22 +42,21 @@ pub fn referenced_field_paths(expr: &Expression, scope: &DType) -> VortexResult< /// column projection). Any other function is opaque—we cannot assume it preserves a field's /// provenance—so its children conservatively re-request the whole scope, which is what keeps an /// expression like `f($).x` reading every field of `$` rather than just `x`. -struct ReferencedFieldPaths<'a> { - scope: &'a DType, +struct ReferencedFieldPaths { field_paths: FieldPathSet, } -impl NodeFolderContext for ReferencedFieldPaths<'_> { - type NodeTy = Expression; +impl NodeFolderContext for ReferencedFieldPaths { + type NodeTy = BoundExpression; type Result = (); type Context = Vec; fn visit_down( &mut self, requested: &Self::Context, - node: &Expression, + node: &BoundExpression, ) -> VortexResult> { - if node.is::() { + if node.is_root() { self.field_paths.extend( requested .iter() @@ -102,7 +65,10 @@ impl NodeFolderContext for ReferencedFieldPaths<'_> { return Ok(FoldDownContext::Skip(())); } - if let Some(field_name) = node.as_opt::() { + if let Some(field_name) = node + .as_scalar() + .and_then(|scalar_fn| scalar_fn.as_opt::()) + { let appended = requested .iter() .map(|path| path.clone().push(Field::Name(field_name.clone()))) @@ -112,9 +78,12 @@ impl NodeFolderContext for ReferencedFieldPaths<'_> { // Keep requested paths whose head is included, expanding a whole-scope request into one // path per included field. - if let Some(selection) = node.as_opt::()) + { + let child_fields = node.children()[0] + .dtype() .as_struct_fields_opt() .ok_or_else(|| vortex_err!("Select child is not a struct"))?; let included_fields = selection.normalize_to_included_fields(child_fields.names())?; @@ -146,7 +115,7 @@ impl NodeFolderContext for ReferencedFieldPaths<'_> { fn visit_up( &mut self, - _node: Expression, + _node: BoundExpression, _requested: &Self::Context, _children: Vec<()>, ) -> VortexResult> { @@ -159,9 +128,11 @@ mod tests { use vortex_utils::aliases::hash_set::HashSet; use super::*; + use crate::dtype::DType; use crate::dtype::Nullability::NonNullable; use crate::dtype::PType::I32; use crate::dtype::StructFields; + use crate::expr::Expression; use crate::expr::get_item; use crate::expr::pack; use crate::expr::root; @@ -183,7 +154,7 @@ mod tests { /// Collects the prefix-minimal field paths referenced by `expr` against [`scope`]. fn referenced(expr: &Expression) -> VortexResult> { - Ok(referenced_field_paths(expr, &scope())? + Ok(referenced_field_paths(&expr.bind(&scope())?)? .into_iter() .collect()) } @@ -259,6 +230,9 @@ mod tests { #[test] fn invalid_get_item_path_returns_error() { - assert!(referenced_field_paths(&get_item("missing", root()), &scope()).is_err()); + let result = get_item("missing", root()) + .bind(&scope()) + .and_then(|expr| referenced_field_paths(&expr)); + assert!(result.is_err()); } } diff --git a/vortex-array/src/expr/bound_expression.rs b/vortex-array/src/expr/bound_expression.rs index 452e4df8377..3a35757068a 100644 --- a/vortex-array/src/expr/bound_expression.rs +++ b/vortex-array/src/expr/bound_expression.rs @@ -57,12 +57,8 @@ pub struct ExactBoundExpr(pub BoundExpression); impl PartialEq for ExactBoundExpr { fn eq(&self, other: &Self) -> bool { - if self.0.dtype != other.0.dtype { - return false; - } - match (&self.0.kind, &other.0.kind) { - (BoundKind::Root, BoundKind::Root) => true, + (BoundKind::Root, BoundKind::Root) => self.0.dtype == other.0.dtype, ( BoundKind::Scalar { scalar_fn: lhs_fn, @@ -72,7 +68,11 @@ impl PartialEq for ExactBoundExpr { scalar_fn: rhs_fn, children: rhs_children, }, - ) => lhs_fn == rhs_fn && Arc::ptr_eq(lhs_children, rhs_children), + ) => { + lhs_fn == rhs_fn + && Arc::ptr_eq(lhs_children, rhs_children) + && self.0.dtype == other.0.dtype + } _ => false, } } @@ -82,7 +82,8 @@ impl Eq for ExactBoundExpr {} impl Hash for ExactBoundExpr { fn hash(&self, state: &mut H) { - self.0.dtype.hash(state); + // DType differences are resolved by equality. Omitting the potentially lazy dtype keeps + // identity-keyed cache lookups from deserializing an entire schema just to compute a hash. match &self.0.kind { BoundKind::Root => state.write_u8(0), BoundKind::Scalar { diff --git a/vortex-array/src/expr/transform/bound_partition.rs b/vortex-array/src/expr/transform/bound_partition.rs index 9f1c59bcd6f..95fe3f4a7ea 100644 --- a/vortex-array/src/expr/transform/bound_partition.rs +++ b/vortex-array/src/expr/transform/bound_partition.rs @@ -102,7 +102,10 @@ where }) } -/// The result of partitioning an expression. +/// The result of partitioning a bound expression. +/// +/// The root and partitions remain bound so callers can cache and reuse their shared tree identity +/// without an unbind/rebind round trip. #[derive(Debug)] pub struct BoundPartitionedExpr { /// The root expression used to re-assemble the results. diff --git a/vortex-file/src/v2/file_stats_reader.rs b/vortex-file/src/v2/file_stats_reader.rs index 034ceb3b4af..20102a2ad66 100644 --- a/vortex-file/src/v2/file_stats_reader.rs +++ b/vortex-file/src/v2/file_stats_reader.rs @@ -15,6 +15,7 @@ use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; use vortex_array::dtype::StructFields; use vortex_array::expr::BoundExpression; +use vortex_array::expr::ExactBoundExpr; use vortex_array::expr::Expression; use vortex_error::VortexResult; use vortex_layout::ArrayFuture; @@ -42,7 +43,7 @@ pub struct FileStatsLayoutReader { file_stats: FileStatistics, struct_fields: StructFields, session: VortexSession, - prune_cache: DashMap, + prune_cache: DashMap, } impl FileStatsLayoutReader { @@ -117,8 +118,10 @@ impl LayoutReader for FileStatsLayoutReader { expr: &BoundExpression, mask: Mask, ) -> VortexResult { + let key = ExactBoundExpr(expr.clone()); + // Check cache first with read-only lock. - if let Some(pruned) = self.prune_cache.get(expr) { + if let Some(pruned) = self.prune_cache.get(&key) { if *pruned { return Ok(MaskFuture::ready(Mask::new_false(mask.len()))); } @@ -128,7 +131,7 @@ impl LayoutReader for FileStatsLayoutReader { // Evaluate and cache. let expression = expr.unbind(); let pruned = self.evaluate_file_stats(&expression)?; - self.prune_cache.insert(expr.clone(), pruned); + self.prune_cache.insert(key, pruned); if pruned { Ok(MaskFuture::ready(Mask::new_false(mask.len()))) diff --git a/vortex-layout/src/layouts/dict/reader.rs b/vortex-layout/src/layouts/dict/reader.rs index 235c61a3307..eefdee6e4c4 100644 --- a/vortex-layout/src/layouts/dict/reader.rs +++ b/vortex-layout/src/layouts/dict/reader.rs @@ -20,14 +20,15 @@ use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; use vortex_array::dtype::Nullability; use vortex_array::expr::BoundExpression; -use vortex_array::expr::Expression; -use vortex_array::expr::direct_annotations; -use vortex_array::expr::is_root; -use vortex_array::expr::label_tree; -use vortex_array::expr::pack; +use vortex_array::expr::ExactBoundExpr; +use vortex_array::expr::direct_bound_annotations; +use vortex_array::expr::label_bound_tree; use vortex_array::expr::root; -use vortex_array::expr::transform::partition_annotations; +use vortex_array::expr::transform::partition_bound_annotations; use vortex_array::optimizer::ArrayOptimizer; +use vortex_array::scalar_fn::ScalarFnVTableExt; +use vortex_array::scalar_fn::fns::pack::Pack; +use vortex_array::scalar_fn::fns::pack::PackOptions; use vortex_array::scalar_fn::is_negative_cost; use vortex_error::VortexError; use vortex_error::VortexExpect; @@ -185,23 +186,29 @@ const PUSHDOWN_ANNOTATION: &str = ""; /// We want to push to the array only if the expression has a negative cost, is infallible, and is /// strict. Strictness ensures dictionary null codes still force a null result after pushdown. fn split_expression_for_pushdown( - expr: Expression, - dtype: &DType, -) -> VortexResult<(Expression, Option)> { - let references_root = label_tree(&expr, is_root, |acc, &child| acc | child); - let annotations = direct_annotations(&expr, |expr| { - let signature = expr.signature(); + expr: &BoundExpression, +) -> VortexResult<(BoundExpression, Option)> { + let references_root = + label_bound_tree(expr, BoundExpression::is_root, |acc, &child| acc | child); + let annotations = direct_bound_annotations(expr, |expr: &BoundExpression| { + let Some(scalar_fn) = expr.as_scalar() else { + return vec![]; + }; + let signature = scalar_fn.signature(); if !signature.is_fallible() && signature.is_strict() - && is_negative_cost(expr.id()) - && references_root.get(&expr).copied().unwrap_or(true) + && is_negative_cost(scalar_fn.id()) + && references_root + .get(&ExactBoundExpr(expr.clone())) + .copied() + .unwrap_or(true) { vec![PUSHDOWN_ANNOTATION] } else { vec![] } }); - let partition = partition_annotations(expr.clone(), dtype, annotations)?; + let partition = partition_bound_annotations(expr.clone(), annotations)?; if partition.partitions.is_empty() { Ok((partition.root, None)) } else { @@ -289,13 +296,18 @@ impl LayoutReader for DictReader { .projection_evaluation(row_range, &codes_root, mask) .map_err(|err| err.with_context("While evaluating projection on codes"))?; - let (expr_outer, expr_inner) = split_expression_for_pushdown(expr.unbind(), self.dtype())?; + let (expr_outer, expr_inner) = split_expression_for_pushdown(expr)?; let values_eval = if let Some(inner) = expr_inner { // "outer" takes a struct field with PUSHDOWN_ANNOTATION name, so // pack inner with this name as well - let inner = pack([(PUSHDOWN_ANNOTATION, inner)], Nullability::NonNullable) - .bind(self.values.dtype())?; + let inner = BoundExpression::try_new( + Pack.bind(PackOptions { + names: [PUSHDOWN_ANNOTATION].into(), + nullability: Nullability::NonNullable, + }), + [inner], + )?; // We can't use values_eval as it uses values_array_uncanonical // which in turn gets populated from self.values. If @@ -333,7 +345,6 @@ impl LayoutReader for DictReader { .into_array() .optimize()?; - let expr_outer = expr_outer.bind(array.dtype())?; array.apply_bound(&expr_outer) } .boxed()) @@ -753,9 +764,18 @@ mod tests { Ok(()) } + fn split_unbound( + expr: Expression, + dtype: &DType, + ) -> VortexResult<(Expression, Option)> { + let bound = expr.bind(dtype)?; + let (outer, inner) = split_expression_for_pushdown(&bound)?; + Ok((outer.unbind(), inner.map(|expr| expr.unbind()))) + } + #[test] fn split_expr_root() { - let (outer, inner) = split_expression_for_pushdown(root(), &DType::Null).unwrap(); + let (outer, inner) = split_unbound(root(), &DType::Null).unwrap(); assert_eq!(outer, root()); assert_eq!(inner, None); } @@ -765,8 +785,7 @@ mod tests { // cast is fallible, thus not pushed let target = DType::Primitive(PType::I64, Nullability::Nullable); let expr = cast(byte_length(root()), target.clone()); - let (outer, inner) = - split_expression_for_pushdown(expr.clone(), &DType::Utf8(false.into()))?; + let (outer, inner) = split_unbound(expr.clone(), &DType::Utf8(false.into()))?; let inner = inner.unwrap(); // [0] = cast([1], dtype) // [1] = byte_length(root) @@ -778,8 +797,7 @@ mod tests { #[test] fn split_expr_full_pushdown() -> VortexResult<()> { let expr = byte_length(root()); - let (outer, inner) = - split_expression_for_pushdown(expr.clone(), &DType::Utf8(false.into()))?; + let (outer, inner) = split_unbound(expr.clone(), &DType::Utf8(false.into()))?; let inner = inner.unwrap(); assert_eq!(outer, pushed_ref(0)); assert_eq!(inner, pushed_inner([byte_length(root())])); @@ -789,9 +807,8 @@ mod tests { #[test] fn split_expr_no_pushdown() { // like is fallible, thus not pushed. lit() does not reference root() - let expr = like(root(), lit(1u64)); - let (outer, inner) = - split_expression_for_pushdown(expr.clone(), &DType::Utf8(true.into())).unwrap(); + let expr = like(root(), lit("abc")); + let (outer, inner) = split_unbound(expr.clone(), &DType::Utf8(true.into())).unwrap(); assert_eq!(outer, expr); assert_eq!(inner, None); } diff --git a/vortex-layout/src/layouts/partitioned.rs b/vortex-layout/src/layouts/partitioned.rs index 09636c4c339..94d1c3f782d 100644 --- a/vortex-layout/src/layouts/partitioned.rs +++ b/vortex-layout/src/layouts/partitioned.rs @@ -13,8 +13,8 @@ use vortex_array::VortexSessionExecute; use vortex_array::arrays::StructArray; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; -use vortex_array::expr::Expression; -use vortex_array::expr::transform::PartitionedExpr; +use vortex_array::expr::BoundExpression; +use vortex_array::expr::transform::BoundPartitionedExpr; use vortex_array::validity::Validity; use vortex_error::VortexError; use vortex_error::VortexResult; @@ -22,28 +22,29 @@ use vortex_session::VortexSession; use crate::ArrayFuture; -pub trait PartitionedExprEval

{ +/// Evaluates cached bound partitions without rebuilding their expression trees per split. +pub(crate) trait BoundPartitionedExprEval

{ fn into_mask_future( self: Arc, mask: MaskFuture, - mask_fn: impl Fn(&P, &Expression, MaskFuture) -> VortexResult, - array_fn: impl Fn(&P, &Expression, MaskFuture) -> VortexResult, + mask_fn: impl Fn(&P, &BoundExpression, MaskFuture) -> VortexResult, + array_fn: impl Fn(&P, &BoundExpression, MaskFuture) -> VortexResult, session: VortexSession, ) -> VortexResult; fn into_array_future( self: Arc, mask: MaskFuture, - array_fn: impl Fn(&P, &Expression, MaskFuture) -> VortexResult, + array_fn: impl Fn(&P, &BoundExpression, MaskFuture) -> VortexResult, ) -> VortexResult; } -impl PartitionedExprEval

for PartitionedExpr

{ +impl BoundPartitionedExprEval

for BoundPartitionedExpr

{ fn into_mask_future( self: Arc, mask: MaskFuture, - mask_fn: impl Fn(&P, &Expression, MaskFuture) -> VortexResult, - array_fn: impl Fn(&P, &Expression, MaskFuture) -> VortexResult, + mask_fn: impl Fn(&P, &BoundExpression, MaskFuture) -> VortexResult, + array_fn: impl Fn(&P, &BoundExpression, MaskFuture) -> VortexResult, session: VortexSession, ) -> VortexResult { // Construct evaluations for each child. @@ -51,22 +52,23 @@ impl PartitionedExprEval

for PartitionedExpr

{ .partition_annotations .iter() .zip_eq(self.partitions.iter()) - .zip_eq(self.partition_dtypes.iter()) - .map(|((annotation, expr), dtype)| { - Ok::<_, VortexError>(if matches!(dtype, DType::Bool(Nullability::NonNullable)) { - // If the partition evaluates to a boolean, we can evaluate it as a mask which - // can often be more efficient since nulls are turned into `false` early on, - // and layouts can perform predicate pruning / indexing. - PartitionEval::Mask(mask_fn(annotation, expr, mask.clone())?) - } else { - // Otherwise, we evaluate the projection as an array, and combine the results - // at the end. - PartitionEval::Array(array_fn( - annotation, - expr, - MaskFuture::new_true(mask.len()), - )?) - }) + .map(|(annotation, expr)| { + Ok::<_, VortexError>( + if matches!(expr.dtype(), DType::Bool(Nullability::NonNullable)) { + // If the partition evaluates to a boolean, we can evaluate it as a mask which + // can often be more efficient since nulls are turned into `false` early on, + // and layouts can perform predicate pruning / indexing. + PartitionEval::Mask(mask_fn(annotation, expr, mask.clone())?) + } else { + // Otherwise, we evaluate the projection as an array, and combine the results + // at the end. + PartitionEval::Array(array_fn( + annotation, + expr, + MaskFuture::new_true(mask.len()), + )?) + }, + ) }) .try_collect()?; @@ -89,9 +91,8 @@ impl PartitionedExprEval

for PartitionedExpr

{ .into_array(); let mut ctx = session.create_execution_ctx(); - let root = self.root.bind(root_scope.dtype())?; let root_mask = root_scope - .apply_bound(&root)? + .apply_bound(&self.root)? .null_as_false() .execute(&mut ctx)?; @@ -104,7 +105,7 @@ impl PartitionedExprEval

for PartitionedExpr

{ fn into_array_future( self: Arc, mask: MaskFuture, - array_fn: impl Fn(&P, &Expression, MaskFuture) -> VortexResult, + array_fn: impl Fn(&P, &BoundExpression, MaskFuture) -> VortexResult, ) -> VortexResult { // Construct evaluations for each child. let field_evals: Vec<_> = self @@ -127,8 +128,7 @@ impl PartitionedExprEval

for PartitionedExpr

{ )? .into_array(); - let root = self.root.bind(root_scope.dtype())?; - root_scope.apply_bound(&root) + root_scope.apply_bound(&self.root) })) } } diff --git a/vortex-layout/src/layouts/row_idx/mod.rs b/vortex-layout/src/layouts/row_idx/mod.rs index f4d2da9ec7a..e7c83ec2950 100644 --- a/vortex-layout/src/layouts/row_idx/mod.rs +++ b/vortex-layout/src/layouts/row_idx/mod.rs @@ -26,12 +26,11 @@ use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::expr::BoundExpression; use vortex_array::expr::ExactBoundExpr; -use vortex_array::expr::Expression; -use vortex_array::expr::is_root; -use vortex_array::expr::root; -use vortex_array::expr::transform::PartitionedExpr; -use vortex_array::expr::transform::partition; -use vortex_array::expr::transform::replace; +use vortex_array::expr::transform::BoundPartitionedExpr; +use vortex_array::expr::transform::partition_bound; +use vortex_array::expr::traversal::NodeExt; +use vortex_array::expr::traversal::Transformed; +use vortex_array::expr::traversal::TraversalOrder; use vortex_array::scalar::PValue; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -45,7 +44,7 @@ use crate::ArrayFuture; use crate::LayoutReader; use crate::RowSplits; use crate::SplitRange; -use crate::layouts::partitioned::PartitionedExprEval; +use crate::layouts::partitioned::BoundPartitionedExprEval; pub struct RowIdxLayoutReader { name: Arc, @@ -76,7 +75,7 @@ impl RowIdxLayoutReader { return Ok(partitioning.clone()); } - let result = self.compute_partitioning(&expr.unbind())?; + let result = self.compute_partitioning(expr)?; self.partition_cache .entry(key) @@ -86,12 +85,15 @@ impl RowIdxLayoutReader { Ok(result) } - fn compute_partitioning(&self, expr: &Expression) -> VortexResult { + fn compute_partitioning(&self, expr: &BoundExpression) -> VortexResult { // Partition the expression into row idx and child expressions. - let mut partitioned = partition(expr.clone(), self.dtype(), |expr| { - if expr.is::() { + let mut partitioned = partition_bound(expr.clone(), |expr: &BoundExpression| { + if expr + .as_scalar() + .is_some_and(|scalar_fn| scalar_fn.is::()) + { vec![Partition::RowIdx] - } else if is_root(expr) { + } else if expr.is_root() { vec![Partition::Child] } else { vec![] @@ -101,19 +103,20 @@ impl RowIdxLayoutReader { // If there's only a single partition, we can directly return the expression. if partitioned.partitions.len() == 1 { return Ok(match &partitioned.partition_annotations[0] { - Partition::RowIdx => { - Partitioning::RowIdx(replace(expr.clone(), &row_idx(), root())) - } + Partition::RowIdx => Partitioning::RowIdx(replace_row_idx(expr.clone())?), Partition::Child => Partitioning::Child(expr.clone()), }); } // Replace the row_idx expression with the root expression in the row_idx partition. - partitioned.partitions = partitioned + let partitions = partitioned .partitions - .into_iter() - .map(|p| replace(p, &row_idx(), root())) - .collect(); + .iter() + .cloned() + .map(replace_row_idx) + .collect::>>()? + .into_boxed_slice(); + partitioned.replace_partitions(partitions)?; Ok(Partitioning::Partitioned(Arc::new(partitioned))) } @@ -122,11 +125,11 @@ impl RowIdxLayoutReader { #[derive(Clone)] enum Partitioning { // An expression that only references the row index (e.g., `row_idx == 5`). - RowIdx(Expression), + RowIdx(BoundExpression), // An expression that does not reference the row index. - Child(Expression), + Child(BoundExpression), // Contains both the RowIdx and Child expressions, (e.g., `row_idx < child.some_field`). - Partitioned(Arc>), + Partitioned(Arc>), } #[derive(Clone, PartialEq, Eq, Hash)] @@ -185,20 +188,14 @@ impl LayoutReader for RowIdxLayoutReader { mask: Mask, ) -> VortexResult { Ok(match &self.partition_expr(expr)? { - Partitioning::RowIdx(expr) => { - let expr = expr.bind(&row_idx_dtype())?; - row_idx_mask_future( - self.row_offset, - row_range, - expr, - MaskFuture::ready(mask), - self.session.clone(), - ) - } - Partitioning::Child(expr) => { - let expr = expr.bind(self.child.dtype())?; - self.child.pruning_evaluation(row_range, &expr, mask)? - } + Partitioning::RowIdx(expr) => row_idx_mask_future( + self.row_offset, + row_range, + expr.clone(), + MaskFuture::ready(mask), + self.session.clone(), + ), + Partitioning::Child(expr) => self.child.pruning_evaluation(row_range, expr, mask)?, Partitioning::Partitioned(..) => MaskFuture::ready(mask), }) } @@ -213,43 +210,28 @@ impl LayoutReader for RowIdxLayoutReader { // Since this is run during pruning, we skip re-evaluating the row index expression // during the filter evaluation. Partitioning::RowIdx(_) => Ok(mask), - Partitioning::Child(expr) => { - let expr = expr.bind(self.child.dtype())?; - self.child.filter_evaluation(row_range, &expr, mask) - } + Partitioning::Child(expr) => self.child.filter_evaluation(row_range, expr, mask), Partitioning::Partitioned(p) => Arc::clone(p).into_mask_future( mask, |annotation, expr, mask| match annotation { - Partition::RowIdx => { - let expr = expr.bind(&row_idx_dtype())?; - Ok(row_idx_mask_future( - self.row_offset, - row_range, - expr, - mask, - self.session.clone(), - )) - } - Partition::Child => { - let expr = expr.bind(self.child.dtype())?; - self.child.filter_evaluation(row_range, &expr, mask) - } + Partition::RowIdx => Ok(row_idx_mask_future( + self.row_offset, + row_range, + expr.clone(), + mask, + self.session.clone(), + )), + Partition::Child => self.child.filter_evaluation(row_range, expr, mask), }, |annotation, expr, mask| match annotation { - Partition::RowIdx => { - let expr = expr.bind(&row_idx_dtype())?; - Ok(row_idx_array_future( - self.row_offset, - row_range, - expr, - mask, - self.session.clone(), - )) - } - Partition::Child => { - let expr = expr.bind(self.child.dtype())?; - self.child.projection_evaluation(row_range, &expr, mask) - } + Partition::RowIdx => Ok(row_idx_array_future( + self.row_offset, + row_range, + expr.clone(), + mask, + self.session.clone(), + )), + Partition::Child => self.child.projection_evaluation(row_range, expr, mask), }, self.session.clone(), ), @@ -263,36 +245,24 @@ impl LayoutReader for RowIdxLayoutReader { mask: MaskFuture, ) -> VortexResult>> { match &self.partition_expr(expr)? { - Partitioning::RowIdx(expr) => { - let expr = expr.bind(&row_idx_dtype())?; - Ok(row_idx_array_future( - self.row_offset, - row_range, - expr, - mask, - self.session.clone(), - )) - } - Partitioning::Child(expr) => { - let expr = expr.bind(self.child.dtype())?; - self.child.projection_evaluation(row_range, &expr, mask) - } + Partitioning::RowIdx(expr) => Ok(row_idx_array_future( + self.row_offset, + row_range, + expr.clone(), + mask, + self.session.clone(), + )), + Partitioning::Child(expr) => self.child.projection_evaluation(row_range, expr, mask), Partitioning::Partitioned(p) => { Arc::clone(p).into_array_future(mask, |annotation, expr, mask| match annotation { - Partition::RowIdx => { - let expr = expr.bind(&row_idx_dtype())?; - Ok(row_idx_array_future( - self.row_offset, - row_range, - expr, - mask, - self.session.clone(), - )) - } - Partition::Child => { - let expr = expr.bind(self.child.dtype())?; - self.child.projection_evaluation(row_range, &expr, mask) - } + Partition::RowIdx => Ok(row_idx_array_future( + self.row_offset, + row_range, + expr.clone(), + mask, + self.session.clone(), + )), + Partition::Child => self.child.projection_evaluation(row_range, expr, mask), }) } } @@ -303,6 +273,25 @@ impl LayoutReader for RowIdxLayoutReader { } } +fn replace_row_idx(expr: BoundExpression) -> VortexResult { + Ok(expr + .transform_down(|node| { + if node + .as_scalar() + .is_some_and(|scalar_fn| scalar_fn.is::()) + { + Ok(Transformed { + value: BoundExpression::new_root(row_idx_dtype()), + changed: true, + order: TraversalOrder::Skip, + }) + } else { + Ok(Transformed::no(node)) + } + })? + .into_inner()) +} + fn row_idx_dtype() -> DType { DType::Primitive(PType::U64, NonNullable) } diff --git a/vortex-layout/src/layouts/struct_/reader.rs b/vortex-layout/src/layouts/struct_/reader.rs index 073cb2bd6ec..1c0a9006f48 100644 --- a/vortex-layout/src/layouts/struct_/reader.rs +++ b/vortex-layout/src/layouts/struct_/reader.rs @@ -17,20 +17,24 @@ use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; use vortex_array::dtype::FieldName; +use vortex_array::dtype::FieldNames; use vortex_array::dtype::Nullability; use vortex_array::dtype::StructFields; use vortex_array::expr::BoundExpression; use vortex_array::expr::ExactBoundExpr; -use vortex_array::expr::Expression; -use vortex_array::expr::col; -use vortex_array::expr::make_free_field_annotator; +use vortex_array::expr::make_bound_free_field_annotator; use vortex_array::expr::root; -use vortex_array::expr::transform::PartitionedExpr; -use vortex_array::expr::transform::partition; -use vortex_array::expr::transform::replace; -use vortex_array::expr::transform::replace_root_fields; +use vortex_array::expr::transform::BoundPartitionedExpr; +use vortex_array::expr::transform::partition_bound; +use vortex_array::expr::traversal::NodeExt; +use vortex_array::expr::traversal::Transformed; +use vortex_array::expr::traversal::TraversalOrder; +use vortex_array::scalar_fn::ScalarFnVTableExt; +use vortex_array::scalar_fn::fns::get_item::GetItem; use vortex_array::scalar_fn::fns::merge::Merge; use vortex_array::scalar_fn::fns::pack::Pack; +use vortex_array::scalar_fn::fns::pack::PackOptions; +use vortex_array::scalar_fn::fns::select::Select; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_err; @@ -45,7 +49,7 @@ use crate::LayoutReaderRef; use crate::LazyReaderChildren; use crate::RowSplits; use crate::SplitRange; -use crate::layouts::partitioned::PartitionedExprEval; +use crate::layouts::partitioned::BoundPartitionedExprEval; use crate::layouts::struct_::StructLayout; use crate::segments::SegmentSource; @@ -57,7 +61,7 @@ pub struct StructReader { /// A `pack` expression that holds each individual field of the root DType. This expansion /// ensures we can correctly partition expressions over the fields of the struct. - expanded_root_expr: Expression, + expanded_root_expr: BoundExpression, field_lookup: Option>, partitioned_expr_cache: DashMap>>, @@ -105,7 +109,7 @@ impl StructReader { ); // Create an expanded root expression that contains all fields of the struct. - let expanded_root_expr = replace_root_fields(root(), struct_dt); + let expanded_root_expr = expanded_struct_root(layout.dtype(), struct_dt)?; // This is where we need to do some complex things with the scan in order to split it into // different scans for different fields. @@ -176,21 +180,22 @@ impl StructReader { if let Some(value) = cell.get() { return Ok(value.clone()); } - let result = self.compute_partitioned_expr(expr.unbind())?; + // Cache the bound trees themselves: rebinding would create fresh child identity and defeat + // downstream `ExactBoundExpr` caches on every split. + let result = self.compute_partitioned_expr(expr)?; Ok(cell.get_or_init(|| result).clone()) } - fn compute_partitioned_expr(&self, expr: Expression) -> VortexResult { + fn compute_partitioned_expr(&self, expr: &BoundExpression) -> VortexResult { // First, we expand the root scope into the fields of the struct to ensure // that partitioning works correctly. - let expr = replace(expr, &root(), self.expanded_root_expr.clone()); - let expr = expr.optimize_recursive(self.dtype())?; + let expr = + expand_struct_root(expr.clone(), &self.expanded_root_expr, self.struct_fields())?; // Partition the expression into expressions that can be evaluated over individual fields - let mut partitioned = partition( + let mut partitioned = partition_bound( expr.clone(), - self.dtype(), - make_free_field_annotator( + make_bound_free_field_annotator( self.dtype() .as_struct_fields_opt() .vortex_expect("We know it's a struct DType"), @@ -202,33 +207,157 @@ impl StructReader { // expression by replacing any `$.a` with `$`. return Ok(Partitioned::Single( partitioned.partition_names[0].clone(), - replace(expr, &col(partitioned.partition_names[0].clone()), root()), + step_into_struct_field( + expr, + &partitioned.partition_names[0], + self.field_reader(&partitioned.partition_names[0])? + .dtype() + .clone(), + )?, )); } // We now need to process the partitioned expressions to rewrite the root scope // to be that of the field, rather than the struct. In other words, "stepping in" // to the field scope. - partitioned.partitions = partitioned + let partitions = partitioned .partitions .iter() .zip_eq(partitioned.partition_names.iter()) - .map(|(e, name)| replace(e.clone(), &col(name.clone()), root())) - .collect(); + .map(|(expr, name)| { + step_into_struct_field(expr.clone(), name, self.field_reader(name)?.dtype().clone()) + }) + .try_collect::<_, Vec<_>, _>()? + .into_boxed_slice(); + partitioned.replace_partitions(partitions)?; Ok(Partitioned::Multi(Arc::new(partitioned))) } } +fn expanded_struct_root( + root_dtype: &DType, + fields: &StructFields, +) -> VortexResult { + let root = BoundExpression::new_root(root_dtype.clone()); + let children = fields + .names() + .iter() + .map(|name| BoundExpression::try_new(GetItem.bind(name.clone()), [root.clone()])) + .try_collect()?; + bound_pack(fields.names().clone(), children) +} + +fn expand_struct_root( + expr: BoundExpression, + expanded_root: &BoundExpression, + fields: &StructFields, +) -> VortexResult { + Ok(expr + .transform_down(|node| { + if node.is_root() { + return Ok(Transformed { + value: expanded_root.clone(), + changed: true, + order: TraversalOrder::Skip, + }); + } + + let Some(scalar_fn) = node.as_scalar() else { + return Ok(Transformed::no(node)); + }; + if !node + .children() + .first() + .is_some_and(BoundExpression::is_root) + { + return Ok(Transformed::no(node)); + } + + if let Some(field_name) = scalar_fn.as_opt::() { + let idx = fields.find(field_name).ok_or_else(|| { + vortex_err!("Field {field_name} not found while expanding struct root") + })?; + return Ok(Transformed { + value: expanded_root.children()[idx].clone(), + changed: true, + order: TraversalOrder::Skip, + }); + } + + if let Some(selection) = scalar_fn.as_opt::