diff --git a/vortex-array/src/expr/analysis/immediate_access.rs b/vortex-array/src/expr/analysis/immediate_access.rs index 8df9e056f7e..fd720607284 100644 --- a/vortex-array/src/expr/analysis/immediate_access.rs +++ b/vortex-array/src/expr/analysis/immediate_access.rs @@ -2,21 +2,16 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_error::VortexExpect; -use vortex_utils::aliases::hash_set::HashSet; use crate::dtype::FieldName; use crate::dtype::StructFields; use crate::expr::BoundExpression; use crate::expr::Expression; use crate::expr::analysis::AnnotationFn; -use crate::expr::analysis::Annotations; -use crate::expr::descendent_annotations; use crate::scalar_fn::fns::get_item::GetItem; use crate::scalar_fn::fns::root::Root; use crate::scalar_fn::fns::select::Select; -pub type FieldAccesses<'a> = Annotations<'a, FieldName>; - /// Returns the "free fields" for this expression node. /// /// A "free field" is a top-level field from the root scope that this expression references—not @@ -92,28 +87,3 @@ pub fn make_bound_free_field_annotator( vec![] } } - -/// For all subexpressions in an expression, find the fields that are accessed directly from the -/// scope, but not any fields in those fields -/// e.g. scope = {a: {b: .., c: ..}, d: ..}, expr = root().a.b + root().d accesses {a,d} (not b). -/// -/// Note: This is a very naive, but simple analysis to find the fields that are accessed directly on an -/// identity node. This is combined to provide an over-approximation of the fields that are accessed -/// by an expression. -pub fn immediate_scope_accesses<'a>( - expr: &'a Expression, - scope: &'a StructFields, -) -> FieldAccesses<'a> { - descendent_annotations(expr, make_free_field_annotator(scope)) -} - -/// This returns the immediate scope_access (as explained `immediate_scope_accesses`) for `expr`. -pub fn immediate_scope_access<'a>( - expr: &'a Expression, - scope: &'a StructFields, -) -> HashSet { - immediate_scope_accesses(expr, scope) - .get(expr) - .vortex_expect("Expression missing from scope accesses, this is a internal bug") - .clone() -} 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-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..20102a2ad66 100644 --- a/vortex-file/src/v2/file_stats_reader.rs +++ b/vortex-file/src/v2/file_stats_reader.rs @@ -14,6 +14,8 @@ 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::ExactBoundExpr; use vortex_array::expr::Expression; use vortex_error::VortexResult; use vortex_layout::ArrayFuture; @@ -41,7 +43,7 @@ pub struct FileStatsLayoutReader { file_stats: FileStatistics, struct_fields: StructFields, session: VortexSession, - prune_cache: DashMap, + prune_cache: DashMap, } impl FileStatsLayoutReader { @@ -113,11 +115,13 @@ impl LayoutReader for FileStatsLayoutReader { fn pruning_evaluation( &self, row_range: &Range, - expr: &Expression, + 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()))); } @@ -125,8 +129,9 @@ impl LayoutReader for FileStatsLayoutReader { } // Evaluate and cache. - let pruned = self.evaluate_file_stats(expr)?; - self.prune_cache.insert(expr.clone(), pruned); + let expression = expr.unbind(); + let pruned = self.evaluate_file_stats(&expression)?; + self.prune_cache.insert(key, pruned); if pruned { Ok(MaskFuture::ready(Mask::new_false(mask.len()))) @@ -138,7 +143,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 +152,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 +265,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 +304,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 +341,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 +397,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 +441,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..eefdee6e4c4 100644 --- a/vortex-layout/src/layouts/dict/reader.rs +++ b/vortex-layout/src/layouts/dict/reader.rs @@ -19,14 +19,16 @@ use vortex_array::arrays::SharedArray; use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; use vortex_array::dtype::Nullability; -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::BoundExpression; +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; @@ -53,7 +55,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 +105,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 +130,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 +147,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 +163,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() @@ -178,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 { @@ -228,7 +242,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 +255,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 +264,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 +286,44 @@ 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)?; 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 = 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 // 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 +345,7 @@ impl LayoutReader for DictReader { .into_array() .optimize()?; - array.apply(&expr_outer) + array.apply_bound(&expr_outer) } .boxed()) } @@ -483,9 +504,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 +592,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 +657,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 +714,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() @@ -735,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); } @@ -747,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) @@ -760,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())])); @@ -771,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/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..f446ad11ab6 100644 --- a/vortex-layout/src/layouts/list/expr.rs +++ b/vortex-layout/src/layouts/list/expr.rs @@ -1,13 +1,15 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_array::expr::Expression; -use vortex_array::expr::is_root; -use vortex_array::expr::not; -use vortex_array::expr::root; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::expr::BoundExpression; +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. @@ -26,53 +28,75 @@ pub(super) enum ListChildrenNeeded { All, } -/// The minimal set of list children needed to evaluate `expr`, where `root()` is a field with list dtype. -pub(super) fn get_necessary_list_children(expr: &Expression) -> ListChildrenNeeded { - if is_null_root(expr) { +/// 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_list_length_root(expr) { + if is_bound_list_length_root(expr) { return ListChildrenNeeded::OffsetsAndValidity; } - if is_root(expr) { + if expr.is_root() { return ListChildrenNeeded::All; } - // Otherwise the requirement is the max over the operands. Childless expressions that never - // touch the list, such as literals, fall back to the cheapest usable child. expr.children() .iter() - .map(get_necessary_list_children) + .map(get_necessary_bound_list_children) .max() .unwrap_or(ListChildrenNeeded::Validity) } -fn is_null_root(expr: &Expression) -> bool { - (expr.is::() || expr.is::()) +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 - && is_root(expr.child(0)) + && expr.children()[0].is_root() } -fn is_list_length_root(expr: &Expression) -> bool { - expr.is::() && expr.children().len() == 1 && is_root(expr.child(0)) +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() } /// 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.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::() && expr.children().len() == 1 && is_root(expr.child(0)) { - return Ok(not(root())); + 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,70 +105,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) } - -#[cfg(test)] -mod tests { - use rstest::rstest; - use vortex_array::dtype::DType; - use vortex_array::dtype::Nullability; - use vortex_array::dtype::PType; - use vortex_array::expr::cast; - use vortex_array::expr::eq; - use vortex_array::expr::gt; - use vortex_array::expr::is_not_null; - use vortex_array::expr::is_null; - use vortex_array::expr::list_length; - use vortex_array::expr::lit; - use vortex_array::expr::not; - use vortex_array::expr::root; - - use super::*; - - /// `get_necessary_list_children` keys off the deepest list child an expression touches; `All` - /// is the always-correct default for anything not specifically recognized. - #[rstest] - // `is_null` / `is_not_null` of the list itself need only validity. - #[case::is_null(is_null(root()), ListChildrenNeeded::Validity)] - #[case::is_not_null(is_not_null(root()), ListChildrenNeeded::Validity)] - // Compound over validity-only operands stays validity. - #[case::not_is_null(not(is_null(root())), ListChildrenNeeded::Validity)] - // A list-independent (constant) expression falls to the cheapest usable child. - #[case::constant(lit(5), ListChildrenNeeded::Validity)] - // `list_length(root())` needs offsets and validity, but not elements. - #[case::list_length(list_length(root()), ListChildrenNeeded::OffsetsAndValidity)] - // Compound over offsets-only operands stays offsets. - #[case::list_length_filter( - gt(list_length(root()), lit(1u64)), - ListChildrenNeeded::OffsetsAndValidity - )] - #[case::cast_list_length( - cast( - list_length(root()), - DType::Primitive(PType::I64, Nullability::Nullable), - ), - ListChildrenNeeded::OffsetsAndValidity - )] - // A bare list reference needs the elements. - #[case::bare_root(root(), ListChildrenNeeded::All)] - // Any other fn over the list needs the elements. - #[case::not_root(not(root()), ListChildrenNeeded::All)] - // `is_null` only short-circuits to validity when its argument is the list itself. - #[case::is_null_of_derived(is_null(not(root())), ListChildrenNeeded::All)] - // Max over operands: validity + elements => elements. - #[case::validity_and_elements(eq(is_null(root()), root()), ListChildrenNeeded::All)] - fn classify_expr_class(#[case] expr: Expression, #[case] expected: ListChildrenNeeded) { - assert_eq!(get_necessary_list_children(&expr), expected); - } -} 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..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()?; @@ -90,7 +92,7 @@ impl PartitionedExprEval

