From cc7d9f477a8a5ffffccec8a03b14f082f88b78b3 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 7 Aug 2026 11:22:08 +0000 Subject: [PATCH] refactor: build known scan expressions as bound expressions directly MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Scan callers that know their projection or filter statically were writing `expr.optimize_recursive(dtype)?.bind(dtype)?`. For these shapes the optimizer pass is a no-op, so the round trip only costs a tree walk and produces a fresh tree identity that defeats the identity-keyed caches the parent PR introduces. Build them with the `bound::*` constructors instead, at the crate doc example, the `vortex` and `vortex-file` tests, and the compress/TPC-H benchmarks. The one exception is `and(gt, lt_eq)` over a single column, which the optimizer folds into a `between`; that site now constructs the `between` directly, which is the form the scan was already receiving. Callers whose expression arrives from outside — Python, DataFusion, scan requests and the fuzz target — still optimize and bind, as do the `vortex-file` tests that deliberately exercise expressions which bind but fail during execution. Add `bound_constructors_match_optimize_then_bind`, which asserts each converted shape equals its `optimize_recursive(..).bind(..)` result so the two cannot drift apart. Signed-off-by: Claude --- benchmarks/compress-bench/src/vortex.rs | 7 +- vortex-array/src/expr/mod.rs | 92 +++++++++++++ vortex-bench/src/datasets/tpch_l_comment.rs | 10 +- vortex-file/src/tests.rs | 138 +++++++++----------- vortex/src/lib.rs | 19 +-- 5 files changed, 168 insertions(+), 98 deletions(-) diff --git a/benchmarks/compress-bench/src/vortex.rs b/benchmarks/compress-bench/src/vortex.rs index 0726886bc4c..96f2108adf1 100644 --- a/benchmarks/compress-bench/src/vortex.rs +++ b/benchmarks/compress-bench/src/vortex.rs @@ -14,8 +14,7 @@ use futures::StreamExt; use futures::pin_mut; use vortex::array::IntoArray; use vortex::dtype::FieldNames; -use vortex::expr::root; -use vortex::expr::select; +use vortex::expr::bound; use vortex::file::OpenOptionsSessionExt; use vortex::file::WriteOptionsSessionExt; use vortex_arrow::ToArrowType; @@ -71,9 +70,7 @@ impl Compressor for VortexCompressor { if let Some(cols) = read_projection(root_columns) { // Columns are named "0".."num_columns-1"; project the given subset. let names: FieldNames = cols.iter().map(|i| i.to_string()).collect(); - let projection = select(names, root()) - .optimize_recursive(&source_dtype)? - .bind(&source_dtype)?; + let projection = bound::select(names, bound::root(source_dtype.clone())); scan = scan.with_projection(projection); } let schema = Arc::new(scan.dtype()?.to_arrow_schema()?); diff --git a/vortex-array/src/expr/mod.rs b/vortex-array/src/expr/mod.rs index 0ddd003e65d..ece6328bfaf 100644 --- a/vortex-array/src/expr/mod.rs +++ b/vortex-array/src/expr/mod.rs @@ -225,9 +225,12 @@ mod tests { use crate::expr::not; use crate::expr::not_eq; use crate::expr::or; + use crate::expr::pack; use crate::expr::select; use crate::expr::select_exclude; use crate::scalar::Scalar; + use crate::scalar_fn::fns::between::BetweenOptions; + use crate::scalar_fn::fns::between::StrictComparison; use crate::scalar_fn::fns::literal::Literal; #[test] @@ -410,4 +413,93 @@ mod tests { let expression = root(); assert!(!expression.contains::().unwrap()); } + + /// Scan callers that know their expression statically build it with the `bound::*` + /// constructors instead of `optimize_recursive(..).bind(..)`. Both must agree, otherwise + /// those call sites would silently ship a different expression to the scan. + #[test] + fn bound_constructors_match_optimize_then_bind() -> vortex_error::VortexResult<()> { + let u64_scope = DType::Primitive(PType::U64, Nullability::NonNullable); + let struct_scope = DType::Struct( + StructFields::from_iter([ + ("name", DType::Utf8(Nullability::Nullable)), + ("age", DType::Primitive(PType::I32, Nullability::Nullable)), + ]), + Nullability::NonNullable, + ); + + let cases: Vec<(DType, Expression, BoundExpression)> = vec![ + ( + u64_scope.clone(), + gt(root(), lit(2u64)), + bound::gt(bound::root(u64_scope), bound::lit(2u64)), + ), + ( + struct_scope.clone(), + select(["name"], root()), + bound::select(["name"], bound::root(struct_scope.clone())), + ), + ( + struct_scope.clone(), + pack([("name", col("name"))], Nullability::NonNullable), + bound::pack( + [("name", bound::col("name", struct_scope.clone()))], + Nullability::NonNullable, + ), + ), + ( + struct_scope.clone(), + eq(get_item("name", root()), lit("Joseph")), + bound::eq( + bound::col("name", struct_scope.clone()), + bound::lit("Joseph"), + ), + ), + ( + struct_scope.clone(), + or( + eq(get_item("name", root()), lit("Angela")), + and( + gt_eq(get_item("age", root()), lit(20)), + lt_eq(get_item("age", root()), lit(30)), + ), + ), + bound::or( + bound::eq( + bound::col("name", struct_scope.clone()), + bound::lit("Angela"), + ), + bound::and( + bound::gt_eq(bound::col("age", struct_scope.clone()), bound::lit(20)), + bound::lt_eq(bound::col("age", struct_scope.clone()), bound::lit(30)), + ), + ), + ), + // `and(gt, lt_eq)` over a single column is folded into a `between`, so the direct + // form is the `between` rather than the conjunction it was written as. + ( + struct_scope.clone(), + and( + gt(get_item("age", root()), lit(21)), + lt_eq(get_item("age", root()), lit(33)), + ), + bound::between( + bound::col("age", struct_scope), + bound::lit(21), + bound::lit(33), + BetweenOptions { + lower_strict: StrictComparison::Strict, + upper_strict: StrictComparison::NonStrict, + }, + ), + ), + ]; + + for (scope, unbound, direct) in cases { + let via_optimizer = unbound.optimize_recursive(&scope)?.bind(&scope)?; + assert_eq!(via_optimizer, direct, "mismatch for {unbound}"); + } + + Ok(()) + } } diff --git a/vortex-bench/src/datasets/tpch_l_comment.rs b/vortex-bench/src/datasets/tpch_l_comment.rs index c57bc91a65d..5541ad2e8b3 100644 --- a/vortex-bench/src/datasets/tpch_l_comment.rs +++ b/vortex-bench/src/datasets/tpch_l_comment.rs @@ -13,8 +13,7 @@ use vortex::array::IntoArray; use vortex::array::arrays::ChunkedArray; use vortex::array::arrays::StructArray; use vortex::dtype::Nullability::NonNullable; -use vortex::expr::col; -use vortex::expr::pack; +use vortex::expr::bound; use vortex::file::OpenOptionsSessionExt; use crate::Format; @@ -66,9 +65,10 @@ impl Dataset for TPCHLCommentChunked { let path = data_dir.join("lineitem.vortex"); let file = SESSION.open_options().open_path(path).await?; - let projection = pack(vec![("l_comment", col("l_comment"))], NonNullable) - .optimize_recursive(file.dtype())? - .bind(file.dtype())?; + let projection = bound::pack( + vec![("l_comment", bound::col("l_comment", file.dtype().clone()))], + NonNullable, + ); let chunks: Vec<_> = file .scan()? .with_projection(projection) diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index f5c177c9cdf..7c26f9aea89 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -39,28 +39,23 @@ use vortex_array::dtype::PType; use vortex_array::dtype::PType::I32; use vortex_array::dtype::StructFields; use vortex_array::expr::BoundExpression; -use vortex_array::expr::Expression; use vortex_array::expr::and; -use vortex_array::expr::cast; +use vortex_array::expr::bound; use vortex_array::expr::col; use vortex_array::expr::eq; use vortex_array::expr::get_item; use vortex_array::expr::gt; -use vortex_array::expr::gt_eq; use vortex_array::expr::lit; use vortex_array::expr::lt; -use vortex_array::expr::lt_eq; use vortex_array::expr::or; use vortex_array::expr::root; -use vortex_array::expr::select; use vortex_array::extension::datetime::TimeUnit; use vortex_array::extension::datetime::Timestamp; use vortex_array::extension::datetime::TimestampOptions; use vortex_array::field_path; use vortex_array::scalar::Scalar; -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::fns::between::BetweenOptions; +use vortex_array::scalar_fn::fns::between::StrictComparison; use vortex_array::stats::PRUNING_STATS; use vortex_array::stream::ArrayStreamAdapter; use vortex_array::stream::ArrayStreamExt; @@ -115,12 +110,6 @@ fn strict_sorted(indices: Buffer) -> StrictSortedBuffer { StrictSortedBuffer::try_new(indices).expect("test indices should be strictly increasing") } -fn bind_scan_expr(file: &VortexFile, expr: Expression) -> BoundExpression { - expr.optimize_recursive(file.dtype()) - .and_then(|expr| expr.bind(file.dtype())) - .vortex_expect("scan expression should bind") -} - #[tokio::test] async fn test_eof_values() { // this test exists as a reminder to think about whether we should increment the version @@ -328,7 +317,10 @@ async fn test_read_projection() { let array = file .scan() .unwrap() - .with_projection(bind_scan_expr(&file, select(["strings"], root()))) + .with_projection(bound::select( + ["strings"], + bound::root(file.dtype().clone()), + )) .into_array_stream() .unwrap() .read_all() @@ -354,7 +346,10 @@ async fn test_read_projection() { let array = file .scan() .unwrap() - .with_projection(bind_scan_expr(&file, select(["numbers"], root()))) + .with_projection(bound::select( + ["numbers"], + bound::root(file.dtype().clone()), + )) .into_array_stream() .unwrap() .read_all() @@ -529,15 +524,12 @@ async fn issue_5385_filter_casted_column() { let result = file .scan() .unwrap() - .with_filter(bind_scan_expr( - &file, - eq( - cast( - get_item("x", root()), - DType::Primitive(PType::U16, Nullability::NonNullable), - ), - lit(1u16), + .with_filter(bound::eq( + bound::cast( + bound::col("x", file.dtype().clone()), + DType::Primitive(PType::U16, Nullability::NonNullable), ), + bound::lit(1u16), )) .into_array_stream() .unwrap() @@ -582,9 +574,9 @@ async fn filter_string() { let result: Vec<_> = file .scan() .unwrap() - .with_filter(bind_scan_expr( - &file, - eq(get_item("name", root()), lit("Joseph")), + .with_filter(bound::eq( + bound::col("name", file.dtype().clone()), + bound::lit("Joseph"), )) .into_array_stream() .unwrap() @@ -643,14 +635,14 @@ async fn filter_or() { let result: Vec<_> = file .scan() .unwrap() - .with_filter(bind_scan_expr( - &file, - or( - eq(get_item("name", root()), lit("Angela")), - and( - gt_eq(get_item("age", root()), lit(20)), - lt_eq(get_item("age", root()), lit(30)), - ), + .with_filter(bound::or( + bound::eq( + bound::col("name", file.dtype().clone()), + bound::lit("Angela"), + ), + bound::and( + bound::gt_eq(bound::col("age", file.dtype().clone()), bound::lit(20)), + bound::lt_eq(bound::col("age", file.dtype().clone()), bound::lit(30)), ), )) .into_array_stream() @@ -712,12 +704,16 @@ async fn filter_and() { let result: Vec<_> = file .scan() .unwrap() - .with_filter(bind_scan_expr( - &file, - and( - gt(get_item("age", root()), lit(21)), - lt_eq(get_item("age", root()), lit(33)), - ), + // `and(gt, lt_eq)` over one column is what the optimizer folds into a `between`, so + // build that form directly rather than round-tripping through the optimizer. + .with_filter(bound::between( + bound::col("age", file.dtype().clone()), + bound::lit(21), + bound::lit(33), + BetweenOptions { + lower_strict: StrictComparison::Strict, + upper_strict: StrictComparison::NonStrict, + }, )) .into_array_stream() .unwrap() @@ -923,9 +919,9 @@ async fn test_with_indices_and_with_row_filter_simple() { let actual_kept_array = file .scan() .unwrap() - .with_filter(bind_scan_expr( - &file, - gt(get_item("numbers", root()), lit(50_i16)), + .with_filter(bound::gt( + bound::col("numbers", file.dtype().clone()), + bound::lit(50_i16), )) .with_row_indices(strict_sorted(Buffer::empty())) .into_array_stream() @@ -944,9 +940,9 @@ async fn test_with_indices_and_with_row_filter_simple() { let actual_kept_array = file .scan() .unwrap() - .with_filter(bind_scan_expr( - &file, - gt(get_item("numbers", root()), lit(50_i16)), + .with_filter(bound::gt( + bound::col("numbers", file.dtype().clone()), + bound::lit(50_i16), )) .with_row_indices(strict_sorted(Buffer::from_iter(kept_indices))) .into_array_stream() @@ -975,9 +971,9 @@ async fn test_with_indices_and_with_row_filter_simple() { let actual_array = file .scan() .unwrap() - .with_filter(bind_scan_expr( - &file, - gt(get_item("numbers", root()), lit(50_i16)), + .with_filter(bound::gt( + bound::col("numbers", file.dtype().clone()), + bound::lit(50_i16), )) .with_row_indices(strict_sorted((0..500).collect::>())) .into_array_stream() @@ -1040,9 +1036,9 @@ async fn filter_string_chunked() { let actual_array = file .scan() .unwrap() - .with_filter(bind_scan_expr( - &file, - eq(get_item("name", root()), lit("Joseph")), + .with_filter(bound::eq( + bound::col("name", file.dtype().clone()), + bound::lit("Joseph"), )) .into_array_stream() .unwrap() @@ -1133,12 +1129,9 @@ async fn test_pruning_with_or() { let actual_array = file .scan() .unwrap() - .with_filter(bind_scan_expr( - &file, - or( - lt_eq(get_item("letter", root()), lit("J")), - lt(get_item("number", root()), lit(25)), - ), + .with_filter(bound::or( + bound::lt_eq(bound::col("letter", file.dtype().clone()), bound::lit("J")), + bound::lt(bound::col("number", file.dtype().clone()), bound::lit(25)), )) .into_array_stream() .unwrap() @@ -1211,9 +1204,9 @@ async fn test_repeated_projection() { let actual = file .scan() .unwrap() - .with_projection(bind_scan_expr( - &file, - select(["strings", "strings"], root()), + .with_projection(bound::select( + ["strings", "strings"], + bound::root(file.dtype().clone()), )) .into_array_stream() .unwrap() @@ -1391,16 +1384,7 @@ async fn write_nullable_nested_struct() -> VortexResult<()> { #[tokio::test] async fn scan_empty_fields() -> VortexResult<()> { let array = (0..10000).collect::(); - let projection = Pack - .new_expr( - PackOptions { - names: Default::default(), - nullability: Nullability::Nullable, - }, - [], - ) - .optimize_recursive(array.dtype())? - .bind(array.dtype())?; + let projection = bound::pack(Vec::<(&str, BoundExpression)>::new(), Nullability::Nullable); let result = round_trip(&array.clone().into_array(), |scan| { Ok(scan.with_projection(projection)) @@ -2375,7 +2359,10 @@ async fn test_large_flat_chunk_scan_subdivides_splits() -> VortexResult<()> { // A filtered scan crossing sub-split boundaries selects exactly the matching rows. let result = file .scan()? - .with_filter(bind_scan_expr(&file, gt(root(), lit(0i32)))) + .with_filter(bound::gt( + bound::root(file.dtype().clone()), + bound::lit(0i32), + )) .into_array_stream()? .read_all() .await?; @@ -2422,7 +2409,10 @@ async fn test_flat_chunk_scan_with_row_count_splits( let result = file .scan()? .with_split_by(SplitBy::RowCount(rows_per_split)) - .with_filter(bind_scan_expr(&file, gt(root(), lit(0i32)))) + .with_filter(bound::gt( + bound::root(file.dtype().clone()), + bound::lit(0i32), + )) .into_array_stream()? .read_all() .await?; diff --git a/vortex/src/lib.rs b/vortex/src/lib.rs index 0bc9fe53f65..4b20867e89b 100644 --- a/vortex/src/lib.rs +++ b/vortex/src/lib.rs @@ -66,7 +66,7 @@ //! use vortex::VortexSessionDefault; //! use vortex::array::{IntoArray, stream::ArrayStreamExt}; //! use vortex::array::arrays::PrimitiveArray; -//! use vortex::array::expr::{gt, lit, root}; +//! use vortex::array::expr::bound; //! use vortex::array::validity::Validity; //! use vortex::buffer::{ByteBufferMut, buffer}; //! use vortex::file::{OpenOptionsSessionExt, WriteOptionsSessionExt}; @@ -85,9 +85,7 @@ //! let file = session //! .open_options() //! .open_buffer(bytes)?; -//! let filter = gt(root(), lit(2u64)) -//! .optimize_recursive(file.dtype())? -//! .bind(file.dtype())?; +//! let filter = bound::gt(bound::root(file.dtype().clone()), bound::lit(2u64)); //! let filtered = file //! .scan()? //! .with_filter(filter) @@ -349,10 +347,7 @@ mod test { use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; use vortex_array::dtype::FieldNames; - use vortex_array::expr::gt; - use vortex_array::expr::lit; - use vortex_array::expr::root; - use vortex_array::expr::select; + use vortex_array::expr::bound; use vortex_array::stream::ArrayStreamExt; use vortex_array::validity::Validity; use vortex_btrblocks::BtrBlocksCompressorBuilder; @@ -443,9 +438,7 @@ mod test { // [read] let file = session.open_options().open_path(path.clone()).await?; - let filter = gt(root(), lit(2u64)) - .optimize_recursive(file.dtype())? - .bind(file.dtype())?; + let filter = bound::gt(bound::root(file.dtype().clone()), bound::lit(2u64)); let array = file .scan()? .with_filter(filter) @@ -543,9 +536,7 @@ mod test { // Read the file back, but project down to just the "value" column. let file = session.open_options().open_path(path.clone()).await?; - let projection = select(["value"], root()) - .optimize_recursive(file.dtype())? - .bind(file.dtype())?; + let projection = bound::select(["value"], bound::root(file.dtype().clone())); let projected = file .scan()? .with_projection(projection)