for PartitionedExpr

{ let mut ctx = session.create_execution_ctx(); let root_mask = root_scope - .apply(&self.root)? + .apply_bound(&self.root)? .null_as_false() .execute(&mut ctx)?; @@ -103,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 @@ -126,7 +128,7 @@ impl PartitionedExprEval

for PartitionedExpr

{ )? .into_array(); - root_scope.apply(&self.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 7eaa820ce2a..e7c83ec2950 100644 --- a/vortex-layout/src/layouts/row_idx/mod.rs +++ b/vortex-layout/src/layouts/row_idx/mod.rs @@ -24,13 +24,13 @@ 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::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::BoundExpression; +use vortex_array::expr::ExactBoundExpr; +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; @@ -44,13 +44,13 @@ 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, row_offset: u64, child: Arc, - partition_cache: DashMap>>, + partition_cache: DashMap>>, session: VortexSession, } @@ -65,8 +65,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) @@ -85,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![] @@ -100,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))) } @@ -121,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)] @@ -180,14 +184,14 @@ 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, + expr.clone(), MaskFuture::ready(mask), self.session.clone(), ), @@ -199,7 +203,7 @@ impl LayoutReader for RowIdxLayoutReader { fn filter_evaluation( &self, row_range: &Range, - expr: &Expression, + expr: &BoundExpression, mask: MaskFuture, ) -> VortexResult { match &self.partition_expr(expr)? { @@ -213,7 +217,7 @@ impl LayoutReader for RowIdxLayoutReader { Partition::RowIdx => Ok(row_idx_mask_future( self.row_offset, row_range, - expr, + expr.clone(), mask, self.session.clone(), )), @@ -223,7 +227,7 @@ impl LayoutReader for RowIdxLayoutReader { Partition::RowIdx => Ok(row_idx_array_future( self.row_offset, row_range, - expr, + expr.clone(), mask, self.session.clone(), )), @@ -237,14 +241,14 @@ 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, + expr.clone(), mask, self.session.clone(), )), @@ -254,7 +258,7 @@ impl LayoutReader for RowIdxLayoutReader { Partition::RowIdx => Ok(row_idx_array_future( self.row_offset, row_range, - expr, + expr.clone(), mask, self.session.clone(), )), @@ -269,6 +273,29 @@ 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) +} + // 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 +312,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 +333,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 +398,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 +445,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 +496,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..1c0a9006f48 100644 --- a/vortex-layout/src/layouts/struct_/reader.rs +++ b/vortex-layout/src/layouts/struct_/reader.rs @@ -17,19 +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::ExactExpr; -use vortex_array::expr::Expression; -use vortex_array::expr::col; -use vortex_array::expr::make_free_field_annotator; +use vortex_array::expr::BoundExpression; +use vortex_array::expr::ExactBoundExpr; +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; @@ -44,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; @@ -56,10 +61,10 @@ 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>>, + partitioned_expr_cache: DashMap>>, } impl StructReader { @@ -104,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. @@ -155,8 +160,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,21 +180,22 @@ impl StructReader { if let Some(value) = cell.get() { return Ok(value.clone()); } + // 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"), @@ -201,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::