diff --git a/.github/workflows/codspeed.yml b/.github/workflows/codspeed.yml index 23f86dde9cd..8222dcb761a 100644 --- a/.github/workflows/codspeed.yml +++ b/.github/workflows/codspeed.yml @@ -56,9 +56,9 @@ jobs: - { shard: 3, name: "Array dictionaries", packages: "vortex-array", array-group: "dictionary", features: "--features _test-harness" } - { shard: 4, name: "Array builders & views", packages: "vortex-array", array-group: "builders", features: "--features _test-harness" } - { shard: 5, name: "Core & encodings", packages: "vortex-buffer vortex-error vortex-mask vortex-compute vortex-file vortex-alp vortex-bytebool vortex-datetime-parts" } - - { shard: 6, name: "Main library & storage", packages: "vortex vortex-btrblocks vortex-compressor vortex-row" } + - { shard: 6, name: "Main library, functions & storage", packages: "vortex vortex-functions vortex-arrow vortex-btrblocks vortex-compressor vortex-row" } - { shard: 7, name: "FastLanes & decimals", packages: "vortex-decimal-byte-parts vortex-fastlanes", features: "--features _test-harness" } - - { shard: 8, name: "Encodings & layout", packages: "vortex-pco vortex-runend vortex-sequence vortex-sparse vortex-zigzag vortex-zstd vortex-layout" } + - { shard: 8, name: "Encodings & layout", packages: "vortex-onpair vortex-pco vortex-runend vortex-sequence vortex-sparse vortex-zigzag vortex-zstd vortex-layout" } - { shard: 9, name: "Strings, tensor & spatial", packages: "vortex-fsst vortex-tensor vortex-spatial", features: "--features _test-harness" } name: "Benchmark with Codspeed (Shard #${{ matrix.shard }} - ${{ matrix.name }})" timeout-minutes: 30 @@ -92,7 +92,7 @@ jobs: groups = { "dictionary": ("dict_", "chunked_dict_", "patches_"), "selection": ("take_", "filter_", "piecewise_sequence_"), - "compute": ("aggregate_", "cast_", "expr_", "binary_ops", "compare", "kleene_bool", "like", "list_sum"), + "compute": ("aggregate_", "cast_", "expr_", "binary_ops", "compare", "kleene_bool"), } manifest = tomllib.loads(pathlib.Path("vortex-array/Cargo.toml").read_text()) benches = [] diff --git a/Cargo.lock b/Cargo.lock index bb85a4861c3..5a154683fff 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11329,6 +11329,7 @@ dependencies = [ "vortex-fastlanes", "vortex-file", "vortex-fsst", + "vortex-functions", "vortex-io", "vortex-ipc", "vortex-layout", @@ -11363,6 +11364,7 @@ dependencies = [ "vortex-buffer", "vortex-error", "vortex-fastlanes", + "vortex-functions", "vortex-mask", "vortex-session", ] @@ -11829,6 +11831,7 @@ dependencies = [ "vortex-buffer", "vortex-compute", "vortex-error", + "vortex-functions", "vortex-mask", "vortex-session", ] @@ -11848,6 +11851,7 @@ dependencies = [ "vortex-bench-support", "vortex-buffer", "vortex-error", + "vortex-functions", "vortex-mask", "vortex-session", ] @@ -11924,6 +11928,7 @@ dependencies = [ "vortex-array", "vortex-buffer", "vortex-error", + "vortex-functions", "vortex-session", "vortex-utils", ] @@ -11949,6 +11954,7 @@ dependencies = [ "vortex-compute", "vortex-error", "vortex-fastlanes", + "vortex-functions", "vortex-mask", "vortex-session", ] @@ -12012,6 +12018,7 @@ dependencies = [ "vortex-expr", "vortex-fastlanes", "vortex-fsst", + "vortex-functions", "vortex-io", "vortex-layout", "vortex-mask", @@ -12044,10 +12051,39 @@ dependencies = [ "vortex-array", "vortex-buffer", "vortex-error", + "vortex-functions", "vortex-mask", "vortex-session", ] +[[package]] +name = "vortex-functions" +version = "0.1.0" +dependencies = [ + "arrow-buffer 59.3.1", + "codspeed-divan-compat", + "insta", + "itertools 0.14.0", + "jiff", + "memchr", + "mimalloc", + "num-traits", + "parking_lot", + "prost 0.14.4", + "rand 0.10.3", + "regex", + "regex-syntax", + "rstest", + "tracing", + "vortex-array", + "vortex-buffer", + "vortex-compute", + "vortex-error", + "vortex-mask", + "vortex-session", + "vortex-utils", +] + [[package]] name = "vortex-fuzz" version = "0.1.0" @@ -12069,6 +12105,7 @@ dependencies = [ "vortex-error", "vortex-file", "vortex-fsst", + "vortex-functions", "vortex-io", "vortex-mask", "vortex-row", @@ -12164,6 +12201,7 @@ dependencies = [ "vortex-arrow", "vortex-edition", "vortex-error", + "vortex-functions", "vortex-session", ] @@ -12214,6 +12252,7 @@ dependencies = [ "vortex-error", "vortex-expr", "vortex-file", + "vortex-functions", "vortex-io", "vortex-mask", "vortex-metrics", @@ -12272,6 +12311,7 @@ dependencies = [ "vortex-array", "vortex-buffer", "vortex-error", + "vortex-functions", "vortex-mask", "vortex-session", ] @@ -12295,6 +12335,7 @@ dependencies = [ "vortex-edition", "vortex-error", "vortex-file", + "vortex-functions", "vortex-io", "vortex-json", "vortex-layout", @@ -12439,6 +12480,7 @@ dependencies = [ "vortex-buffer", "vortex-compute", "vortex-error", + "vortex-functions", "vortex-mask", "vortex-session", ] @@ -12468,6 +12510,7 @@ dependencies = [ "vortex-arrow", "vortex-buffer", "vortex-error", + "vortex-functions", "vortex-mask", "vortex-session", ] @@ -12493,6 +12536,7 @@ dependencies = [ "vortex-edition", "vortex-error", "vortex-expr", + "vortex-functions", "vortex-layout", "vortex-mask", "vortex-session", diff --git a/Cargo.toml b/Cargo.toml index 31c8a368200..0f3b75271c0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,6 +13,7 @@ members = [ "vortex-io", "vortex-cloud", "vortex-array", + "vortex-functions", "vortex-expr", "vortex-arrow", "vortex-row", @@ -318,6 +319,7 @@ vortex-expr = { version = "0.1.0", path = "./vortex-expr", default-features = fa vortex-fastlanes = { version = "0.1.0", path = "./encodings/fastlanes", default-features = false } vortex-file = { version = "0.1.0", path = "./vortex-file", default-features = false } vortex-fsst = { version = "0.1.0", path = "./encodings/fsst", default-features = false } +vortex-functions = { version = "0.1.0", path = "./vortex-functions", default-features = false } vortex-io = { version = "0.1.0", path = "./vortex-io", default-features = false } vortex-ipc = { version = "0.1.0", path = "./vortex-ipc", default-features = false } vortex-json = { version = "0.1.0", path = "./vortex-json", default-features = false } diff --git a/encodings/alp/Cargo.toml b/encodings/alp/Cargo.toml index 8907e85aec8..46daee1cbe0 100644 --- a/encodings/alp/Cargo.toml +++ b/encodings/alp/Cargo.toml @@ -24,6 +24,7 @@ vortex-array = { workspace = true } vortex-buffer = { workspace = true } vortex-error = { workspace = true } vortex-fastlanes = { workspace = true } +vortex-functions = { workspace = true } vortex-mask = { workspace = true } vortex-session = { workspace = true } diff --git a/encodings/alp/src/alp/compute/between.rs b/encodings/alp/src/alp/compute/between.rs index 48e2e7fb5ee..79c00d8cf05 100644 --- a/encodings/alp/src/alp/compute/between.rs +++ b/encodings/alp/src/alp/compute/between.rs @@ -8,16 +8,16 @@ use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::IntoArray; use vortex_array::arrays::ConstantArray; -use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::NativeDType; use vortex_array::dtype::NativePType; use vortex_array::dtype::Nullability; use vortex_array::scalar::PValue; use vortex_array::scalar::Scalar; -use vortex_array::scalar_fn::fns::between::BetweenOptions; -use vortex_array::scalar_fn::fns::between::BetweenReduce; -use vortex_array::scalar_fn::fns::between::StrictComparison; use vortex_error::VortexResult; +use vortex_functions::between::BetweenOptions; +use vortex_functions::between::BetweenReduce; +use vortex_functions::between::StrictComparison; +use vortex_functions::builtins::ArrayFunctions; use crate::ALP; use crate::ALPFloat; @@ -145,8 +145,8 @@ mod tests { use vortex_array::assert_arrays_eq; use vortex_array::dtype::Nullability; use vortex_array::scalar::Scalar; - use vortex_array::scalar_fn::fns::between::BetweenOptions; - use vortex_array::scalar_fn::fns::between::StrictComparison; + use vortex_functions::between::BetweenOptions; + use vortex_functions::between::StrictComparison; use vortex_session::VortexSession; use crate::ALPArray; diff --git a/encodings/alp/src/alp/compute/cast.rs b/encodings/alp/src/alp/compute/cast.rs index 6b05ed7d949..198fcdaaafc 100644 --- a/encodings/alp/src/alp/compute/cast.rs +++ b/encodings/alp/src/alp/compute/cast.rs @@ -6,7 +6,7 @@ use vortex_array::ArrayView; use vortex_array::IntoArray; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; -use vortex_array::scalar_fn::fns::cast::CastReduce; +use vortex_array::scalar_fn::builtin::cast::CastReduce; use vortex_error::VortexResult; use crate::ALPArrayExt; diff --git a/encodings/alp/src/alp/compute/compare.rs b/encodings/alp/src/alp/compute/compare.rs index 457aacee172..55b1d5843e1 100644 --- a/encodings/alp/src/alp/compute/compare.rs +++ b/encodings/alp/src/alp/compute/compare.rs @@ -11,9 +11,9 @@ use vortex_array::arrays::ConstantArray; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::NativePType; use vortex_array::scalar::Scalar; -use vortex_array::scalar_fn::fns::binary::CompareKernel; -use vortex_array::scalar_fn::fns::operators::CompareOperator; -use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::scalar_fn::builtin::binary::CompareKernel; +use vortex_array::scalar_fn::builtin::operators::CompareOperator; +use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; @@ -165,8 +165,8 @@ mod tests { use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::scalar::Scalar; - use vortex_array::scalar_fn::fns::operators::CompareOperator; - use vortex_array::scalar_fn::fns::operators::Operator; + use vortex_array::scalar_fn::builtin::operators::CompareOperator; + use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_session::VortexSession; use super::*; diff --git a/encodings/alp/src/alp/compute/mask.rs b/encodings/alp/src/alp/compute/mask.rs index b67848ca6a4..890147f0ff1 100644 --- a/encodings/alp/src/alp/compute/mask.rs +++ b/encodings/alp/src/alp/compute/mask.rs @@ -6,8 +6,8 @@ use vortex_array::ArrayView; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::builtins::ArrayBuiltins; -use vortex_array::scalar_fn::fns::mask::MaskKernel; -use vortex_array::scalar_fn::fns::mask::MaskReduce; +use vortex_array::scalar_fn::builtin::mask::MaskKernel; +use vortex_array::scalar_fn::builtin::mask::MaskReduce; use vortex_array::validity::Validity; use vortex_error::VortexResult; @@ -58,7 +58,7 @@ mod test { use vortex_array::arrays::PrimitiveArray; use vortex_array::compute::conformance::mask::test_mask_conformance; use vortex_array::dtype::Nullability; - use vortex_array::scalar_fn::fns::mask::MaskKernel; + use vortex_array::scalar_fn::builtin::mask::MaskKernel; use vortex_buffer::buffer; use vortex_session::VortexSession; diff --git a/encodings/alp/src/alp/rules.rs b/encodings/alp/src/alp/rules.rs index 3b888ae8269..e401e594369 100644 --- a/encodings/alp/src/alp/rules.rs +++ b/encodings/alp/src/alp/rules.rs @@ -11,13 +11,13 @@ use vortex_array::arrays::slice::SliceExecuteAdaptor; use vortex_array::optimizer::kernels::ArrayKernelsExt; use vortex_array::optimizer::rules::ParentRuleSet; use vortex_array::scalar_fn::ScalarFnVTable; -use vortex_array::scalar_fn::fns::between::BetweenReduceAdaptor; -use vortex_array::scalar_fn::fns::binary::Binary; -use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; -use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor; -use vortex_array::scalar_fn::fns::mask::Mask; -use vortex_array::scalar_fn::fns::mask::MaskExecuteAdaptor; -use vortex_array::scalar_fn::fns::mask::MaskReduceAdaptor; +use vortex_array::scalar_fn::builtin::binary::Binary; +use vortex_array::scalar_fn::builtin::binary::CompareExecuteAdaptor; +use vortex_array::scalar_fn::builtin::cast::CastReduceAdaptor; +use vortex_array::scalar_fn::builtin::mask::Mask; +use vortex_array::scalar_fn::builtin::mask::MaskExecuteAdaptor; +use vortex_array::scalar_fn::builtin::mask::MaskReduceAdaptor; +use vortex_functions::between::BetweenReduceAdaptor; use vortex_session::VortexSession; use crate::ALP; diff --git a/encodings/alp/src/alp_rd/compute/cast.rs b/encodings/alp/src/alp_rd/compute/cast.rs index 70e1affeee9..c88817e8bc4 100644 --- a/encodings/alp/src/alp_rd/compute/cast.rs +++ b/encodings/alp/src/alp_rd/compute/cast.rs @@ -6,7 +6,7 @@ use vortex_array::ArrayView; use vortex_array::IntoArray; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; -use vortex_array::scalar_fn::fns::cast::CastReduce; +use vortex_array::scalar_fn::builtin::cast::CastReduce; use vortex_error::VortexResult; use crate::ALPRDArrayExt; diff --git a/encodings/alp/src/alp_rd/compute/mask.rs b/encodings/alp/src/alp_rd/compute/mask.rs index 9ffcfa06375..8c504f14d7b 100644 --- a/encodings/alp/src/alp_rd/compute/mask.rs +++ b/encodings/alp/src/alp_rd/compute/mask.rs @@ -4,8 +4,8 @@ use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::IntoArray; -use vortex_array::scalar_fn::fns::mask::Mask as MaskExpr; -use vortex_array::scalar_fn::fns::mask::MaskReduce; +use vortex_array::scalar_fn::builtin::mask::Mask as MaskExpr; +use vortex_array::scalar_fn::builtin::mask::MaskReduce; use vortex_error::VortexResult; use crate::ALPRD; diff --git a/encodings/alp/src/alp_rd/rules.rs b/encodings/alp/src/alp_rd/rules.rs index be5eda6e7f3..34b7a8272ce 100644 --- a/encodings/alp/src/alp_rd/rules.rs +++ b/encodings/alp/src/alp_rd/rules.rs @@ -2,8 +2,8 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_array::optimizer::rules::ParentRuleSet; -use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor; -use vortex_array::scalar_fn::fns::mask::MaskReduceAdaptor; +use vortex_array::scalar_fn::builtin::cast::CastReduceAdaptor; +use vortex_array::scalar_fn::builtin::mask::MaskReduceAdaptor; use crate::alp_rd::ALPRD; diff --git a/encodings/bytebool/src/compute.rs b/encodings/bytebool/src/compute.rs index 6a92a039ec8..e5d993282b9 100644 --- a/encodings/bytebool/src/compute.rs +++ b/encodings/bytebool/src/compute.rs @@ -12,13 +12,13 @@ use vortex_array::arrays::dict::TakeExecute; use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; use vortex_array::match_each_integer_ptype; -use vortex_array::scalar_fn::fns::binary::BooleanKernel; -use vortex_array::scalar_fn::fns::binary::kleene_boolean_buffer_scalar; -use vortex_array::scalar_fn::fns::binary::kleene_boolean_buffers; -use vortex_array::scalar_fn::fns::cast::CastKernel; -use vortex_array::scalar_fn::fns::cast::CastReduce; -use vortex_array::scalar_fn::fns::mask::MaskReduce; -use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::scalar_fn::builtin::binary::BooleanKernel; +use vortex_array::scalar_fn::builtin::binary::kleene_boolean_buffer_scalar; +use vortex_array::scalar_fn::builtin::binary::kleene_boolean_buffers; +use vortex_array::scalar_fn::builtin::cast::CastKernel; +use vortex_array::scalar_fn::builtin::cast::CastReduce; +use vortex_array::scalar_fn::builtin::mask::MaskReduce; +use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_array::validity::Validity; use vortex_buffer::BitBuffer; use vortex_buffer::ByteBuffer; @@ -182,7 +182,7 @@ mod tests { use vortex_array::compute::conformance::filter::test_filter_conformance; use vortex_array::compute::conformance::mask::test_mask_conformance; use vortex_array::compute::conformance::take::test_take_conformance; - use vortex_array::scalar_fn::fns::operators::Operator; + use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_error::vortex_err; use vortex_session::VortexSession; diff --git a/encodings/bytebool/src/kernel.rs b/encodings/bytebool/src/kernel.rs index 678da8642a7..c480f68184d 100644 --- a/encodings/bytebool/src/kernel.rs +++ b/encodings/bytebool/src/kernel.rs @@ -6,10 +6,10 @@ use vortex_array::arrays::Dict; use vortex_array::arrays::dict::TakeExecuteAdaptor; use vortex_array::optimizer::kernels::ArrayKernelsExt; use vortex_array::scalar_fn::ScalarFnVTable; -use vortex_array::scalar_fn::fns::binary::Binary; -use vortex_array::scalar_fn::fns::binary::BooleanExecuteAdaptor; -use vortex_array::scalar_fn::fns::cast::Cast; -use vortex_array::scalar_fn::fns::cast::CastExecuteAdaptor; +use vortex_array::scalar_fn::builtin::binary::Binary; +use vortex_array::scalar_fn::builtin::binary::BooleanExecuteAdaptor; +use vortex_array::scalar_fn::builtin::cast::Cast; +use vortex_array::scalar_fn::builtin::cast::CastExecuteAdaptor; use vortex_session::VortexSession; use crate::ByteBool; diff --git a/encodings/bytebool/src/rules.rs b/encodings/bytebool/src/rules.rs index f67d3567326..a29efa0ea1a 100644 --- a/encodings/bytebool/src/rules.rs +++ b/encodings/bytebool/src/rules.rs @@ -3,8 +3,8 @@ use vortex_array::arrays::slice::SliceReduceAdaptor; use vortex_array::optimizer::rules::ParentRuleSet; -use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor; -use vortex_array::scalar_fn::fns::mask::MaskReduceAdaptor; +use vortex_array::scalar_fn::builtin::cast::CastReduceAdaptor; +use vortex_array::scalar_fn::builtin::mask::MaskReduceAdaptor; use crate::ByteBool; diff --git a/encodings/datetime-parts/Cargo.toml b/encodings/datetime-parts/Cargo.toml index 5f5657c6e3d..c8eb48fee0c 100644 --- a/encodings/datetime-parts/Cargo.toml +++ b/encodings/datetime-parts/Cargo.toml @@ -22,6 +22,7 @@ vortex-array = { workspace = true } vortex-buffer = { workspace = true } vortex-compute = { workspace = true } vortex-error = { workspace = true } +vortex-functions = { workspace = true } vortex-mask = { workspace = true } vortex-session = { workspace = true } diff --git a/encodings/datetime-parts/src/compute/cast.rs b/encodings/datetime-parts/src/compute/cast.rs index 8d5b0cd2e2c..7354871c6c7 100644 --- a/encodings/datetime-parts/src/compute/cast.rs +++ b/encodings/datetime-parts/src/compute/cast.rs @@ -6,7 +6,7 @@ use vortex_array::ArrayView; use vortex_array::IntoArray; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; -use vortex_array::scalar_fn::fns::cast::CastReduce; +use vortex_array::scalar_fn::builtin::cast::CastReduce; use vortex_error::VortexResult; use crate::DateTimeParts; diff --git a/encodings/datetime-parts/src/compute/compare.rs b/encodings/datetime-parts/src/compute/compare.rs index 96885800424..41b6cf3eefe 100644 --- a/encodings/datetime-parts/src/compute/compare.rs +++ b/encodings/datetime-parts/src/compute/compare.rs @@ -13,9 +13,9 @@ use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::extension::datetime::Timestamp; use vortex_array::scalar::Scalar; -use vortex_array::scalar_fn::fns::binary::CompareKernel; -use vortex_array::scalar_fn::fns::operators::CompareOperator; -use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::scalar_fn::builtin::binary::CompareKernel; +use vortex_array::scalar_fn::builtin::operators::CompareOperator; +use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_array::validity::Validity; use vortex_error::VortexResult; use vortex_mask::AllOr; @@ -302,7 +302,7 @@ mod test { use vortex_array::extension::datetime::Timestamp; use vortex_array::extension::datetime::TimestampOptions; use vortex_array::scalar::Scalar; - use vortex_array::scalar_fn::fns::binary::CompareKernel; + use vortex_array::scalar_fn::builtin::binary::CompareKernel; use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_session::VortexSession; diff --git a/encodings/datetime-parts/src/compute/kernel.rs b/encodings/datetime-parts/src/compute/kernel.rs index 4bd0cf2e1b3..46cc7b20b3a 100644 --- a/encodings/datetime-parts/src/compute/kernel.rs +++ b/encodings/datetime-parts/src/compute/kernel.rs @@ -6,8 +6,8 @@ use vortex_array::arrays::Dict; use vortex_array::arrays::dict::TakeExecuteAdaptor; use vortex_array::optimizer::kernels::ArrayKernelsExt; use vortex_array::scalar_fn::ScalarFnVTable; -use vortex_array::scalar_fn::fns::binary::Binary; -use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; +use vortex_array::scalar_fn::builtin::binary::Binary; +use vortex_array::scalar_fn::builtin::binary::CompareExecuteAdaptor; use vortex_session::VortexSession; use crate::DateTimeParts; diff --git a/encodings/datetime-parts/src/compute/mask.rs b/encodings/datetime-parts/src/compute/mask.rs index bd81605c446..67132e7dbb4 100644 --- a/encodings/datetime-parts/src/compute/mask.rs +++ b/encodings/datetime-parts/src/compute/mask.rs @@ -5,7 +5,7 @@ use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::IntoArray; use vortex_array::builtins::ArrayBuiltins; -use vortex_array::scalar_fn::fns::mask::MaskReduce; +use vortex_array::scalar_fn::builtin::mask::MaskReduce; use vortex_error::VortexResult; use crate::DateTimeParts; diff --git a/encodings/datetime-parts/src/compute/rules.rs b/encodings/datetime-parts/src/compute/rules.rs index 5712b94961d..34d42354bd7 100644 --- a/encodings/datetime-parts/src/compute/rules.rs +++ b/encodings/datetime-parts/src/compute/rules.rs @@ -16,15 +16,14 @@ use vortex_array::arrays::slice::SliceReduceAdaptor; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::extension::datetime::Timestamp; -use vortex_array::optimizer::ArrayOptimizer; use vortex_array::optimizer::rules::ArrayParentReduceRule; use vortex_array::optimizer::rules::ParentRuleSet; -use vortex_array::scalar_fn::fns::between::Between; -use vortex_array::scalar_fn::fns::binary::Binary; -use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor; -use vortex_array::scalar_fn::fns::mask::MaskReduceAdaptor; +use vortex_array::scalar_fn::builtin::binary::Binary; +use vortex_array::scalar_fn::builtin::cast::CastReduceAdaptor; +use vortex_array::scalar_fn::builtin::mask::MaskReduceAdaptor; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_functions::between::Between; use crate::DateTimeParts; use crate::array::DateTimePartsArraySlotsExt; @@ -135,9 +134,7 @@ impl ArrayParentReduceRule for DTPComparisonPushDownRule { } } - let result = ScalarFnArray::try_new(parent.scalar_fn().clone(), new_children)? - .into_array() - .optimize()?; + let result = ScalarFnArray::try_new(parent.scalar_fn().clone(), new_children)?.into_array(); Ok(Some(result)) } @@ -192,11 +189,11 @@ mod tests { use vortex_array::extension::datetime::TimestampOptions; use vortex_array::optimizer::ArrayOptimizer; use vortex_array::scalar::Scalar; - use vortex_array::scalar_fn::fns::between::BetweenOptions; - use vortex_array::scalar_fn::fns::between::StrictComparison; - use vortex_array::scalar_fn::fns::operators::Operator; + use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_array::validity::Validity; use vortex_buffer::Buffer; + use vortex_functions::between::BetweenOptions; + use vortex_functions::between::StrictComparison; use vortex_session::VortexSession; use super::*; @@ -283,7 +280,7 @@ mod tests { .into_array(); // Optimize should push down to days - let optimized = comparison.optimize().unwrap(); + let optimized = comparison.optimize(&SESSION).unwrap(); // Verify correctness: days [0, 1, 2] <= 1 should give [true, true, false] let mut ctx = SESSION.create_execution_ctx(); @@ -317,7 +314,7 @@ mod tests { .into_array(); // Optimize should push down to days - let optimized = between.optimize().unwrap(); + let optimized = between.optimize(&SESSION).unwrap(); // Verify correctness: days [0, 1, 2, 3, 4] between 1 and 3 should give [false, true, true, true, false] let mut ctx = SESSION.create_execution_ctx(); @@ -354,7 +351,7 @@ mod tests { let constant = ConstantArray::new(scalar, dtp.len()).into_array(); let comparison = Binary::try_new(dtp, constant, Operator::Lte)?.into_array(); - let optimized = comparison.optimize()?; + let optimized = comparison.optimize(&SESSION)?; assert_eq!(optimized.dtype(), comparison.dtype()); let mut ctx = SESSION.create_execution_ctx(); @@ -379,7 +376,7 @@ mod tests { .into_array(); // Optimize should NOT push down (constant has non-zero seconds) - let optimized = comparison.optimize().unwrap(); + let optimized = comparison.optimize(&SESSION).unwrap(); // days [0, 1, 2] at midnight <= day 1 at noon: [true, true, false] let mut ctx = SESSION.create_execution_ctx(); @@ -415,7 +412,7 @@ mod tests { .into_array(); // Should still compute correctly (just not optimized via pushdown) - let optimized = comparison.optimize().unwrap(); + let optimized = comparison.optimize(&SESSION).unwrap(); // timestamps at 1am on days [0, 1, 2] <= day 1 midnight: [true, false, false] assert_arrays_eq!( optimized, diff --git a/encodings/decimal-byte-parts/Cargo.toml b/encodings/decimal-byte-parts/Cargo.toml index dcccabf001a..be283bf96ba 100644 --- a/encodings/decimal-byte-parts/Cargo.toml +++ b/encodings/decimal-byte-parts/Cargo.toml @@ -21,6 +21,7 @@ prost = { workspace = true } vortex-array = { workspace = true } vortex-buffer = { workspace = true } vortex-error = { workspace = true } +vortex-functions = { workspace = true } vortex-mask = { workspace = true } vortex-session = { workspace = true } diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/between.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/between.rs index 8ef4d81daf3..2fe075465b1 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/between.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/between.rs @@ -7,17 +7,17 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ConstantArray; -use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::scalar::PValue; use vortex_array::scalar::Scalar; use vortex_array::scalar::ScalarValue; -use vortex_array::scalar_fn::fns::between::BetweenKernel; -use vortex_array::scalar_fn::fns::between::BetweenOptions; -use vortex_array::scalar_fn::fns::between::StrictComparison; use vortex_buffer::BitBuffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_functions::between::BetweenKernel; +use vortex_functions::between::BetweenOptions; +use vortex_functions::between::StrictComparison; +use vortex_functions::builtins::ArrayFunctions; use crate::DecimalByteParts; use crate::decimal_byte_parts::DecimalBytePartsArraySlotsExt; @@ -116,17 +116,17 @@ mod tests { use vortex_array::arrays::DecimalArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; - use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DecimalDType; use vortex_array::dtype::Nullability; use vortex_array::scalar::DecimalValue; use vortex_array::scalar::Scalar; - use vortex_array::scalar_fn::fns::between::BetweenOptions; - use vortex_array::scalar_fn::fns::between::StrictComparison; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::buffer; use vortex_error::VortexResult; + use vortex_functions::between::BetweenOptions; + use vortex_functions::between::StrictComparison; + use vortex_functions::builtins::ArrayFunctions; use vortex_session::VortexSession; use crate::DecimalByteParts; diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/cast.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/cast.rs index c700c6f2515..7b0080ed854 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/cast.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/cast.rs @@ -6,7 +6,7 @@ use vortex_array::ArrayView; use vortex_array::IntoArray; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; -use vortex_array::scalar_fn::fns::cast::CastReduce; +use vortex_array::scalar_fn::builtin::cast::CastReduce; use vortex_error::VortexResult; use crate::DecimalByteParts; diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/compare.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/compare.rs index 4626b325215..6069b1c2b79 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/compare.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/compare.rs @@ -18,9 +18,9 @@ use vortex_array::match_each_integer_ptype; use vortex_array::scalar::DecimalValue; use vortex_array::scalar::Scalar; use vortex_array::scalar::ScalarValue; -use vortex_array::scalar_fn::fns::binary::CompareKernel; -use vortex_array::scalar_fn::fns::operators::CompareOperator; -use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::scalar_fn::builtin::binary::CompareKernel; +use vortex_array::scalar_fn::builtin::operators::CompareOperator; +use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -165,7 +165,7 @@ mod tests { use vortex_array::dtype::Nullability; use vortex_array::scalar::DecimalValue; use vortex_array::scalar::Scalar; - use vortex_array::scalar_fn::fns::operators::Operator; + use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexExpect; diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/kernel.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/kernel.rs index 9f26305a94d..fc85a00931c 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/kernel.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/kernel.rs @@ -3,10 +3,10 @@ use vortex_array::optimizer::kernels::ArrayKernelsExt; use vortex_array::scalar_fn::ScalarFnVTable; -use vortex_array::scalar_fn::fns::between::Between; -use vortex_array::scalar_fn::fns::between::BetweenExecuteAdaptor; -use vortex_array::scalar_fn::fns::binary::Binary; -use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; +use vortex_array::scalar_fn::builtin::binary::Binary; +use vortex_array::scalar_fn::builtin::binary::CompareExecuteAdaptor; +use vortex_functions::between::Between; +use vortex_functions::between::BetweenExecuteAdaptor; use vortex_session::VortexSession; use crate::DecimalByteParts; diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mask.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mask.rs index 9a022ef34ce..a898bbe35e7 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mask.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mask.rs @@ -4,8 +4,8 @@ use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::IntoArray; -use vortex_array::scalar_fn::fns::mask::Mask as MaskExpr; -use vortex_array::scalar_fn::fns::mask::MaskReduce; +use vortex_array::scalar_fn::builtin::mask::Mask as MaskExpr; +use vortex_array::scalar_fn::builtin::mask::MaskReduce; use vortex_error::VortexResult; use crate::DecimalByteParts; diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/take.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/take.rs index 6a2ca24f390..6918ae4100d 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/take.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/take.rs @@ -40,6 +40,7 @@ mod tests { use vortex_array::assert_arrays_eq; use vortex_array::dtype::DecimalDType; use vortex_array::dtype::i256; + use vortex_array::optimizer::ArrayOptimizer; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::buffer; @@ -55,8 +56,8 @@ mod tests { }); /// Taking pushes down into the parts during optimization, with no execution context in - /// play: `ArrayRef::take` wraps the array in a `Dict` and optimizes, and the reduce rule - /// must rewrite that into a `DecimalByteParts` of taken parts. + /// play: `ArrayRef::take` wraps the array in a `Dict`, and optimizing it must rewrite that + /// into a `DecimalByteParts` of taken parts. #[test] fn take_pushes_down_without_executing() -> VortexResult<()> { let decimal = DecimalArray::new( @@ -67,7 +68,8 @@ mod tests { let indices = buffer![0u64, 2].into_array(); let taken = DecimalByteParts::encode(&decimal, &mut SESSION.create_execution_ctx())? .into_array() - .take(indices)?; + .take(indices)? + .optimize(&SESSION)?; assert!( taken.is::(), diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/rules.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/rules.rs index 28503d5d8af..3c4c9ecf6d7 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/rules.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/rules.rs @@ -5,8 +5,8 @@ use vortex_array::arrays::dict::TakeReduceAdaptor; use vortex_array::arrays::filter::FilterReduceAdaptor; use vortex_array::arrays::slice::SliceReduceAdaptor; use vortex_array::optimizer::rules::ParentRuleSet; -use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor; -use vortex_array::scalar_fn::fns::mask::MaskReduceAdaptor; +use vortex_array::scalar_fn::builtin::cast::CastReduceAdaptor; +use vortex_array::scalar_fn::builtin::mask::MaskReduceAdaptor; use crate::DecimalByteParts; diff --git a/encodings/fastlanes/Cargo.toml b/encodings/fastlanes/Cargo.toml index c1261092b61..4b69c99259a 100644 --- a/encodings/fastlanes/Cargo.toml +++ b/encodings/fastlanes/Cargo.toml @@ -29,6 +29,7 @@ vortex-array = { workspace = true } vortex-buffer = { workspace = true } vortex-compute = { workspace = true } vortex-error = { workspace = true } +vortex-functions = { workspace = true } vortex-mask = { workspace = true } vortex-session = { workspace = true } diff --git a/encodings/fastlanes/benches/bitpack_compare.rs b/encodings/fastlanes/benches/bitpack_compare.rs index acc8c89f2e8..674ebbc342f 100644 --- a/encodings/fastlanes/benches/bitpack_compare.rs +++ b/encodings/fastlanes/benches/bitpack_compare.rs @@ -26,7 +26,7 @@ use vortex_array::arrays::PrimitiveArray; use vortex_array::builtins::ArrayBuiltins; use vortex_array::memory::MemorySession; use vortex_array::optimizer::kernels::KernelSession; -use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::Alignment; diff --git a/encodings/fastlanes/benches/bitpack_compare_sweep.rs b/encodings/fastlanes/benches/bitpack_compare_sweep.rs index 3dcb3f50a11..f7ae88f341d 100644 --- a/encodings/fastlanes/benches/bitpack_compare_sweep.rs +++ b/encodings/fastlanes/benches/bitpack_compare_sweep.rs @@ -29,7 +29,7 @@ use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::NativePType; use vortex_array::optimizer::kernels::KernelSession; use vortex_array::scalar::Scalar; -use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::Alignment; diff --git a/encodings/fastlanes/benches/compute_between.rs b/encodings/fastlanes/benches/compute_between.rs index d3e85f61e3d..080043a5f1f 100644 --- a/encodings/fastlanes/benches/compute_between.rs +++ b/encodings/fastlanes/benches/compute_between.rs @@ -99,11 +99,11 @@ mod primitive { use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; use vortex_array::arrays::ConstantArray; - use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::NativePType; - use vortex_array::scalar_fn::fns::between::BetweenOptions; - use vortex_array::scalar_fn::fns::between::StrictComparison::NonStrict; use vortex_error::VortexExpect; + use vortex_functions::between::BetweenOptions; + use vortex_functions::between::StrictComparison::NonStrict; + use vortex_functions::builtins::ArrayFunctions; use crate::BENCH_ARGS; use crate::SESSION; @@ -152,11 +152,11 @@ mod bitpack { use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; use vortex_array::arrays::ConstantArray; - use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::NativePType; - use vortex_array::scalar_fn::fns::between::BetweenOptions; - use vortex_array::scalar_fn::fns::between::StrictComparison::NonStrict; use vortex_error::VortexExpect; + use vortex_functions::between::BetweenOptions; + use vortex_functions::between::StrictComparison::NonStrict; + use vortex_functions::builtins::ArrayFunctions; use crate::BENCH_ARGS; use crate::SESSION; @@ -204,11 +204,11 @@ mod alp { use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; use vortex_array::arrays::ConstantArray; - use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::NativePType; - use vortex_array::scalar_fn::fns::between::BetweenOptions; - use vortex_array::scalar_fn::fns::between::StrictComparison::NonStrict; use vortex_error::VortexExpect; + use vortex_functions::between::BetweenOptions; + use vortex_functions::between::StrictComparison::NonStrict; + use vortex_functions::builtins::ArrayFunctions; use crate::BENCH_ARGS; use crate::SESSION; diff --git a/encodings/fastlanes/src/bitpacking/array/tests.rs b/encodings/fastlanes/src/bitpacking/array/tests.rs index 94a29adff73..0364c068499 100644 --- a/encodings/fastlanes/src/bitpacking/array/tests.rs +++ b/encodings/fastlanes/src/bitpacking/array/tests.rs @@ -21,8 +21,8 @@ use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::match_each_unsigned_integer_ptype; -use vortex_array::scalar_fn::fns::cast::CastKernel; -use vortex_array::scalar_fn::fns::cast::CastReduce; +use vortex_array::scalar_fn::builtin::cast::CastKernel; +use vortex_array::scalar_fn::builtin::cast::CastReduce; use vortex_array::validity::Validity; use vortex_buffer::ByteBuffer; use vortex_buffer::buffer; diff --git a/encodings/fastlanes/src/bitpacking/compute/between.rs b/encodings/fastlanes/src/bitpacking/compute/between.rs index a878fcb3c05..9734da3f119 100644 --- a/encodings/fastlanes/src/bitpacking/compute/between.rs +++ b/encodings/fastlanes/src/bitpacking/compute/between.rs @@ -21,11 +21,11 @@ use vortex_array::dtype::NativePType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PhysicalPType; use vortex_array::match_each_integer_ptype; -use vortex_array::scalar_fn::fns::between::BetweenKernel; -use vortex_array::scalar_fn::fns::between::BetweenOptions; -use vortex_array::scalar_fn::fns::between::StrictComparison; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_functions::between::BetweenKernel; +use vortex_functions::between::BetweenOptions; +use vortex_functions::between::StrictComparison; use crate::BitPacked; use crate::BitPackedArrayExt; @@ -224,13 +224,13 @@ mod tests { use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; - use vortex_array::builtins::ArrayBuiltins; - use vortex_array::scalar_fn::fns::between::BetweenKernel; - use vortex_array::scalar_fn::fns::between::BetweenOptions; - use vortex_array::scalar_fn::fns::between::StrictComparison; use vortex_array::validity::Validity; use vortex_buffer::BufferMut; use vortex_error::VortexResult; + use vortex_functions::between::BetweenKernel; + use vortex_functions::between::BetweenOptions; + use vortex_functions::between::StrictComparison; + use vortex_functions::builtins::ArrayFunctions; use vortex_session::VortexSession; use crate::BitPacked; diff --git a/encodings/fastlanes/src/bitpacking/compute/cast.rs b/encodings/fastlanes/src/bitpacking/compute/cast.rs index 8203e68a6b7..88d03270002 100644 --- a/encodings/fastlanes/src/bitpacking/compute/cast.rs +++ b/encodings/fastlanes/src/bitpacking/compute/cast.rs @@ -11,8 +11,8 @@ use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::PType; use vortex_array::match_each_integer_ptype; -use vortex_array::scalar_fn::fns::cast::CastKernel; -use vortex_array::scalar_fn::fns::cast::CastReduce; +use vortex_array::scalar_fn::builtin::cast::CastKernel; +use vortex_array::scalar_fn::builtin::cast::CastReduce; use vortex_array::validity::Validity; use vortex_error::VortexResult; diff --git a/encodings/fastlanes/src/bitpacking/compute/compare.rs b/encodings/fastlanes/src/bitpacking/compute/compare.rs index 68389b6eb2b..a3333c28f5a 100644 --- a/encodings/fastlanes/src/bitpacking/compute/compare.rs +++ b/encodings/fastlanes/src/bitpacking/compute/compare.rs @@ -21,8 +21,8 @@ use vortex_array::dtype::NativePType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PhysicalPType; use vortex_array::match_each_integer_ptype; -use vortex_array::scalar_fn::fns::binary::CompareKernel; -use vortex_array::scalar_fn::fns::operators::CompareOperator; +use vortex_array::scalar_fn::builtin::binary::CompareKernel; +use vortex_array::scalar_fn::builtin::operators::CompareOperator; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -203,9 +203,10 @@ mod tests { use vortex_array::arrays::slice::SliceKernel; use vortex_array::assert_arrays_eq; use vortex_array::builtins::ArrayBuiltins; - use vortex_array::scalar_fn::fns::binary::CompareKernel; - use vortex_array::scalar_fn::fns::operators::CompareOperator; - use vortex_array::scalar_fn::fns::operators::Operator; + use vortex_array::optimizer::ArrayOptimizer; + use vortex_array::scalar_fn::builtin::binary::CompareKernel; + use vortex_array::scalar_fn::builtin::operators::CompareOperator; + use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_error::VortexResult; use vortex_session::VortexSession; @@ -317,7 +318,10 @@ mod tests { let prim = PrimitiveArray::from_iter(values); let packed = BitPackedData::encode(&prim.clone().into_array(), 7, &mut ctx)?; - let sliced = packed.into_array().slice(start..start + slice_len)?; + let sliced = packed + .into_array() + .slice(start..start + slice_len)? + .optimize(&SESSION)?; let rhs = ConstantArray::new(50u32, slice_len).into_array(); for op in [ CompareOperator::Eq, diff --git a/encodings/fastlanes/src/bitpacking/compute/filter.rs b/encodings/fastlanes/src/bitpacking/compute/filter.rs index 0b1fde8621f..da9c0f2fccc 100644 --- a/encodings/fastlanes/src/bitpacking/compute/filter.rs +++ b/encodings/fastlanes/src/bitpacking/compute/filter.rs @@ -148,6 +148,7 @@ mod tests { use vortex_array::assert_arrays_eq; use vortex_array::compute::conformance::filter::test_filter_conformance; use vortex_array::dtype::NativePType; + use vortex_array::optimizer::ArrayOptimizer; use vortex_array::validity::Validity; use vortex_buffer::BitBuffer; use vortex_buffer::Buffer; @@ -497,7 +498,8 @@ mod tests { PrimitiveArray::from_iter((0..5000u32).map(|i| >::from((i % 100) as u8))); let bitpacked = BitPackedData::encode(&unpacked.into_array(), 7, &mut ctx)? .into_array() - .slice(range.clone())?; + .slice(range.clone())? + .optimize(&SESSION)?; let bitpacked = bitpacked.as_::(); let bits = selection(pattern, range.len()); diff --git a/encodings/fastlanes/src/bitpacking/plugin/tests.rs b/encodings/fastlanes/src/bitpacking/plugin/tests.rs index 127695c6d26..9fe5c1d6a94 100644 --- a/encodings/fastlanes/src/bitpacking/plugin/tests.rs +++ b/encodings/fastlanes/src/bitpacking/plugin/tests.rs @@ -20,6 +20,7 @@ use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; +use vortex_array::optimizer::ArrayOptimizer; use vortex_array::serde::SerializeOptions; use vortex_array::serde::SerializedArray; use vortex_array::session::ArraySessionExt; @@ -224,7 +225,8 @@ fn serde_roundtrip( let mut ctx = PLUGIN_SESSION.create_execution_ctx(); let array = BitPackedData::encode(&values, bit_width, &mut ctx)? .into_array() - .slice(range.clone())?; + .slice(range.clone())? + .optimize(&PLUGIN_SESSION)?; let view = array.as_::(); let serialization = PLUGIN_SESSION .array_serialize(&array)? diff --git a/encodings/fastlanes/src/bitpacking/vtable/kernels.rs b/encodings/fastlanes/src/bitpacking/vtable/kernels.rs index eb0dd9b7a23..33bce64ff86 100644 --- a/encodings/fastlanes/src/bitpacking/vtable/kernels.rs +++ b/encodings/fastlanes/src/bitpacking/vtable/kernels.rs @@ -10,12 +10,12 @@ use vortex_array::arrays::filter::FilterExecuteAdaptor; use vortex_array::arrays::slice::SliceExecuteAdaptor; use vortex_array::optimizer::kernels::ArrayKernelsExt; use vortex_array::scalar_fn::ScalarFnVTable; -use vortex_array::scalar_fn::fns::between::Between; -use vortex_array::scalar_fn::fns::between::BetweenExecuteAdaptor; -use vortex_array::scalar_fn::fns::binary::Binary; -use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; -use vortex_array::scalar_fn::fns::cast::Cast; -use vortex_array::scalar_fn::fns::cast::CastExecuteAdaptor; +use vortex_array::scalar_fn::builtin::binary::Binary; +use vortex_array::scalar_fn::builtin::binary::CompareExecuteAdaptor; +use vortex_array::scalar_fn::builtin::cast::Cast; +use vortex_array::scalar_fn::builtin::cast::CastExecuteAdaptor; +use vortex_functions::between::Between; +use vortex_functions::between::BetweenExecuteAdaptor; use vortex_session::VortexSession; use crate::BitPacked; diff --git a/encodings/fastlanes/src/bitpacking/vtable/rules.rs b/encodings/fastlanes/src/bitpacking/vtable/rules.rs index a26989569b2..f0e101ddc8b 100644 --- a/encodings/fastlanes/src/bitpacking/vtable/rules.rs +++ b/encodings/fastlanes/src/bitpacking/vtable/rules.rs @@ -3,7 +3,7 @@ use vortex_array::arrays::slice::SliceReduceAdaptor; use vortex_array::optimizer::rules::ParentRuleSet; -use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor; +use vortex_array::scalar_fn::builtin::cast::CastReduceAdaptor; use crate::BitPacked; diff --git a/encodings/fastlanes/src/delta/compute/cast.rs b/encodings/fastlanes/src/delta/compute/cast.rs index 4342fd883bd..d9db6375314 100644 --- a/encodings/fastlanes/src/delta/compute/cast.rs +++ b/encodings/fastlanes/src/delta/compute/cast.rs @@ -6,7 +6,7 @@ use vortex_array::ArrayView; use vortex_array::IntoArray; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability::NonNullable; -use vortex_array::scalar_fn::fns::cast::CastReduce; +use vortex_array::scalar_fn::builtin::cast::CastReduce; use vortex_error::VortexResult; use crate::delta::Delta; @@ -63,6 +63,7 @@ mod tests { use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; + use vortex_array::optimizer::ArrayOptimizer; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_error::vortex_err; @@ -129,7 +130,8 @@ mod tests { let casted = array .into_array() - .cast(DType::Primitive(PType::U32, Nullability::Nullable))?; + .cast(DType::Primitive(PType::U32, Nullability::Nullable))? + .optimize(&SESSION)?; assert_eq!( casted.dtype(), diff --git a/encodings/fastlanes/src/delta/vtable/rules.rs b/encodings/fastlanes/src/delta/vtable/rules.rs index d6892897ab5..c3cc19ef892 100644 --- a/encodings/fastlanes/src/delta/vtable/rules.rs +++ b/encodings/fastlanes/src/delta/vtable/rules.rs @@ -3,7 +3,7 @@ use vortex_array::arrays::slice::SliceReduceAdaptor; use vortex_array::optimizer::rules::ParentRuleSet; -use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor; +use vortex_array::scalar_fn::builtin::cast::CastReduceAdaptor; use crate::delta::vtable::Delta; diff --git a/encodings/fastlanes/src/for_/compute/cast.rs b/encodings/fastlanes/src/for_/compute/cast.rs index af1ac50af3f..18f78d0c6c1 100644 --- a/encodings/fastlanes/src/for_/compute/cast.rs +++ b/encodings/fastlanes/src/for_/compute/cast.rs @@ -6,7 +6,7 @@ use vortex_array::ArrayView; use vortex_array::IntoArray; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; -use vortex_array::scalar_fn::fns::cast::CastReduce; +use vortex_array::scalar_fn::builtin::cast::CastReduce; use vortex_error::VortexResult; use crate::FoRReferences; diff --git a/encodings/fastlanes/src/for_/compute/compare.rs b/encodings/fastlanes/src/for_/compute/compare.rs index d9dfcc6d9f8..d5b761e6592 100644 --- a/encodings/fastlanes/src/for_/compute/compare.rs +++ b/encodings/fastlanes/src/for_/compute/compare.rs @@ -15,9 +15,9 @@ use vortex_array::dtype::Nullability; use vortex_array::match_each_integer_ptype; use vortex_array::scalar::PValue; use vortex_array::scalar::Scalar; -use vortex_array::scalar_fn::fns::binary::CompareKernel; -use vortex_array::scalar_fn::fns::operators::CompareOperator; -use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::scalar_fn::builtin::binary::CompareKernel; +use vortex_array::scalar_fn::builtin::operators::CompareOperator; +use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_error::VortexError; use vortex_error::VortexExpect as _; use vortex_error::VortexResult; diff --git a/encodings/fastlanes/src/for_/plugin/tests.rs b/encodings/fastlanes/src/for_/plugin/tests.rs index 16278e070df..a88c3f22c4d 100644 --- a/encodings/fastlanes/src/for_/plugin/tests.rs +++ b/encodings/fastlanes/src/for_/plugin/tests.rs @@ -10,6 +10,7 @@ use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; +use vortex_array::optimizer::ArrayOptimizer; use vortex_array::scalar::ScalarValue; use vortex_array::serde::SerializeOptions; use vortex_array::serde::SerializedArray; @@ -81,7 +82,10 @@ fn chunked() -> VortexResult { #[case::whole(0..5000)] #[case::sliced(1500..2600)] fn v2_roundtrip(#[case] range: std::ops::Range) -> VortexResult<()> { - let array = chunked()?.into_array().slice(range.clone())?; + let array = chunked()? + .into_array() + .slice(range.clone())? + .optimize(&SESSION)?; let serialization = SESSION .array_serialize(&array)? .ok_or_else(|| vortex_err!("FoR must serialize"))?; diff --git a/encodings/fastlanes/src/for_/tests.rs b/encodings/fastlanes/src/for_/tests.rs index 63116924f5e..17c4d94c8e9 100644 --- a/encodings/fastlanes/src/for_/tests.rs +++ b/encodings/fastlanes/src/for_/tests.rs @@ -19,6 +19,7 @@ use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::compute::conformance::consistency::test_array_consistency; use vortex_array::dtype::NativePType; +use vortex_array::optimizer::ArrayOptimizer; use vortex_array::scalar::Scalar; use vortex_array::session::ArraySessionExt; use vortex_array::smallvec::smallvec; @@ -161,7 +162,7 @@ fn global_reference_survives_slicing() -> VortexResult<()> { )?; assert_eq!(global(&array)?, 7u32.into()); assert!(array.blocked_references().is_none()); - let sliced = array.into_array().slice(1500..1600)?; + let sliced = array.into_array().slice(1500..1600)?.optimize(&SESSION)?; assert_eq!(global(&sliced.as_::().into_owned())?, 7u32.into()); Ok(()) } @@ -328,7 +329,7 @@ fn fused_decode_sliced( ) -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); let (array, expected) = fused(bit_width, signed)?; - let sliced = array.into_array().slice(start..end)?; + let sliced = array.into_array().slice(start..end)?.optimize(&SESSION)?; let sliced_for = sliced.as_::(); // A BitPacked child with patches slices lazily, which takes the unfused path. assert_eq!(sliced_for.encoded().is::(), bit_width == 3); diff --git a/encodings/fastlanes/src/for_/vtable/kernels.rs b/encodings/fastlanes/src/for_/vtable/kernels.rs index 6c63292bb9f..7687ede7cf1 100644 --- a/encodings/fastlanes/src/for_/vtable/kernels.rs +++ b/encodings/fastlanes/src/for_/vtable/kernels.rs @@ -6,8 +6,8 @@ use vortex_array::arrays::Dict; use vortex_array::arrays::dict::TakeExecuteAdaptor; use vortex_array::optimizer::kernels::ArrayKernelsExt; use vortex_array::scalar_fn::ScalarFnVTable; -use vortex_array::scalar_fn::fns::binary::Binary; -use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; +use vortex_array::scalar_fn::builtin::binary::Binary; +use vortex_array::scalar_fn::builtin::binary::CompareExecuteAdaptor; use vortex_session::VortexSession; use crate::FoR; diff --git a/encodings/fastlanes/src/for_/vtable/rules.rs b/encodings/fastlanes/src/for_/vtable/rules.rs index 541232865e3..5e22819e619 100644 --- a/encodings/fastlanes/src/for_/vtable/rules.rs +++ b/encodings/fastlanes/src/for_/vtable/rules.rs @@ -9,7 +9,7 @@ use vortex_array::arrays::filter::FilterReduceAdaptor; use vortex_array::arrays::slice::SliceReduceAdaptor; use vortex_array::optimizer::rules::ArrayParentReduceRule; use vortex_array::optimizer::rules::ParentRuleSet; -use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor; +use vortex_array::scalar_fn::builtin::cast::CastReduceAdaptor; use vortex_error::VortexResult; use crate::FoR; diff --git a/encodings/fastlanes/src/rle/compute/cast.rs b/encodings/fastlanes/src/rle/compute/cast.rs index 194c6c36c3e..994a8f0c4d1 100644 --- a/encodings/fastlanes/src/rle/compute/cast.rs +++ b/encodings/fastlanes/src/rle/compute/cast.rs @@ -7,7 +7,7 @@ use vortex_array::IntoArray; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; -use vortex_array::scalar_fn::fns::cast::CastReduce; +use vortex_array::scalar_fn::builtin::cast::CastReduce; use vortex_error::VortexResult; use crate::rle::RLE; diff --git a/encodings/fastlanes/src/rle/probe/tests.rs b/encodings/fastlanes/src/rle/probe/tests.rs index f8523270095..43fa87164a1 100644 --- a/encodings/fastlanes/src/rle/probe/tests.rs +++ b/encodings/fastlanes/src/rle/probe/tests.rs @@ -13,7 +13,7 @@ use vortex_array::builders::builder_with_capacity_in; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::Nullability; use vortex_array::scalar::Scalar; -use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_array::validity::Validity; use vortex_error::VortexResult; diff --git a/encodings/fastlanes/src/rle/vtable/rules.rs b/encodings/fastlanes/src/rle/vtable/rules.rs index 50dd1f16a5a..fb605cabe74 100644 --- a/encodings/fastlanes/src/rle/vtable/rules.rs +++ b/encodings/fastlanes/src/rle/vtable/rules.rs @@ -2,7 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_array::optimizer::rules::ParentRuleSet; -use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor; +use vortex_array::scalar_fn::builtin::cast::CastReduceAdaptor; use crate::RLE; diff --git a/encodings/fsst/Cargo.toml b/encodings/fsst/Cargo.toml index 5b15dca85fb..3d6313686dd 100644 --- a/encodings/fsst/Cargo.toml +++ b/encodings/fsst/Cargo.toml @@ -22,6 +22,7 @@ rand = { workspace = true, optional = true } vortex-array = { workspace = true } vortex-buffer = { workspace = true } vortex-error = { workspace = true } +vortex-functions = { workspace = true } vortex-mask = { workspace = true } vortex-session = { workspace = true } diff --git a/encodings/fsst/benches/fsst_compress.rs b/encodings/fsst/benches/fsst_compress.rs index 3a0a4e682e6..69c76be10b3 100644 --- a/encodings/fsst/benches/fsst_compress.rs +++ b/encodings/fsst/benches/fsst_compress.rs @@ -26,7 +26,7 @@ use vortex_array::dtype::Nullability; use vortex_array::memory::MemorySession; use vortex_array::optimizer::kernels::KernelSession; use vortex_array::scalar::Scalar; -use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_array::session::ArraySession; use vortex_fsst::fsst_compress; use vortex_fsst::fsst_train_compressor; diff --git a/encodings/fsst/benches/fsst_like.rs b/encodings/fsst/benches/fsst_like.rs index a6336292ef2..fbbab52d57b 100644 --- a/encodings/fsst/benches/fsst_like.rs +++ b/encodings/fsst/benches/fsst_like.rs @@ -14,8 +14,6 @@ use vortex_array::VortexSessionExecute; use vortex_array::arrays::ConstantArray; use vortex_array::memory::MemorySession; use vortex_array::optimizer::kernels::KernelSession; -use vortex_array::scalar_fn::fns::like::Like; -use vortex_array::scalar_fn::fns::like::LikeOptions; use vortex_array::session::ArraySession; use vortex_fsst::FSSTArray; use vortex_fsst::test_utils::NUM_STRINGS; @@ -26,6 +24,8 @@ use vortex_fsst::test_utils::make_fsst_json_strings; use vortex_fsst::test_utils::make_fsst_log_lines; use vortex_fsst::test_utils::make_fsst_rare_match; use vortex_fsst::test_utils::make_fsst_short_urls; +use vortex_functions::like::Like; +use vortex_functions::like::LikeOptions; use vortex_session::VortexSession; #[global_allocator] diff --git a/encodings/fsst/benches/fsst_url_compare.rs b/encodings/fsst/benches/fsst_url_compare.rs index f62af505f3d..f4f3983ecb3 100644 --- a/encodings/fsst/benches/fsst_url_compare.rs +++ b/encodings/fsst/benches/fsst_url_compare.rs @@ -16,13 +16,12 @@ use vortex_array::arrays::VarBinArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::varbin::VarBinArrayExt; use vortex_array::builtins::ArrayBuiltins; -use vortex_array::expr::like; use vortex_array::expr::lit; use vortex_array::expr::root; use vortex_array::memory::MemorySession; use vortex_array::optimizer::kernels::KernelSession; use vortex_array::scalar::Scalar; -use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_array::session::ArraySession; use vortex_fsst::fsst_compress; use vortex_fsst::fsst_train_compressor; @@ -30,6 +29,7 @@ use vortex_fsst::test_utils::HIGH_MATCH_DOMAIN; use vortex_fsst::test_utils::LOW_MATCH_DOMAIN; use vortex_fsst::test_utils::NUM_STRINGS; use vortex_fsst::test_utils::generate_url_data; +use vortex_functions::expr::like; use vortex_session::VortexSession; #[global_allocator] diff --git a/encodings/fsst/src/array.rs b/encodings/fsst/src/array.rs index c8cb108adf3..cd5f868d0de 100644 --- a/encodings/fsst/src/array.rs +++ b/encodings/fsst/src/array.rs @@ -1062,6 +1062,7 @@ mod test { use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; + use vortex_array::optimizer::ArrayOptimizer; use vortex_array::test_harness::check_metadata; use vortex_array::vtable::VTable as _; use vortex_buffer::Buffer; @@ -1098,6 +1099,7 @@ mod test { let compressor_ptr = fsst_array.compressor() as *const Compressor; let sliced = fsst_array .slice(1..3)? + .optimize(&SESSION)? .try_downcast::() .map_err(|_| vortex_err!("slice must return an FSST array"))?; let sliced_compressor_ptr = sliced.compressor() as *const Compressor; diff --git a/encodings/fsst/src/compute/byte_length.rs b/encodings/fsst/src/compute/byte_length.rs index 05505fe95e5..668fad8e718 100644 --- a/encodings/fsst/src/compute/byte_length.rs +++ b/encodings/fsst/src/compute/byte_length.rs @@ -10,9 +10,9 @@ use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::PType; use vortex_array::scalar::Scalar; -use vortex_array::scalar_fn::fns::byte_length::ByteLengthKernel; use vortex_array::validity::Validity; use vortex_error::VortexResult; +use vortex_functions::byte_length::ByteLengthKernel; use crate::FSST; use crate::array::FSSTArraySlotsExt; diff --git a/encodings/fsst/src/compute/cast.rs b/encodings/fsst/src/compute/cast.rs index 63513cd4ba3..a9c7cc6bd57 100644 --- a/encodings/fsst/src/compute/cast.rs +++ b/encodings/fsst/src/compute/cast.rs @@ -8,8 +8,8 @@ use vortex_array::IntoArray; use vortex_array::arrays::VarBinArray; use vortex_array::arrays::varbin::VarBinArraySlotsExt; use vortex_array::dtype::DType; -use vortex_array::scalar_fn::fns::cast::CastKernel; -use vortex_array::scalar_fn::fns::cast::CastReduce; +use vortex_array::scalar_fn::builtin::cast::CastKernel; +use vortex_array::scalar_fn::builtin::cast::CastReduce; use vortex_array::validity::Validity; use vortex_error::VortexResult; diff --git a/encodings/fsst/src/compute/compare.rs b/encodings/fsst/src/compute/compare.rs index 53422acebee..dc6699bf26c 100644 --- a/encodings/fsst/src/compute/compare.rs +++ b/encodings/fsst/src/compute/compare.rs @@ -10,9 +10,9 @@ use vortex_array::arrays::ConstantArray; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::scalar::Scalar; -use vortex_array::scalar_fn::fns::binary::CompareKernel; -use vortex_array::scalar_fn::fns::operators::CompareOperator; -use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::scalar_fn::builtin::binary::CompareKernel; +use vortex_array::scalar_fn::builtin::operators::CompareOperator; +use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_buffer::BitBuffer; use vortex_buffer::ByteBuffer; use vortex_error::VortexExpect; @@ -135,7 +135,7 @@ mod tests { use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::scalar::Scalar; - use vortex_array::scalar_fn::fns::operators::Operator; + use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_error::VortexResult; use vortex_session::VortexSession; diff --git a/encodings/fsst/src/compute/like.rs b/encodings/fsst/src/compute/like.rs index 380de2f82cc..328675cf777 100644 --- a/encodings/fsst/src/compute/like.rs +++ b/encodings/fsst/src/compute/like.rs @@ -9,9 +9,9 @@ use vortex_array::arrays::BoolArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::varbin::VarBinArraySlotsExt; use vortex_array::match_each_integer_ptype; -use vortex_array::scalar_fn::fns::like::LikeKernel; -use vortex_array::scalar_fn::fns::like::LikeOptions; use vortex_error::VortexResult; +use vortex_functions::like::LikeKernel; +use vortex_functions::like::LikeOptions; use crate::FSST; use crate::FSSTArrayExt; @@ -89,10 +89,10 @@ mod tests { use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; - use vortex_array::scalar_fn::fns::like::Like; - use vortex_array::scalar_fn::fns::like::LikeKernel; - use vortex_array::scalar_fn::fns::like::LikeOptions; use vortex_error::VortexResult; + use vortex_functions::like::Like; + use vortex_functions::like::LikeKernel; + use vortex_functions::like::LikeOptions; use vortex_session::VortexSession; use crate::FSST; diff --git a/encodings/fsst/src/compute/mod.rs b/encodings/fsst/src/compute/mod.rs index 01b36149a3b..cf821ff4b88 100644 --- a/encodings/fsst/src/compute/mod.rs +++ b/encodings/fsst/src/compute/mod.rs @@ -175,6 +175,7 @@ mod tests { use vortex_array::dtype::Nullability; use vortex_array::expr::stats::Precision; use vortex_array::expr::stats::Stat; + use vortex_array::optimizer::ArrayOptimizer; use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; @@ -286,7 +287,7 @@ mod tests { let array = fsst_compress(&input, &compressor, &mut ctx)?.into_array(); let indices = buffer![0u64].into_array(); for (range, gathers) in [(0..8, true), (4..16, false)] { - let sliced = array.slice(range.clone())?; + let sliced = array.slice(range.clone())?.optimize(&SESSION)?; assert_eq!( FSST::take(sliced.as_::(), &indices, &mut ctx)?.is_some(), gathers diff --git a/encodings/fsst/src/dfa/tests.rs b/encodings/fsst/src/dfa/tests.rs index c97682e27e3..b53a6a53385 100644 --- a/encodings/fsst/src/dfa/tests.rs +++ b/encodings/fsst/src/dfa/tests.rs @@ -16,9 +16,9 @@ use vortex_array::arrays::VarBinArray; use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; -use vortex_array::scalar_fn::fns::like::Like; -use vortex_array::scalar_fn::fns::like::LikeOptions; use vortex_error::VortexResult; +use vortex_functions::like::Like; +use vortex_functions::like::LikeOptions; use vortex_session::VortexSession; use super::FsstMatcher; diff --git a/encodings/fsst/src/kernel.rs b/encodings/fsst/src/kernel.rs index 44d9c7b2531..840594d7b17 100644 --- a/encodings/fsst/src/kernel.rs +++ b/encodings/fsst/src/kernel.rs @@ -8,14 +8,14 @@ use vortex_array::arrays::dict::TakeExecuteAdaptor; use vortex_array::arrays::filter::FilterExecuteAdaptor; use vortex_array::optimizer::kernels::ArrayKernelsExt; use vortex_array::scalar_fn::ScalarFnVTable; -use vortex_array::scalar_fn::fns::binary::Binary; -use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; -use vortex_array::scalar_fn::fns::byte_length::ByteLength; -use vortex_array::scalar_fn::fns::byte_length::ByteLengthExecuteAdaptor; -use vortex_array::scalar_fn::fns::cast::Cast; -use vortex_array::scalar_fn::fns::cast::CastExecuteAdaptor; -use vortex_array::scalar_fn::fns::like::Like; -use vortex_array::scalar_fn::fns::like::LikeExecuteAdaptor; +use vortex_array::scalar_fn::builtin::binary::Binary; +use vortex_array::scalar_fn::builtin::binary::CompareExecuteAdaptor; +use vortex_array::scalar_fn::builtin::cast::Cast; +use vortex_array::scalar_fn::builtin::cast::CastExecuteAdaptor; +use vortex_functions::byte_length::ByteLength; +use vortex_functions::byte_length::ByteLengthExecuteAdaptor; +use vortex_functions::like::Like; +use vortex_functions::like::LikeExecuteAdaptor; use vortex_session::VortexSession; use crate::FSST; @@ -44,9 +44,9 @@ mod tests { use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; - use vortex_array::expr::byte_length; use vortex_array::expr::root; use vortex_error::VortexResult; + use vortex_functions::expr::byte_length; use vortex_mask::Mask; use vortex_session::VortexSession; diff --git a/encodings/fsst/src/rules.rs b/encodings/fsst/src/rules.rs index bec81c0fd3a..eb68e392981 100644 --- a/encodings/fsst/src/rules.rs +++ b/encodings/fsst/src/rules.rs @@ -3,7 +3,7 @@ use vortex_array::arrays::slice::SliceReduceAdaptor; use vortex_array::optimizer::rules::ParentRuleSet; -use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor; +use vortex_array::scalar_fn::builtin::cast::CastReduceAdaptor; use crate::FSST; diff --git a/encodings/fsst/src/slice.rs b/encodings/fsst/src/slice.rs index c2b2e320d31..1a4d921ff90 100644 --- a/encodings/fsst/src/slice.rs +++ b/encodings/fsst/src/slice.rs @@ -23,9 +23,7 @@ impl SliceReduce for FSST { FSST::new_unchecked_with_symbol_table( array.dtype().clone(), array.symbol_table(), - array - .codes() - .slice(range.clone())? + VarBin::_slice(array.codes().as_view(), range.clone())? .try_downcast::() .map_err(|_| vortex_err!("cannot fail conversion"))?, array.uncompressed_lengths().slice(range)?, diff --git a/encodings/fsst/src/tests.rs b/encodings/fsst/src/tests.rs index dbd9de56dcb..cb78aa4737c 100644 --- a/encodings/fsst/src/tests.rs +++ b/encodings/fsst/src/tests.rs @@ -14,6 +14,7 @@ use vortex_array::assert_arrays_eq; use vortex_array::assert_nth_scalar; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; +use vortex_array::optimizer::ArrayOptimizer; use vortex_buffer::buffer; use vortex_mask::Mask; use vortex_session::VortexSession; @@ -73,7 +74,7 @@ fn test_fsst_array_ops() { ); // test slice - let fsst_sliced = fsst_array.slice(1..3).unwrap(); + let fsst_sliced = fsst_array.slice(1..3).unwrap().optimize(&SESSION).unwrap(); assert!(fsst_sliced.is::()); assert_eq!(fsst_sliced.len(), 2); assert_nth_scalar!( diff --git a/encodings/onpair/Cargo.toml b/encodings/onpair/Cargo.toml index 4c82ba08aba..f703ca60fb9 100644 --- a/encodings/onpair/Cargo.toml +++ b/encodings/onpair/Cargo.toml @@ -21,6 +21,7 @@ prost = { workspace = true } vortex-array = { workspace = true } vortex-buffer = { workspace = true } vortex-error = { workspace = true } +vortex-functions = { workspace = true } vortex-mask = { workspace = true } vortex-session = { workspace = true } diff --git a/encodings/onpair/src/compute/byte_length.rs b/encodings/onpair/src/compute/byte_length.rs index 933e8bccd85..3ee181ad551 100644 --- a/encodings/onpair/src/compute/byte_length.rs +++ b/encodings/onpair/src/compute/byte_length.rs @@ -8,8 +8,8 @@ use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::PType; use vortex_array::scalar::Scalar; -use vortex_array::scalar_fn::fns::byte_length::ByteLengthKernel; use vortex_array::validity::Validity; +use vortex_functions::byte_length::ByteLengthKernel; use crate::OnPair; use crate::OnPairArraySlotsExt; diff --git a/encodings/onpair/src/compute/cast.rs b/encodings/onpair/src/compute/cast.rs index 737b6933f8f..fd9968e8678 100644 --- a/encodings/onpair/src/compute/cast.rs +++ b/encodings/onpair/src/compute/cast.rs @@ -6,8 +6,8 @@ use vortex_array::ArrayView; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::dtype::DType; -use vortex_array::scalar_fn::fns::cast::CastKernel; -use vortex_array::scalar_fn::fns::cast::CastReduce; +use vortex_array::scalar_fn::builtin::cast::CastKernel; +use vortex_array::scalar_fn::builtin::cast::CastReduce; use vortex_array::validity::Validity; use vortex_error::VortexResult; diff --git a/encodings/onpair/src/compute/compare.rs b/encodings/onpair/src/compute/compare.rs index fe8e6a0cec3..28f8af1c4f6 100644 --- a/encodings/onpair/src/compute/compare.rs +++ b/encodings/onpair/src/compute/compare.rs @@ -11,8 +11,8 @@ use vortex_array::arrays::ConstantArray; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::scalar::Scalar; -use vortex_array::scalar_fn::fns::binary::CompareKernel; -use vortex_array::scalar_fn::fns::operators::CompareOperator; +use vortex_array::scalar_fn::builtin::binary::CompareKernel; +use vortex_array::scalar_fn::builtin::operators::CompareOperator; use vortex_buffer::BitBuffer; use vortex_error::VortexResult; @@ -101,10 +101,11 @@ mod tests { use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; + use vortex_array::optimizer::ArrayOptimizer; use vortex_array::scalar::Scalar; - use vortex_array::scalar_fn::fns::binary::CompareKernel; - use vortex_array::scalar_fn::fns::operators::CompareOperator; - use vortex_array::scalar_fn::fns::operators::Operator; + use vortex_array::scalar_fn::builtin::binary::CompareKernel; + use vortex_array::scalar_fn::builtin::operators::CompareOperator; + use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_error::VortexResult; use vortex_error::vortex_err; use vortex_session::VortexSession; @@ -250,7 +251,7 @@ mod tests { ); let mut ctx = SESSION.create_execution_ctx(); let arr = onpair_compress(input.as_array(), DEFAULT_CONFIG, &mut ctx)?.into_array(); - let sliced = arr.slice(1..4)?; + let sliced = arr.slice(1..4)?.optimize(&SESSION)?; assert!(sliced.is::(), "slice dropped OnPair encoding"); let sliced = sliced .try_downcast::() diff --git a/encodings/onpair/src/kernel.rs b/encodings/onpair/src/kernel.rs index 6a76b1a5215..6e5c24cb36f 100644 --- a/encodings/onpair/src/kernel.rs +++ b/encodings/onpair/src/kernel.rs @@ -8,12 +8,12 @@ use vortex_array::arrays::dict::TakeExecuteAdaptor; use vortex_array::arrays::filter::FilterExecuteAdaptor; use vortex_array::optimizer::kernels::ArrayKernelsExt; use vortex_array::scalar_fn::ScalarFnVTable; -use vortex_array::scalar_fn::fns::binary::Binary; -use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; -use vortex_array::scalar_fn::fns::byte_length::ByteLength; -use vortex_array::scalar_fn::fns::byte_length::ByteLengthExecuteAdaptor; -use vortex_array::scalar_fn::fns::cast::Cast; -use vortex_array::scalar_fn::fns::cast::CastExecuteAdaptor; +use vortex_array::scalar_fn::builtin::binary::Binary; +use vortex_array::scalar_fn::builtin::binary::CompareExecuteAdaptor; +use vortex_array::scalar_fn::builtin::cast::Cast; +use vortex_array::scalar_fn::builtin::cast::CastExecuteAdaptor; +use vortex_functions::byte_length::ByteLength; +use vortex_functions::byte_length::ByteLengthExecuteAdaptor; use vortex_session::VortexSession; use crate::OnPair; diff --git a/encodings/onpair/src/rules.rs b/encodings/onpair/src/rules.rs index 279c160c1eb..ab18f256604 100644 --- a/encodings/onpair/src/rules.rs +++ b/encodings/onpair/src/rules.rs @@ -3,7 +3,7 @@ use vortex_array::arrays::slice::SliceReduceAdaptor; use vortex_array::optimizer::rules::ParentRuleSet; -use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor; +use vortex_array::scalar_fn::builtin::cast::CastReduceAdaptor; use crate::OnPair; diff --git a/encodings/onpair/src/tests.rs b/encodings/onpair/src/tests.rs index aa735b0524a..0759f204afe 100644 --- a/encodings/onpair/src/tests.rs +++ b/encodings/onpair/src/tests.rs @@ -25,7 +25,8 @@ use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::match_each_integer_ptype; -use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::optimizer::ArrayOptimizer; +use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_array::test_harness::check_metadata; use vortex_array::validity::Validity; use vortex_buffer::BufferMut; @@ -284,7 +285,7 @@ fn test_onpair_scalar_at_window() -> vortex_error::VortexResult<()> { // Sliced array: `codes_offsets` is narrowed (first boundary > 0), so the // point lookup must resolve indices relative to the slice. let (start, end) = (700usize, 1300usize); - let sliced = arr.slice(start..end)?; + let sliced = arr.slice(start..end)?.optimize(&SESSION)?; assert!(sliced.is::(), "slice dropped OnPair encoding"); for &j in &[0usize, 1, 300, end - start - 1] { let got = sliced.execute_scalar(j, &mut ctx)?; @@ -549,7 +550,7 @@ fn test_onpair_slice_canonicalize() -> vortex_error::VortexResult<()> { // interior (start>0, end0, // end=n), and a near-full window. for (start, end) in [(1234usize, 1240usize), (0, 7), (4993, n), (1, n - 1)] { - let sliced = arr.clone().slice(start..end)?; + let sliced = arr.clone().slice(start..end)?.optimize(&SESSION)?; assert_eq!(sliced.len(), end - start); assert!( sliced.is::(), diff --git a/encodings/onpair/tests/big_data.rs b/encodings/onpair/tests/big_data.rs index 80cbaeb18d0..37b39ba6e46 100644 --- a/encodings/onpair/tests/big_data.rs +++ b/encodings/onpair/tests/big_data.rs @@ -26,7 +26,7 @@ use vortex_array::assert_arrays_eq; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; -use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_onpair::DEFAULT_CONFIG; use vortex_onpair::onpair_compress; use vortex_session::VortexSession; diff --git a/encodings/parquet-variant/Cargo.toml b/encodings/parquet-variant/Cargo.toml index bb7e786b5ef..903610bfadf 100644 --- a/encodings/parquet-variant/Cargo.toml +++ b/encodings/parquet-variant/Cargo.toml @@ -27,6 +27,7 @@ vortex-array = { workspace = true } vortex-arrow = { workspace = true } vortex-buffer = { workspace = true } vortex-error = { workspace = true } +vortex-functions = { workspace = true } vortex-json = { workspace = true } vortex-mask = { workspace = true } vortex-session = { workspace = true } diff --git a/encodings/parquet-variant/src/json_to_variant_tests.rs b/encodings/parquet-variant/src/json_to_variant_tests.rs index e66ed56c67d..16febc9d8ce 100644 --- a/encodings/parquet-variant/src/json_to_variant_tests.rs +++ b/encodings/parquet-variant/src/json_to_variant_tests.rs @@ -27,11 +27,11 @@ use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::expr::root; -use vortex_array::expr::variant_get; -use vortex_array::scalar_fn::fns::variant_get::VariantPath; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; +use vortex_functions::expr::variant_get; +use vortex_functions::variant_get::VariantPath; use vortex_json::Json; use vortex_session::VortexSession; diff --git a/encodings/parquet-variant/src/kernel.rs b/encodings/parquet-variant/src/kernel.rs index b218d9404f4..d0c1d2d4e1f 100644 --- a/encodings/parquet-variant/src/kernel.rs +++ b/encodings/parquet-variant/src/kernel.rs @@ -40,15 +40,15 @@ use vortex_array::dtype::DType; use vortex_array::kernel::ExecuteParentKernel; use vortex_array::optimizer::kernels::ArrayKernelsExt; use vortex_array::scalar_fn::ScalarFnVTable; -use vortex_array::scalar_fn::fns::variant_get::VariantGet; -use vortex_array::scalar_fn::fns::variant_get::VariantPath; -use vortex_array::scalar_fn::fns::variant_get::VariantPathElement; use vortex_arrow::ArrowExportOptions; use vortex_arrow::ArrowSession; use vortex_arrow::ArrowSessionExt; use vortex_error::VortexResult; use vortex_error::vortex_ensure_eq; use vortex_error::vortex_err; +use vortex_functions::variant_get::VariantGet; +use vortex_functions::variant_get::VariantPath; +use vortex_functions::variant_get::VariantPathElement; use vortex_json::Json; use vortex_json::JsonToVariant; use vortex_mask::Mask; @@ -326,15 +326,16 @@ mod tests { use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::expr::root; - use vortex_array::expr::variant_get; - use vortex_array::scalar_fn::fns::variant_get::VariantPath; - use vortex_array::scalar_fn::fns::variant_get::VariantPathElement; + use vortex_array::optimizer::ArrayOptimizer; use vortex_array::validity::Validity; use vortex_arrow::ArrowSessionExt; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_err; + use vortex_functions::expr::variant_get; + use vortex_functions::variant_get::VariantPath; + use vortex_functions::variant_get::VariantPathElement; use vortex_mask::Mask; use vortex_session::VortexSession; @@ -994,7 +995,7 @@ mod tests { expected: impl IntoIterator>, ) -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); - let executed = array.clone().execute::(&mut ctx)?; + let executed = array.optimize(&SESSION)?.execute::(&mut ctx)?; let executed_variant = executed.as_::(); let typed_value = executed_variant .typed_value() diff --git a/encodings/pco/src/compute/cast.rs b/encodings/pco/src/compute/cast.rs index af1d003d7f7..4d5e800ff1a 100644 --- a/encodings/pco/src/compute/cast.rs +++ b/encodings/pco/src/compute/cast.rs @@ -5,7 +5,7 @@ use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::IntoArray; use vortex_array::dtype::DType; -use vortex_array::scalar_fn::fns::cast::CastReduce; +use vortex_array::scalar_fn::builtin::cast::CastReduce; use vortex_error::VortexResult; use crate::Pco; diff --git a/encodings/pco/src/probe/tests.rs b/encodings/pco/src/probe/tests.rs index c5386f1aa00..beacb432f25 100644 --- a/encodings/pco/src/probe/tests.rs +++ b/encodings/pco/src/probe/tests.rs @@ -12,6 +12,7 @@ use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::builders::builder_with_capacity_in; +use vortex_array::optimizer::ArrayOptimizer; use vortex_array::validity::Validity; use vortex_error::VortexResult; use vortex_session::VortexSession; @@ -47,7 +48,7 @@ fn random_access( }; let encoded = Pco::from_primitive(input.as_view(), 3, 128, &mut ctx)?.into_array(); let range = if sliced { 777..3333 } else { 0..4096 }; - let source = encoded.slice(range.clone())?; + let source = encoded.slice(range.clone())?.optimize(&SESSION)?; assert!(source.is::()); let input = input.slice(range)?; let indices = [0u32, 1, 127, 128, 512, 2048, 129, 2, 1024, 0]; diff --git a/encodings/pco/src/rules.rs b/encodings/pco/src/rules.rs index 4d27392dc01..c7bd63a123b 100644 --- a/encodings/pco/src/rules.rs +++ b/encodings/pco/src/rules.rs @@ -3,7 +3,7 @@ use vortex_array::arrays::slice::SliceReduceAdaptor; use vortex_array::optimizer::rules::ParentRuleSet; -use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor; +use vortex_array::scalar_fn::builtin::cast::CastReduceAdaptor; use crate::Pco; diff --git a/encodings/runend/src/compute/cast.rs b/encodings/runend/src/compute/cast.rs index 0619ae1614f..43fcf66294a 100644 --- a/encodings/runend/src/compute/cast.rs +++ b/encodings/runend/src/compute/cast.rs @@ -6,7 +6,7 @@ use vortex_array::ArrayView; use vortex_array::IntoArray; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; -use vortex_array::scalar_fn::fns::cast::CastReduce; +use vortex_array::scalar_fn::builtin::cast::CastReduce; use vortex_error::VortexResult; use crate::RunEnd; diff --git a/encodings/runend/src/compute/compare.rs b/encodings/runend/src/compute/compare.rs index 00ce5a904e4..af6a84f9f9e 100644 --- a/encodings/runend/src/compute/compare.rs +++ b/encodings/runend/src/compute/compare.rs @@ -9,9 +9,9 @@ use vortex_array::arrays::BoolArray; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::builtins::ArrayBuiltins; -use vortex_array::scalar_fn::fns::binary::CompareKernel; -use vortex_array::scalar_fn::fns::operators::CompareOperator; -use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::scalar_fn::builtin::binary::CompareKernel; +use vortex_array::scalar_fn::builtin::operators::CompareOperator; +use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_error::VortexResult; use crate::RunEnd; @@ -59,7 +59,7 @@ mod test { use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::builtins::ArrayBuiltins; - use vortex_array::scalar_fn::fns::operators::Operator; + use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_session::VortexSession; use crate::RunEnd; diff --git a/encodings/runend/src/compute/fill_null.rs b/encodings/runend/src/compute/fill_null.rs index 90284be3e51..046d4154300 100644 --- a/encodings/runend/src/compute/fill_null.rs +++ b/encodings/runend/src/compute/fill_null.rs @@ -6,7 +6,7 @@ use vortex_array::ArrayView; use vortex_array::IntoArray; use vortex_array::builtins::ArrayBuiltins; use vortex_array::scalar::Scalar; -use vortex_array::scalar_fn::fns::fill_null::FillNullReduce; +use vortex_array::scalar_fn::builtin::fill_null::FillNullReduce; use vortex_error::VortexResult; use crate::RunEnd; diff --git a/encodings/runend/src/kernel.rs b/encodings/runend/src/kernel.rs index 2eed964e3a9..673826f2478 100644 --- a/encodings/runend/src/kernel.rs +++ b/encodings/runend/src/kernel.rs @@ -17,8 +17,8 @@ use vortex_array::arrays::filter::FilterExecuteAdaptor; use vortex_array::kernel::ExecuteParentKernel; use vortex_array::optimizer::kernels::ArrayKernelsExt; use vortex_array::scalar_fn::ScalarFnVTable; -use vortex_array::scalar_fn::fns::binary::Binary; -use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; +use vortex_array::scalar_fn::builtin::binary::Binary; +use vortex_array::scalar_fn::builtin::binary::CompareExecuteAdaptor; use vortex_error::VortexResult; use vortex_session::VortexSession; diff --git a/encodings/runend/src/probe/tests.rs b/encodings/runend/src/probe/tests.rs index cedfe8af16d..0e4d7986cb8 100644 --- a/encodings/runend/src/probe/tests.rs +++ b/encodings/runend/src/probe/tests.rs @@ -15,7 +15,7 @@ use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::scalar::Scalar; -use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::buffer; diff --git a/encodings/runend/src/rules.rs b/encodings/runend/src/rules.rs index a6f3d6a4ae3..9c32a36b57b 100644 --- a/encodings/runend/src/rules.rs +++ b/encodings/runend/src/rules.rs @@ -13,8 +13,8 @@ use vortex_array::arrays::scalar_fn::ScalarFnArrayExt; use vortex_array::dtype::DType; use vortex_array::optimizer::rules::ArrayParentReduceRule; use vortex_array::optimizer::rules::ParentRuleSet; -use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor; -use vortex_array::scalar_fn::fns::fill_null::FillNullReduceAdaptor; +use vortex_array::scalar_fn::builtin::cast::CastReduceAdaptor; +use vortex_array::scalar_fn::builtin::fill_null::FillNullReduceAdaptor; use vortex_error::VortexResult; use crate::RunEnd; diff --git a/encodings/runend/src/trace_tests.rs b/encodings/runend/src/trace_tests.rs index 302fc9871fd..f6ef73cd6f2 100644 --- a/encodings/runend/src/trace_tests.rs +++ b/encodings/runend/src/trace_tests.rs @@ -3,6 +3,8 @@ //! Snapshot tests tracing how reductions and executions flow through run-end arrays. +use std::sync::LazyLock; + use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -15,8 +17,8 @@ use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::optimizer::ArrayOptimizer; use vortex_array::scalar::Scalar; -use vortex_array::scalar_fn::fns::binary::Binary; -use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::scalar_fn::builtin::binary::Binary; +use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_array::session::ArraySession; use vortex_array::test_harness::trace::trace_op; use vortex_error::VortexResult; @@ -25,6 +27,8 @@ use vortex_session::VortexSession; use crate::RunEnd; +static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + fn execution_ctx() -> ExecutionCtx { let session = VortexSession::empty().with::(); // Execute-parent kernels live in the session registry; without registering RunEnd's own @@ -48,9 +52,9 @@ fn trace_compare_on_runend() -> VortexResult<()> { let rhs = ConstantArray::new(Scalar::from(2i32), runend.len()).into_array(); let compared = Binary::try_new(runend, rhs, Operator::Eq)?.into_array(); - let traced = trace_op(|| compared.optimize())?; + let traced = trace_op(|| compared.optimize(&SESSION))?; insta::assert_snapshot!(traced.trace.to_string(), @" - optimize root=vortex.binary(bool, len=9) session=false + optimize root=vortex.binary(bool, len=9) reduce_parent static:RunEndScalarFnRule slot=0 parent=vortex.binary(bool, len=9) child=vortex.runend(i32, len=9) -> vortex.runend(bool, len=9) done output=vortex.runend(bool, len=9) "); @@ -68,6 +72,11 @@ fn trace_compare_on_runend() -> VortexResult<()> { &mut execution_ctx() ); insta::assert_snapshot!(traced.trace.to_string(), @" + optimize_recursive root=vortex.runend(bool, len=9) + optimize_recursive root=vortex.primitive(u8, len=3) + optimize_recursive root=vortex.binary(bool, len=3) + optimize_recursive root=vortex.primitive(i32, len=3) + optimize_recursive root=vortex.constant(i32, len=3) execute_until target=AnyCanonical root=vortex.runend(bool, len=9) iter 0 current=vortex.runend(bool, len=9) builder_active=false ExecuteSlot slot=1 parent=vortex.runend(bool, len=9) child=vortex.binary(bool, len=3) @@ -78,6 +87,7 @@ fn trace_compare_on_runend() -> VortexResult<()> { iter 3 current=vortex.runend(bool, len=9) builder_active=false Done array=vortex.bool(bool, len=9) iter 4 current=vortex.bool(bool, len=9) builder_active=false + optimize_recursive root=vortex.bool(bool, len=9) return output=vortex.bool(bool, len=9) "); @@ -105,12 +115,17 @@ fn trace_filter_on_runend() -> VortexResult<()> { &mut execution_ctx() ); insta::assert_snapshot!(traced.trace.to_string(), @" + optimize_recursive root=vortex.filter(i32, len=4) + optimize_recursive root=vortex.runend(i32, len=9) + optimize_recursive root=vortex.primitive(u8, len=3) + optimize_recursive root=vortex.primitive(i32, len=3) execute_until target=AnyCanonical root=vortex.filter(i32, len=4) iter 0 current=vortex.filter(i32, len=4) builder_active=false child_execute_parent session[0]:execute_parent_fn slot=0 parent=vortex.filter(i32, len=4) child=vortex.runend(i32, len=9) -> vortex.dict(i32, len=4) iter 1 current=vortex.dict(i32, len=4) builder_active=false child_execute_parent session[0]:execute_parent_fn slot=1 parent=vortex.dict(i32, len=4) child=vortex.primitive(i32, len=3) -> vortex.primitive(i32, len=4) iter 2 current=vortex.primitive(i32, len=4) builder_active=false + optimize_recursive root=vortex.primitive(i32, len=4) return output=vortex.primitive(i32, len=4) "); @@ -135,12 +150,18 @@ fn trace_take_on_runend() -> VortexResult<()> { &mut execution_ctx() ); insta::assert_snapshot!(traced.trace.to_string(), @" + optimize_recursive root=vortex.dict(i32, len=4) + optimize_recursive root=vortex.primitive(u64, len=4) + optimize_recursive root=vortex.runend(i32, len=9) + optimize_recursive root=vortex.primitive(u8, len=3) + optimize_recursive root=vortex.primitive(i32, len=3) execute_until target=AnyCanonical root=vortex.dict(i32, len=4) iter 0 current=vortex.dict(i32, len=4) builder_active=false child_execute_parent session[0]:execute_parent_fn slot=1 parent=vortex.dict(i32, len=4) child=vortex.runend(i32, len=9) -> vortex.dict(i32, len=4) iter 1 current=vortex.dict(i32, len=4) builder_active=false child_execute_parent session[0]:execute_parent_fn slot=1 parent=vortex.dict(i32, len=4) child=vortex.primitive(i32, len=3) -> vortex.primitive(i32, len=4) iter 2 current=vortex.primitive(i32, len=4) builder_active=false + optimize_recursive root=vortex.primitive(i32, len=4) return output=vortex.primitive(i32, len=4) "); diff --git a/encodings/sequence/Cargo.toml b/encodings/sequence/Cargo.toml index a6233c27a7b..e9b98fdfccb 100644 --- a/encodings/sequence/Cargo.toml +++ b/encodings/sequence/Cargo.toml @@ -20,6 +20,7 @@ vortex-array = { workspace = true } vortex-buffer = { workspace = true } vortex-compute = { workspace = true } vortex-error = { workspace = true } +vortex-functions = { workspace = true } vortex-mask = { workspace = true } vortex-session = { workspace = true } diff --git a/encodings/sequence/src/array.rs b/encodings/sequence/src/array.rs index 6b088ff57aa..db9ebae5cbf 100644 --- a/encodings/sequence/src/array.rs +++ b/encodings/sequence/src/array.rs @@ -586,6 +586,7 @@ mod tests { use vortex_array::expr::stats::Precision as StatPrecision; use vortex_array::expr::stats::Stat; use vortex_array::expr::stats::StatsProviderExt; + use vortex_array::optimizer::ArrayOptimizer; use vortex_array::scalar::PValue; use vortex_array::scalar::Scalar; use vortex_array::serde::SerializeOptions; @@ -711,7 +712,8 @@ mod tests { )?; let sliced = Sequence::try_new_typed(10i64, multiplier, Nullability::NonNullable, 4)? .into_array() - .slice(1..2)?; + .slice(1..2)? + .optimize(&SESSION)?; assert!(sliced.is::()); for array in [typed.into_array(), untyped.into_array(), sliced] { diff --git a/encodings/sequence/src/compute/cast.rs b/encodings/sequence/src/compute/cast.rs index 8eca1229f3d..104be45c5db 100644 --- a/encodings/sequence/src/compute/cast.rs +++ b/encodings/sequence/src/compute/cast.rs @@ -5,7 +5,7 @@ use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::IntoArray; use vortex_array::dtype::DType; -use vortex_array::scalar_fn::fns::cast::CastReduce; +use vortex_array::scalar_fn::builtin::cast::CastReduce; use vortex_error::VortexResult; use crate::Sequence; @@ -54,6 +54,7 @@ mod tests { use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; + use vortex_array::optimizer::ArrayOptimizer; use vortex_array::scalar::Scalar; use vortex_array::scalar::ScalarValue; use vortex_error::VortexResult; @@ -156,7 +157,8 @@ mod tests { fn test_cast_sequence_narrows_to_output_dtype() -> VortexResult<()> { let casted = Sequence::try_new_typed(100i32, -10i32, Nullability::NonNullable, 5)? .into_array() - .cast(DType::Primitive(PType::U8, Nullability::NonNullable))?; + .cast(DType::Primitive(PType::U8, Nullability::NonNullable))? + .optimize(&SESSION)?; let sequence = casted .as_typed::() diff --git a/encodings/sequence/src/compute/compare.rs b/encodings/sequence/src/compute/compare.rs index e2f3a0186be..cac8da99a47 100644 --- a/encodings/sequence/src/compute/compare.rs +++ b/encodings/sequence/src/compute/compare.rs @@ -10,8 +10,8 @@ use vortex_array::arrays::ConstantArray; use vortex_array::dtype::Nullability; use vortex_array::scalar::PValue; use vortex_array::scalar::Scalar; -use vortex_array::scalar_fn::fns::binary::CompareKernel; -use vortex_array::scalar_fn::fns::operators::CompareOperator; +use vortex_array::scalar_fn::builtin::binary::CompareKernel; +use vortex_array::scalar_fn::builtin::operators::CompareOperator; use vortex_buffer::BitBufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -135,7 +135,7 @@ mod tests { use vortex_array::dtype::Nullability::Nullable; use vortex_array::dtype::PType; use vortex_array::scalar::PValue; - use vortex_array::scalar_fn::fns::operators::Operator; + use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_error::VortexResult; use vortex_session::VortexSession; diff --git a/encodings/sequence/src/compute/list_contains.rs b/encodings/sequence/src/compute/list_contains.rs index ea966049ea4..8c7cd01c773 100644 --- a/encodings/sequence/src/compute/list_contains.rs +++ b/encodings/sequence/src/compute/list_contains.rs @@ -8,13 +8,13 @@ use vortex_array::arrays::BoolArray; use vortex_array::arrays::Constant; use vortex_array::arrays::ConstantArray; use vortex_array::scalar::Scalar; -use vortex_array::scalar_fn::fns::list_contains::ListContainsElementReduce; -use vortex_array::scalar_fn::fns::list_contains::ListContainsOptions; -use vortex_array::scalar_fn::fns::list_contains::PreparedSet; -use vortex_array::scalar_fn::fns::list_contains::PreparedSetData; use vortex_array::validity::Validity; use vortex_buffer::BitBufferMut; use vortex_error::VortexResult; +use vortex_functions::list_contains::ListContainsElementReduce; +use vortex_functions::list_contains::ListContainsOptions; +use vortex_functions::list_contains::PreparedSet; +use vortex_functions::list_contains::PreparedSetData; use crate::array::Sequence; use crate::compute::compare::Intersection; @@ -97,13 +97,14 @@ mod tests { use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType::I32; - use vortex_array::expr::list_contains; - use vortex_array::expr::list_contains_opts; use vortex_array::expr::lit; use vortex_array::expr::root; + use vortex_array::optimizer::ArrayOptimizer; use vortex_array::scalar::Scalar; - use vortex_array::scalar_fn::fns::list_contains::ListContainsOptions; use vortex_error::VortexResult; + use vortex_functions::expr::list_contains; + use vortex_functions::expr::list_contains_opts; + use vortex_functions::list_contains::ListContainsOptions; use vortex_session::VortexSession; use crate::Sequence; @@ -174,7 +175,7 @@ mod tests { ); let array = Sequence::try_new_typed(1, 1, Nullability::NonNullable, 3)?.into_array(); - let result = apply_list_contains(array, list_scalar, optimize)?; + let result = apply_list_contains(array, list_scalar, optimize)?.optimize(&SESSION)?; assert!(!result.is::()); let expected = BoolArray::from_iter([true, false, true]); diff --git a/encodings/sequence/src/kernel.rs b/encodings/sequence/src/kernel.rs index 37521a261e7..045d34eb0f9 100644 --- a/encodings/sequence/src/kernel.rs +++ b/encodings/sequence/src/kernel.rs @@ -8,8 +8,8 @@ use vortex_array::arrays::dict::TakeExecuteAdaptor; use vortex_array::arrays::filter::FilterExecuteAdaptor; use vortex_array::optimizer::kernels::ArrayKernelsExt; use vortex_array::scalar_fn::ScalarFnVTable; -use vortex_array::scalar_fn::fns::binary::Binary; -use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; +use vortex_array::scalar_fn::builtin::binary::Binary; +use vortex_array::scalar_fn::builtin::binary::CompareExecuteAdaptor; use vortex_session::VortexSession; use crate::Sequence; diff --git a/encodings/sequence/src/model_tests.rs b/encodings/sequence/src/model_tests.rs index 9bfb4624bf0..8c2da4b806d 100644 --- a/encodings/sequence/src/model_tests.rs +++ b/encodings/sequence/src/model_tests.rs @@ -26,7 +26,7 @@ use vortex_array::match_each_integer_ptype; use vortex_array::scalar::PValue; use vortex_array::scalar::Scalar; use vortex_array::scalar::ScalarValue; -use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_error::VortexResult; use vortex_error::vortex_err; use vortex_mask::Mask; diff --git a/encodings/sequence/src/rules.rs b/encodings/sequence/src/rules.rs index f87c08e8109..924c06fc5a0 100644 --- a/encodings/sequence/src/rules.rs +++ b/encodings/sequence/src/rules.rs @@ -3,8 +3,8 @@ use vortex_array::arrays::slice::SliceReduceAdaptor; use vortex_array::optimizer::rules::ParentRuleSet; -use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor; -use vortex_array::scalar_fn::fns::list_contains::ListContainsElementReduceAdaptor; +use vortex_array::scalar_fn::builtin::cast::CastReduceAdaptor; +use vortex_functions::list_contains::ListContainsElementReduceAdaptor; use crate::Sequence; diff --git a/encodings/sparse/Cargo.toml b/encodings/sparse/Cargo.toml index 78f8b9faeb0..84a3eac4756 100644 --- a/encodings/sparse/Cargo.toml +++ b/encodings/sparse/Cargo.toml @@ -22,6 +22,7 @@ prost = { workspace = true } vortex-array = { workspace = true } vortex-buffer = { workspace = true } vortex-error = { workspace = true } +vortex-functions = { workspace = true } vortex-mask = { workspace = true } vortex-session = { workspace = true } diff --git a/encodings/sparse/benches/sparse_pushdown.rs b/encodings/sparse/benches/sparse_pushdown.rs index f3e6133dccc..1692db3d846 100644 --- a/encodings/sparse/benches/sparse_pushdown.rs +++ b/encodings/sparse/benches/sparse_pushdown.rs @@ -35,7 +35,7 @@ use vortex_array::dtype::PType; use vortex_array::memory::MemorySession; use vortex_array::optimizer::kernels::KernelSession; use vortex_array::scalar::Scalar; -use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_array::session::ArraySession; use vortex_buffer::Buffer; use vortex_error::VortexExpect; diff --git a/encodings/sparse/src/compute/between.rs b/encodings/sparse/src/compute/between.rs index 1a01958ac44..bff24fdab71 100644 --- a/encodings/sparse/src/compute/between.rs +++ b/encodings/sparse/src/compute/between.rs @@ -6,10 +6,10 @@ use vortex_array::ArrayView; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::ConstantArray; -use vortex_array::builtins::ArrayBuiltins; -use vortex_array::scalar_fn::fns::between::BetweenKernel; -use vortex_array::scalar_fn::fns::between::BetweenOptions; use vortex_error::VortexResult; +use vortex_functions::between::BetweenKernel; +use vortex_functions::between::BetweenOptions; +use vortex_functions::builtins::ArrayFunctions; use crate::Sparse; use crate::SparseExt as _; @@ -70,11 +70,11 @@ mod tests { use vortex_array::VortexSessionExecute; use vortex_array::arrays::ConstantArray; use vortex_array::assert_arrays_eq; - use vortex_array::builtins::ArrayBuiltins; use vortex_array::scalar::Scalar; - use vortex_array::scalar_fn::fns::between::BetweenOptions; - use vortex_array::scalar_fn::fns::between::StrictComparison; use vortex_buffer::buffer; + use vortex_functions::between::BetweenOptions; + use vortex_functions::between::StrictComparison; + use vortex_functions::builtins::ArrayFunctions; use vortex_session::VortexSession; use crate::Sparse; diff --git a/encodings/sparse/src/compute/cast.rs b/encodings/sparse/src/compute/cast.rs index 88b07e0d9f4..74b8ad652d9 100644 --- a/encodings/sparse/src/compute/cast.rs +++ b/encodings/sparse/src/compute/cast.rs @@ -7,7 +7,7 @@ use vortex_array::IntoArray; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::scalar::Scalar; -use vortex_array::scalar_fn::fns::cast::CastReduce; +use vortex_array::scalar_fn::builtin::cast::CastReduce; use vortex_error::VortexResult; use crate::Sparse; diff --git a/encodings/sparse/src/compute/compare.rs b/encodings/sparse/src/compute/compare.rs index ad1fad781eb..8e6b5c3ac91 100644 --- a/encodings/sparse/src/compute/compare.rs +++ b/encodings/sparse/src/compute/compare.rs @@ -7,10 +7,10 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::ConstantArray; use vortex_array::builtins::ArrayBuiltins; -use vortex_array::scalar_fn::fns::binary::CompareKernel; -use vortex_array::scalar_fn::fns::binary::scalar_cmp; -use vortex_array::scalar_fn::fns::operators::CompareOperator; -use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::scalar_fn::builtin::binary::CompareKernel; +use vortex_array::scalar_fn::builtin::binary::scalar_cmp; +use vortex_array::scalar_fn::builtin::operators::CompareOperator; +use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_error::VortexResult; use crate::Sparse; @@ -64,7 +64,7 @@ mod tests { use vortex_array::assert_arrays_eq; use vortex_array::builtins::ArrayBuiltins; use vortex_array::scalar::Scalar; - use vortex_array::scalar_fn::fns::operators::Operator; + use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_buffer::buffer; use vortex_session::VortexSession; diff --git a/encodings/sparse/src/compute/fill_null.rs b/encodings/sparse/src/compute/fill_null.rs index 1ae877068d2..66ab0a3617c 100644 --- a/encodings/sparse/src/compute/fill_null.rs +++ b/encodings/sparse/src/compute/fill_null.rs @@ -7,7 +7,7 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::builtins::ArrayBuiltins; use vortex_array::scalar::Scalar; -use vortex_array::scalar_fn::fns::fill_null::FillNullKernel; +use vortex_array::scalar_fn::builtin::fill_null::FillNullKernel; use vortex_error::VortexResult; use crate::Sparse; diff --git a/encodings/sparse/src/compute/like.rs b/encodings/sparse/src/compute/like.rs index a4cb23a41b3..bc53335512c 100644 --- a/encodings/sparse/src/compute/like.rs +++ b/encodings/sparse/src/compute/like.rs @@ -6,10 +6,10 @@ use vortex_array::ArrayView; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::ConstantArray; -use vortex_array::scalar_fn::fns::like::Like; -use vortex_array::scalar_fn::fns::like::LikeKernel; -use vortex_array::scalar_fn::fns::like::LikeOptions; use vortex_error::VortexResult; +use vortex_functions::like::Like; +use vortex_functions::like::LikeKernel; +use vortex_functions::like::LikeOptions; use crate::Sparse; use crate::SparseExt as _; @@ -71,10 +71,10 @@ mod tests { use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::scalar::Scalar; - use vortex_array::scalar_fn::fns::like::Like; - use vortex_array::scalar_fn::fns::like::LikeOptions; use vortex_buffer::buffer; use vortex_error::VortexResult; + use vortex_functions::like::Like; + use vortex_functions::like::LikeOptions; use vortex_session::VortexSession; use crate::Sparse; diff --git a/encodings/sparse/src/kernel.rs b/encodings/sparse/src/kernel.rs index 1265a27d678..2b09d2a3758 100644 --- a/encodings/sparse/src/kernel.rs +++ b/encodings/sparse/src/kernel.rs @@ -10,14 +10,14 @@ use vortex_array::arrays::filter::FilterExecuteAdaptor; use vortex_array::arrays::slice::SliceExecuteAdaptor; use vortex_array::optimizer::kernels::ArrayKernelsExt; use vortex_array::scalar_fn::ScalarFnVTable; -use vortex_array::scalar_fn::fns::between::Between; -use vortex_array::scalar_fn::fns::between::BetweenExecuteAdaptor; -use vortex_array::scalar_fn::fns::binary::Binary; -use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; -use vortex_array::scalar_fn::fns::fill_null::FillNull; -use vortex_array::scalar_fn::fns::fill_null::FillNullExecuteAdaptor; -use vortex_array::scalar_fn::fns::like::Like; -use vortex_array::scalar_fn::fns::like::LikeExecuteAdaptor; +use vortex_array::scalar_fn::builtin::binary::Binary; +use vortex_array::scalar_fn::builtin::binary::CompareExecuteAdaptor; +use vortex_array::scalar_fn::builtin::fill_null::FillNull; +use vortex_array::scalar_fn::builtin::fill_null::FillNullExecuteAdaptor; +use vortex_functions::between::Between; +use vortex_functions::between::BetweenExecuteAdaptor; +use vortex_functions::like::Like; +use vortex_functions::like::LikeExecuteAdaptor; use vortex_session::VortexSession; use crate::Sparse; diff --git a/encodings/sparse/src/lib.rs b/encodings/sparse/src/lib.rs index 55e5ed883e2..d2ecb52dda8 100644 --- a/encodings/sparse/src/lib.rs +++ b/encodings/sparse/src/lib.rs @@ -43,7 +43,7 @@ use vortex_array::require_child; use vortex_array::require_opt_child; use vortex_array::scalar::Scalar; use vortex_array::scalar::ScalarValue; -use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_array::serde::ArrayChildren; use vortex_array::validity::Validity; use vortex_array::vtable::VTable; diff --git a/encodings/sparse/src/rules.rs b/encodings/sparse/src/rules.rs index d888d64a389..a001a875077 100644 --- a/encodings/sparse/src/rules.rs +++ b/encodings/sparse/src/rules.rs @@ -6,9 +6,9 @@ use vortex_array::ArrayView; use vortex_array::IntoArray; use vortex_array::builtins::ArrayBuiltins; use vortex_array::optimizer::rules::ParentRuleSet; -use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor; -use vortex_array::scalar_fn::fns::not::NotReduce; -use vortex_array::scalar_fn::fns::not::NotReduceAdaptor; +use vortex_array::scalar_fn::builtin::cast::CastReduceAdaptor; +use vortex_array::scalar_fn::builtin::not::NotReduce; +use vortex_array::scalar_fn::builtin::not::NotReduceAdaptor; use vortex_error::VortexResult; use crate::Sparse; diff --git a/encodings/zigzag/src/array.rs b/encodings/zigzag/src/array.rs index 25fed1c159e..8252344daf3 100644 --- a/encodings/zigzag/src/array.rs +++ b/encodings/zigzag/src/array.rs @@ -285,6 +285,7 @@ mod test { use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; + use vortex_array::optimizer::ArrayOptimizer; use vortex_array::scalar::Scalar; use vortex_buffer::buffer; use vortex_session::VortexSession; @@ -315,7 +316,7 @@ mod test { array.statistics().compute_is_constant(&mut ctx) ); - let sliced = zigzag.slice(0..2)?; + let sliced = zigzag.slice(0..2)?.optimize(&SESSION)?; let sliced = sliced.as_::(); assert_eq!( sliced.array().execute_scalar(sliced.len() - 1, &mut ctx,)?, diff --git a/encodings/zigzag/src/compute/cast.rs b/encodings/zigzag/src/compute/cast.rs index 1bd582f77f6..4e4a38e5e0a 100644 --- a/encodings/zigzag/src/compute/cast.rs +++ b/encodings/zigzag/src/compute/cast.rs @@ -6,7 +6,7 @@ use vortex_array::ArrayView; use vortex_array::IntoArray; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; -use vortex_array::scalar_fn::fns::cast::CastReduce; +use vortex_array::scalar_fn::builtin::cast::CastReduce; use vortex_error::VortexResult; use crate::ZigZag; @@ -38,6 +38,7 @@ mod tests { use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; + use vortex_array::optimizer::ArrayOptimizer; use vortex_session::VortexSession; use crate::ZigZagArray; @@ -57,6 +58,8 @@ mod tests { let casted = zigzag .into_array() .cast(DType::Primitive(PType::I64, Nullability::NonNullable)) + .unwrap() + .optimize(&SESSION) .unwrap(); assert_eq!( casted.dtype(), @@ -87,6 +90,8 @@ mod tests { let casted = zigzag .into_array() .cast(DType::Primitive(PType::I16, Nullability::NonNullable)) + .unwrap() + .optimize(&SESSION) .unwrap(); assert_eq!( casted.encoding_id().as_ref(), @@ -107,6 +112,8 @@ mod tests { let casted64 = zigzag16 .into_array() .cast(DType::Primitive(PType::I64, Nullability::NonNullable)) + .unwrap() + .optimize(&SESSION) .unwrap(); assert_eq!( casted64.encoding_id().as_ref(), diff --git a/encodings/zigzag/src/compute/mod.rs b/encodings/zigzag/src/compute/mod.rs index 8e301febcc1..0c8fad015f2 100644 --- a/encodings/zigzag/src/compute/mod.rs +++ b/encodings/zigzag/src/compute/mod.rs @@ -9,8 +9,8 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::dict::TakeExecute; use vortex_array::arrays::filter::FilterReduce; -use vortex_array::scalar_fn::fns::mask::Mask as MaskExpr; -use vortex_array::scalar_fn::fns::mask::MaskReduce; +use vortex_array::scalar_fn::builtin::mask::Mask as MaskExpr; +use vortex_array::scalar_fn::builtin::mask::MaskReduce; use vortex_error::VortexResult; use vortex_mask::Mask; diff --git a/encodings/zigzag/src/rules.rs b/encodings/zigzag/src/rules.rs index c3b612d101d..18f106ddf9a 100644 --- a/encodings/zigzag/src/rules.rs +++ b/encodings/zigzag/src/rules.rs @@ -4,8 +4,8 @@ use vortex_array::arrays::filter::FilterReduceAdaptor; use vortex_array::arrays::slice::SliceReduceAdaptor; use vortex_array::optimizer::rules::ParentRuleSet; -use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor; -use vortex_array::scalar_fn::fns::mask::MaskReduceAdaptor; +use vortex_array::scalar_fn::builtin::cast::CastReduceAdaptor; +use vortex_array::scalar_fn::builtin::mask::MaskReduceAdaptor; use crate::ZigZag; diff --git a/encodings/zstd/src/compute/cast.rs b/encodings/zstd/src/compute/cast.rs index f1bb4ebdcda..e1cd9fcd02c 100644 --- a/encodings/zstd/src/compute/cast.rs +++ b/encodings/zstd/src/compute/cast.rs @@ -6,7 +6,7 @@ use vortex_array::ArrayView; use vortex_array::IntoArray; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; -use vortex_array::scalar_fn::fns::cast::CastReduce; +use vortex_array::scalar_fn::builtin::cast::CastReduce; use vortex_array::vtable::child_to_validity; use vortex_error::VortexResult; diff --git a/encodings/zstd/src/rules.rs b/encodings/zstd/src/rules.rs index cfbf93b09e8..5b547144515 100644 --- a/encodings/zstd/src/rules.rs +++ b/encodings/zstd/src/rules.rs @@ -3,7 +3,7 @@ use vortex_array::arrays::slice::SliceReduceAdaptor; use vortex_array::optimizer::rules::ParentRuleSet; -use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor; +use vortex_array::scalar_fn::builtin::cast::CastReduceAdaptor; use crate::Zstd; diff --git a/fuzz/Cargo.toml b/fuzz/Cargo.toml index eed6fb69495..24e104cc785 100644 --- a/fuzz/Cargo.toml +++ b/fuzz/Cargo.toml @@ -36,6 +36,7 @@ vortex-btrblocks = { workspace = true } vortex-buffer = { workspace = true } vortex-error = { workspace = true } vortex-fsst = { workspace = true, features = ["_test-harness"] } +vortex-functions = { workspace = true } vortex-io = { workspace = true, features = ["wasm-bindgen"] } vortex-mask = { workspace = true } vortex-row = { workspace = true } diff --git a/fuzz/src/array/compare.rs b/fuzz/src/array/compare.rs index 03ce50a7170..ac44571552a 100644 --- a/fuzz/src/array/compare.rs +++ b/fuzz/src/array/compare.rs @@ -15,8 +15,8 @@ use vortex_array::dtype::Nullability; use vortex_array::match_each_decimal_value_type; use vortex_array::match_each_native_ptype; use vortex_array::scalar::Scalar; -use vortex_array::scalar_fn::fns::binary::scalar_cmp; -use vortex_array::scalar_fn::fns::operators::CompareOperator; +use vortex_array::scalar_fn::builtin::binary::scalar_cmp; +use vortex_array::scalar_fn::builtin::operators::CompareOperator; use vortex_array::validity::Validity; use vortex_buffer::BitBuffer; use vortex_buffer::ByteBuffer; diff --git a/fuzz/src/array/mod.rs b/fuzz/src/array/mod.rs index 638bd1559cf..bdfa0cffdce 100644 --- a/fuzz/src/array/mod.rs +++ b/fuzz/src/array/mod.rs @@ -59,8 +59,8 @@ use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::scalar::Scalar; use vortex_array::scalar::arbitrary::random_scalar; -use vortex_array::scalar_fn::fns::operators::CompareOperator; -use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::scalar_fn::builtin::operators::CompareOperator; +use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_array::search_sorted::SearchResult; use vortex_array::search_sorted::SearchSorted; use vortex_array::search_sorted::SearchSortedArray; diff --git a/fuzz/src/fsst_like.rs b/fuzz/src/fsst_like.rs index 0cd9b4c0cae..8ab863616a6 100644 --- a/fuzz/src/fsst_like.rs +++ b/fuzz/src/fsst_like.rs @@ -19,12 +19,12 @@ use vortex_array::arrays::VarBinArray; use vortex_array::arrays::bool::BoolArrayExt; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; -use vortex_array::scalar_fn::fns::like::Like; -use vortex_array::scalar_fn::fns::like::LikeOptions; use vortex_error::VortexResult; use vortex_fsst::FSSTArray; use vortex_fsst::fsst_compress; use vortex_fsst::fsst_train_compressor; +use vortex_functions::like::Like; +use vortex_functions::like::LikeOptions; use vortex_session::VortexSession; use crate::error::Backtrace; diff --git a/vortex-array/Cargo.toml b/vortex-array/Cargo.toml index fbbc3b42575..5238b5cd07f 100644 --- a/vortex-array/Cargo.toml +++ b/vortex-array/Cargo.toml @@ -155,14 +155,6 @@ harness = false name = "kleene_bool" harness = false -[[bench]] -name = "like" -harness = false - -[[bench]] -name = "list_contains_set" -harness = false - [[bench]] name = "interleave" harness = false @@ -200,11 +192,6 @@ name = "expr_large_struct_pack" path = "benches/expr/large_struct_pack.rs" harness = false -[[bench]] -name = "expr_case_when" -path = "benches/expr/case_when_bench.rs" -harness = false - [[bench]] name = "expr_optimize" path = "benches/expr/optimize_bench.rs" @@ -241,10 +228,6 @@ harness = false name = "dict_unreferenced_mask" harness = false -[[bench]] -name = "varbinview_zip" -harness = false - [[bench]] name = "varbinview_try_new" harness = false @@ -253,18 +236,6 @@ harness = false name = "varbinview_fill_null" harness = false -[[bench]] -name = "bool_zip" -harness = false - -[[bench]] -name = "primitive_zip" -harness = false - -[[bench]] -name = "listview_zip" -harness = false - [[bench]] name = "take_primitive" harness = false @@ -309,14 +280,6 @@ harness = false name = "filter_list" harness = false -[[bench]] -name = "list_length" -harness = false - -[[bench]] -name = "list_sum" -harness = false - [[bench]] name = "listview_rebuild" harness = false diff --git a/vortex-array/benches/binary_ops.rs b/vortex-array/benches/binary_ops.rs index 1103b6571fb..55dfdb73dbb 100644 --- a/vortex-array/benches/binary_ops.rs +++ b/vortex-array/benches/binary_ops.rs @@ -36,7 +36,7 @@ use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DecimalDType; use vortex_array::memory::MemorySession; use vortex_array::optimizer::kernels::KernelSession; -use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_session::VortexSession; #[global_allocator] diff --git a/vortex-array/benches/compare.rs b/vortex-array/benches/compare.rs index ddf6343e9ab..e98eb4a91ee 100644 --- a/vortex-array/benches/compare.rs +++ b/vortex-array/benches/compare.rs @@ -46,7 +46,7 @@ use vortex_array::dtype::DecimalDType; use vortex_array::memory::MemorySession; use vortex_array::optimizer::kernels::KernelSession; use vortex_array::scalar::Scalar; -use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_session::VortexSession; diff --git a/vortex-array/benches/compare_sorted.rs b/vortex-array/benches/compare_sorted.rs index a50a6666949..60175b012cb 100644 --- a/vortex-array/benches/compare_sorted.rs +++ b/vortex-array/benches/compare_sorted.rs @@ -28,7 +28,7 @@ use vortex_array::memory::MemorySession; use vortex_array::optimizer::kernels::KernelSession; use vortex_array::scalar::Scalar; use vortex_array::scalar::ScalarValue; -use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_buffer::Buffer; use vortex_session::VortexSession; diff --git a/vortex-array/benches/dict_compare.rs b/vortex-array/benches/dict_compare.rs index f97a12908d2..f541d054001 100644 --- a/vortex-array/benches/dict_compare.rs +++ b/vortex-array/benches/dict_compare.rs @@ -24,7 +24,7 @@ use vortex_array::expr::lit; use vortex_array::expr::root; use vortex_array::memory::MemorySession; use vortex_array::optimizer::kernels::KernelSession; -use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_session::VortexSession; #[global_allocator] diff --git a/vortex-array/benches/kleene_bool.rs b/vortex-array/benches/kleene_bool.rs index b6bfb25e442..970b8344fe4 100644 --- a/vortex-array/benches/kleene_bool.rs +++ b/vortex-array/benches/kleene_bool.rs @@ -20,7 +20,7 @@ use vortex_array::dtype::Nullability; use vortex_array::memory::MemorySession; use vortex_array::optimizer::kernels::KernelSession; use vortex_array::scalar::Scalar; -use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_session::VortexSession; #[global_allocator] diff --git a/vortex-array/benches/scalar_fn_probe.rs b/vortex-array/benches/scalar_fn_probe.rs index 4ce01520d17..cdf9c19e1d3 100644 --- a/vortex-array/benches/scalar_fn_probe.rs +++ b/vortex-array/benches/scalar_fn_probe.rs @@ -20,8 +20,8 @@ use vortex_array::arrays::ScalarFnArray; use vortex_array::memory::MemorySession; use vortex_array::optimizer::kernels::KernelSession; use vortex_array::scalar_fn::TypedScalarFnInstance; -use vortex_array::scalar_fn::fns::binary::Binary; -use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::scalar_fn::builtin::binary::Binary; +use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_session::VortexSession; #[global_allocator] diff --git a/vortex-array/benches/scalar_subtract.rs b/vortex-array/benches/scalar_subtract.rs index c882fc9fb76..cf6f0a5a43b 100644 --- a/vortex-array/benches/scalar_subtract.rs +++ b/vortex-array/benches/scalar_subtract.rs @@ -20,7 +20,7 @@ use vortex_array::builtins::ArrayBuiltins; use vortex_array::memory::MemorySession; use vortex_array::optimizer::kernels::KernelSession; use vortex_array::scalar::Scalar; -use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_buffer::Buffer; use vortex_session::VortexSession; diff --git a/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs b/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs index ff8e526402b..82091fa8772 100644 --- a/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs @@ -49,7 +49,7 @@ use crate::expr::stats::Stat; use crate::expr::stats::StatsProvider; use crate::expr::stats::StatsProviderExt; use crate::scalar::Scalar; -use crate::scalar_fn::fns::operators::Operator; +use crate::scalar_fn::builtin::operators::Operator; /// Check if two arrays of the same length have equal values at every position (null-safe). /// diff --git a/vortex-array/src/aggregate_fn/fns/mean/mod.rs b/vortex-array/src/aggregate_fn/fns/mean/mod.rs index b2dcf09fab4..c5e0c2bbb9c 100644 --- a/vortex-array/src/aggregate_fn/fns/mean/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/mean/mod.rs @@ -32,7 +32,7 @@ use crate::dtype::PType; use crate::dtype::i256; use crate::scalar::DecimalValue; use crate::scalar::Scalar; -use crate::scalar_fn::fns::operators::Operator; +use crate::scalar_fn::builtin::operators::Operator; /// Compute the arithmetic mean of an array. /// diff --git a/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs b/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs index 46fe681ab0a..9fe62f99a63 100644 --- a/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum_v2/mod.rs @@ -50,7 +50,7 @@ use crate::expr::stats::Stat; use crate::expr::stats::StatsProviderExt; use crate::scalar::Scalar; use crate::scalar::ScalarValue; -use crate::scalar_fn::fns::operators::Operator; +use crate::scalar_fn::builtin::operators::Operator; use crate::validity::Validity; const SUM_FIELD: &str = "sum"; diff --git a/vortex-array/src/array/erased.rs b/vortex-array/src/array/erased.rs index f630d42a4f8..98308fb7109 100644 --- a/vortex-array/src/array/erased.rs +++ b/vortex-array/src/array/erased.rs @@ -210,6 +210,9 @@ impl ArrayRef { } /// Performs a constant-time slice of the array. + /// + /// The result is a lazy [`SliceArray`] unless the slice is empty or covers the whole array; + /// optimizing it with a session pushes the slice into the child where a rule allows. pub fn slice(&self, range: Range) -> VortexResult { let len = self.len(); let start = range.start; @@ -228,9 +231,7 @@ impl ArrayRef { return Ok(Canonical::empty(self.dtype()).into_array()); } - let sliced = SliceArray::try_new(self.clone(), range)? - .into_array() - .optimize()?; + let sliced = SliceArray::try_new(self.clone(), range)?.into_array(); // Propagate some stats from the original array to the sliced array. if !sliced.is::() { @@ -252,16 +253,12 @@ impl ArrayRef { /// Wraps the array in a [`FilterArray`] such that it is logically filtered by the given mask. pub fn filter(&self, mask: Mask) -> VortexResult { - FilterArray::try_new(self.clone(), mask)? - .into_array() - .optimize() + Ok(FilterArray::try_new(self.clone(), mask)?.into_array()) } /// Wraps the array in a [`DictArray`] such that it is logically taken by the given indices. pub fn take(&self, indices: ArrayRef) -> VortexResult { - DictArray::try_new(indices, self.clone())? - .into_array() - .optimize() + Ok(DictArray::try_new(indices, self.clone())?.into_array()) } /// Execute the array to extract a scalar at the given index. @@ -390,12 +387,16 @@ impl ArrayRef { } /// Writes the array into the canonical builder. + /// + /// The array is optimized first, so a lazy wrapper such as a slice reaches the encoding's own + /// builder path instead of being executed to canonical form. pub fn append_to_builder( &self, builder: &mut dyn ArrayBuilder, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { - self.0.data.append_to_builder(self, builder, ctx) + let array = self.optimize(ctx.session())?; + array.0.data.append_to_builder(&array, builder, ctx) } /// Returns the statistics of the array. diff --git a/vortex-array/src/array/mod.rs b/vortex-array/src/array/mod.rs index fae796e1714..4ffee3201cc 100644 --- a/vortex-array/src/array/mod.rs +++ b/vortex-array/src/array/mod.rs @@ -32,8 +32,6 @@ pub use erased::*; mod plugin; pub use plugin::*; -mod haystack; -pub(crate) use haystack::ArrayHaystack; mod probe; pub use probe::*; diff --git a/vortex-array/src/arrays/bool/compute/boolean.rs b/vortex-array/src/arrays/bool/compute/boolean.rs index 28d1bbb7f52..3980911f2eb 100644 --- a/vortex-array/src/arrays/bool/compute/boolean.rs +++ b/vortex-array/src/arrays/bool/compute/boolean.rs @@ -10,10 +10,10 @@ use crate::array::ArrayView; use crate::arrays::Bool; use crate::arrays::Constant; use crate::arrays::bool::BoolArrayExt; -use crate::scalar_fn::fns::binary::BooleanKernel; -use crate::scalar_fn::fns::binary::kleene_boolean_buffer_scalar; -use crate::scalar_fn::fns::binary::kleene_boolean_buffers; -use crate::scalar_fn::fns::operators::Operator; +use crate::scalar_fn::builtin::binary::BooleanKernel; +use crate::scalar_fn::builtin::binary::kleene_boolean_buffer_scalar; +use crate::scalar_fn::builtin::binary::kleene_boolean_buffers; +use crate::scalar_fn::builtin::operators::Operator; impl BooleanKernel for Bool { fn boolean( diff --git a/vortex-array/src/arrays/bool/compute/cast.rs b/vortex-array/src/arrays/bool/compute/cast.rs index af98c78909c..846e042ea39 100644 --- a/vortex-array/src/arrays/bool/compute/cast.rs +++ b/vortex-array/src/arrays/bool/compute/cast.rs @@ -16,8 +16,8 @@ use crate::arrays::PrimitiveArray; use crate::arrays::bool::BoolArrayExt; use crate::dtype::DType; use crate::match_each_native_ptype; -use crate::scalar_fn::fns::cast::CastKernel; -use crate::scalar_fn::fns::cast::CastReduce; +use crate::scalar_fn::builtin::cast::CastKernel; +use crate::scalar_fn::builtin::cast::CastReduce; impl CastReduce for Bool { fn cast(array: ArrayView<'_, Bool>, dtype: &DType) -> VortexResult> { diff --git a/vortex-array/src/arrays/bool/compute/fill_null.rs b/vortex-array/src/arrays/bool/compute/fill_null.rs index 94966f2e829..71358a32590 100644 --- a/vortex-array/src/arrays/bool/compute/fill_null.rs +++ b/vortex-array/src/arrays/bool/compute/fill_null.rs @@ -12,7 +12,7 @@ use crate::arrays::Bool; use crate::arrays::BoolArray; use crate::arrays::bool::BoolArrayExt; use crate::scalar::Scalar; -use crate::scalar_fn::fns::fill_null::FillNullKernel; +use crate::scalar_fn::builtin::fill_null::FillNullKernel; use crate::validity::Validity; impl FillNullKernel for Bool { diff --git a/vortex-array/src/arrays/bool/compute/mask.rs b/vortex-array/src/arrays/bool/compute/mask.rs index 818dbaebca5..14eea59b24f 100644 --- a/vortex-array/src/arrays/bool/compute/mask.rs +++ b/vortex-array/src/arrays/bool/compute/mask.rs @@ -8,7 +8,7 @@ use crate::IntoArray; use crate::array::ArrayView; use crate::arrays::Bool; use crate::arrays::BoolArray; -use crate::scalar_fn::fns::mask::MaskReduce; +use crate::scalar_fn::builtin::mask::MaskReduce; use crate::validity::Validity; impl MaskReduce for Bool { @@ -52,7 +52,7 @@ mod tests { use crate::buffer::BufferHandle; use crate::buffer::DeviceBuffer; use crate::compute::conformance::mask::test_mask_conformance; - use crate::scalar_fn::fns::mask::MaskReduce; + use crate::scalar_fn::builtin::mask::MaskReduce; use crate::validity::Validity; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); diff --git a/vortex-array/src/arrays/bool/compute/mod.rs b/vortex-array/src/arrays/bool/compute/mod.rs index da5fa766d85..fc0f25eadfc 100644 --- a/vortex-array/src/arrays/bool/compute/mod.rs +++ b/vortex-array/src/arrays/bool/compute/mod.rs @@ -9,7 +9,6 @@ mod mask; pub mod rules; mod slice; mod take; -mod zip; #[cfg(test)] mod tests { diff --git a/vortex-array/src/arrays/bool/compute/rules.rs b/vortex-array/src/arrays/bool/compute/rules.rs index 43f811785ea..8bfdd6f9146 100644 --- a/vortex-array/src/arrays/bool/compute/rules.rs +++ b/vortex-array/src/arrays/bool/compute/rules.rs @@ -14,8 +14,8 @@ use crate::arrays::filter::FilterReduceAdaptor; use crate::arrays::slice::SliceReduceAdaptor; use crate::optimizer::rules::ArrayParentReduceRule; use crate::optimizer::rules::ParentRuleSet; -use crate::scalar_fn::fns::cast::CastReduceAdaptor; -use crate::scalar_fn::fns::mask::MaskReduceAdaptor; +use crate::scalar_fn::builtin::cast::CastReduceAdaptor; +use crate::scalar_fn::builtin::mask::MaskReduceAdaptor; pub(crate) const RULES: ParentRuleSet = ParentRuleSet::new(&[ ParentRuleSet::lift(&BoolMaskedValidityRule), diff --git a/vortex-array/src/arrays/bool/vtable/kernel.rs b/vortex-array/src/arrays/bool/vtable/kernel.rs index c19d00d5118..4bf52ca0435 100644 --- a/vortex-array/src/arrays/bool/vtable/kernel.rs +++ b/vortex-array/src/arrays/bool/vtable/kernel.rs @@ -9,14 +9,12 @@ use crate::arrays::Dict; use crate::arrays::dict::TakeExecuteAdaptor; use crate::optimizer::kernels::ArrayKernelsExt; use crate::scalar_fn::ScalarFnVTable; -use crate::scalar_fn::fns::binary::Binary; -use crate::scalar_fn::fns::binary::BooleanExecuteAdaptor; -use crate::scalar_fn::fns::cast::Cast; -use crate::scalar_fn::fns::cast::CastExecuteAdaptor; -use crate::scalar_fn::fns::fill_null::FillNull; -use crate::scalar_fn::fns::fill_null::FillNullExecuteAdaptor; -use crate::scalar_fn::fns::zip::Zip; -use crate::scalar_fn::fns::zip::ZipExecuteAdaptor; +use crate::scalar_fn::builtin::binary::Binary; +use crate::scalar_fn::builtin::binary::BooleanExecuteAdaptor; +use crate::scalar_fn::builtin::cast::Cast; +use crate::scalar_fn::builtin::cast::CastExecuteAdaptor; +use crate::scalar_fn::builtin::fill_null::FillNull; +use crate::scalar_fn::builtin::fill_null::FillNullExecuteAdaptor; pub(crate) fn initialize(session: &VortexSession) { let kernels = session.kernels(); @@ -24,5 +22,4 @@ pub(crate) fn initialize(session: &VortexSession) { kernels.register_execute_parent_kernel(Cast.id(), Bool, CastExecuteAdaptor(Bool)); kernels.register_execute_parent_kernel(FillNull.id(), Bool, FillNullExecuteAdaptor(Bool)); kernels.register_execute_parent_kernel(Dict.id(), Bool, TakeExecuteAdaptor(Bool)); - kernels.register_execute_parent_kernel(Zip.id(), Bool, ZipExecuteAdaptor(Bool)); } diff --git a/vortex-array/src/arrays/chunked/compute/cast.rs b/vortex-array/src/arrays/chunked/compute/cast.rs index 6f52c792f91..de718143fe1 100644 --- a/vortex-array/src/arrays/chunked/compute/cast.rs +++ b/vortex-array/src/arrays/chunked/compute/cast.rs @@ -11,7 +11,7 @@ use crate::arrays::ChunkedArray; use crate::arrays::chunked::ChunkedArrayExt; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; -use crate::scalar_fn::fns::cast::CastReduce; +use crate::scalar_fn::builtin::cast::CastReduce; impl CastReduce for Chunked { fn cast(array: ArrayView<'_, Chunked>, dtype: &DType) -> VortexResult> { diff --git a/vortex-array/src/arrays/chunked/compute/fill_null.rs b/vortex-array/src/arrays/chunked/compute/fill_null.rs index 7d05a2914e6..9d2680f2701 100644 --- a/vortex-array/src/arrays/chunked/compute/fill_null.rs +++ b/vortex-array/src/arrays/chunked/compute/fill_null.rs @@ -11,7 +11,7 @@ use crate::arrays::ChunkedArray; use crate::arrays::chunked::ChunkedArrayExt; use crate::builtins::ArrayBuiltins; use crate::scalar::Scalar; -use crate::scalar_fn::fns::fill_null::FillNullReduce; +use crate::scalar_fn::builtin::fill_null::FillNullReduce; impl FillNullReduce for Chunked { fn fill_null( diff --git a/vortex-array/src/arrays/chunked/compute/kernel.rs b/vortex-array/src/arrays/chunked/compute/kernel.rs index db0042105cd..dc074dc4775 100644 --- a/vortex-array/src/arrays/chunked/compute/kernel.rs +++ b/vortex-array/src/arrays/chunked/compute/kernel.rs @@ -13,10 +13,8 @@ use crate::arrays::filter::FilterExecuteAdaptor; use crate::arrays::slice::SliceExecuteAdaptor; use crate::optimizer::kernels::ArrayKernelsExt; use crate::scalar_fn::ScalarFnVTable; -use crate::scalar_fn::fns::mask::Mask; -use crate::scalar_fn::fns::mask::MaskExecuteAdaptor; -use crate::scalar_fn::fns::zip::Zip; -use crate::scalar_fn::fns::zip::ZipExecuteAdaptor; +use crate::scalar_fn::builtin::mask::Mask; +use crate::scalar_fn::builtin::mask::MaskExecuteAdaptor; pub(crate) fn initialize(session: &VortexSession) { let kernels = session.kernels(); @@ -24,5 +22,4 @@ pub(crate) fn initialize(session: &VortexSession) { kernels.register_execute_parent_kernel(Mask.id(), Chunked, MaskExecuteAdaptor(Chunked)); kernels.register_execute_parent_kernel(Slice.id(), Chunked, SliceExecuteAdaptor(Chunked)); kernels.register_execute_parent_kernel(Dict.id(), Chunked, TakeExecuteAdaptor(Chunked)); - kernels.register_execute_parent_kernel(Zip.id(), Chunked, ZipExecuteAdaptor(Chunked)); } diff --git a/vortex-array/src/arrays/chunked/compute/mask.rs b/vortex-array/src/arrays/chunked/compute/mask.rs index f99d4254780..215d27bd9c0 100644 --- a/vortex-array/src/arrays/chunked/compute/mask.rs +++ b/vortex-array/src/arrays/chunked/compute/mask.rs @@ -10,8 +10,8 @@ use crate::array::ArrayView; use crate::arrays::Chunked; use crate::arrays::ChunkedArray; use crate::arrays::chunked::ChunkedArrayExt; -use crate::scalar_fn::fns::mask::Mask as MaskExpr; -use crate::scalar_fn::fns::mask::MaskKernel; +use crate::scalar_fn::builtin::mask::Mask as MaskExpr; +use crate::scalar_fn::builtin::mask::MaskKernel; impl MaskKernel for Chunked { fn mask( diff --git a/vortex-array/src/arrays/chunked/compute/mod.rs b/vortex-array/src/arrays/chunked/compute/mod.rs index e7c27e00361..cf7c79425fe 100644 --- a/vortex-array/src/arrays/chunked/compute/mod.rs +++ b/vortex-array/src/arrays/chunked/compute/mod.rs @@ -6,12 +6,10 @@ mod cast; mod fill_null; mod filter; pub(crate) mod kernel; -mod list_contains; mod mask; pub(crate) mod rules; mod slice; mod take; -mod zip; #[cfg(test)] mod tests { diff --git a/vortex-array/src/arrays/chunked/compute/rules.rs b/vortex-array/src/arrays/chunked/compute/rules.rs index c5f904910f2..74f4d3a28cb 100644 --- a/vortex-array/src/arrays/chunked/compute/rules.rs +++ b/vortex-array/src/arrays/chunked/compute/rules.rs @@ -15,17 +15,14 @@ use crate::arrays::ScalarFnArray; use crate::arrays::chunked::ChunkedArrayExt; use crate::arrays::scalar_fn::AnyScalarFn; use crate::arrays::scalar_fn::ScalarFnArrayExt; -use crate::optimizer::ArrayOptimizer; use crate::optimizer::rules::ArrayParentReduceRule; use crate::optimizer::rules::ParentRuleSet; -use crate::scalar_fn::fns::cast::CastReduceAdaptor; -use crate::scalar_fn::fns::fill_null::FillNullReduceAdaptor; -use crate::scalar_fn::fns::list_contains::ListContainsElementReduceAdaptor; +use crate::scalar_fn::builtin::cast::CastReduceAdaptor; +use crate::scalar_fn::builtin::fill_null::FillNullReduceAdaptor; pub(crate) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ ParentRuleSet::lift(&CastReduceAdaptor(Chunked)), ParentRuleSet::lift(&ChunkedUnaryScalarFnPushDownRule), - ParentRuleSet::lift(&ListContainsElementReduceAdaptor(Chunked)), ParentRuleSet::lift(&ChunkedConstantScalarFnPushDownRule), ParentRuleSet::lift(&FillNullReduceAdaptor(Chunked)), ]); @@ -48,10 +45,11 @@ impl ArrayParentReduceRule for ChunkedUnaryScalarFnPushDownRule { let new_chunks: Vec<_> = array .iter_chunks() - .map(|chunk| { - ScalarFnArray::try_new(parent.scalar_fn().clone(), vec![chunk.clone()])? - .into_array() - .optimize() + .map(|chunk| -> VortexResult { + Ok( + ScalarFnArray::try_new(parent.scalar_fn().clone(), vec![chunk.clone()])? + .into_array(), + ) }) .try_collect()?; @@ -85,7 +83,7 @@ impl ArrayParentReduceRule for ChunkedConstantScalarFnPushDownRule { let mut offset = 0; let new_chunks: Vec<_> = array .iter_chunks() - .map(|chunk| { + .map(|chunk| -> VortexResult { let range = offset..offset + chunk.len(); offset += chunk.len(); // Each sibling is sliced to the chunk's rows, so that every chunk shares what the @@ -102,9 +100,7 @@ impl ArrayParentReduceRule for ChunkedConstantScalarFnPushDownRule { }) .try_collect()?; - ScalarFnArray::try_new(parent.scalar_fn().clone(), new_children)? - .into_array() - .optimize() + Ok(ScalarFnArray::try_new(parent.scalar_fn().clone(), new_children)?.into_array()) }) .try_collect()?; diff --git a/vortex-array/src/arrays/chunked/mod.rs b/vortex-array/src/arrays/chunked/mod.rs index 3bd361d66a1..3f0cd0b2c74 100644 --- a/vortex-array/src/arrays/chunked/mod.rs +++ b/vortex-array/src/arrays/chunked/mod.rs @@ -9,7 +9,6 @@ pub use array::ChunkedSlotsView; pub use vtable::ChunkedArray; pub(crate) mod compute; -pub(crate) mod paired_chunks; mod vtable; pub use vtable::Chunked; diff --git a/vortex-array/src/arrays/chunked/vtable/mod.rs b/vortex-array/src/arrays/chunked/vtable/mod.rs index cfd8b94cb90..af6709b3f2b 100644 --- a/vortex-array/src/arrays/chunked/vtable/mod.rs +++ b/vortex-array/src/arrays/chunked/vtable/mod.rs @@ -279,6 +279,8 @@ impl VTable for Chunked { fn reduce(array: ArrayView<'_, Self>) -> VortexResult> { Ok(match array.nchunks() { + // Variant has no empty canonical form; canonicalization handles empty variant chunks. + 0 if matches!(array.dtype(), DType::Variant(_)) => None, 0 => Some(Canonical::empty(array.dtype()).into_array()), 1 => Some(array.chunk(0).clone()), _ => None, diff --git a/vortex-array/src/arrays/constant/compute/cast.rs b/vortex-array/src/arrays/constant/compute/cast.rs index ea26f632bb8..60c7ecd8d3c 100644 --- a/vortex-array/src/arrays/constant/compute/cast.rs +++ b/vortex-array/src/arrays/constant/compute/cast.rs @@ -9,7 +9,7 @@ use crate::array::ArrayView; use crate::arrays::Constant; use crate::arrays::ConstantArray; use crate::dtype::DType; -use crate::scalar_fn::fns::cast::CastReduce; +use crate::scalar_fn::builtin::cast::CastReduce; impl CastReduce for Constant { fn cast(array: ArrayView<'_, Constant>, dtype: &DType) -> VortexResult> { diff --git a/vortex-array/src/arrays/constant/compute/fill_null.rs b/vortex-array/src/arrays/constant/compute/fill_null.rs index a559bbb8a70..ee34bec0dc2 100644 --- a/vortex-array/src/arrays/constant/compute/fill_null.rs +++ b/vortex-array/src/arrays/constant/compute/fill_null.rs @@ -7,8 +7,8 @@ use crate::ArrayRef; use crate::array::ArrayView; use crate::arrays::Constant; use crate::scalar::Scalar; -use crate::scalar_fn::fns::fill_null::FillNullReduce; -use crate::scalar_fn::fns::fill_null::fill_null_constant; +use crate::scalar_fn::builtin::fill_null::FillNullReduce; +use crate::scalar_fn::builtin::fill_null::fill_null_constant; impl FillNullReduce for Constant { fn fill_null( diff --git a/vortex-array/src/arrays/constant/compute/mod.rs b/vortex-array/src/arrays/constant/compute/mod.rs index 2a6f0bb0ce5..1303992bea0 100644 --- a/vortex-array/src/arrays/constant/compute/mod.rs +++ b/vortex-array/src/arrays/constant/compute/mod.rs @@ -1,7 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -mod between; mod cast; mod fill_null; mod filter; diff --git a/vortex-array/src/arrays/constant/compute/not.rs b/vortex-array/src/arrays/constant/compute/not.rs index 1eb912d0a91..3e068a56735 100644 --- a/vortex-array/src/arrays/constant/compute/not.rs +++ b/vortex-array/src/arrays/constant/compute/not.rs @@ -9,7 +9,7 @@ use crate::array::ArrayView; use crate::arrays::Constant; use crate::arrays::ConstantArray; use crate::scalar::Scalar; -use crate::scalar_fn::fns::not::NotReduce; +use crate::scalar_fn::builtin::not::NotReduce; impl NotReduce for Constant { fn invert(array: ArrayView<'_, Constant>) -> VortexResult> { diff --git a/vortex-array/src/arrays/constant/compute/rules.rs b/vortex-array/src/arrays/constant/compute/rules.rs index f97e6c596c7..b7424dd493b 100644 --- a/vortex-array/src/arrays/constant/compute/rules.rs +++ b/vortex-array/src/arrays/constant/compute/rules.rs @@ -14,13 +14,11 @@ use crate::arrays::filter::FilterReduceAdaptor; use crate::arrays::slice::SliceReduceAdaptor; use crate::optimizer::rules::ArrayParentReduceRule; use crate::optimizer::rules::ParentRuleSet; -use crate::scalar_fn::fns::between::BetweenReduceAdaptor; -use crate::scalar_fn::fns::cast::CastReduceAdaptor; -use crate::scalar_fn::fns::fill_null::FillNullReduceAdaptor; -use crate::scalar_fn::fns::not::NotReduceAdaptor; +use crate::scalar_fn::builtin::cast::CastReduceAdaptor; +use crate::scalar_fn::builtin::fill_null::FillNullReduceAdaptor; +use crate::scalar_fn::builtin::not::NotReduceAdaptor; pub(crate) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ - ParentRuleSet::lift(&BetweenReduceAdaptor(Constant)), ParentRuleSet::lift(&CastReduceAdaptor(Constant)), ParentRuleSet::lift(&ConstantFilterRule), ParentRuleSet::lift(&FillNullReduceAdaptor(Constant)), diff --git a/vortex-array/src/arrays/constant/vtable/mod.rs b/vortex-array/src/arrays/constant/vtable/mod.rs index 3b55071ffb8..244ffdcbc3e 100644 --- a/vortex-array/src/arrays/constant/vtable/mod.rs +++ b/vortex-array/src/arrays/constant/vtable/mod.rs @@ -8,6 +8,7 @@ use std::hash::Hasher; use itertools::Itertools; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use vortex_error::vortex_ensure_eq; use vortex_error::vortex_panic; use vortex_session::VortexSession; @@ -35,6 +36,7 @@ use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::builders::BoolBuilder; use crate::builders::DecimalBuilder; +use crate::builders::ExtensionBuilder; use crate::builders::FixedSizeListBuilder; use crate::builders::ListViewBuilder; use crate::builders::NullBuilder; @@ -264,10 +266,19 @@ impl VTable for Constant { // is a constant storage array, which stays constant-encoded in the builder. // Canonicalizing instead would materialize the storage: see the note in // `constant_canonicalize` about `ExtensionConstantRule`. + // Append straight into the extension builder: going through + // `ArrayRef::append_to_builder` would optimize the extension array back into + // this constant. let storage = ConstantArray::new(scalar.as_extension().to_storage_scalar(), n); - ExtensionArray::new(ext_dtype.clone(), storage.into_array()) - .into_array() - .append_to_builder(builder, ctx)? + let Some(builder) = builder.as_any_mut().downcast_mut::() else { + vortex_bail!( + "append_to_builder for an extension constant requires an ExtensionBuilder" + ); + }; + builder.append_extension_array( + &ExtensionArray::new(ext_dtype.clone(), storage.into_array()), + ctx, + )? } DType::FixedSizeList(..) => { append_constant_fixed_size_list_run(array, n, builder, ctx)? diff --git a/vortex-array/src/arrays/decimal/compute/cast.rs b/vortex-array/src/arrays/decimal/compute/cast.rs index 29b38ce9406..ac61556451a 100644 --- a/vortex-array/src/arrays/decimal/compute/cast.rs +++ b/vortex-array/src/arrays/decimal/compute/cast.rs @@ -38,8 +38,8 @@ use crate::scalar::DecimalToFloatCast; use crate::scalar::DecimalToIntegerCast; use crate::scalar::DecimalValue; use crate::scalar::Scalar; -use crate::scalar_fn::fns::cast::CastKernel; -use crate::scalar_fn::fns::cast::CastReduce; +use crate::scalar_fn::builtin::cast::CastKernel; +use crate::scalar_fn::builtin::cast::CastReduce; use crate::validity::Validity; impl CastReduce for Decimal { diff --git a/vortex-array/src/arrays/decimal/compute/fill_null.rs b/vortex-array/src/arrays/decimal/compute/fill_null.rs index 3755f6c6d1f..d817bcf8706 100644 --- a/vortex-array/src/arrays/decimal/compute/fill_null.rs +++ b/vortex-array/src/arrays/decimal/compute/fill_null.rs @@ -20,7 +20,7 @@ use crate::dtype::NativeDecimalType; use crate::match_each_decimal_value_type; use crate::scalar::DecimalValue; use crate::scalar::Scalar; -use crate::scalar_fn::fns::fill_null::FillNullKernel; +use crate::scalar_fn::builtin::fill_null::FillNullKernel; use crate::validity::Validity; impl FillNullKernel for Decimal { diff --git a/vortex-array/src/arrays/decimal/compute/mask.rs b/vortex-array/src/arrays/decimal/compute/mask.rs index 91490a288ad..dc51edb9e31 100644 --- a/vortex-array/src/arrays/decimal/compute/mask.rs +++ b/vortex-array/src/arrays/decimal/compute/mask.rs @@ -8,7 +8,7 @@ use crate::IntoArray; use crate::array::ArrayView; use crate::arrays::Decimal; use crate::arrays::DecimalArray; -use crate::scalar_fn::fns::mask::MaskReduce; +use crate::scalar_fn::builtin::mask::MaskReduce; use crate::validity::Validity; impl MaskReduce for Decimal { diff --git a/vortex-array/src/arrays/decimal/compute/mod.rs b/vortex-array/src/arrays/decimal/compute/mod.rs index 88a8713c2dd..62fe5f496cf 100644 --- a/vortex-array/src/arrays/decimal/compute/mod.rs +++ b/vortex-array/src/arrays/decimal/compute/mod.rs @@ -1,7 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -mod between; mod cast; mod fill_null; mod fixed_width; diff --git a/vortex-array/src/arrays/decimal/compute/rules.rs b/vortex-array/src/arrays/decimal/compute/rules.rs index 8f1d09ea92a..1b79c9e6e5c 100644 --- a/vortex-array/src/arrays/decimal/compute/rules.rs +++ b/vortex-array/src/arrays/decimal/compute/rules.rs @@ -16,8 +16,8 @@ use crate::arrays::slice::SliceReduceAdaptor; use crate::match_each_decimal_value_type; use crate::optimizer::rules::ArrayParentReduceRule; use crate::optimizer::rules::ParentRuleSet; -use crate::scalar_fn::fns::cast::CastReduceAdaptor; -use crate::scalar_fn::fns::mask::MaskReduceAdaptor; +use crate::scalar_fn::builtin::cast::CastReduceAdaptor; +use crate::scalar_fn::builtin::mask::MaskReduceAdaptor; pub(crate) static RULES: ParentRuleSet = ParentRuleSet::new(&[ ParentRuleSet::lift(&DecimalMaskedValidityRule), @@ -111,6 +111,7 @@ mod tests { use crate::dtype::DecimalDType; use crate::dtype::DecimalType; use crate::match_each_decimal_value_type; + use crate::optimizer::ArrayOptimizer; use crate::validity::Validity; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); @@ -183,7 +184,7 @@ mod tests { let sliced = array.slice(1..5)?; let range = 1..3; - let nested = sliced.slice(range.clone())?; + let nested = sliced.slice(range.clone())?.optimize(&SESSION)?; let decimal = nested.as_::(); let byte_width = values_type.byte_width(); let original_range = range.start + 1..range.end + 1; diff --git a/vortex-array/src/arrays/decimal/vtable/kernel.rs b/vortex-array/src/arrays/decimal/vtable/kernel.rs index 19bb30d8c95..0f2ea03d7be 100644 --- a/vortex-array/src/arrays/decimal/vtable/kernel.rs +++ b/vortex-array/src/arrays/decimal/vtable/kernel.rs @@ -9,16 +9,13 @@ use crate::arrays::Dict; use crate::arrays::dict::TakeExecuteAdaptor; use crate::optimizer::kernels::ArrayKernelsExt; use crate::scalar_fn::ScalarFnVTable; -use crate::scalar_fn::fns::between::Between; -use crate::scalar_fn::fns::between::BetweenExecuteAdaptor; -use crate::scalar_fn::fns::cast::Cast; -use crate::scalar_fn::fns::cast::CastExecuteAdaptor; -use crate::scalar_fn::fns::fill_null::FillNull; -use crate::scalar_fn::fns::fill_null::FillNullExecuteAdaptor; +use crate::scalar_fn::builtin::cast::Cast; +use crate::scalar_fn::builtin::cast::CastExecuteAdaptor; +use crate::scalar_fn::builtin::fill_null::FillNull; +use crate::scalar_fn::builtin::fill_null::FillNullExecuteAdaptor; pub(crate) fn initialize(session: &VortexSession) { let kernels = session.kernels(); - kernels.register_execute_parent_kernel(Between.id(), Decimal, BetweenExecuteAdaptor(Decimal)); kernels.register_execute_parent_kernel(Cast.id(), Decimal, CastExecuteAdaptor(Decimal)); kernels.register_execute_parent_kernel(FillNull.id(), Decimal, FillNullExecuteAdaptor(Decimal)); kernels.register_execute_parent_kernel(Dict.id(), Decimal, TakeExecuteAdaptor(Decimal)); diff --git a/vortex-array/src/arrays/decimal/vtable/operations.rs b/vortex-array/src/arrays/decimal/vtable/operations.rs index 09d1d31a29c..d311c4edb2b 100644 --- a/vortex-array/src/arrays/decimal/vtable/operations.rs +++ b/vortex-array/src/arrays/decimal/vtable/operations.rs @@ -43,6 +43,7 @@ mod tests { use crate::assert_arrays_eq; use crate::dtype::DecimalDType; use crate::dtype::Nullability; + use crate::optimizer::ArrayOptimizer; use crate::scalar::DecimalValue; use crate::scalar::Scalar; use crate::validity::Validity; @@ -58,7 +59,7 @@ mod tests { ) .into_array(); - let sliced = array.slice(1..3).unwrap(); + let sliced = array.slice(1..3).unwrap().optimize(&SESSION).unwrap(); assert_eq!(sliced.len(), 2); let decimal = sliced.as_::(); diff --git a/vortex-array/src/arrays/dict/compute/cast.rs b/vortex-array/src/arrays/dict/compute/cast.rs index 2c705464d53..4bea1fb67f9 100644 --- a/vortex-array/src/arrays/dict/compute/cast.rs +++ b/vortex-array/src/arrays/dict/compute/cast.rs @@ -12,7 +12,7 @@ use crate::arrays::dict::DictArrayExt; use crate::arrays::dict::DictArraySlotsExt; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; -use crate::scalar_fn::fns::cast::CastReduce; +use crate::scalar_fn::builtin::cast::CastReduce; impl CastReduce for Dict { fn cast(array: ArrayView<'_, Dict>, dtype: &DType) -> VortexResult> { @@ -67,6 +67,7 @@ mod tests { use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; + use crate::optimizer::ArrayOptimizer; use crate::validity::Validity; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); @@ -128,6 +129,8 @@ mod tests { let non_nullable = dict .into_array() .cast(DType::Primitive(PType::I32, Nullability::NonNullable)) + .unwrap() + .optimize(&SESSION) .unwrap(); assert_eq!( non_nullable.dtype(), @@ -148,6 +151,8 @@ mod tests { // Cast to Nullable let nullable = non_nullable .cast(DType::Primitive(PType::I32, Nullability::Nullable)) + .unwrap() + .optimize(&SESSION) .unwrap(); assert_eq!( nullable.dtype(), @@ -168,6 +173,8 @@ mod tests { // Cast back to NonNullable let back_to_non_nullable = nullable .cast(DType::Primitive(PType::I32, Nullability::NonNullable)) + .unwrap() + .optimize(&SESSION) .unwrap(); assert_eq!( back_to_non_nullable.dtype(), diff --git a/vortex-array/src/arrays/dict/compute/compare.rs b/vortex-array/src/arrays/dict/compute/compare.rs index d59be4bf963..7fdd149d7ea 100644 --- a/vortex-array/src/arrays/dict/compute/compare.rs +++ b/vortex-array/src/arrays/dict/compute/compare.rs @@ -14,9 +14,9 @@ use crate::arrays::ConstantArray; use crate::arrays::dict::DictArrayExt; use crate::arrays::dict::DictArraySlotsExt; use crate::builtins::ArrayBuiltins; -use crate::scalar_fn::fns::binary::CompareKernel; -use crate::scalar_fn::fns::operators::CompareOperator; -use crate::scalar_fn::fns::operators::Operator; +use crate::scalar_fn::builtin::binary::CompareKernel; +use crate::scalar_fn::builtin::operators::CompareOperator; +use crate::scalar_fn::builtin::operators::Operator; impl CompareKernel for Dict { fn compare( diff --git a/vortex-array/src/arrays/dict/compute/fill_null.rs b/vortex-array/src/arrays/dict/compute/fill_null.rs index 077ee865bda..3e65be51599 100644 --- a/vortex-array/src/arrays/dict/compute/fill_null.rs +++ b/vortex-array/src/arrays/dict/compute/fill_null.rs @@ -19,8 +19,8 @@ use crate::builtins::ArrayBuiltins; use crate::match_each_integer_ptype; use crate::scalar::Scalar; use crate::scalar::ScalarValue; -use crate::scalar_fn::fns::fill_null::FillNullKernel; -use crate::scalar_fn::fns::operators::Operator; +use crate::scalar_fn::builtin::fill_null::FillNullKernel; +use crate::scalar_fn::builtin::operators::Operator; impl FillNullKernel for Dict { fn fill_null( diff --git a/vortex-array/src/arrays/dict/compute/filter.rs b/vortex-array/src/arrays/dict/compute/filter.rs index 46387684a5c..952841bbe8f 100644 --- a/vortex-array/src/arrays/dict/compute/filter.rs +++ b/vortex-array/src/arrays/dict/compute/filter.rs @@ -56,6 +56,7 @@ mod tests { use crate::expr::stats::Precision; use crate::expr::stats::Stat; use crate::expr::stats::StatsProviderExt; + use crate::optimizer::ArrayOptimizer; use crate::scalar::ScalarValue; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); @@ -67,7 +68,9 @@ mod tests { ConstantArray::new(1u8, 5).into_array(), buffer![10i32, 20, 30].into_array(), ); - let filtered = dict.filter(Mask::from_iter([true, false, true, true, false]))?; + let filtered = dict + .filter(Mask::from_iter([true, false, true, true, false]))? + .optimize(&SESSION)?; let filtered_dict = filtered.as_::(); assert_arrays_eq!( @@ -94,7 +97,9 @@ mod tests { .set(Stat::IsConstant, Precision::exact(ScalarValue::Bool(false))); let dict = DictArray::new(codes, buffer![10i32, 20, 30].into_array()); - let filtered = dict.filter(Mask::from_iter([true, false, true, true]))?; + let filtered = dict + .filter(Mask::from_iter([true, false, true, true]))? + .optimize(&SESSION)?; let filtered_codes = filtered.as_::().codes().clone(); assert_eq!( diff --git a/vortex-array/src/arrays/dict/compute/mask.rs b/vortex-array/src/arrays/dict/compute/mask.rs index b129fdd19c5..f56e1ae2ddc 100644 --- a/vortex-array/src/arrays/dict/compute/mask.rs +++ b/vortex-array/src/arrays/dict/compute/mask.rs @@ -9,8 +9,8 @@ use crate::array::ArrayView; use crate::arrays::Dict; use crate::arrays::DictArray; use crate::arrays::dict::DictArraySlotsExt; -use crate::scalar_fn::fns::mask::Mask as MaskExpr; -use crate::scalar_fn::fns::mask::MaskReduce; +use crate::scalar_fn::builtin::mask::Mask as MaskExpr; +use crate::scalar_fn::builtin::mask::MaskReduce; impl MaskReduce for Dict { fn mask(array: ArrayView<'_, Dict>, mask: &ArrayRef) -> VortexResult> { diff --git a/vortex-array/src/arrays/dict/compute/mod.rs b/vortex-array/src/arrays/dict/compute/mod.rs index 25eb0c5e442..163da7f63de 100644 --- a/vortex-array/src/arrays/dict/compute/mod.rs +++ b/vortex-array/src/arrays/dict/compute/mod.rs @@ -7,7 +7,6 @@ mod fill_null; mod filter; pub(crate) mod is_constant; pub(crate) mod is_sorted; -mod like; mod mask; pub(crate) mod min_max; pub(crate) mod rules; @@ -65,7 +64,7 @@ mod test { use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType::I32; - use crate::scalar_fn::fns::operators::Operator; + use crate::scalar_fn::builtin::operators::Operator; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); diff --git a/vortex-array/src/arrays/dict/compute/rules.rs b/vortex-array/src/arrays/dict/compute/rules.rs index 273d2b8562d..7b2bd7269df 100644 --- a/vortex-array/src/arrays/dict/compute/rules.rs +++ b/vortex-array/src/arrays/dict/compute/rules.rs @@ -25,21 +25,18 @@ use crate::arrays::scalar_fn::AnyScalarFn; use crate::arrays::scalar_fn::ScalarFnArrayExt; use crate::arrays::slice::SliceReduceAdaptor; use crate::builtins::ArrayBuiltins; -use crate::optimizer::ArrayOptimizer; use crate::optimizer::rules::ArrayParentReduceRule; use crate::optimizer::rules::ParentRuleSet; -use crate::scalar_fn::fns::cast::Cast; -use crate::scalar_fn::fns::cast::CastReduceAdaptor; -use crate::scalar_fn::fns::like::LikeReduceAdaptor; -use crate::scalar_fn::fns::mask::MaskReduceAdaptor; -use crate::scalar_fn::fns::pack::Pack; +use crate::scalar_fn::builtin::cast::Cast; +use crate::scalar_fn::builtin::cast::CastReduceAdaptor; +use crate::scalar_fn::builtin::mask::MaskReduceAdaptor; +use crate::scalar_fn::builtin::pack::Pack; use crate::validity::Validity; pub(crate) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ ParentRuleSet::lift(&FilterReduceAdaptor(Dict)), ParentRuleSet::lift(&CastReduceAdaptor(Dict)), ParentRuleSet::lift(&MaskReduceAdaptor(Dict)), - ParentRuleSet::lift(&LikeReduceAdaptor(Dict)), ParentRuleSet::lift(&DictionaryChunkedValuesPullUpRule), ParentRuleSet::lift(&DictionaryScalarFnValuesPushDownRule), ParentRuleSet::lift(&DictionaryScalarFnCodesPullUpRule), @@ -174,9 +171,8 @@ impl ArrayParentReduceRule for DictionaryScalarFnValuesPushDownRule { } } - let transformed_values = ScalarFnArray::try_new(scalar_fn.clone(), value_children)? - .into_array() - .optimize()?; + let transformed_values = + ScalarFnArray::try_new(scalar_fn.clone(), value_children)?.into_array(); // A non-strict function reaches this point only when the codes are all valid, but their // dtype may still be nullable. Remove that declared nullability while rebuilding the @@ -252,9 +248,8 @@ impl ArrayParentReduceRule for DictionaryScalarFnCodesPullUpRule { } } - let new_values = ScalarFnArray::try_new(parent.scalar_fn().clone(), new_children)? - .into_array() - .optimize()?; + let new_values = + ScalarFnArray::try_new(parent.scalar_fn().clone(), new_children)?.into_array(); let new_dict = unsafe { DictArray::new_unchecked(array.codes().clone(), new_values) }.into_array(); @@ -288,9 +283,9 @@ mod tests { use crate::executor::VortexSessionExecute; use crate::optimizer::ArrayOptimizer; use crate::scalar::Scalar; - use crate::scalar_fn::fns::binary::Binary; - use crate::scalar_fn::fns::not::Not; - use crate::scalar_fn::fns::operators::Operator; + use crate::scalar_fn::builtin::binary::Binary; + use crate::scalar_fn::builtin::not::Not; + use crate::scalar_fn::builtin::operators::Operator; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); @@ -304,7 +299,7 @@ mod tests { let array = ChunkedArray::try_new(vec![chunk0, chunk1], values.dtype().clone())?.into_array(); - let optimized = array.optimize()?; + let optimized = array.optimize(&SESSION)?; let dict = optimized.as_::(); let codes = dict.codes().as_::(); @@ -331,7 +326,7 @@ mod tests { let array = ChunkedArray::try_new(vec![chunk0, chunk1], values0.dtype().clone())?.into_array(); - let optimized = array.optimize()?; + let optimized = array.optimize(&SESSION)?; assert!(optimized.is::()); let mut ctx = SESSION.create_execution_ctx(); @@ -355,7 +350,7 @@ mod tests { } .into_array(); - let result = Not::try_new(dict)?.into_array().optimize()?; + let result = Not::try_new(dict)?.into_array().optimize(&SESSION)?; let result = result.as_::(); assert!(result.has_all_values_referenced()); @@ -380,7 +375,7 @@ mod tests { let mut ctx = SESSION.create_execution_ctx(); // Kleene AND: null AND false == false let expected = BoolArray::from_iter([false, false, false]).into_array(); - let optimized = expr.optimize()?; + let optimized = expr.optimize(&SESSION)?; assert_arrays_eq!(optimized, expected, &mut ctx); Ok(()) } @@ -397,7 +392,7 @@ mod tests { let mut ctx = SESSION.create_execution_ctx(); // Kleene OR: null OR true == true let expected = BoolArray::from_iter([true, true, true]).into_array(); - let optimized = expr.optimize()?; + let optimized = expr.optimize(&SESSION)?; assert_arrays_eq!(optimized, expected, &mut ctx); Ok(()) } diff --git a/vortex-array/src/arrays/dict/vtable/kernel.rs b/vortex-array/src/arrays/dict/vtable/kernel.rs index 71786f63f66..77838761180 100644 --- a/vortex-array/src/arrays/dict/vtable/kernel.rs +++ b/vortex-array/src/arrays/dict/vtable/kernel.rs @@ -11,10 +11,10 @@ use crate::arrays::dict::TakeExecuteAdaptor; use crate::arrays::dict::take::DictConstantCodesKernel; use crate::optimizer::kernels::ArrayKernelsExt; use crate::scalar_fn::ScalarFnVTable; -use crate::scalar_fn::fns::binary::Binary; -use crate::scalar_fn::fns::binary::CompareExecuteAdaptor; -use crate::scalar_fn::fns::fill_null::FillNull; -use crate::scalar_fn::fns::fill_null::FillNullExecuteAdaptor; +use crate::scalar_fn::builtin::binary::Binary; +use crate::scalar_fn::builtin::binary::CompareExecuteAdaptor; +use crate::scalar_fn::builtin::fill_null::FillNull; +use crate::scalar_fn::builtin::fill_null::FillNullExecuteAdaptor; pub(crate) fn initialize(session: &VortexSession) { let kernels = session.kernels(); diff --git a/vortex-array/src/arrays/extension/compute/cast.rs b/vortex-array/src/arrays/extension/compute/cast.rs index 5328e897225..2fc73d875e9 100644 --- a/vortex-array/src/arrays/extension/compute/cast.rs +++ b/vortex-array/src/arrays/extension/compute/cast.rs @@ -11,7 +11,7 @@ use crate::arrays::ExtensionArray; use crate::arrays::extension::ExtensionArrayExt; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; -use crate::scalar_fn::fns::cast::CastReduce; +use crate::scalar_fn::builtin::cast::CastReduce; impl CastReduce for Extension { fn cast(array: ArrayView<'_, Extension>, dtype: &DType) -> VortexResult> { diff --git a/vortex-array/src/arrays/extension/compute/compare.rs b/vortex-array/src/arrays/extension/compute/compare.rs index b3f978ef503..48d4b7d6c8c 100644 --- a/vortex-array/src/arrays/extension/compute/compare.rs +++ b/vortex-array/src/arrays/extension/compute/compare.rs @@ -11,9 +11,9 @@ use crate::arrays::ConstantArray; use crate::arrays::Extension; use crate::arrays::extension::ExtensionArrayExt; use crate::builtins::ArrayBuiltins; -use crate::scalar_fn::fns::binary::CompareKernel; -use crate::scalar_fn::fns::operators::CompareOperator; -use crate::scalar_fn::fns::operators::Operator; +use crate::scalar_fn::builtin::binary::CompareKernel; +use crate::scalar_fn::builtin::operators::CompareOperator; +use crate::scalar_fn::builtin::operators::Operator; impl CompareKernel for Extension { fn compare( diff --git a/vortex-array/src/arrays/extension/compute/mask.rs b/vortex-array/src/arrays/extension/compute/mask.rs index 10ae5e904bc..16e0db13a90 100644 --- a/vortex-array/src/arrays/extension/compute/mask.rs +++ b/vortex-array/src/arrays/extension/compute/mask.rs @@ -9,8 +9,8 @@ use crate::array::ArrayView; use crate::arrays::Extension; use crate::arrays::ExtensionArray; use crate::arrays::extension::ExtensionArrayExt; -use crate::scalar_fn::fns::mask::Mask as MaskExpr; -use crate::scalar_fn::fns::mask::MaskReduce; +use crate::scalar_fn::builtin::mask::Mask as MaskExpr; +use crate::scalar_fn::builtin::mask::MaskReduce; impl MaskReduce for Extension { fn mask(array: ArrayView<'_, Extension>, mask: &ArrayRef) -> VortexResult> { diff --git a/vortex-array/src/arrays/extension/compute/rules.rs b/vortex-array/src/arrays/extension/compute/rules.rs index fd92940ac6d..d4819d8f6d5 100644 --- a/vortex-array/src/arrays/extension/compute/rules.rs +++ b/vortex-array/src/arrays/extension/compute/rules.rs @@ -20,8 +20,8 @@ use crate::optimizer::rules::ArrayReduceRule; use crate::optimizer::rules::ParentRuleSet; use crate::optimizer::rules::ReduceRuleSet; use crate::scalar::Scalar; -use crate::scalar_fn::fns::cast::CastReduceAdaptor; -use crate::scalar_fn::fns::mask::MaskReduceAdaptor; +use crate::scalar_fn::builtin::cast::CastReduceAdaptor; +use crate::scalar_fn::builtin::mask::MaskReduceAdaptor; pub(crate) const RULES: ReduceRuleSet = ReduceRuleSet::new(&[&ExtensionConstantRule]); @@ -110,8 +110,8 @@ mod tests { use crate::optimizer::ArrayOptimizer; use crate::scalar::Scalar; use crate::scalar::ScalarValue; - use crate::scalar_fn::fns::binary::Binary; - use crate::scalar_fn::fns::operators::Operator; + use crate::scalar_fn::builtin::binary::Binary; + use crate::scalar_fn::builtin::operators::Operator; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); @@ -166,7 +166,7 @@ mod tests { let filter_array = FilterArray::new(ext_array, mask).into_array(); // Optimize should push the filter into the storage - let optimized = filter_array.optimize().unwrap(); + let optimized = filter_array.optimize(&SESSION).unwrap(); // The result should be an ExtensionArray, not a FilterArray assert!( @@ -204,7 +204,7 @@ mod tests { let mask = Mask::from_iter([true, true, false, false, true]); let filter_array = FilterArray::new(ext_array, mask).into_array(); - let optimized = filter_array.optimize().unwrap(); + let optimized = filter_array.optimize(&SESSION).unwrap(); assert!(optimized.as_opt::().is_some()); let ext_result = optimized.as_::(); @@ -297,7 +297,7 @@ mod tests { .unwrap() .into_array(); - let optimized = scalar_fn_array.optimize().unwrap(); + let optimized = scalar_fn_array.optimize(&SESSION).unwrap(); // The first child should still be an ExtensionArray (no pushdown happened) let scalar_fn = optimized.as_opt::().unwrap(); @@ -322,7 +322,7 @@ mod tests { .unwrap() .into_array(); - let optimized = scalar_fn_array.optimize().unwrap(); + let optimized = scalar_fn_array.optimize(&SESSION).unwrap(); // No pushdown should happen because sibling is not a constant let scalar_fn = optimized.as_opt::().unwrap(); @@ -345,7 +345,7 @@ mod tests { .unwrap() .into_array(); - let optimized = scalar_fn_array.optimize().unwrap(); + let optimized = scalar_fn_array.optimize(&SESSION).unwrap(); // No pushdown should happen because constant is not an extension scalar let scalar_fn = optimized.as_opt::().unwrap(); diff --git a/vortex-array/src/arrays/extension/vtable/kernel.rs b/vortex-array/src/arrays/extension/vtable/kernel.rs index c90458f1d5f..c88f321025a 100644 --- a/vortex-array/src/arrays/extension/vtable/kernel.rs +++ b/vortex-array/src/arrays/extension/vtable/kernel.rs @@ -9,8 +9,8 @@ use crate::arrays::Extension; use crate::arrays::dict::TakeExecuteAdaptor; use crate::optimizer::kernels::ArrayKernelsExt; use crate::scalar_fn::ScalarFnVTable; -use crate::scalar_fn::fns::binary::Binary; -use crate::scalar_fn::fns::binary::CompareExecuteAdaptor; +use crate::scalar_fn::builtin::binary::Binary; +use crate::scalar_fn::builtin::binary::CompareExecuteAdaptor; pub(crate) fn initialize(session: &VortexSession) { let kernels = session.kernels(); diff --git a/vortex-array/src/arrays/filter/execute/mod.rs b/vortex-array/src/arrays/filter/execute/mod.rs index 2ed36ab128e..a633a87af67 100644 --- a/vortex-array/src/arrays/filter/execute/mod.rs +++ b/vortex-array/src/arrays/filter/execute/mod.rs @@ -197,6 +197,7 @@ mod tests { use crate::VortexSessionExecute; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; + use crate::optimizer::ArrayOptimizer; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); @@ -211,7 +212,7 @@ mod tests { values.cached_indices().is_none() && values.cached_slices().is_none() })); - let filtered = array.into_array().filter(mask)?; + let filtered = array.into_array().filter(mask)?.optimize(&SESSION)?; assert!(filtered.is::()); let filtered = filtered.execute::(&mut SESSION.create_execution_ctx())?; let filtered_values = filtered.to_buffer::(); diff --git a/vortex-array/src/arrays/filter/rules.rs b/vortex-array/src/arrays/filter/rules.rs index 5535d2cb52e..6fea1d672d3 100644 --- a/vortex-array/src/arrays/filter/rules.rs +++ b/vortex-array/src/arrays/filter/rules.rs @@ -24,7 +24,7 @@ use crate::optimizer::rules::ArrayParentReduceRule; use crate::optimizer::rules::ArrayReduceRule; use crate::optimizer::rules::ParentRuleSet; use crate::optimizer::rules::ReduceRuleSet; -use crate::scalar_fn::fns::get_item::GetItem; +use crate::scalar_fn::builtin::get_item::GetItem; pub(super) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ ParentRuleSet::lift(&FilterReduceAdaptor(Filter)), diff --git a/vortex-array/src/arrays/fixed_size_list/compute/cast.rs b/vortex-array/src/arrays/fixed_size_list/compute/cast.rs index 8e813209a5b..d11e75263a6 100644 --- a/vortex-array/src/arrays/fixed_size_list/compute/cast.rs +++ b/vortex-array/src/arrays/fixed_size_list/compute/cast.rs @@ -13,8 +13,8 @@ use crate::arrays::fixed_size_list::FixedSizeListArrayExt; use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; -use crate::scalar_fn::fns::cast::CastKernel; -use crate::scalar_fn::fns::cast::CastReduce; +use crate::scalar_fn::builtin::cast::CastKernel; +use crate::scalar_fn::builtin::cast::CastReduce; use crate::validity::Validity; fn build_with_validity( diff --git a/vortex-array/src/arrays/fixed_size_list/compute/mask.rs b/vortex-array/src/arrays/fixed_size_list/compute/mask.rs index 983c848c85a..9fc948dc3d4 100644 --- a/vortex-array/src/arrays/fixed_size_list/compute/mask.rs +++ b/vortex-array/src/arrays/fixed_size_list/compute/mask.rs @@ -10,7 +10,7 @@ use crate::arrays::FixedSizeList; use crate::arrays::FixedSizeListArray; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; -use crate::scalar_fn::fns::mask::MaskReduce; +use crate::scalar_fn::builtin::mask::MaskReduce; use crate::validity::Validity; impl MaskReduce for FixedSizeList { diff --git a/vortex-array/src/arrays/fixed_size_list/compute/rules.rs b/vortex-array/src/arrays/fixed_size_list/compute/rules.rs index da1d91423e2..69173dde685 100644 --- a/vortex-array/src/arrays/fixed_size_list/compute/rules.rs +++ b/vortex-array/src/arrays/fixed_size_list/compute/rules.rs @@ -4,8 +4,8 @@ use crate::arrays::FixedSizeList; use crate::arrays::slice::SliceReduceAdaptor; use crate::optimizer::rules::ParentRuleSet; -use crate::scalar_fn::fns::cast::CastReduceAdaptor; -use crate::scalar_fn::fns::mask::MaskReduceAdaptor; +use crate::scalar_fn::builtin::cast::CastReduceAdaptor; +use crate::scalar_fn::builtin::mask::MaskReduceAdaptor; pub(crate) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ ParentRuleSet::lift(&CastReduceAdaptor(FixedSizeList)), diff --git a/vortex-array/src/arrays/fixed_size_list/compute/take.rs b/vortex-array/src/arrays/fixed_size_list/compute/take.rs index acc05cca73d..d1b85c1c1fa 100644 --- a/vortex-array/src/arrays/fixed_size_list/compute/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/compute/take.rs @@ -91,7 +91,7 @@ fn take_empty_fsl( FixedSizeListArray::new_unchecked(new_elements, array.list_size(), new_validity, new_len) } .into_array() - .optimize_ctx(ctx.session()) + .optimize(ctx.session()) } fn take_non_empty_fsl( @@ -220,7 +220,7 @@ fn take_piecewise_fsl( FixedSizeListArray::new_unchecked(new_elements, array.list_size(), new_validity, new_len) } .into_array() - .optimize_ctx(ctx.session()) + .optimize(ctx.session()) .map(Some) } @@ -254,7 +254,7 @@ fn take_non_empty_degenerate_fsl( ) } .into_array() - .optimize_ctx(ctx.session()) + .optimize(ctx.session()) } fn take_non_empty_non_degenerate_fsl( @@ -276,7 +276,7 @@ fn take_non_empty_non_degenerate_fsl( FixedSizeListArray::new_unchecked(new_elements, array.list_size(), new_validity, new_len) } .into_array() - .optimize_ctx(ctx.session()) + .optimize(ctx.session()) } fn take_non_empty_non_degenerate_elements( diff --git a/vortex-array/src/arrays/fixed_size_list/vtable/kernel.rs b/vortex-array/src/arrays/fixed_size_list/vtable/kernel.rs index 47dc3ec0651..2095334a940 100644 --- a/vortex-array/src/arrays/fixed_size_list/vtable/kernel.rs +++ b/vortex-array/src/arrays/fixed_size_list/vtable/kernel.rs @@ -9,8 +9,8 @@ use crate::arrays::FixedSizeList; use crate::arrays::dict::TakeExecuteAdaptor; use crate::optimizer::kernels::ArrayKernelsExt; use crate::scalar_fn::ScalarFnVTable; -use crate::scalar_fn::fns::cast::Cast; -use crate::scalar_fn::fns::cast::CastExecuteAdaptor; +use crate::scalar_fn::builtin::cast::Cast; +use crate::scalar_fn::builtin::cast::CastExecuteAdaptor; pub(crate) fn initialize(session: &VortexSession) { let kernels = session.kernels(); diff --git a/vortex-array/src/arrays/list/array.rs b/vortex-array/src/arrays/list/array.rs index 377499faea6..e5c9f3e47b5 100644 --- a/vortex-array/src/arrays/list/array.rs +++ b/vortex-array/src/arrays/list/array.rs @@ -35,7 +35,8 @@ use crate::dtype::DType; use crate::dtype::NativePType; use crate::match_each_integer_ptype; use crate::match_each_native_ptype; -use crate::scalar_fn::fns::operators::Operator; +use crate::optimizer::ArrayOptimizer; +use crate::scalar_fn::builtin::operators::Operator; use crate::validity::Validity; #[array_slots(List)] @@ -321,13 +322,13 @@ pub trait ListArrayExt: ListArraySlotsExt { fn list_elements_at(&self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult { let start = self.offset_at(index, ctx)?; let end = self.offset_at(index + 1, ctx)?; - self.elements().slice(start..end) + self.elements().slice(start..end)?.optimize(ctx.session()) } fn sliced_elements(&self, ctx: &mut ExecutionCtx) -> VortexResult { let start = self.offset_at(0, ctx)?; let end = self.offset_at(self.len(), ctx)?; - self.elements().slice(start..end) + self.elements().slice(start..end)?.optimize(ctx.session()) } fn element_dtype(&self) -> &DType { diff --git a/vortex-array/src/arrays/list/compute/cast.rs b/vortex-array/src/arrays/list/compute/cast.rs index 8ecd998dbd4..2ec444dcafb 100644 --- a/vortex-array/src/arrays/list/compute/cast.rs +++ b/vortex-array/src/arrays/list/compute/cast.rs @@ -12,8 +12,8 @@ use crate::arrays::ListArray; use crate::arrays::list::ListArraySlotsExt; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; -use crate::scalar_fn::fns::cast::CastKernel; -use crate::scalar_fn::fns::cast::CastReduce; +use crate::scalar_fn::builtin::cast::CastKernel; +use crate::scalar_fn::builtin::cast::CastReduce; impl CastReduce for List { fn cast(array: ArrayView<'_, List>, dtype: &DType) -> VortexResult> { diff --git a/vortex-array/src/arrays/list/compute/kernels.rs b/vortex-array/src/arrays/list/compute/kernels.rs index 6af858dbf49..426657a7099 100644 --- a/vortex-array/src/arrays/list/compute/kernels.rs +++ b/vortex-array/src/arrays/list/compute/kernels.rs @@ -11,8 +11,8 @@ use crate::arrays::dict::TakeExecuteAdaptor; use crate::arrays::filter::FilterExecuteAdaptor; use crate::optimizer::kernels::ArrayKernelsExt; use crate::scalar_fn::ScalarFnVTable; -use crate::scalar_fn::fns::cast::Cast; -use crate::scalar_fn::fns::cast::CastExecuteAdaptor; +use crate::scalar_fn::builtin::cast::Cast; +use crate::scalar_fn::builtin::cast::CastExecuteAdaptor; pub(crate) fn initialize(session: &VortexSession) { let kernels = session.kernels(); diff --git a/vortex-array/src/arrays/list/compute/mask.rs b/vortex-array/src/arrays/list/compute/mask.rs index 40737eb6583..0ad7bddf2f2 100644 --- a/vortex-array/src/arrays/list/compute/mask.rs +++ b/vortex-array/src/arrays/list/compute/mask.rs @@ -9,7 +9,7 @@ use crate::array::ArrayView; use crate::arrays::List; use crate::arrays::ListArray; use crate::arrays::list::ListArraySlotsExt; -use crate::scalar_fn::fns::mask::MaskReduce; +use crate::scalar_fn::builtin::mask::MaskReduce; use crate::validity::Validity; impl MaskReduce for List { diff --git a/vortex-array/src/arrays/list/compute/rules.rs b/vortex-array/src/arrays/list/compute/rules.rs index d2aca066f7f..afc2aaf1c80 100644 --- a/vortex-array/src/arrays/list/compute/rules.rs +++ b/vortex-array/src/arrays/list/compute/rules.rs @@ -4,8 +4,8 @@ use crate::arrays::List; use crate::arrays::slice::SliceReduceAdaptor; use crate::optimizer::rules::ParentRuleSet; -use crate::scalar_fn::fns::cast::CastReduceAdaptor; -use crate::scalar_fn::fns::mask::MaskReduceAdaptor; +use crate::scalar_fn::builtin::cast::CastReduceAdaptor; +use crate::scalar_fn::builtin::mask::MaskReduceAdaptor; pub(crate) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ ParentRuleSet::lift(&CastReduceAdaptor(List)), diff --git a/vortex-array/src/arrays/list/tests.rs b/vortex-array/src/arrays/list/tests.rs index 08792fc251a..26a7b863026 100644 --- a/vortex-array/src/arrays/list/tests.rs +++ b/vortex-array/src/arrays/list/tests.rs @@ -23,6 +23,7 @@ use crate::builders::ListBuilder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType::I32; +use crate::optimizer::ArrayOptimizer; use crate::scalar::Scalar; use crate::validity::Validity; @@ -480,7 +481,12 @@ fn test_offset_to_0() { .as_list(), ) .vortex_expect("operation should succeed in test"); - let list = builder.finish().slice(2..4).unwrap(); + let list = builder + .finish() + .slice(2..4) + .unwrap() + .optimize(&SESSION) + .unwrap(); // The sliced list should be a ListArray since we built it with ListBuilder // and slice doesn't change the encoding @@ -704,7 +710,11 @@ fn test_list_of_lists() { assert_eq!(list_scalar.len(), 2); // Test slicing. - let sliced = list_of_lists.slice(1..3).unwrap(); + let sliced = list_of_lists + .slice(1..3) + .unwrap() + .optimize(&SESSION) + .unwrap(); let sliced_list = sliced.as_::(); assert_eq!(sliced_list.len(), 2); @@ -990,7 +1000,7 @@ fn test_recursive_compact_list_of_lists() { let original = create_list_of_lists_nullable(nested_data); // Slice to remove prefix - creates wasted space since offsets no longer reference early elements - let sliced = original.slice(1..3).unwrap(); + let sliced = original.slice(1..3).unwrap().optimize(&SESSION).unwrap(); let sliced_list = sliced.as_::(); // Test non-recursive compaction: only resets outer list offsets diff --git a/vortex-array/src/arrays/listview/array.rs b/vortex-array/src/arrays/listview/array.rs index bebd5df30fe..23821602dc5 100644 --- a/vortex-array/src/arrays/listview/array.rs +++ b/vortex-array/src/arrays/listview/array.rs @@ -37,7 +37,8 @@ use crate::dtype::PType; use crate::expr::stats::Stat; use crate::match_each_integer_ptype; use crate::match_each_unsigned_integer_ptype; -use crate::scalar_fn::fns::operators::Operator; +use crate::optimizer::ArrayOptimizer; +use crate::scalar_fn::builtin::operators::Operator; use crate::validity::Validity; #[array_slots(ListView)] @@ -409,7 +410,9 @@ pub trait ListViewArrayExt: ListViewArraySlotsExt { fn list_elements_at(&self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult { let offset = self.offset_at(index, ctx); let size = self.size_at(index, ctx); - self.elements().slice(offset..offset + size) + self.elements() + .slice(offset..offset + size)? + .optimize(ctx.session()) } /// Returns a [`Mask`] of length `elements.len()` where each bit is set iff that diff --git a/vortex-array/src/arrays/listview/compute/cast.rs b/vortex-array/src/arrays/listview/compute/cast.rs index 1c5e1fdfd84..b0f555c8ac9 100644 --- a/vortex-array/src/arrays/listview/compute/cast.rs +++ b/vortex-array/src/arrays/listview/compute/cast.rs @@ -12,8 +12,8 @@ use crate::arrays::ListViewArray; use crate::arrays::listview::ListViewArraySlotsExt; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; -use crate::scalar_fn::fns::cast::CastKernel; -use crate::scalar_fn::fns::cast::CastReduce; +use crate::scalar_fn::builtin::cast::CastKernel; +use crate::scalar_fn::builtin::cast::CastReduce; use crate::validity::Validity; fn build_with_validity( diff --git a/vortex-array/src/arrays/listview/compute/mask.rs b/vortex-array/src/arrays/listview/compute/mask.rs index b9d12f39772..ba44cfb2f78 100644 --- a/vortex-array/src/arrays/listview/compute/mask.rs +++ b/vortex-array/src/arrays/listview/compute/mask.rs @@ -11,7 +11,7 @@ use crate::arrays::ListView; use crate::arrays::ListViewArray; use crate::arrays::listview::ListViewArraySlotsExt; use crate::arrays::listview::ListViewData; -use crate::scalar_fn::fns::mask::MaskReduce; +use crate::scalar_fn::builtin::mask::MaskReduce; use crate::validity::Validity; impl MaskReduce for ListView { diff --git a/vortex-array/src/arrays/listview/compute/mod.rs b/vortex-array/src/arrays/listview/compute/mod.rs index 87587495f7a..9a43503c4b5 100644 --- a/vortex-array/src/arrays/listview/compute/mod.rs +++ b/vortex-array/src/arrays/listview/compute/mod.rs @@ -6,4 +6,3 @@ mod mask; pub(crate) mod rules; mod slice; mod take; -mod zip; diff --git a/vortex-array/src/arrays/listview/compute/rules.rs b/vortex-array/src/arrays/listview/compute/rules.rs index 233533202a8..8997d00d255 100644 --- a/vortex-array/src/arrays/listview/compute/rules.rs +++ b/vortex-array/src/arrays/listview/compute/rules.rs @@ -14,8 +14,8 @@ use crate::arrays::listview::ListViewArraySlotsExt; use crate::arrays::slice::SliceReduceAdaptor; use crate::optimizer::rules::ArrayParentReduceRule; use crate::optimizer::rules::ParentRuleSet; -use crate::scalar_fn::fns::cast::CastReduceAdaptor; -use crate::scalar_fn::fns::mask::MaskReduceAdaptor; +use crate::scalar_fn::builtin::cast::CastReduceAdaptor; +use crate::scalar_fn::builtin::mask::MaskReduceAdaptor; pub(crate) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ ParentRuleSet::lift(&ListViewFilterPushDown), diff --git a/vortex-array/src/arrays/listview/rebuild.rs b/vortex-array/src/arrays/listview/rebuild.rs index 0fe73d4f0ad..a2297ea02e0 100644 --- a/vortex-array/src/arrays/listview/rebuild.rs +++ b/vortex-array/src/arrays/listview/rebuild.rs @@ -26,7 +26,7 @@ use crate::dtype::PType; use crate::match_each_integer_ptype; use crate::match_each_unsigned_integer_ptype; use crate::scalar::Scalar; -use crate::scalar_fn::fns::operators::Operator; +use crate::scalar_fn::builtin::operators::Operator; use crate::validity::Validity; /// The widened offset type used for a rebuilt, zero-copy-to-list array. diff --git a/vortex-array/src/arrays/listview/tests/nested.rs b/vortex-array/src/arrays/listview/tests/nested.rs index c3f1c5a53b9..d73ef7386da 100644 --- a/vortex-array/src/arrays/listview/tests/nested.rs +++ b/vortex-array/src/arrays/listview/tests/nested.rs @@ -19,6 +19,7 @@ use crate::dtype::FieldNames; use crate::dtype::Nullability; use crate::dtype::PType; use crate::dtype::StructFields; +use crate::optimizer::ArrayOptimizer; use crate::validity::Validity; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); @@ -123,7 +124,11 @@ fn test_listview_of_listview_with_overlapping() { ); // Test slicing the outer ListView. - let sliced = outer_listview.slice(1..2).unwrap(); + let sliced = outer_listview + .slice(1..2) + .unwrap() + .optimize(&SESSION) + .unwrap(); assert_eq!(sliced.len(), 1); let sliced_lv = sliced.as_::(); let inner_after_slice = sliced_lv.list_elements_at(0, &mut ctx).unwrap(); @@ -257,7 +262,11 @@ fn test_mixed_offset_size_types() { assert_eq!(first_outer.len(), 3); // Test slicing with mixed types. - let sliced = outer_listview.slice(1..3).unwrap(); + let sliced = outer_listview + .slice(1..3) + .unwrap() + .optimize(&SESSION) + .unwrap(); assert_eq!(sliced.len(), 2); let sliced_lv = sliced.as_::(); diff --git a/vortex-array/src/arrays/listview/tests/operations.rs b/vortex-array/src/arrays/listview/tests/operations.rs index 45cace9c945..f6ae9567661 100644 --- a/vortex-array/src/arrays/listview/tests/operations.rs +++ b/vortex-array/src/arrays/listview/tests/operations.rs @@ -6,14 +6,12 @@ use std::sync::LazyLock; use rstest::rstest; use vortex_buffer::buffer; -use vortex_error::VortexResult; use vortex_mask::Mask; use vortex_session::VortexSession; use super::common::create_basic_listview; use super::common::create_large_listview; use super::common::create_nullable_listview; -use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; use crate::aggregate_fn::fns::is_constant::is_constant; @@ -30,6 +28,7 @@ use crate::compute::conformance::mask::test_mask_conformance; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; +use crate::optimizer::ArrayOptimizer; use crate::validity::Validity; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); @@ -51,7 +50,7 @@ fn test_slice_comprehensive() { ListViewArray::new(elements, offsets, sizes, Validity::NonNullable, &mut ctx).into_array(); // Test basic slice [1..3] - middle portion. - let sliced = listview.slice(1..3).unwrap(); + let sliced = listview.slice(1..3).unwrap().optimize(&SESSION).unwrap(); let sliced_list = sliced.as_::(); assert_eq!(sliced_list.len(), 2, "Wrong slice length"); assert_eq!( @@ -76,7 +75,7 @@ fn test_slice_comprehensive() { ); // Test full array slice [0..4]. - let full = listview.slice(0..4).unwrap(); + let full = listview.slice(0..4).unwrap().optimize(&SESSION).unwrap(); let full_list = full.as_::(); assert_eq!(full_list.len(), 4, "Full slice should preserve length"); for i in 0..4 { @@ -95,7 +94,7 @@ fn test_slice_comprehensive() { } // Test single element slice [2..3]. - let single = listview.slice(2..3).unwrap(); + let single = listview.slice(2..3).unwrap().optimize(&SESSION).unwrap(); let single_list = single.as_::(); assert_eq!(single_list.len(), 1, "Single element slice failed"); assert_eq!( @@ -123,7 +122,7 @@ fn test_slice_out_of_order() { ListViewArray::new(elements, offsets, sizes, Validity::NonNullable, &mut ctx).into_array(); // Slice [1..4] should maintain the out-of-order offsets. - let sliced = listview.slice(1..4).unwrap(); + let sliced = listview.slice(1..4).unwrap().optimize(&SESSION).unwrap(); let sliced_list = sliced.as_::(); assert_eq!( @@ -198,7 +197,7 @@ fn test_slice_with_nulls() { .into_array(); // Slice [1..3] should preserve nulls. - let sliced = listview.slice(1..3).unwrap(); + let sliced = listview.slice(1..3).unwrap().optimize(&SESSION).unwrap(); let sliced_list = sliced.as_::(); assert_eq!(sliced_list.len(), 2); @@ -507,108 +506,6 @@ fn test_cast_large_dataset() { assert_arrays_eq!(result_list, expected, &mut SESSION.create_execution_ctx()); } -//////////////////////////////////////////////////////////////////////////////////////////////////// -// Zip tests -//////////////////////////////////////////////////////////////////////////////////////////////////// - -#[test] -fn test_zip_widens_false_element_nullability() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - // [[1, 2], [3], [4]] - let if_true = ListViewArray::new( - buffer![1i32, 2, 3, 4].into_array(), - buffer![0u32, 2, 3].into_array(), - buffer![2u32, 1, 1].into_array(), - Validity::NonNullable, - &mut ctx, - ) - .into_array(); - // [[10, null], [30], [40]] - let if_false = ListViewArray::new( - PrimitiveArray::from_option_iter([Some(10i32), None, Some(30), Some(40)]).into_array(), - buffer![0u32, 2, 3].into_array(), - buffer![2u32, 1, 1].into_array(), - Validity::NonNullable, - &mut ctx, - ) - .into_array(); - let mask = Mask::from_iter([false, true, false]); - - let result = mask - .into_array() - .zip(if_true, if_false)? - .execute::(&mut SESSION.create_execution_ctx())?; - assert!(result.is::()); - assert_eq!( - result.dtype(), - &DType::List( - Arc::new(DType::Primitive(PType::I32, Nullability::Nullable)), - Nullability::NonNullable, - ) - ); - - // [[10, null], [3], [40]] - let expected = ListViewArray::new( - PrimitiveArray::from_option_iter([Some(10i32), None, Some(3), Some(40)]).into_array(), - buffer![0u32, 2, 3].into_array(), - buffer![2u32, 1, 1].into_array(), - Validity::NonNullable, - &mut ctx, - ) - .into_array(); - assert_arrays_eq!(result, expected, &mut ctx); - Ok(()) -} - -#[test] -fn test_zip_widens_true_element_nullability() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - // [[1, null], [3], [4]] - let if_true = ListViewArray::new( - PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), Some(4)]).into_array(), - buffer![0u32, 2, 3].into_array(), - buffer![2u32, 1, 1].into_array(), - Validity::NonNullable, - &mut ctx, - ) - .into_array(); - // [[10], [20], [30]] - let if_false = ListViewArray::new( - buffer![10i32, 20, 30].into_array(), - buffer![0u32, 1, 2].into_array(), - buffer![1u32, 1, 1].into_array(), - Validity::NonNullable, - &mut ctx, - ) - .into_array(); - let mask = Mask::from_iter([true, false, true]); - - let result = mask - .into_array() - .zip(if_true, if_false)? - .execute::(&mut SESSION.create_execution_ctx())?; - assert!(result.is::()); - assert_eq!( - result.dtype(), - &DType::List( - Arc::new(DType::Primitive(PType::I32, Nullability::Nullable)), - Nullability::NonNullable, - ) - ); - - // [[1, null], [20], [4]] - let expected = ListViewArray::new( - PrimitiveArray::from_option_iter([Some(1i32), None, Some(20), Some(4)]).into_array(), - buffer![0u32, 2, 3].into_array(), - buffer![2u32, 1, 1].into_array(), - Validity::NonNullable, - &mut ctx, - ) - .into_array(); - assert_arrays_eq!(result, expected, &mut ctx); - Ok(()) -} - //////////////////////////////////////////////////////////////////////////////////////////////////// // Constant tests //////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/vortex-array/src/arrays/listview/vtable/kernel.rs b/vortex-array/src/arrays/listview/vtable/kernel.rs index 5621b23e79e..80107f2c769 100644 --- a/vortex-array/src/arrays/listview/vtable/kernel.rs +++ b/vortex-array/src/arrays/listview/vtable/kernel.rs @@ -6,13 +6,10 @@ use vortex_session::VortexSession; use crate::arrays::ListView; use crate::optimizer::kernels::ArrayKernelsExt; use crate::scalar_fn::ScalarFnVTable; -use crate::scalar_fn::fns::cast::Cast; -use crate::scalar_fn::fns::cast::CastExecuteAdaptor; -use crate::scalar_fn::fns::zip::Zip; -use crate::scalar_fn::fns::zip::ZipExecuteAdaptor; +use crate::scalar_fn::builtin::cast::Cast; +use crate::scalar_fn::builtin::cast::CastExecuteAdaptor; pub(crate) fn initialize(session: &VortexSession) { let kernels = session.kernels(); kernels.register_execute_parent_kernel(Cast.id(), ListView, CastExecuteAdaptor(ListView)); - kernels.register_execute_parent_kernel(Zip.id(), ListView, ZipExecuteAdaptor(ListView)); } diff --git a/vortex-array/src/arrays/map/compute/cast.rs b/vortex-array/src/arrays/map/compute/cast.rs index bfbab6e17e9..0fb3622d990 100644 --- a/vortex-array/src/arrays/map/compute/cast.rs +++ b/vortex-array/src/arrays/map/compute/cast.rs @@ -16,8 +16,8 @@ use crate::arrays::map::compute::rebuild_map_from_array; use crate::dtype::DType; use crate::dtype::MapDType; use crate::executor::ExecutionCtx; -use crate::scalar_fn::fns::cast::CastKernel; -use crate::scalar_fn::fns::cast::CastReduce; +use crate::scalar_fn::builtin::cast::CastKernel; +use crate::scalar_fn::builtin::cast::CastReduce; fn prepare_map_cast_target( array: ArrayView<'_, Map>, diff --git a/vortex-array/src/arrays/map/compute/mask.rs b/vortex-array/src/arrays/map/compute/mask.rs index 97218d08454..a56c453ffd3 100644 --- a/vortex-array/src/arrays/map/compute/mask.rs +++ b/vortex-array/src/arrays/map/compute/mask.rs @@ -11,8 +11,8 @@ use crate::arrays::map::MapArrayExt; use crate::arrays::map::MapArraySlotsExt; use crate::arrays::map::compute::rebuild_map_from_array; use crate::executor::ExecutionCtx; -use crate::scalar_fn::fns::mask::MaskKernel; -use crate::scalar_fn::fns::mask::MaskReduce; +use crate::scalar_fn::builtin::mask::MaskKernel; +use crate::scalar_fn::builtin::mask::MaskReduce; impl MaskReduce for Map { const VALIDITY_IS_METADATA_ONLY: bool = true; diff --git a/vortex-array/src/arrays/map/compute/rules.rs b/vortex-array/src/arrays/map/compute/rules.rs index 27622ff1c5f..518d00bf2f6 100644 --- a/vortex-array/src/arrays/map/compute/rules.rs +++ b/vortex-array/src/arrays/map/compute/rules.rs @@ -6,8 +6,8 @@ use crate::arrays::dict::TakeReduceAdaptor; use crate::arrays::filter::FilterReduceAdaptor; use crate::arrays::slice::SliceReduceAdaptor; use crate::optimizer::rules::ParentRuleSet; -use crate::scalar_fn::fns::cast::CastReduceAdaptor; -use crate::scalar_fn::fns::mask::MaskReduceAdaptor; +use crate::scalar_fn::builtin::cast::CastReduceAdaptor; +use crate::scalar_fn::builtin::mask::MaskReduceAdaptor; pub(crate) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ ParentRuleSet::lift(&FilterReduceAdaptor(Map)), diff --git a/vortex-array/src/arrays/map/tests.rs b/vortex-array/src/arrays/map/tests.rs index f9f60ec624a..914a292fea2 100644 --- a/vortex-array/src/arrays/map/tests.rs +++ b/vortex-array/src/arrays/map/tests.rs @@ -39,6 +39,7 @@ use crate::dtype::DType; use crate::dtype::MapDType; use crate::dtype::Nullability; use crate::dtype::PType; +use crate::optimizer::ArrayOptimizer; use crate::scalar::Scalar; use crate::serde::SerializeOptions; use crate::serde::SerializedArray; @@ -299,7 +300,7 @@ fn slice_preserves_map_rows() -> VortexResult<()> { Some(vec![]), ], )?; - let sliced = source.slice(0..2)?; + let sliced = source.slice(0..2)?.optimize(&SESSION)?; let mut ctx = SESSION.create_execution_ctx(); assert!(sliced.is::()); @@ -314,11 +315,15 @@ fn filter_handles_all_none_and_mixed_maps() -> VortexResult<()> { let map_dtype = map_dtype()?; let mut ctx = SESSION.create_execution_ctx(); - let all = source.filter(Mask::from_iter([true, true, true]))?; + let all = source + .filter(Mask::from_iter([true, true, true]))? + .optimize(&SESSION)?; assert!(all.is::()); assert_arrays_eq!(all, sample_array()?, &mut ctx); - let none = source.filter(Mask::from_iter([false, false, false]))?; + let none = source + .filter(Mask::from_iter([false, false, false]))? + .optimize(&SESSION)?; let expected_none = map_array_from_rows( map_dtype.clone(), Nullability::Nullable, @@ -327,7 +332,9 @@ fn filter_handles_all_none_and_mixed_maps() -> VortexResult<()> { assert!(none.is::()); assert_arrays_eq!(none, expected_none, &mut ctx); - let mixed = source.filter(Mask::from_iter([true, false, true]))?; + let mixed = source + .filter(Mask::from_iter([true, false, true]))? + .optimize(&SESSION)?; let expected_mixed = map_array_from_rows( map_dtype, Nullability::Nullable, @@ -375,7 +382,7 @@ fn filter_dropping_nonempty_middle_row_clears_zero_copy_flag() -> VortexResult<( let mask = Mask::from_iter([true, false, true]); let mut ctx = SESSION.create_execution_ctx(); - let reduced = source.filter(mask.clone())?; + let reduced = source.filter(mask.clone())?.optimize(&SESSION)?; assert!(reduced.is::()); // The dropped middle row leaves a gap in the referenced entry elements, so the filtered // entries must not claim zero-copy-to-list convertibility. @@ -410,7 +417,9 @@ fn take_supports_reordered_duplicate_and_nullable_indices() -> VortexResult<()> let map_dtype = map_dtype()?; let mut ctx = SESSION.create_execution_ctx(); - let taken = source.take(PrimitiveArray::from_iter([2u64, 0, 0]).into_array())?; + let taken = source + .take(PrimitiveArray::from_iter([2u64, 0, 0]).into_array())? + .optimize(&SESSION)?; let expected_taken = map_array_from_rows( map_dtype.clone(), Nullability::Nullable, @@ -442,7 +451,8 @@ fn take_supports_reordered_duplicate_and_nullable_indices() -> VortexResult<()> )? .into_array(); let nullable_taken = nonnullable_source - .take(PrimitiveArray::from_option_iter([Some(1u64), None, Some(0)]).into_array())?; + .take(PrimitiveArray::from_option_iter([Some(1u64), None, Some(0)]).into_array())? + .optimize(&SESSION)?; let expected_nullable_taken = map_array_from_rows( map_dtype, Nullability::Nullable, @@ -459,7 +469,7 @@ fn take_supports_reordered_duplicate_and_nullable_indices() -> VortexResult<()> fn mask_combines_with_existing_map_validity() -> VortexResult<()> { let source = sample_array()?.into_array(); let mask = BoolArray::from_iter([true, false, true]).into_array(); - let masked = source.mask(mask)?; + let masked = source.mask(mask)?.optimize(&SESSION)?; let expected = map_array_from_rows( map_dtype()?, Nullability::Nullable, @@ -516,7 +526,7 @@ fn cast_widens_map_key_value_and_outer_nullability() -> VortexResult<()> { .as_map_opt() .vortex_expect("target dtype is map") .clone(); - let cast = source.cast(target_dtype)?; + let cast = source.cast(target_dtype)?.optimize(&SESSION)?; let expected = map_array_from_rows( target_map_dtype, Nullability::Nullable, @@ -551,7 +561,7 @@ fn cast_can_drop_but_not_create_sortedness_assertion() -> VortexResult<()> { false, Nullability::Nullable, )?; - let cast = sorted_source.cast(unsorted_target)?; + let cast = sorted_source.cast(unsorted_target)?.optimize(&SESSION)?; let expected = map_array_from_rows( MapDType::try_new( DType::Primitive(PType::I32, Nullability::NonNullable), @@ -589,7 +599,13 @@ fn cast_can_drop_but_not_create_sortedness_assertion() -> VortexResult<()> { true, Nullability::Nullable, )?; - assert!(unsorted_source.cast(sorted_target).is_err()); + // The cast is rejected when it executes. + assert!( + unsorted_source + .cast(sorted_target)? + .execute::(&mut SESSION.create_execution_ctx()) + .is_err() + ); Ok(()) } @@ -619,7 +635,12 @@ fn null_map_cast_cannot_create_sortedness_assertion() -> VortexResult<()> { let all_null = map_array_from_rows(unsorted_map_dtype, Nullability::Nullable, [None, None])?.into_array(); - assert!(all_null.cast(sorted_dtype).is_err()); + assert!( + all_null + .cast(sorted_dtype)? + .execute::(&mut ctx) + .is_err() + ); Ok(()) } diff --git a/vortex-array/src/arrays/map/vtable/kernel.rs b/vortex-array/src/arrays/map/vtable/kernel.rs index 860496834d6..d2c1665c542 100644 --- a/vortex-array/src/arrays/map/vtable/kernel.rs +++ b/vortex-array/src/arrays/map/vtable/kernel.rs @@ -13,10 +13,10 @@ use crate::arrays::filter::FilterExecuteAdaptor; use crate::arrays::slice::SliceExecuteAdaptor; use crate::optimizer::kernels::ArrayKernelsExt; use crate::scalar_fn::ScalarFnVTable; -use crate::scalar_fn::fns::cast::Cast; -use crate::scalar_fn::fns::cast::CastExecuteAdaptor; -use crate::scalar_fn::fns::mask::Mask; -use crate::scalar_fn::fns::mask::MaskExecuteAdaptor; +use crate::scalar_fn::builtin::cast::Cast; +use crate::scalar_fn::builtin::cast::CastExecuteAdaptor; +use crate::scalar_fn::builtin::mask::Mask; +use crate::scalar_fn::builtin::mask::MaskExecuteAdaptor; pub(crate) fn initialize(session: &VortexSession) { let kernels = session.kernels(); diff --git a/vortex-array/src/arrays/masked/compute/mask.rs b/vortex-array/src/arrays/masked/compute/mask.rs index 0de2dea5cd9..b92b9c4f921 100644 --- a/vortex-array/src/arrays/masked/compute/mask.rs +++ b/vortex-array/src/arrays/masked/compute/mask.rs @@ -8,8 +8,8 @@ use crate::IntoArray; use crate::array::ArrayView; use crate::arrays::Masked; use crate::arrays::masked::MaskedArraySlotsExt; -use crate::scalar_fn::fns::mask::Mask as MaskExpr; -use crate::scalar_fn::fns::mask::MaskReduce; +use crate::scalar_fn::builtin::mask::Mask as MaskExpr; +use crate::scalar_fn::builtin::mask::MaskReduce; use crate::validity::Validity; impl MaskReduce for Masked { diff --git a/vortex-array/src/arrays/masked/compute/rules.rs b/vortex-array/src/arrays/masked/compute/rules.rs index ac2123aedc8..e213452d38e 100644 --- a/vortex-array/src/arrays/masked/compute/rules.rs +++ b/vortex-array/src/arrays/masked/compute/rules.rs @@ -16,8 +16,8 @@ use crate::arrays::slice::SliceReduceAdaptor; use crate::builtins::ArrayBuiltins; use crate::optimizer::rules::ArrayParentReduceRule; use crate::optimizer::rules::ParentRuleSet; -use crate::scalar_fn::fns::get_item::GetItem; -use crate::scalar_fn::fns::mask::MaskReduceAdaptor; +use crate::scalar_fn::builtin::get_item::GetItem; +use crate::scalar_fn::builtin::mask::MaskReduceAdaptor; pub(crate) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ ParentRuleSet::lift(&FilterReduceAdaptor(Masked)), diff --git a/vortex-array/src/arrays/mod.rs b/vortex-array/src/arrays/mod.rs index ed6e272b2b1..9208f83a863 100644 --- a/vortex-array/src/arrays/mod.rs +++ b/vortex-array/src/arrays/mod.rs @@ -148,7 +148,6 @@ pub(crate) fn initialize(session: &VortexSession) { struct_::initialize(session); varbin::initialize(session); varbinview::initialize(session); - variant::initialize(session); } #[cfg(feature = "arbitrary")] diff --git a/vortex-array/src/arrays/null/compute/cast.rs b/vortex-array/src/arrays/null/compute/cast.rs index be9da91131e..f85e0ba81f6 100644 --- a/vortex-array/src/arrays/null/compute/cast.rs +++ b/vortex-array/src/arrays/null/compute/cast.rs @@ -11,7 +11,7 @@ use crate::arrays::ConstantArray; use crate::arrays::Null; use crate::dtype::DType; use crate::scalar::Scalar; -use crate::scalar_fn::fns::cast::CastReduce; +use crate::scalar_fn::builtin::cast::CastReduce; impl CastReduce for Null { fn cast(array: ArrayView<'_, Null>, dtype: &DType) -> VortexResult> { @@ -34,6 +34,7 @@ mod tests { use rstest::rstest; use vortex_session::VortexSession; + use crate::Canonical; use crate::IntoArray; use crate::VortexSessionExecute; use crate::arrays::NullArray; @@ -82,9 +83,11 @@ mod tests { #[test] fn test_cast_null_to_non_nullable_fails() { let null_array = NullArray::new(5); + // The cast is rejected when it executes. let result = null_array .into_array() - .cast(DType::Primitive(PType::I32, Nullability::NonNullable)); + .cast(DType::Primitive(PType::I32, Nullability::NonNullable)) + .and_then(|cast| cast.execute::(&mut SESSION.create_execution_ctx())); assert!(result.is_err()); } diff --git a/vortex-array/src/arrays/null/compute/mask.rs b/vortex-array/src/arrays/null/compute/mask.rs index dbb3afaad1c..8205ee6dbb0 100644 --- a/vortex-array/src/arrays/null/compute/mask.rs +++ b/vortex-array/src/arrays/null/compute/mask.rs @@ -6,7 +6,7 @@ use vortex_error::VortexResult; use crate::ArrayRef; use crate::array::ArrayView; use crate::arrays::Null; -use crate::scalar_fn::fns::mask::MaskReduce; +use crate::scalar_fn::builtin::mask::MaskReduce; impl MaskReduce for Null { fn mask(array: ArrayView<'_, Null>, _mask: &ArrayRef) -> VortexResult> { diff --git a/vortex-array/src/arrays/null/compute/rules.rs b/vortex-array/src/arrays/null/compute/rules.rs index b2ee44a19dd..996d7d869e6 100644 --- a/vortex-array/src/arrays/null/compute/rules.rs +++ b/vortex-array/src/arrays/null/compute/rules.rs @@ -5,8 +5,8 @@ use crate::arrays::Null; use crate::arrays::filter::FilterReduceAdaptor; use crate::arrays::slice::SliceReduceAdaptor; use crate::optimizer::rules::ParentRuleSet; -use crate::scalar_fn::fns::cast::CastReduceAdaptor; -use crate::scalar_fn::fns::mask::MaskReduceAdaptor; +use crate::scalar_fn::builtin::cast::CastReduceAdaptor; +use crate::scalar_fn::builtin::mask::MaskReduceAdaptor; pub(crate) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ ParentRuleSet::lift(&FilterReduceAdaptor(Null)), diff --git a/vortex-array/src/arrays/patched/compute/compare.rs b/vortex-array/src/arrays/patched/compute/compare.rs index 800783b11c5..6adfd994f55 100644 --- a/vortex-array/src/arrays/patched/compute/compare.rs +++ b/vortex-array/src/arrays/patched/compute/compare.rs @@ -22,8 +22,8 @@ use crate::arrays::primitive::NativeValue; use crate::builtins::ArrayBuiltins; use crate::dtype::NativePType; use crate::match_each_native_ptype; -use crate::scalar_fn::fns::binary::CompareKernel; -use crate::scalar_fn::fns::operators::CompareOperator; +use crate::scalar_fn::builtin::binary::CompareKernel; +use crate::scalar_fn::builtin::operators::CompareOperator; impl CompareKernel for Patched { fn compare( @@ -175,8 +175,8 @@ mod tests { use crate::assert_arrays_eq; use crate::optimizer::ArrayOptimizer; use crate::patches::Patches; - use crate::scalar_fn::fns::binary::CompareKernel; - use crate::scalar_fn::fns::operators::CompareOperator; + use crate::scalar_fn::builtin::binary::CompareKernel; + use crate::scalar_fn::builtin::operators::CompareOperator; use crate::validity::Validity; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); @@ -234,7 +234,12 @@ mod tests { let lhs = Patched::from_array_and_patches(lhs, &patches, &mut ctx).unwrap(); // Slice the array so that the first patch should be skipped. - let lhs_ref = lhs.into_array().slice(10..512).unwrap().optimize().unwrap(); + let lhs_ref = lhs + .into_array() + .slice(10..512) + .unwrap() + .optimize(&SESSION) + .unwrap(); let lhs = lhs_ref.try_downcast::().unwrap(); assert_eq!(lhs.len(), 502); diff --git a/vortex-array/src/arrays/patched/compute/filter.rs b/vortex-array/src/arrays/patched/compute/filter.rs index e8621ff0eab..12dd6078b7d 100644 --- a/vortex-array/src/arrays/patched/compute/filter.rs +++ b/vortex-array/src/arrays/patched/compute/filter.rs @@ -94,7 +94,7 @@ mod tests { let array = Patched::from_array_and_patches(array, &patches, &mut ctx)?.into_array(); let mask = Mask::from_iter([true, false, false, false, true]); - let filtered = array.filter(mask)?.optimize()?; + let filtered = array.filter(mask)?.optimize(&SESSION)?; // Values at indices 0 and 4: MIN and MAX. let expected = PrimitiveArray::from_iter([u16::MIN, u16::MAX]); @@ -117,11 +117,11 @@ mod tests { let array = Patched::from_array_and_patches(array, &patches, &mut ctx)?.into_array(); - let sliced = array.slice(5..4096)?.optimize()?; + let sliced = array.slice(5..4096)?.optimize(&SESSION)?; // Filter that touches only the first 2 chunks. let mask = Mask::from_indices(4091, vec![0, 1, 2, 1025]); - let filtered = sliced.filter(mask)?.optimize()?; + let filtered = sliced.filter(mask)?.optimize(&SESSION)?; let expected = PrimitiveArray::from_iter([u16::MAX, u16::MIN, u16::MIN, u16::MAX]); assert_arrays_eq!(filtered, expected, &mut ctx); @@ -144,7 +144,7 @@ mod tests { // Filter that only touches the middle 2 chunks. let mask = Mask::from_indices(4096, vec![1024, 1025, 3000]); - let filtered = array.filter(mask)?.optimize()?; + let filtered = array.filter(mask)?.optimize(&SESSION)?; let expected = PrimitiveArray::from_iter([u16::MAX, u16::MAX, u16::MIN]); @@ -168,7 +168,7 @@ mod tests { // Filter that only touches the middle 2 chunks. let mask = Mask::from_indices(4096, vec![1024, 1025, 3000]); - let filtered = array.filter(mask)?.optimize()?; + let filtered = array.filter(mask)?.optimize(&SESSION)?; let expected = PrimitiveArray::from_iter([u16::MAX, u16::MIN, u16::MIN]); @@ -202,7 +202,7 @@ mod tests { // Indices 1024 and 1025 fall in chunk 1, and 3000 falls in chunk 2. let mask = Mask::from_indices(4096, vec![1024, 1025, 3000]); - let filtered = sliced.filter(mask)?.optimize()?; + let filtered = sliced.filter(mask)?.optimize(&SESSION)?; let expected = PrimitiveArray::from_iter([u16::MAX, u16::MAX, u16::MIN]); @@ -230,13 +230,13 @@ mod tests { // Slice to create offset > 0. // After slice(5..4096), logical position 0 = original position 5 (patched to 9999). - let sliced = array.slice(5..4096)?.optimize()?; + let sliced = array.slice(5..4096)?.optimize(&SESSION)?; assert_eq!(sliced.len(), 4091); // Filter that touches the first 2 chunks. // Logical indices: 0 (was 5, patched), 1 (was 6, value 6), 1025 (was 1030, patched). let mask = Mask::from_indices(4091, vec![0, 1, 1025]); - let filtered = sliced.filter(mask)?.optimize()?; + let filtered = sliced.filter(mask)?.optimize(&SESSION)?; // Expected: 9999 (patched at logical 0), 6 (original at logical 1), 8888 (patched at logical 1025). let expected = PrimitiveArray::from_iter([9999u16, 6, 8888]); @@ -264,7 +264,7 @@ mod tests { // Slice at chunk boundary to create offset > 0. // Slice [1024..6144] gives us 5120 elements (5 chunks). // Original patches at 5000 and 6000 become relative indices 3976 and 4976. - let sliced = patched.into_array().slice(1024..6144)?.optimize()?; + let sliced = patched.into_array().slice(1024..6144)?.optimize(&SESSION)?; assert_eq!(sliced.len(), 5120); // Filter that touches only the last 2 chunks (chunks 3 and 4). @@ -272,7 +272,7 @@ mod tests { // Patch at 3976 is in chunk 3, patch at 4976 is in chunk 4. let mask = Mask::from_indices(5120, vec![3976, 4976, 5119]); - let filtered = sliced.filter(mask)?.optimize()?; + let filtered = sliced.filter(mask)?.optimize(&SESSION)?; // Expected: patch at 3976 (was 5000), patch at 4976 (was 6000), and MIN at 5119. let expected = PrimitiveArray::from_iter([u16::MAX, u16::MAX, u16::MIN]); diff --git a/vortex-array/src/arrays/patched/vtable/kernels.rs b/vortex-array/src/arrays/patched/vtable/kernels.rs index 51f53d15625..6ce98942c61 100644 --- a/vortex-array/src/arrays/patched/vtable/kernels.rs +++ b/vortex-array/src/arrays/patched/vtable/kernels.rs @@ -9,8 +9,8 @@ use crate::arrays::Patched; use crate::arrays::dict::TakeExecuteAdaptor; use crate::optimizer::kernels::ArrayKernelsExt; use crate::scalar_fn::ScalarFnVTable; -use crate::scalar_fn::fns::binary::Binary; -use crate::scalar_fn::fns::binary::CompareExecuteAdaptor; +use crate::scalar_fn::builtin::binary::Binary; +use crate::scalar_fn::builtin::binary::CompareExecuteAdaptor; pub(crate) fn initialize(session: &VortexSession) { let kernels = session.kernels(); diff --git a/vortex-array/src/arrays/patched/vtable/operations.rs b/vortex-array/src/arrays/patched/vtable/operations.rs index af3a6d253d8..a222957dff9 100644 --- a/vortex-array/src/arrays/patched/vtable/operations.rs +++ b/vortex-array/src/arrays/patched/vtable/operations.rs @@ -38,7 +38,7 @@ impl OperationsVTable for Patched { let patch_indices = array .patch_indices() .slice(range.clone())? - .optimize()? + .optimize(ctx.session())? .execute::(ctx)?; // NOTE: we do linear scan as lane has <= 32 patches, binary search would likely @@ -174,7 +174,7 @@ mod tests { .into_array() .slice(3..4096) .unwrap() - .optimize() + .optimize(&SESSION) .unwrap(); assert!(array.is::()); diff --git a/vortex-array/src/arrays/patched/vtable/slice.rs b/vortex-array/src/arrays/patched/vtable/slice.rs index 6c59f2b528c..f250e548f11 100644 --- a/vortex-array/src/arrays/patched/vtable/slice.rs +++ b/vortex-array/src/arrays/patched/vtable/slice.rs @@ -73,6 +73,7 @@ mod tests { use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::dtype::NativePType; + use crate::optimizer::ArrayOptimizer; use crate::patches::Patches; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); @@ -89,7 +90,10 @@ mod tests { let patched_array = Patched::from_array_and_patches(values, &patches, &mut ctx)?; - let sliced = patched_array.into_array().slice(1..10)?; + let sliced = patched_array + .into_array() + .slice(1..10)? + .optimize(&SESSION)?; insta::assert_snapshot!( sliced.display_tree_encodings_only(), diff --git a/vortex-array/src/arrays/primitive/compute/cast.rs b/vortex-array/src/arrays/primitive/compute/cast.rs index c937d68d611..dad36d912a0 100644 --- a/vortex-array/src/arrays/primitive/compute/cast.rs +++ b/vortex-array/src/arrays/primitive/compute/cast.rs @@ -44,8 +44,8 @@ use crate::match_each_native_ptype; use crate::match_each_signed_integer_ptype; use crate::scalar::DecimalValue; use crate::scalar::Scalar; -use crate::scalar_fn::fns::cast::CastKernel; -use crate::scalar_fn::fns::cast::CastReduce; +use crate::scalar_fn::builtin::cast::CastKernel; +use crate::scalar_fn::builtin::cast::CastReduce; use crate::validity::Validity; impl CastReduce for Primitive { diff --git a/vortex-array/src/arrays/primitive/compute/fill_null.rs b/vortex-array/src/arrays/primitive/compute/fill_null.rs index ea657448327..cc105683e7c 100644 --- a/vortex-array/src/arrays/primitive/compute/fill_null.rs +++ b/vortex-array/src/arrays/primitive/compute/fill_null.rs @@ -15,7 +15,7 @@ use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::match_each_native_ptype; use crate::scalar::Scalar; -use crate::scalar_fn::fns::fill_null::FillNullKernel; +use crate::scalar_fn::builtin::fill_null::FillNullKernel; use crate::validity::Validity; impl FillNullKernel for Primitive { diff --git a/vortex-array/src/arrays/primitive/compute/mask.rs b/vortex-array/src/arrays/primitive/compute/mask.rs index 4cd307f0909..1965997abfb 100644 --- a/vortex-array/src/arrays/primitive/compute/mask.rs +++ b/vortex-array/src/arrays/primitive/compute/mask.rs @@ -8,7 +8,7 @@ use crate::IntoArray; use crate::array::ArrayView; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; -use crate::scalar_fn::fns::mask::MaskReduce; +use crate::scalar_fn::builtin::mask::MaskReduce; use crate::validity::Validity; impl MaskReduce for Primitive { diff --git a/vortex-array/src/arrays/primitive/compute/mod.rs b/vortex-array/src/arrays/primitive/compute/mod.rs index c12771d3da1..142d42e6887 100644 --- a/vortex-array/src/arrays/primitive/compute/mod.rs +++ b/vortex-array/src/arrays/primitive/compute/mod.rs @@ -1,14 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -mod between; mod cast; mod fill_null; mod fixed_width; mod mask; pub(crate) mod rules; mod slice; -mod zip; #[cfg(test)] mod tests { diff --git a/vortex-array/src/arrays/primitive/compute/rules.rs b/vortex-array/src/arrays/primitive/compute/rules.rs index 528e5e91520..cf8192f8a22 100644 --- a/vortex-array/src/arrays/primitive/compute/rules.rs +++ b/vortex-array/src/arrays/primitive/compute/rules.rs @@ -12,8 +12,8 @@ use crate::arrays::PrimitiveArray; use crate::arrays::slice::SliceReduceAdaptor; use crate::optimizer::rules::ArrayParentReduceRule; use crate::optimizer::rules::ParentRuleSet; -use crate::scalar_fn::fns::cast::CastReduceAdaptor; -use crate::scalar_fn::fns::mask::MaskReduceAdaptor; +use crate::scalar_fn::builtin::cast::CastReduceAdaptor; +use crate::scalar_fn::builtin::mask::MaskReduceAdaptor; pub(crate) const RULES: ParentRuleSet = ParentRuleSet::new(&[ ParentRuleSet::lift(&PrimitiveMaskedValidityRule), diff --git a/vortex-array/src/arrays/primitive/vtable/kernel.rs b/vortex-array/src/arrays/primitive/vtable/kernel.rs index 6382ea73794..26589e42c08 100644 --- a/vortex-array/src/arrays/primitive/vtable/kernel.rs +++ b/vortex-array/src/arrays/primitive/vtable/kernel.rs @@ -9,22 +9,13 @@ use crate::arrays::Primitive; use crate::arrays::dict::TakeExecuteAdaptor; use crate::optimizer::kernels::ArrayKernelsExt; use crate::scalar_fn::ScalarFnVTable; -use crate::scalar_fn::fns::between::Between; -use crate::scalar_fn::fns::between::BetweenExecuteAdaptor; -use crate::scalar_fn::fns::cast::Cast; -use crate::scalar_fn::fns::cast::CastExecuteAdaptor; -use crate::scalar_fn::fns::fill_null::FillNull; -use crate::scalar_fn::fns::fill_null::FillNullExecuteAdaptor; -use crate::scalar_fn::fns::zip::Zip; -use crate::scalar_fn::fns::zip::ZipExecuteAdaptor; +use crate::scalar_fn::builtin::cast::Cast; +use crate::scalar_fn::builtin::cast::CastExecuteAdaptor; +use crate::scalar_fn::builtin::fill_null::FillNull; +use crate::scalar_fn::builtin::fill_null::FillNullExecuteAdaptor; pub(crate) fn initialize(session: &VortexSession) { let kernels = session.kernels(); - kernels.register_execute_parent_kernel( - Between.id(), - Primitive, - BetweenExecuteAdaptor(Primitive), - ); kernels.register_execute_parent_kernel(Cast.id(), Primitive, CastExecuteAdaptor(Primitive)); kernels.register_execute_parent_kernel( FillNull.id(), @@ -32,5 +23,4 @@ pub(crate) fn initialize(session: &VortexSession) { FillNullExecuteAdaptor(Primitive), ); kernels.register_execute_parent_kernel(Dict.id(), Primitive, TakeExecuteAdaptor(Primitive)); - kernels.register_execute_parent_kernel(Zip.id(), Primitive, ZipExecuteAdaptor(Primitive)); } diff --git a/vortex-array/src/arrays/scalar_fn/array.rs b/vortex-array/src/arrays/scalar_fn/array.rs index 391ea9a03cd..c36f39e929e 100644 --- a/vortex-array/src/arrays/scalar_fn/array.rs +++ b/vortex-array/src/arrays/scalar_fn/array.rs @@ -18,7 +18,7 @@ use crate::array::TypedArrayRef; use crate::arrays::ScalarFn; use crate::dtype::DType; use crate::scalar_fn::ScalarFnRef; -use crate::scalar_fn::fns::literal::Literal; +use crate::scalar_fn::builtin::literal::Literal; // ScalarFnArray has a variable number of slots (one per child) diff --git a/vortex-array/src/arrays/scalar_fn/rules.rs b/vortex-array/src/arrays/scalar_fn/rules.rs index 5383a614bd6..4d70966f871 100644 --- a/vortex-array/src/arrays/scalar_fn/rules.rs +++ b/vortex-array/src/arrays/scalar_fn/rules.rs @@ -21,7 +21,7 @@ use crate::optimizer::rules::ArrayReduceRule; use crate::optimizer::rules::ParentRuleSet; use crate::optimizer::rules::ReduceRuleSet; use crate::scalar_fn::ArrayReduceNode; -use crate::scalar_fn::fns::pack::Pack; +use crate::scalar_fn::builtin::pack::Pack; use crate::validity::Validity; pub(super) const RULES: ReduceRuleSet = @@ -170,11 +170,9 @@ mod tests { use crate::expr::is_null; use crate::expr::root; use crate::optimizer::rules::ArrayParentReduceRule; - use crate::scalar::Scalar; use crate::scalar_fn::TypedScalarFnInstance; - use crate::scalar_fn::fns::binary::Binary; - use crate::scalar_fn::fns::list_contains::PreparedSet; - use crate::scalar_fn::fns::operators::Operator; + use crate::scalar_fn::builtin::binary::Binary; + use crate::scalar_fn::builtin::operators::Operator; use crate::validity::Validity; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); @@ -227,30 +225,6 @@ mod tests { Ok(()) } - #[test] - fn test_filter_pushdown_without_children() -> VortexResult<()> { - let list = Scalar::list( - DType::Primitive(PType::I32, Nullability::NonNullable), - vec![Scalar::from(7i32)], - Nullability::NonNullable, - ); - let array = PreparedSet::try_new_array(list.clone(), 3)?; - let parent = FilterArray::try_new(array.clone(), Mask::from_iter([true, false, true]))?; - - let result = ScalarFilterPushdownRule - .reduce_parent(array.as_::(), parent.as_view(), 0)? - .ok_or_else(|| vortex_err!("expected filter pushdown"))?; - - assert_eq!(result.as_::().nchildren(), 0); - assert_eq!(result.len(), 2); - assert_arrays_eq!( - result, - ConstantArray::new(list, 2), - &mut SESSION.create_execution_ctx() - ); - Ok(()) - } - #[test] fn test_empty_constants() { let array = ChunkedArray::try_new( diff --git a/vortex-array/src/arrays/scalar_fn/vtable/operations.rs b/vortex-array/src/arrays/scalar_fn/vtable/operations.rs index 943d7a2a041..c38213bad02 100644 --- a/vortex-array/src/arrays/scalar_fn/vtable/operations.rs +++ b/vortex-array/src/arrays/scalar_fn/vtable/operations.rs @@ -80,23 +80,17 @@ mod tests { use crate::array::Array; use crate::array::ArrayParts; use crate::arrays::BoolArray; - use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; use crate::arrays::ScalarFnArray; - use crate::arrays::scalar_fn::ScalarFnArrayExt; use crate::arrays::scalar_fn::array::ScalarFnData; use crate::arrays::scalar_fn::vtable::ScalarFn; use crate::assert_arrays_eq; - use crate::dtype::DType; - use crate::dtype::Nullability; use crate::dtype::PType; use crate::scalar::Scalar; use crate::scalar_fn::TypedScalarFnInstance; - use crate::scalar_fn::fns::binary::Binary; - use crate::scalar_fn::fns::list_contains::PreparedSet; - use crate::scalar_fn::fns::list_contains::PreparedSetData; - use crate::scalar_fn::fns::literal::Literal; - use crate::scalar_fn::fns::operators::Operator; + use crate::scalar_fn::builtin::binary::Binary; + use crate::scalar_fn::builtin::literal::Literal; + use crate::scalar_fn::builtin::operators::Operator; use crate::validity::Validity; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); @@ -177,37 +171,6 @@ mod tests { ); } - #[test] - fn test_scalar_fn_without_children_requires_explicit_len() -> VortexResult<()> { - let list = Scalar::list( - DType::Primitive(PType::I32, Nullability::NonNullable), - vec![Scalar::from(1i32)], - Nullability::NonNullable, - ); - let scalar_fn = - TypedScalarFnInstance::new(PreparedSet, PreparedSetData::try_new(list.clone())?) - .erased(); - - let Err(err) = ScalarFnArray::try_new(scalar_fn.clone(), vec![]) else { - panic!("ScalarFnArray::try_new should reject zero children"); - }; - assert!( - err.to_string() - .contains("ScalarFnArray length cannot be inferred without children") - ); - - let scalar_fn_array = ScalarFnArray::try_new_with_len(scalar_fn, vec![], 3)?; - assert_eq!(scalar_fn_array.len(), 3); - assert_eq!(scalar_fn_array.child_count(), 0); - assert_arrays_eq!( - scalar_fn_array, - ConstantArray::new(list, 3), - &mut SESSION.create_execution_ctx() - ); - - Ok(()) - } - #[test] fn rejects_literal() { let scalar_fn = TypedScalarFnInstance::new(Literal, Scalar::from(1i32)).erased(); diff --git a/vortex-array/src/arrays/slice/rules.rs b/vortex-array/src/arrays/slice/rules.rs index 62ed8892541..25e96b5c15b 100644 --- a/vortex-array/src/arrays/slice/rules.rs +++ b/vortex-array/src/arrays/slice/rules.rs @@ -13,7 +13,7 @@ use crate::arrays::slice::SliceReduceAdaptor; use crate::builtins::ArrayBuiltins; use crate::optimizer::rules::ArrayParentReduceRule; use crate::optimizer::rules::ParentRuleSet; -use crate::scalar_fn::fns::get_item::GetItem; +use crate::scalar_fn::builtin::get_item::GetItem; pub(super) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ ParentRuleSet::lift(&SliceReduceAdaptor(Slice)), diff --git a/vortex-array/src/arrays/slice/vtable.rs b/vortex-array/src/arrays/slice/vtable.rs index 3b2351a065c..043869059b4 100644 --- a/vortex-array/src/arrays/slice/vtable.rs +++ b/vortex-array/src/arrays/slice/vtable.rs @@ -39,6 +39,7 @@ use crate::buffer::BufferHandle; use crate::dtype::DType; use crate::executor::ExecutionCtx; use crate::executor::ExecutionResult; +use crate::optimizer::ArrayOptimizer; use crate::require_child; use crate::scalar::Scalar; use crate::serde::ArrayChildren; @@ -150,14 +151,15 @@ impl VTable for Slice { vortex_bail!("Slice array is not serializable") } - fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { + fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { let array = require_child!(array, array.child(), SliceSlots::CHILD => AnyCanonical); debug_assert!(array.child().is_canonical()); // TODO(ngates): we should inline canonical slice logic here. array .child() - .slice(array.range.clone()) + .slice(array.range.clone())? + .optimize(ctx.session()) .map(ExecutionResult::done) } diff --git a/vortex-array/src/arrays/struct_/compute/cast.rs b/vortex-array/src/arrays/struct_/compute/cast.rs index e73ea26446b..8daf31428ef 100644 --- a/vortex-array/src/arrays/struct_/compute/cast.rs +++ b/vortex-array/src/arrays/struct_/compute/cast.rs @@ -23,7 +23,7 @@ use crate::kernel::ExecuteParentKernel; use crate::optimizer::kernels::ArrayKernelsExt; use crate::scalar::Scalar; use crate::scalar_fn::ScalarFnVTable; -use crate::scalar_fn::fns::cast::Cast; +use crate::scalar_fn::builtin::cast::Cast; pub(crate) fn initialize(session: &VortexSession) { let kernels = session.kernels(); @@ -153,7 +153,7 @@ mod tests { use crate::optimizer::kernels::ExecuteParentFn; use crate::optimizer::kernels::KernelSession; use crate::scalar::Scalar; - use crate::scalar_fn::fns::cast::Cast; + use crate::scalar_fn::builtin::cast::Cast; use crate::validity::Validity; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); diff --git a/vortex-array/src/arrays/struct_/compute/mask.rs b/vortex-array/src/arrays/struct_/compute/mask.rs index 68db6166455..758ede1bbf0 100644 --- a/vortex-array/src/arrays/struct_/compute/mask.rs +++ b/vortex-array/src/arrays/struct_/compute/mask.rs @@ -9,7 +9,7 @@ use crate::array::ArrayView; use crate::arrays::Struct; use crate::arrays::StructArray; use crate::arrays::struct_::StructArrayExt; -use crate::scalar_fn::fns::mask::MaskReduce; +use crate::scalar_fn::builtin::mask::MaskReduce; use crate::validity::Validity; impl MaskReduce for Struct { diff --git a/vortex-array/src/arrays/struct_/compute/mod.rs b/vortex-array/src/arrays/struct_/compute/mod.rs index 1d528448ad3..6edc2c92e19 100644 --- a/vortex-array/src/arrays/struct_/compute/mod.rs +++ b/vortex-array/src/arrays/struct_/compute/mod.rs @@ -6,7 +6,6 @@ mod mask; pub(crate) mod rules; mod slice; mod take; -mod zip; #[cfg(test)] mod tests { diff --git a/vortex-array/src/arrays/struct_/compute/rules.rs b/vortex-array/src/arrays/struct_/compute/rules.rs index f7bcd750030..a44ea316399 100644 --- a/vortex-array/src/arrays/struct_/compute/rules.rs +++ b/vortex-array/src/arrays/struct_/compute/rules.rs @@ -21,9 +21,9 @@ use crate::matcher::Matcher; use crate::optimizer::rules::ArrayParentReduceRule; use crate::optimizer::rules::ParentRuleSet; use crate::scalar::Scalar; -use crate::scalar_fn::fns::cast::Cast; -use crate::scalar_fn::fns::get_item::GetItem; -use crate::scalar_fn::fns::mask::MaskReduceAdaptor; +use crate::scalar_fn::builtin::cast::Cast; +use crate::scalar_fn::builtin::get_item::GetItem; +use crate::scalar_fn::builtin::mask::MaskReduceAdaptor; use crate::validity::Validity; pub(crate) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ @@ -222,8 +222,10 @@ mod tests { Nullability::NonNullable, ); + // Without the session-registered rule the cast stays a scalar function. + let session = VortexSession::empty().with_some(KernelSession::empty()); let cast = source.cast(target).unwrap(); - let optimized = cast.optimize().unwrap(); + let optimized = cast.optimize(&session).unwrap(); assert!(optimized.is::()); } @@ -248,10 +250,10 @@ mod tests { ); let cast = source.cast(target).unwrap(); - assert!(cast.optimize_ctx(&SESSION).unwrap().is::()); + assert!(cast.optimize(&SESSION).unwrap().is::()); let session = VortexSession::empty().with_some(KernelSession::empty()); - let optimized = cast.optimize_ctx(&session).unwrap(); + let optimized = cast.optimize(&session).unwrap(); assert!(optimized.is::()); } @@ -273,7 +275,7 @@ mod tests { let result = source .into_array() .cast(DType::Primitive(PType::I32, Nullability::NonNullable)) - .and_then(|arr| arr.optimize_ctx(&SESSION)); + .and_then(|arr| arr.optimize(&SESSION)); // Whether this errors or succeeds depends on execution, but the key invariant is that the // optimizer rule does not panic. @@ -393,7 +395,7 @@ mod tests { let arr = source .into_array() .cast(target) - .and_then(|arr| arr.optimize_ctx(&SESSION)); + .and_then(|arr| arr.optimize(&SESSION)); assert!(arr.is_err()); } } diff --git a/vortex-array/src/arrays/struct_/mod.rs b/vortex-array/src/arrays/struct_/mod.rs index ab533dd8cdb..faaf03dc23c 100644 --- a/vortex-array/src/arrays/struct_/mod.rs +++ b/vortex-array/src/arrays/struct_/mod.rs @@ -15,7 +15,6 @@ pub use vtable::Struct; pub(crate) fn initialize(session: &vortex_session::VortexSession) { compute::cast::initialize(session); - vtable::initialize(session); } #[cfg(test)] diff --git a/vortex-array/src/arrays/struct_/vtable/kernel.rs b/vortex-array/src/arrays/struct_/vtable/kernel.rs deleted file mode 100644 index 1e58da2d0d0..00000000000 --- a/vortex-array/src/arrays/struct_/vtable/kernel.rs +++ /dev/null @@ -1,15 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use vortex_session::VortexSession; - -use crate::arrays::Struct; -use crate::optimizer::kernels::ArrayKernelsExt; -use crate::scalar_fn::ScalarFnVTable; -use crate::scalar_fn::fns::zip::Zip; -use crate::scalar_fn::fns::zip::ZipExecuteAdaptor; - -pub(crate) fn initialize(session: &VortexSession) { - let kernels = session.kernels(); - kernels.register_execute_parent_kernel(Zip.id(), Struct, ZipExecuteAdaptor(Struct)); -} diff --git a/vortex-array/src/arrays/struct_/vtable/mod.rs b/vortex-array/src/arrays/struct_/vtable/mod.rs index 15cf128b8e9..398b5db53a4 100644 --- a/vortex-array/src/arrays/struct_/vtable/mod.rs +++ b/vortex-array/src/arrays/struct_/vtable/mod.rs @@ -26,7 +26,6 @@ use crate::builders::StructBuilder; use crate::dtype::DType; use crate::serde::ArrayChildren; use crate::validity::Validity; -mod kernel; mod operations; mod validity; @@ -37,10 +36,6 @@ use crate::array::ArrayId; /// A [`Struct`]-encoded Vortex array. pub type StructArray = Array; -pub(crate) fn initialize(session: &VortexSession) { - kernel::initialize(session); -} - impl VTable for Struct { type TypedArrayData = EmptyArrayData; diff --git a/vortex-array/src/arrays/union/compute/mask.rs b/vortex-array/src/arrays/union/compute/mask.rs index 63bb9ae60f7..18f74b5daf8 100644 --- a/vortex-array/src/arrays/union/compute/mask.rs +++ b/vortex-array/src/arrays/union/compute/mask.rs @@ -11,7 +11,7 @@ use crate::arrays::UnionArray; use crate::arrays::union::UnionArrayExt; use crate::arrays::union::UnionArraySlotsExt; use crate::builtins::ArrayBuiltins; -use crate::scalar_fn::fns::mask::MaskReduce; +use crate::scalar_fn::builtin::mask::MaskReduce; impl MaskReduce for Union { fn mask(array: ArrayView<'_, Union>, mask: &ArrayRef) -> VortexResult> { diff --git a/vortex-array/src/arrays/union/compute/rules.rs b/vortex-array/src/arrays/union/compute/rules.rs index 921d66a24d3..33123cc7793 100644 --- a/vortex-array/src/arrays/union/compute/rules.rs +++ b/vortex-array/src/arrays/union/compute/rules.rs @@ -5,7 +5,7 @@ use crate::arrays::Union; use crate::arrays::dict::TakeReduceAdaptor; use crate::arrays::slice::SliceReduceAdaptor; use crate::optimizer::rules::ParentRuleSet; -use crate::scalar_fn::fns::mask::MaskReduceAdaptor; +use crate::scalar_fn::builtin::mask::MaskReduceAdaptor; pub(crate) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ ParentRuleSet::lift(&MaskReduceAdaptor(Union)), diff --git a/vortex-array/src/arrays/union/tests/take.rs b/vortex-array/src/arrays/union/tests/take.rs index ba0de72c01f..ea478d92f88 100644 --- a/vortex-array/src/arrays/union/tests/take.rs +++ b/vortex-array/src/arrays/union/tests/take.rs @@ -25,6 +25,7 @@ use crate::compute::conformance::take::test_take_conformance; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; +use crate::optimizer::ArrayOptimizer; use crate::scalar::Scalar; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); @@ -36,6 +37,7 @@ fn take(array: &UnionArray, indices: ArrayRef) -> VortexResult { .clone() .into_array() .take(indices)? + .optimize(&SESSION)? .as_::() .into_owned()) } diff --git a/vortex-array/src/arrays/varbin/compute/cast.rs b/vortex-array/src/arrays/varbin/compute/cast.rs index 09b68449103..c8449d8b2be 100644 --- a/vortex-array/src/arrays/varbin/compute/cast.rs +++ b/vortex-array/src/arrays/varbin/compute/cast.rs @@ -11,8 +11,8 @@ use crate::arrays::VarBin; use crate::arrays::VarBinArray; use crate::arrays::varbin::VarBinArraySlotsExt; use crate::dtype::DType; -use crate::scalar_fn::fns::cast::CastKernel; -use crate::scalar_fn::fns::cast::CastReduce; +use crate::scalar_fn::builtin::cast::CastKernel; +use crate::scalar_fn::builtin::cast::CastReduce; use crate::validity::Validity; fn build_with_validity( diff --git a/vortex-array/src/arrays/varbin/compute/compare.rs b/vortex-array/src/arrays/varbin/compute/compare.rs index ab2abce91c9..7559b1a30b3 100644 --- a/vortex-array/src/arrays/varbin/compute/compare.rs +++ b/vortex-array/src/arrays/varbin/compute/compare.rs @@ -19,8 +19,8 @@ use crate::arrays::varbin::VarBinArraySlotsExt; use crate::dtype::DType; use crate::dtype::IntegerPType; use crate::match_each_integer_ptype; -use crate::scalar_fn::fns::binary::CompareKernel; -use crate::scalar_fn::fns::operators::CompareOperator; +use crate::scalar_fn::builtin::binary::CompareKernel; +use crate::scalar_fn::builtin::operators::CompareOperator; // This implementation exists so we can compare against a constant in encoded space, without // canonicalizing the VarBin array to VarBinView. @@ -189,7 +189,7 @@ mod test { use crate::dtype::DType; use crate::dtype::Nullability; use crate::scalar::Scalar; - use crate::scalar_fn::fns::operators::Operator; + use crate::scalar_fn::builtin::operators::Operator; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); @@ -261,7 +261,7 @@ mod tests { use crate::dtype::DType; use crate::dtype::Nullability; use crate::scalar::Scalar; - use crate::scalar_fn::fns::operators::Operator; + use crate::scalar_fn::builtin::operators::Operator; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); diff --git a/vortex-array/src/arrays/varbin/compute/mask.rs b/vortex-array/src/arrays/varbin/compute/mask.rs index 5f75ccc3839..14a48924f71 100644 --- a/vortex-array/src/arrays/varbin/compute/mask.rs +++ b/vortex-array/src/arrays/varbin/compute/mask.rs @@ -9,7 +9,7 @@ use crate::array::ArrayView; use crate::arrays::VarBin; use crate::arrays::VarBinArray; use crate::arrays::varbin::VarBinArraySlotsExt; -use crate::scalar_fn::fns::mask::MaskReduce; +use crate::scalar_fn::builtin::mask::MaskReduce; use crate::validity::Validity; impl MaskReduce for VarBin { diff --git a/vortex-array/src/arrays/varbin/compute/rules.rs b/vortex-array/src/arrays/varbin/compute/rules.rs index 89e22ce0d20..75c3bd6086a 100644 --- a/vortex-array/src/arrays/varbin/compute/rules.rs +++ b/vortex-array/src/arrays/varbin/compute/rules.rs @@ -4,8 +4,8 @@ use crate::arrays::VarBin; use crate::arrays::slice::SliceReduceAdaptor; use crate::optimizer::rules::ParentRuleSet; -use crate::scalar_fn::fns::cast::CastReduceAdaptor; -use crate::scalar_fn::fns::mask::MaskReduceAdaptor; +use crate::scalar_fn::builtin::cast::CastReduceAdaptor; +use crate::scalar_fn::builtin::mask::MaskReduceAdaptor; pub(crate) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ ParentRuleSet::lift(&CastReduceAdaptor(VarBin)), diff --git a/vortex-array/src/arrays/varbin/vtable/kernel.rs b/vortex-array/src/arrays/varbin/vtable/kernel.rs index 9e80abd1037..c23d5e2af69 100644 --- a/vortex-array/src/arrays/varbin/vtable/kernel.rs +++ b/vortex-array/src/arrays/varbin/vtable/kernel.rs @@ -11,10 +11,10 @@ use crate::arrays::dict::TakeExecuteAdaptor; use crate::arrays::filter::FilterExecuteAdaptor; use crate::optimizer::kernels::ArrayKernelsExt; use crate::scalar_fn::ScalarFnVTable; -use crate::scalar_fn::fns::binary::Binary; -use crate::scalar_fn::fns::binary::CompareExecuteAdaptor; -use crate::scalar_fn::fns::cast::Cast; -use crate::scalar_fn::fns::cast::CastExecuteAdaptor; +use crate::scalar_fn::builtin::binary::Binary; +use crate::scalar_fn::builtin::binary::CompareExecuteAdaptor; +use crate::scalar_fn::builtin::cast::Cast; +use crate::scalar_fn::builtin::cast::CastExecuteAdaptor; pub(crate) fn initialize(session: &VortexSession) { let kernels = session.kernels(); diff --git a/vortex-array/src/arrays/varbinview/compute/cast.rs b/vortex-array/src/arrays/varbinview/compute/cast.rs index 19b669f2eac..3ee2f80e3a3 100644 --- a/vortex-array/src/arrays/varbinview/compute/cast.rs +++ b/vortex-array/src/arrays/varbinview/compute/cast.rs @@ -12,8 +12,8 @@ use crate::array::ArrayView; use crate::arrays::VarBinView; use crate::arrays::VarBinViewArray; use crate::dtype::DType; -use crate::scalar_fn::fns::cast::CastKernel; -use crate::scalar_fn::fns::cast::CastReduce; +use crate::scalar_fn::builtin::cast::CastKernel; +use crate::scalar_fn::builtin::cast::CastReduce; use crate::validity::Validity; fn build_with_validity( diff --git a/vortex-array/src/arrays/varbinview/compute/fill_null.rs b/vortex-array/src/arrays/varbinview/compute/fill_null.rs index cd3c9bfd10e..322f68cd0ff 100644 --- a/vortex-array/src/arrays/varbinview/compute/fill_null.rs +++ b/vortex-array/src/arrays/varbinview/compute/fill_null.rs @@ -20,7 +20,7 @@ use crate::arrays::VarBinViewArray; use crate::arrays::varbinview::BinaryView; use crate::dtype::DType; use crate::scalar::Scalar; -use crate::scalar_fn::fns::fill_null::FillNullKernel; +use crate::scalar_fn::builtin::fill_null::FillNullKernel; use crate::validity::Validity; // Points every null row at one view of the fill value. The data buffers are reused as they are, diff --git a/vortex-array/src/arrays/varbinview/compute/mask.rs b/vortex-array/src/arrays/varbinview/compute/mask.rs index 8869f816600..d71b2b3b48c 100644 --- a/vortex-array/src/arrays/varbinview/compute/mask.rs +++ b/vortex-array/src/arrays/varbinview/compute/mask.rs @@ -10,7 +10,7 @@ use crate::IntoArray; use crate::array::ArrayView; use crate::arrays::VarBinView; use crate::arrays::VarBinViewArray; -use crate::scalar_fn::fns::mask::MaskReduce; +use crate::scalar_fn::builtin::mask::MaskReduce; use crate::validity::Validity; impl MaskReduce for VarBinView { diff --git a/vortex-array/src/arrays/varbinview/compute/mod.rs b/vortex-array/src/arrays/varbinview/compute/mod.rs index dab35219b5c..d55c77f8e4a 100644 --- a/vortex-array/src/arrays/varbinview/compute/mod.rs +++ b/vortex-array/src/arrays/varbinview/compute/mod.rs @@ -7,7 +7,6 @@ mod mask; pub(crate) mod rules; mod slice; mod take; -mod zip; #[cfg(test)] mod tests { diff --git a/vortex-array/src/arrays/varbinview/compute/rules.rs b/vortex-array/src/arrays/varbinview/compute/rules.rs index 5ec24dca7de..8aa16def538 100644 --- a/vortex-array/src/arrays/varbinview/compute/rules.rs +++ b/vortex-array/src/arrays/varbinview/compute/rules.rs @@ -3,8 +3,8 @@ use crate::arrays::VarBinView; use crate::arrays::slice::SliceReduceAdaptor; use crate::optimizer::rules::ParentRuleSet; -use crate::scalar_fn::fns::cast::CastReduceAdaptor; -use crate::scalar_fn::fns::mask::MaskReduceAdaptor; +use crate::scalar_fn::builtin::cast::CastReduceAdaptor; +use crate::scalar_fn::builtin::mask::MaskReduceAdaptor; pub(crate) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ ParentRuleSet::lift(&CastReduceAdaptor(VarBinView)), diff --git a/vortex-array/src/arrays/varbinview/vtable/kernel.rs b/vortex-array/src/arrays/varbinview/vtable/kernel.rs index 44dfcf5c028..0b22dbe73c3 100644 --- a/vortex-array/src/arrays/varbinview/vtable/kernel.rs +++ b/vortex-array/src/arrays/varbinview/vtable/kernel.rs @@ -9,12 +9,10 @@ use crate::arrays::VarBinView; use crate::arrays::dict::TakeExecuteAdaptor; use crate::optimizer::kernels::ArrayKernelsExt; use crate::scalar_fn::ScalarFnVTable; -use crate::scalar_fn::fns::cast::Cast; -use crate::scalar_fn::fns::cast::CastExecuteAdaptor; -use crate::scalar_fn::fns::fill_null::FillNull; -use crate::scalar_fn::fns::fill_null::FillNullExecuteAdaptor; -use crate::scalar_fn::fns::zip::Zip; -use crate::scalar_fn::fns::zip::ZipExecuteAdaptor; +use crate::scalar_fn::builtin::cast::Cast; +use crate::scalar_fn::builtin::cast::CastExecuteAdaptor; +use crate::scalar_fn::builtin::fill_null::FillNull; +use crate::scalar_fn::builtin::fill_null::FillNullExecuteAdaptor; pub(crate) fn initialize(session: &VortexSession) { let kernels = session.kernels(); @@ -25,5 +23,4 @@ pub(crate) fn initialize(session: &VortexSession) { FillNullExecuteAdaptor(VarBinView), ); kernels.register_execute_parent_kernel(Dict.id(), VarBinView, TakeExecuteAdaptor(VarBinView)); - kernels.register_execute_parent_kernel(Zip.id(), VarBinView, ZipExecuteAdaptor(VarBinView)); } diff --git a/vortex-array/src/arrays/variant/mod.rs b/vortex-array/src/arrays/variant/mod.rs index 3a44a11019d..e92449c15bf 100644 --- a/vortex-array/src/arrays/variant/mod.rs +++ b/vortex-array/src/arrays/variant/mod.rs @@ -10,11 +10,7 @@ use vortex_error::vortex_ensure; pub use self::vtable::Variant; pub use self::vtable::VariantArray; - -pub(crate) fn initialize(session: &vortex_session::VortexSession) { - vtable::initialize(session); -} - +pub use self::vtable::merge_typed_scalar_as_variant; use crate::ArrayRef; use crate::array::Array; use crate::array::ArrayParts; @@ -92,18 +88,13 @@ mod tests { use crate::arrays::ChunkedArray; use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; - use crate::arrays::StructArray; use crate::arrays::VariantArray; use crate::arrays::variant::VariantArraySlotsExt; use crate::assert_arrays_eq; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::Nullability; - use crate::dtype::PType; - use crate::expr::root; - use crate::expr::variant_get; use crate::scalar::Scalar; - use crate::scalar_fn::fns::variant_get::VariantPath; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); @@ -376,71 +367,4 @@ mod tests { assert_variant_core_rows(&masked, &[Some(1), None, None, Some(4)]) } - - #[test] - fn variant_get_keeps_valid_shredded_rows_for_matching_dtype() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let core_storage = row_storage([1, 2, 3])?; - let shredded = StructArray::try_from_iter([( - "a", - PrimitiveArray::from_iter([10i32, 20, 30]).into_array(), - )])?; - let variant = VariantArray::try_new(core_storage, Some(shredded.into_array()))?; - let expr = variant_get( - root(), - VariantPath::field("a"), - Some(DType::Primitive(PType::I32, Nullability::NonNullable)), - ); - - let result = variant.into_array().apply(&expr)?; - - assert_arrays_eq!( - result, - PrimitiveArray::from_option_iter([Some(10i32), Some(20), Some(30)]), - &mut ctx - ); - Ok(()) - } - - #[test] - fn variant_get_treats_value_and_typed_value_as_logical_field_names() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let core_storage = row_storage([1, 2, 3])?; - let shredded = StructArray::try_from_iter([ - ( - "value", - PrimitiveArray::from_iter([10i32, 20, 30]).into_array(), - ), - ( - "typed_value", - PrimitiveArray::from_iter([40i32, 50, 60]).into_array(), - ), - ])?; - let variant = VariantArray::try_new(core_storage, Some(shredded.into_array()))?; - - let value_expr = variant_get( - root(), - VariantPath::field("value"), - Some(DType::Primitive(PType::I32, Nullability::NonNullable)), - ); - let value_result = variant.clone().into_array().apply(&value_expr)?; - assert_arrays_eq!( - value_result, - PrimitiveArray::from_option_iter([Some(10i32), Some(20), Some(30)]), - &mut ctx - ); - - let typed_value_expr = variant_get( - root(), - VariantPath::field("typed_value"), - Some(DType::Primitive(PType::I32, Nullability::NonNullable)), - ); - let typed_value_result = variant.into_array().apply(&typed_value_expr)?; - assert_arrays_eq!( - typed_value_result, - PrimitiveArray::from_option_iter([Some(40i32), Some(50), Some(60)]), - &mut ctx - ); - Ok(()) - } } diff --git a/vortex-array/src/arrays/variant/vtable/mod.rs b/vortex-array/src/arrays/variant/vtable/mod.rs index 9e07169012a..a1f7724915e 100644 --- a/vortex-array/src/arrays/variant/vtable/mod.rs +++ b/vortex-array/src/arrays/variant/vtable/mod.rs @@ -1,7 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -mod kernel; mod operations; mod validity; @@ -41,10 +40,6 @@ use crate::serde::ArrayChildren; /// A [`Variant`]-encoded Vortex array. pub type VariantArray = Array; -pub(crate) fn initialize(session: &VortexSession) { - kernel::initialize(session); -} - #[derive(Clone, Debug)] pub struct Variant; @@ -198,7 +193,9 @@ impl VTable for Variant { } } -fn merge_typed_scalar_as_variant( +/// Merges a shredded `typed_scalar` with the unshredded `fallback_scalar` of the same row into a +/// variant scalar of `dtype`. +pub fn merge_typed_scalar_as_variant( typed_scalar: Scalar, fallback_scalar: Option, dtype: &DType, diff --git a/vortex-array/src/builders/lazy_null_builder.rs b/vortex-array/src/builders/lazy_null_builder.rs index d46d731cf25..59684a4d101 100644 --- a/vortex-array/src/builders/lazy_null_builder.rs +++ b/vortex-array/src/builders/lazy_null_builder.rs @@ -106,12 +106,18 @@ impl LazyBitBufferBuilder { } /// Returns the current length of the builder. + /// Whether no entries have been appended. + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + pub fn len(&self) -> usize { // self.len is the length of the builder if the inner buffer is not materialized self.inner.as_ref().map(|i| i.len()).unwrap_or(self.len) } - fn finish(&mut self) -> Option { + /// The validity bits appended so far, or `None` when every entry is valid. + pub fn finish(&mut self) -> Option { self.len = 0; self.inner.take().map(|b| b.freeze()) } diff --git a/vortex-array/src/builders/mod.rs b/vortex-array/src/builders/mod.rs index 480c3f6a701..9590ae6fe56 100644 --- a/vortex-array/src/builders/mod.rs +++ b/vortex-array/src/builders/mod.rs @@ -54,7 +54,7 @@ use crate::memory::BufferAllocatorRef; use crate::scalar::Scalar; mod lazy_null_builder; -pub(crate) use lazy_null_builder::LazyBitBufferBuilder; +pub use lazy_null_builder::LazyBitBufferBuilder; mod bool; mod child; diff --git a/vortex-array/src/builders/varbinview.rs b/vortex-array/src/builders/varbinview.rs index 8751802aeb0..475fdc1497a 100644 --- a/vortex-array/src/builders/varbinview.rs +++ b/vortex-array/src/builders/varbinview.rs @@ -1074,14 +1074,13 @@ impl DeduplicatedBuffers { .collect() } - pub(crate) fn extend_from_iter( - &mut self, - buffers: impl Iterator, - ) -> Vec { + /// Pushes each of `buffers`, returning the index of each in the deduplicated buffers. + pub fn extend_from_iter(&mut self, buffers: impl Iterator) -> Vec { buffers.map(|buffer| self.push(buffer)).collect() } - pub(crate) fn finish(self) -> Arc<[ByteBuffer]> { + /// The deduplicated buffers, in the order they were first pushed. + pub fn finish(self) -> Arc<[ByteBuffer]> { Arc::from(self.buffers) } } diff --git a/vortex-array/src/builtins.rs b/vortex-array/src/builtins.rs index 42b696c2f47..1b164180865 100644 --- a/vortex-array/src/builtins.rs +++ b/vortex-array/src/builtins.rs @@ -18,24 +18,18 @@ use crate::arrays::InterleaveArray; use crate::dtype::DType; use crate::dtype::FieldName; use crate::expr::Expression; -use crate::optimizer::ArrayOptimizer; use crate::scalar::Scalar; use crate::scalar_fn::EmptyOptions; use crate::scalar_fn::ScalarFnVTableExt; -use crate::scalar_fn::fns::between::Between; -use crate::scalar_fn::fns::between::BetweenOptions; -use crate::scalar_fn::fns::binary::Binary; -use crate::scalar_fn::fns::cast::Cast; -use crate::scalar_fn::fns::fill_null::FillNull; -use crate::scalar_fn::fns::get_item::GetItem; -use crate::scalar_fn::fns::is_not_null::IsNotNull; -use crate::scalar_fn::fns::is_null::IsNull; -use crate::scalar_fn::fns::list_contains::ListContains; -use crate::scalar_fn::fns::list_contains::ListContainsOptions; -use crate::scalar_fn::fns::mask::Mask; -use crate::scalar_fn::fns::not::Not; -use crate::scalar_fn::fns::operators::Operator; -use crate::scalar_fn::fns::zip::Zip; +use crate::scalar_fn::builtin::binary::Binary; +use crate::scalar_fn::builtin::cast::Cast; +use crate::scalar_fn::builtin::fill_null::FillNull; +use crate::scalar_fn::builtin::get_item::GetItem; +use crate::scalar_fn::builtin::is_not_null::IsNotNull; +use crate::scalar_fn::builtin::is_null::IsNull; +use crate::scalar_fn::builtin::mask::Mask; +use crate::scalar_fn::builtin::not::Not; +use crate::scalar_fn::builtin::operators::Operator; /// A collection of built-in scalar functions that can be applied to expressions or arrays. pub trait ExprBuiltins: Sized { @@ -62,12 +56,6 @@ pub trait ExprBuiltins: Sized { /// Boolean negation. fn not(&self) -> VortexResult; - /// Check if a list contains a value. - fn list_contains(&self, value: Expression) -> VortexResult; - - /// Conditional selection: `result[i] = if mask[i] then if_true[i] else if_false[i]`. - fn zip(&self, if_true: Expression, if_false: Expression) -> VortexResult; - // TODO(joe): add an `interleave` expression builtin mirroring `ArrayBuiltins::interleave`. /// Apply a binary operator to this expression and another. @@ -103,14 +91,6 @@ impl ExprBuiltins for Expression { Not.try_new_expr(EmptyOptions, [self.clone()]) } - fn list_contains(&self, value: Expression) -> VortexResult { - ListContains.try_new_expr(ListContainsOptions::default(), [self.clone(), value]) - } - - fn zip(&self, if_true: Expression, if_false: Expression) -> VortexResult { - Zip.try_new_expr(EmptyOptions, [if_true, if_false, self.clone()]) - } - fn binary(&self, rhs: Expression, op: Operator) -> VortexResult { Binary.try_new_expr(op, [self.clone(), rhs]) } @@ -140,9 +120,6 @@ pub trait ArrayBuiltins: Sized { /// Boolean negation. fn not(&self) -> VortexResult; - /// Conditional selection: `result[i] = if mask[i] then if_true[i] else if_false[i]`. - fn zip(&self, if_true: ArrayRef, if_false: ArrayRef) -> VortexResult; - /// Random-access gather by `(array_index, row_index)`: output row `i` is taken from /// `values[array_indices[i]][row_indices[i]]`, where `self` is the (non-nullable) /// `array_indices` selector and `row_indices` names the position within the selected value. @@ -153,19 +130,8 @@ pub trait ArrayBuiltins: Sized { row_indices: ArrayRef, ) -> VortexResult; - /// Check if a list contains a value. - fn list_contains(&self, value: ArrayRef) -> VortexResult; - /// Apply a binary operator to this array and another. fn binary(&self, rhs: ArrayRef, op: Operator) -> VortexResult; - - /// Compare a values between lower VortexResult; } impl ArrayBuiltins for ArrayRef { @@ -173,7 +139,7 @@ impl ArrayBuiltins for ArrayRef { if self.dtype() == &dtype { return Ok(self.clone()); } - Cast::new(self.clone(), dtype).into_array().optimize() + Ok(Cast::new(self.clone(), dtype).into_array()) } fn fill_null(&self, fill_value: impl Into) -> VortexResult { @@ -181,38 +147,31 @@ impl ArrayBuiltins for ArrayRef { if !self.dtype().is_nullable() { return self.cast(fill_value.dtype().clone()); } - FillNull::try_new( + Ok(FillNull::try_new( self.clone(), ConstantArray::new(fill_value, self.len()).into_array(), )? - .into_array() - .optimize() + .into_array()) } fn get_item(&self, field_name: impl Into) -> VortexResult { - GetItem::try_new(self.clone(), field_name)? - .into_array() - .optimize() + Ok(GetItem::try_new(self.clone(), field_name)?.into_array()) } fn is_null(&self) -> VortexResult { - IsNull::new(self.clone()).into_array().optimize() + Ok(IsNull::new(self.clone()).into_array()) } fn is_not_null(&self) -> VortexResult { - IsNotNull::new(self.clone()).into_array().optimize() + Ok(IsNotNull::new(self.clone()).into_array()) } fn mask(self, mask: ArrayRef) -> VortexResult { - Mask::try_new(self, mask)?.into_array().optimize() + Ok(Mask::try_new(self, mask)?.into_array()) } fn not(&self) -> VortexResult { - Not::try_new(self.clone())?.into_array().optimize() - } - - fn zip(&self, if_true: ArrayRef, if_false: ArrayRef) -> VortexResult { - Ok(Zip::try_new(if_true, if_false, self.clone())?.into_array()) + Ok(Not::try_new(self.clone())?.into_array()) } fn interleave( @@ -226,26 +185,7 @@ impl ArrayBuiltins for ArrayRef { ) } - fn list_contains(&self, value: ArrayRef) -> VortexResult { - ListContains::try_new(self.clone(), value)? - .into_array() - .optimize() - } - fn binary(&self, rhs: ArrayRef, op: Operator) -> VortexResult { - Binary::try_new(self.clone(), rhs, op)? - .into_array() - .optimize() - } - - fn between( - self, - lower: ArrayRef, - upper: ArrayRef, - options: BetweenOptions, - ) -> VortexResult { - Between::try_new(self, lower, upper, options)? - .into_array() - .optimize() + Ok(Binary::try_new(self.clone(), rhs, op)?.into_array()) } } diff --git a/vortex-array/src/compute/conformance/binary_numeric.rs b/vortex-array/src/compute/conformance/binary_numeric.rs index 29f3c3a6531..aa6847902e9 100644 --- a/vortex-array/src/compute/conformance/binary_numeric.rs +++ b/vortex-array/src/compute/conformance/binary_numeric.rs @@ -52,7 +52,7 @@ use crate::scalar::DecimalValue; use crate::scalar::NumericOperator; use crate::scalar::PrimitiveScalar; use crate::scalar::Scalar; -use crate::scalar_fn::fns::binary::numeric_op_result_decimal_dtype; +use crate::scalar_fn::builtin::binary::numeric_op_result_decimal_dtype; fn to_vec_of_scalar(array: &ArrayRef, ctx: &mut ExecutionCtx) -> Vec { // Not fast, but obviously correct diff --git a/vortex-array/src/compute/conformance/consistency.rs b/vortex-array/src/compute/conformance/consistency.rs index 6d1b919e51d..ca8099f8011 100644 --- a/vortex-array/src/compute/conformance/consistency.rs +++ b/vortex-array/src/compute/conformance/consistency.rs @@ -43,7 +43,7 @@ use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; -use crate::scalar_fn::fns::operators::Operator; +use crate::scalar_fn::builtin::operators::Operator; /// Tests that filter and take operations produce consistent results. /// diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index afc5464b03e..f2ac24eab49 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -171,7 +171,8 @@ impl ArrayRef { self, ctx: &mut ExecutionCtx, ) -> VortexResult { - let mut current_array = self; + // Constructors build lazy arrays, so optimize the whole tree before executing it. + let mut current_array = self.optimize_recursive(ctx.session())?; let mut current_builder: Option> = None; let mut stack: Vec = Vec::new(); let execute_parent_kernels = Arc::clone(&ctx.execute_parent_kernels); @@ -180,9 +181,9 @@ impl ArrayRef { trace_op!(record_execute_until_start::(¤t_array)); - for _iteration in 0..max_iterations { + for iteration in 0..max_iterations { trace_op!(record_execute_until_iteration( - _iteration, + iteration, ¤t_array, stack .last() @@ -206,6 +207,11 @@ impl ArrayRef { current_builder.is_none(), "root activation should not retain a builder" ); + // Kernels build their results with lazy constructors too, such as a + // filtered validity, so optimize whatever execution produced. + if iteration > 0 { + current_array = current_array.optimize_recursive(ctx.session())?; + } trace_op!(record_execute_until_return(¤t_array)); return Ok(current_array); } @@ -242,7 +248,7 @@ impl ArrayRef { } { let frame = stack.pop().vortex_expect("just peeked"); - let optimized = result.optimize_ctx(ctx.session())?; + let optimized = result.optimize(ctx.session())?; trace_op!(record_execute_optimized(&result, &optimized)); current_array = optimized; current_builder = frame.parent_builder; @@ -259,7 +265,7 @@ impl ArrayRef { if current_builder.is_none() && let Some(rewritten) = try_execute_parent(¤t_array, kernels, ctx)? { - let optimized = rewritten.optimize_ctx(ctx.session())?; + let optimized = rewritten.optimize(ctx.session())?; trace_op!(record_execute_optimized(&rewritten, &optimized)); current_array = optimized; continue; @@ -485,34 +491,13 @@ impl Executable for ArrayRef { } trace_op!(record_single_step_phase_none("canonical", &array)); - if let Some(reduced) = array.reduce()? { - reduced.statistics().inherit_from(array.statistics()); - trace_op!(record_single_step_applied("reduce", &array, &reduced)); - return Ok(reduced); - } - trace_op!(record_single_step_phase_none("reduce", &array)); - - for (slot_idx, slot) in array.slots().iter().enumerate() { - let Some(child) = slot else { continue }; - if let Some(reduced_parent) = child.reduce_parent(&array, slot_idx)? { - log_parent_rewrite( - ctx, - "reduce_parent", - slot_idx, - &array, - child, - &reduced_parent, - ); - reduced_parent.statistics().inherit_from(array.statistics()); - trace_op!(record_single_step_applied( - "reduce_parent", - &array, - &reduced_parent, - )); - return Ok(reduced_parent); - } + let optimized = array.optimize(ctx.session())?; + if !ArrayRef::ptr_eq(&optimized, &array) { + optimized.statistics().inherit_from(array.statistics()); + trace_op!(record_single_step_applied("optimize", &array, &optimized)); + return Ok(optimized); } - trace_op!(record_single_step_phase_none("reduce_parent", &array)); + trace_op!(record_single_step_phase_none("optimize", &array)); let execute_parent_kernels = Arc::clone(&ctx.execute_parent_kernels); let kernels = execute_parent_kernels.as_ref(); diff --git a/vortex-array/src/expr/arbitrary.rs b/vortex-array/src/expr/arbitrary.rs index f9dd1ce2f25..4da265a3416 100644 --- a/vortex-array/src/expr/arbitrary.rs +++ b/vortex-array/src/expr/arbitrary.rs @@ -15,8 +15,8 @@ use crate::expr::lit; use crate::expr::pack; use crate::scalar::arbitrary::random_scalar; use crate::scalar_fn::ScalarFnVTableExt; -use crate::scalar_fn::fns::binary::Binary; -use crate::scalar_fn::fns::operators::Operator; +use crate::scalar_fn::builtin::binary::Binary; +use crate::scalar_fn::builtin::operators::Operator; pub fn projection_expr(u: &mut Unstructured<'_>, dtype: &DType) -> AResult> { let Some(struct_dtype) = dtype.as_struct_fields_opt() else { diff --git a/vortex-array/src/expr/bound_expression.rs b/vortex-array/src/expr/bound_expression.rs index 035ca91f404..ece19d92729 100644 --- a/vortex-array/src/expr/bound_expression.rs +++ b/vortex-array/src/expr/bound_expression.rs @@ -27,7 +27,7 @@ use crate::scalar_fn::ReduceNodeValidity; use crate::scalar_fn::ScalarFnRef; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::ScalarFnVTableExt; -use crate::scalar_fn::fns::is_not_null::IsNotNull; +use crate::scalar_fn::builtin::is_not_null::IsNotNull; /// An [`Expression`] that has been type-checked against a [`Scope`]. /// @@ -348,7 +348,7 @@ mod tests { use crate::expr::lit; use crate::expr::root; use crate::expr::test_harness::struct_dtype; - use crate::scalar_fn::fns::literal::Literal; + use crate::scalar_fn::builtin::literal::Literal; fn scope() -> Scope { Scope::new(struct_dtype()) diff --git a/vortex-array/src/expr/display.rs b/vortex-array/src/expr/display.rs index 8de2a7d112d..74dcfd26a30 100644 --- a/vortex-array/src/expr/display.rs +++ b/vortex-array/src/expr/display.rs @@ -143,19 +143,13 @@ mod tests { use crate::dtype::Nullability; use crate::dtype::PType; use crate::expr::and; - use crate::expr::between; use crate::expr::cast; use crate::expr::eq; use crate::expr::get_item; use crate::expr::gt; use crate::expr::lit; use crate::expr::not; - use crate::expr::pack; use crate::expr::root; - use crate::expr::select; - use crate::expr::select_exclude; - use crate::scalar_fn::fns::between::BetweenOptions; - use crate::scalar_fn::fns::between::StrictComparison; #[test] fn test_display_tree_root() { @@ -213,26 +207,6 @@ mod tests { "#); } - #[test] - fn test_display_tree_select() { - use insta::assert_snapshot; - let select_expr = select(["name", "age"], root()); - assert_snapshot!(select_expr.display_tree().to_string(), @r" - vortex.select({name, age}) - └── child: vortex.root() - "); - } - - #[test] - fn test_display_tree_select_exclude() { - use insta::assert_snapshot; - let select_exclude_expr = select_exclude(["internal_id", "metadata"], root()); - assert_snapshot!(select_exclude_expr.display_tree().to_string(), @r" - vortex.select(~{internal_id, metadata}) - └── child: vortex.root() - "); - } - #[test] fn test_display_tree_cast() { use insta::assert_snapshot; @@ -259,80 +233,4 @@ mod tests { └── rhs: vortex.literal(true) "); } - - #[test] - fn test_display_tree_between() { - use insta::assert_snapshot; - let between_expr = between( - get_item("score", root()), - lit(0), - lit(100), - BetweenOptions { - lower_strict: StrictComparison::NonStrict, - upper_strict: StrictComparison::NonStrict, - }, - ); - assert_snapshot!(between_expr.display_tree().to_string(), @r" - vortex.between(lower_strict: <=, upper_strict: <=) - ├── array: vortex.get_item(score) - │ └── input: vortex.root() - ├── lower: vortex.literal(0i32) - └── upper: vortex.literal(100i32) - "); - } - - #[test] - fn test_display_tree_nested() { - use insta::assert_snapshot; - let nested_expr = select( - ["result"], - cast( - between( - get_item("score", root()), - lit(50), - lit(100), - BetweenOptions { - lower_strict: StrictComparison::Strict, - upper_strict: StrictComparison::NonStrict, - }, - ), - DType::Bool(Nullability::NonNullable), - ), - ); - assert_snapshot!(nested_expr.display_tree().to_string(), @r" - vortex.select({result}) - └── child: vortex.cast(bool) - └── input: vortex.between(lower_strict: <, upper_strict: <=) - ├── array: vortex.get_item(score) - │ └── input: vortex.root() - ├── lower: vortex.literal(50i32) - └── upper: vortex.literal(100i32) - "); - } - - #[test] - fn test_display_tree_pack() { - use insta::assert_snapshot; - let select_from_pack_expr = select( - ["fizz", "buzz"], - pack( - [ - ("fizz", root()), - ("bar", lit(5)), - ("buzz", eq(lit(42), get_item("answer", root()))), - ], - Nullability::Nullable, - ), - ); - assert_snapshot!(select_from_pack_expr.display_tree().to_string(), @r" - vortex.select({fizz, buzz}) - └── child: vortex.pack(names: [fizz, bar, buzz], nullability: Nullable) - ├── fizz: vortex.root() - ├── bar: vortex.literal(5i32) - └── buzz: vortex.binary(=) - ├── lhs: vortex.literal(42i32) - └── rhs: vortex.get_item(answer) - └── input: vortex.root() - "); - } } diff --git a/vortex-array/src/expr/expression.rs b/vortex-array/src/expr/expression.rs index ded8e2c87d7..2a1debf7611 100644 --- a/vortex-array/src/expr/expression.rs +++ b/vortex-array/src/expr/expression.rs @@ -197,23 +197,18 @@ impl Expression { /// # Example /// /// ```rust - /// # use vortex_array::dtype::{DType, Nullability, PType}; - /// # use vortex_array::scalar_fn::fns::like::{Like, LikeOptions}; - /// # use vortex_array::scalar_fn::ScalarFnVTableExt; - /// # use vortex_array::expr::{and, cast, eq, get_item, gt, lit, not, root, select}; + /// # use vortex_array::dtype::Nullability; + /// # use vortex_array::expr::{and, eq, get_item, gt, is_not_null, lit, not, pack, root}; /// // Build a complex nested expression - /// let complex_expr = select( - /// ["result"], - /// and( + /// let complex_expr = pack( + /// [("result", and( /// not(eq(get_item("status", root()), lit("inactive"))), /// and( - /// Like.new_expr(LikeOptions::default(), [get_item("name", root()), lit("%admin%")]), - /// gt( - /// cast(get_item("score", root()), DType::Primitive(PType::F64, Nullability::NonNullable)), - /// lit(75.0) - /// ) + /// is_not_null(get_item("name", root())), + /// gt(get_item("score", root()), lit(75.0)), /// ) - /// ) + /// ))], + /// Nullability::NonNullable, /// ); /// /// println!("{}", complex_expr.display_tree()); @@ -230,14 +225,12 @@ impl Expression { /// │ │ └── Root /// │ └── rhs: Literal(value: "inactive", dtype: utf8) /// └── rhs: Binary(and) - /// ├── lhs: Like - /// │ ├── child: GetItem(name) - /// │ │ └── Root - /// │ └── pattern: Literal(value: "%admin%", dtype: utf8) + /// ├── lhs: IsNotNull + /// │ └── GetItem(name) + /// │ └── Root /// └── rhs: Binary(>) - /// ├── lhs: Cast(target: f64) - /// │ └── GetItem(score) - /// │ └── Root + /// ├── lhs: GetItem(score) + /// │ └── Root /// └── rhs: Literal(value: 75f64, dtype: f64) /// ``` pub fn display_tree(&self) -> impl Display { @@ -249,7 +242,7 @@ impl Expression { /// # Example /// /// ```rust - /// # use vortex_array::scalar_fn::fns::literal::Literal; + /// # use vortex_array::scalar_fn::builtin::literal::Literal; /// # use vortex_array::expr::{eq, lit, root}; /// let expression = &eq(root(), lit(3u64)); /// assert!(expression.contains::().unwrap()); diff --git a/vortex-array/src/expr/exprs.rs b/vortex-array/src/expr/exprs.rs index 8634ae901b3..db4ec230c13 100644 --- a/vortex-array/src/expr/exprs.rs +++ b/vortex-array/src/expr/exprs.rs @@ -3,65 +3,29 @@ //! Factory functions for creating [`Expression`]s from scalar function vtables. -use std::sync::Arc; - use vortex_error::VortexExpect; -use vortex_error::VortexResult; -use vortex_error::vortex_ensure; -use vortex_error::vortex_panic; use vortex_utils::iter::ReduceBalancedIterExt; -use crate::aggregate_fn::NumericalAggregateOpts; use crate::dtype::DType; use crate::dtype::FieldName; -use crate::dtype::FieldNames; use crate::dtype::Nullability; use crate::expr::BoundExpression; use crate::expr::Expression; use crate::scalar::Scalar; -use crate::scalar::ScalarValue; use crate::scalar_fn::EmptyOptions; use crate::scalar_fn::ScalarFnVTableExt; -use crate::scalar_fn::fns::between::Between; -use crate::scalar_fn::fns::between::BetweenOptions; -use crate::scalar_fn::fns::between::StrictComparison; -use crate::scalar_fn::fns::binary::Binary; -use crate::scalar_fn::fns::byte_length::ByteLength; -use crate::scalar_fn::fns::case_when::CaseWhen; -use crate::scalar_fn::fns::case_when::CaseWhenOptions; -use crate::scalar_fn::fns::cast::Cast; -use crate::scalar_fn::fns::dynamic::DynamicComparison; -use crate::scalar_fn::fns::dynamic::DynamicComparisonExpr; -use crate::scalar_fn::fns::dynamic::Rhs; -use crate::scalar_fn::fns::ext_storage::ExtStorage; -use crate::scalar_fn::fns::fill_null::FillNull; -use crate::scalar_fn::fns::get_item::GetItem; -use crate::scalar_fn::fns::is_nan::IsNan; -use crate::scalar_fn::fns::is_not_null::IsNotNull; -use crate::scalar_fn::fns::is_null::IsNull; -use crate::scalar_fn::fns::like::Like; -use crate::scalar_fn::fns::like::LikeOptions; -use crate::scalar_fn::fns::list_contains::ListContains; -use crate::scalar_fn::fns::list_contains::ListContainsOptions; -use crate::scalar_fn::fns::list_length::ListLength; -use crate::scalar_fn::fns::list_sum::ListSum; -use crate::scalar_fn::fns::literal::Literal; -use crate::scalar_fn::fns::mask::Mask; -use crate::scalar_fn::fns::merge::DuplicateHandling; -use crate::scalar_fn::fns::merge::Merge; -use crate::scalar_fn::fns::not::Not; -use crate::scalar_fn::fns::operators::CompareOperator; -use crate::scalar_fn::fns::operators::Operator; -use crate::scalar_fn::fns::pack::Pack; -use crate::scalar_fn::fns::pack::PackOptions; -use crate::scalar_fn::fns::replace_time_zone::ReplaceTimeZone; -use crate::scalar_fn::fns::replace_time_zone::ReplaceTimeZoneOptions; -use crate::scalar_fn::fns::select::FieldSelection; -use crate::scalar_fn::fns::select::Select; -use crate::scalar_fn::fns::variant_get::VariantGet; -use crate::scalar_fn::fns::variant_get::VariantGetOptions; -use crate::scalar_fn::fns::variant_get::VariantPath; -use crate::scalar_fn::fns::zip::Zip; +use crate::scalar_fn::builtin::binary::Binary; +use crate::scalar_fn::builtin::cast::Cast; +use crate::scalar_fn::builtin::fill_null::FillNull; +use crate::scalar_fn::builtin::get_item::GetItem; +use crate::scalar_fn::builtin::is_not_null::IsNotNull; +use crate::scalar_fn::builtin::is_null::IsNull; +use crate::scalar_fn::builtin::literal::Literal; +use crate::scalar_fn::builtin::mask::Mask; +use crate::scalar_fn::builtin::not::Not; +use crate::scalar_fn::builtin::operators::Operator; +use crate::scalar_fn::builtin::pack::Pack; +use crate::scalar_fn::builtin::pack::PackOptions; /// Creates an expression that references the root scope. /// @@ -92,7 +56,7 @@ pub fn is_root(expr: &Expression) -> bool { /// use vortex_array::arrays::PrimitiveArray; /// use vortex_array::dtype::Nullability; /// use vortex_array::expr::lit; -/// use vortex_array::scalar_fn::fns::literal::Literal; +/// use vortex_array::scalar_fn::builtin::literal::Literal; /// use vortex_array::scalar::Scalar; /// /// let number = lit(34i32); @@ -168,146 +132,6 @@ pub fn var(name: impl Into) -> Expression { Expression::Variable(name.into()) } -// ---- VariantGet ---- - -/// Creates an expression that extracts a path from a Variant expression. -/// -/// Missing paths, traversal mismatches, and failed casts return null. When `dtype` is `None`, -/// results are nullable Variant values; otherwise results are nullable values of `dtype`. -pub fn variant_get( - child: Expression, - path: impl Into, - dtype: Option, -) -> Expression { - VariantGet.new_expr(VariantGetOptions::new(path.into(), dtype), vec![child]) -} - -/// Creates a bound expression that extracts a path from a Variant expression. -pub fn bound_variant_get( - child: BoundExpression, - path: impl Into, - dtype: Option, -) -> BoundExpression { - VariantGet - .try_new_bound_expr(VariantGetOptions::new(path.into(), dtype), [child]) - .vortex_expect("variant-get expressions require a Variant child") -} - -// ---- CaseWhen ---- - -/// Creates a CASE WHEN expression with one WHEN/THEN pair and an ELSE value. -pub fn case_when( - condition: Expression, - then_value: Expression, - else_value: Expression, -) -> Expression { - let options = CaseWhenOptions { - num_when_then_pairs: 1, - has_else: true, - }; - CaseWhen.new_expr(options, [condition, then_value, else_value]) -} - -/// Creates a bound CASE WHEN expression with one WHEN/THEN pair and an ELSE value. -pub fn bound_case_when( - condition: BoundExpression, - then_value: BoundExpression, - else_value: BoundExpression, -) -> BoundExpression { - let options = CaseWhenOptions { - num_when_then_pairs: 1, - has_else: true, - }; - CaseWhen - .try_new_bound_expr(options, [condition, then_value, else_value]) - .vortex_expect("case expressions must have boolean conditions and matching branch dtypes") -} - -/// Creates a CASE WHEN expression with one WHEN/THEN pair and no ELSE value. -pub fn case_when_no_else(condition: Expression, then_value: Expression) -> Expression { - let options = CaseWhenOptions { - num_when_then_pairs: 1, - has_else: false, - }; - CaseWhen.new_expr(options, [condition, then_value]) -} - -/// Creates a bound CASE WHEN expression with one WHEN/THEN pair and no ELSE value. -pub fn bound_case_when_no_else( - condition: BoundExpression, - then_value: BoundExpression, -) -> BoundExpression { - let options = CaseWhenOptions { - num_when_then_pairs: 1, - has_else: false, - }; - CaseWhen - .try_new_bound_expr(options, [condition, then_value]) - .vortex_expect("case expressions must have boolean conditions") -} - -/// Creates an n-ary CASE WHEN expression from WHEN/THEN pairs and an optional ELSE value. -pub fn nested_case_when( - when_then_pairs: Vec<(Expression, Expression)>, - else_value: Option, -) -> Expression { - assert!( - !when_then_pairs.is_empty(), - "nested_case_when requires at least one when/then pair" - ); - - let has_else = else_value.is_some(); - let mut children = Vec::with_capacity(when_then_pairs.len() * 2 + usize::from(has_else)); - for (condition, then_value) in &when_then_pairs { - children.push(condition.clone()); - children.push(then_value.clone()); - } - if let Some(else_expr) = else_value { - children.push(else_expr); - } - - let Ok(num_when_then_pairs) = u32::try_from(when_then_pairs.len()) else { - vortex_panic!("nested_case_when has too many when/then pairs"); - }; - let options = CaseWhenOptions { - num_when_then_pairs, - has_else, - }; - CaseWhen.new_expr(options, children) -} - -/// Creates a bound n-ary CASE WHEN expression from WHEN/THEN pairs and an optional ELSE value. -pub fn bound_nested_case_when( - when_then_pairs: Vec<(BoundExpression, BoundExpression)>, - else_value: Option, -) -> BoundExpression { - assert!( - !when_then_pairs.is_empty(), - "nested_case_when requires at least one when/then pair" - ); - - let Ok(num_when_then_pairs) = u32::try_from(when_then_pairs.len()) else { - vortex_panic!("nested_case_when has too many when/then pairs"); - }; - let has_else = else_value.is_some(); - let mut children = Vec::with_capacity(when_then_pairs.len() * 2 + usize::from(has_else)); - for (condition, then_value) in when_then_pairs { - children.push(condition); - children.push(then_value); - } - if let Some(else_expr) = else_value { - children.push(else_expr); - } - - let options = CaseWhenOptions { - num_when_then_pairs, - has_else, - }; - CaseWhen - .try_new_bound_expr(options, children) - .vortex_expect("case expressions must have boolean conditions and matching branch dtypes") -} - // ---- Binary operators ---- /// Creates a binary expression with the given operator. @@ -680,92 +504,6 @@ pub fn bound_not(operand: BoundExpression) -> BoundExpression { .vortex_expect("not expressions require a boolean operand") } -// ---- Between ---- - -/// Creates an expression that checks if values are between two bounds. -/// -/// Returns a boolean array indicating which values fall within the specified range. -/// The comparison strictness is controlled by the options parameter. -/// -/// ```rust -/// # use vortex_array::scalar_fn::fns::between::BetweenOptions; -/// # use vortex_array::scalar_fn::fns::between::StrictComparison; -/// # use vortex_array::expr::{between, lit, root}; -/// let opts = BetweenOptions { -/// lower_strict: StrictComparison::NonStrict, -/// upper_strict: StrictComparison::NonStrict, -/// }; -/// let expr = between(root(), lit(10), lit(20), opts); -/// ``` -pub fn between( - arr: Expression, - lower: Expression, - upper: Expression, - options: BetweenOptions, -) -> Expression { - Between - .try_new_expr(options, [arr, lower, upper]) - .vortex_expect("Failed to create Between expression") -} - -/// Creates a bound expression that checks if values are between two bounds. -pub fn bound_between( - arr: BoundExpression, - lower: BoundExpression, - upper: BoundExpression, - options: BetweenOptions, -) -> BoundExpression { - Between - .try_new_bound_expr(options, [arr, lower, upper]) - .vortex_expect("between expressions require compatible operand dtypes") -} - -// ---- Select ---- - -/// Creates an expression that selects (includes) specific fields from an array. -/// -/// Projects only the specified fields from the child expression, which must be of DType struct. -/// ```rust -/// # use vortex_array::expr::{select, root}; -/// let expr = select(["name", "age"], root()); -/// ``` -pub fn select(field_names: impl Into, child: Expression) -> Expression { - Select - .try_new_expr(FieldSelection::Include(field_names.into()), [child]) - .vortex_expect("Failed to create Select expression") -} - -/// Creates a bound expression that selects specific fields from a struct expression. -pub fn bound_select(field_names: impl Into, child: BoundExpression) -> BoundExpression { - Select - .try_new_bound_expr(FieldSelection::Include(field_names.into()), [child]) - .vortex_expect("select expressions require fields from a struct child") -} - -/// Creates an expression that excludes specific fields from an array. -/// -/// Projects all fields except the specified ones from the input struct expression. -/// -/// ```rust -/// # use vortex_array::expr::{select_exclude, root}; -/// let expr = select_exclude(["internal_id", "metadata"], root()); -/// ``` -pub fn select_exclude(fields: impl Into, child: Expression) -> Expression { - Select - .try_new_expr(FieldSelection::Exclude(fields.into()), [child]) - .vortex_expect("Failed to create Select expression") -} - -/// Creates a bound expression that excludes specific fields from a struct expression. -pub fn bound_select_exclude( - fields: impl Into, - child: BoundExpression, -) -> BoundExpression { - Select - .try_new_bound_expr(FieldSelection::Exclude(fields.into()), [child]) - .vortex_expect("select expressions require fields from a struct child") -} - // ---- Pack ---- /// Creates an expression that packs values into a struct with named fields. @@ -873,28 +611,6 @@ pub fn bound_is_null(child: BoundExpression) -> BoundExpression { .vortex_expect("is-null expressions are always well-typed") } -// ---- IsNan ---- - -/// Creates an expression that checks for NaN values. -/// -/// The expression is strict: null inputs produce null outputs, so the output nullability -/// follows the input. Only primitive float inputs are supported. -/// -/// ```rust -/// # use vortex_array::expr::{is_nan, root}; -/// let expr = is_nan(root()); -/// ``` -pub fn is_nan(child: Expression) -> Expression { - IsNan.new_expr(EmptyOptions, vec![child]) -} - -/// Creates a bound expression that checks for NaN values. -pub fn bound_is_nan(child: BoundExpression) -> BoundExpression { - IsNan - .try_new_bound_expr(EmptyOptions, [child]) - .vortex_expect("is-nan expressions are always well-typed") -} - // ---- IsNotNull ---- /// Creates an expression that checks for non-null values. @@ -916,88 +632,6 @@ pub fn bound_is_not_null(child: BoundExpression) -> BoundExpression { .vortex_expect("is-not-null expressions are always well-typed") } -// ---- Like ---- - -/// Creates a SQL LIKE expression. -pub fn like(child: Expression, pattern: Expression) -> Expression { - Like.new_expr( - LikeOptions { - negated: false, - case_insensitive: false, - }, - [child, pattern], - ) -} - -/// Creates a bound SQL LIKE expression. -pub fn bound_like(child: BoundExpression, pattern: BoundExpression) -> BoundExpression { - bound_like_with_options(child, pattern, false, false) -} - -/// Creates a case-insensitive SQL ILIKE expression. -pub fn ilike(child: Expression, pattern: Expression) -> Expression { - Like.new_expr( - LikeOptions { - negated: false, - case_insensitive: true, - }, - [child, pattern], - ) -} - -/// Creates a bound case-insensitive SQL ILIKE expression. -pub fn bound_ilike(child: BoundExpression, pattern: BoundExpression) -> BoundExpression { - bound_like_with_options(child, pattern, false, true) -} - -/// Creates a negated SQL NOT LIKE expression. -pub fn not_like(child: Expression, pattern: Expression) -> Expression { - Like.new_expr( - LikeOptions { - negated: true, - case_insensitive: false, - }, - [child, pattern], - ) -} - -/// Creates a bound negated SQL NOT LIKE expression. -pub fn bound_not_like(child: BoundExpression, pattern: BoundExpression) -> BoundExpression { - bound_like_with_options(child, pattern, true, false) -} - -/// Creates a negated case-insensitive SQL NOT ILIKE expression. -pub fn not_ilike(child: Expression, pattern: Expression) -> Expression { - Like.new_expr( - LikeOptions { - negated: true, - case_insensitive: true, - }, - [child, pattern], - ) -} - -/// Creates a bound negated case-insensitive SQL NOT ILIKE expression. -pub fn bound_not_ilike(child: BoundExpression, pattern: BoundExpression) -> BoundExpression { - bound_like_with_options(child, pattern, true, true) -} - -fn bound_like_with_options( - child: BoundExpression, - pattern: BoundExpression, - negated: bool, - case_insensitive: bool, -) -> BoundExpression { - Like.try_new_bound_expr( - LikeOptions { - negated, - case_insensitive, - }, - [child, pattern], - ) - .vortex_expect("like expressions require UTF-8 or binary operands") -} - // ---- Mask ---- /// Creates a mask expression that applies the given boolean mask to the input array. @@ -1011,409 +645,6 @@ pub fn bound_mask(array: BoundExpression, mask: BoundExpression) -> BoundExpress .vortex_expect("mask expressions require a boolean mask") } -// ---- Merge ---- - -/// Creates an expression that merges struct expressions into a single struct. -/// -/// Combines fields from all input expressions. If field names are duplicated, -/// later expressions win. Fields are not recursively merged. -/// -/// ```rust -/// # use vortex_array::dtype::Nullability; -/// # use vortex_array::expr::{merge, get_item, root}; -/// let expr = merge([get_item("a", root()), get_item("b", root())]); -/// ``` -pub fn merge(elements: impl IntoIterator>) -> Expression { - use itertools::Itertools as _; - let values = elements.into_iter().map(|value| value.into()).collect_vec(); - Merge.new_expr(DuplicateHandling::default(), values) -} - -/// Creates a bound expression that merges struct expressions. -pub fn bound_merge(elements: impl IntoIterator) -> BoundExpression { - bound_merge_opts(elements, DuplicateHandling::default()) -} - -/// Creates a merge expression with explicit duplicate handling. -pub fn merge_opts( - elements: impl IntoIterator>, - duplicate_handling: DuplicateHandling, -) -> Expression { - use itertools::Itertools as _; - let values = elements.into_iter().map(|value| value.into()).collect_vec(); - Merge.new_expr(duplicate_handling, values) -} - -/// Creates a bound merge expression with explicit duplicate handling. -pub fn bound_merge_opts( - elements: impl IntoIterator, - duplicate_handling: DuplicateHandling, -) -> BoundExpression { - Merge - .try_new_bound_expr(duplicate_handling, elements) - .vortex_expect("merge expressions require non-nullable struct children") -} - -// ---- Zip ---- - -/// Creates a zip expression that conditionally selects between two arrays. -/// -/// ```rust -/// # use vortex_array::expr::{zip_expr, root, lit}; -/// let expr = zip_expr(lit(true), root(), lit(0i32)); -/// ``` -pub fn zip_expr(mask: Expression, if_true: Expression, if_false: Expression) -> Expression { - Zip.new_expr(EmptyOptions, [if_true, if_false, mask]) -} - -/// Creates a bound zip expression that conditionally selects between two arrays. -pub fn bound_zip_expr( - mask: BoundExpression, - if_true: BoundExpression, - if_false: BoundExpression, -) -> BoundExpression { - Zip.try_new_bound_expr(EmptyOptions, [if_true, if_false, mask]) - .vortex_expect("zip expressions require a boolean mask and compatible value dtypes") -} - -// ---- Dynamic ---- - -/// Creates a dynamic comparison expression from its complete options. -pub fn dynamic_with_options(options: DynamicComparisonExpr, lhs: Expression) -> Expression { - DynamicComparison.new_expr(options, [lhs]) -} - -/// Creates a bound dynamic comparison expression from its complete options. -pub fn bound_dynamic_with_options( - options: DynamicComparisonExpr, - lhs: BoundExpression, -) -> BoundExpression { - DynamicComparison - .try_new_bound_expr(options, [lhs]) - .vortex_expect("dynamic comparisons require a compatible left-hand dtype") -} - -/// Creates a dynamic comparison expression. -pub fn dynamic( - operator: CompareOperator, - rhs_value: impl Fn() -> Option + Send + Sync + 'static, - rhs_dtype: DType, - default: bool, - lhs: Expression, -) -> Expression { - dynamic_with_options( - DynamicComparisonExpr { - operator, - rhs: Arc::new(Rhs { - value: Arc::new(rhs_value), - dtype: rhs_dtype, - }), - default, - }, - lhs, - ) -} - -/// Creates a bound dynamic comparison expression. -pub fn bound_dynamic( - operator: CompareOperator, - rhs_value: impl Fn() -> Option + Send + Sync + 'static, - rhs_dtype: DType, - default: bool, - lhs: BoundExpression, -) -> BoundExpression { - bound_dynamic_with_options( - DynamicComparisonExpr { - operator, - rhs: Arc::new(Rhs { - value: Arc::new(rhs_value), - dtype: rhs_dtype, - }), - default, - }, - lhs, - ) -} - -// ---- ListContains ---- - -/// Creates an expression that checks whether a list contains a value. -/// -/// A null element never matches anything: a needle that matches no element yields `false`. For -/// SQL's three-valued `IN`, where a null element makes that `null`, use [`in_list`]. -/// -/// ```rust -/// # use vortex_array::expr::{list_contains, lit, root}; -/// let expr = list_contains(root(), lit(42)); -/// ``` -pub fn list_contains(list: Expression, value: Expression) -> Expression { - list_contains_opts(list, value, ListContainsOptions::default()) -} - -/// Creates an expression that checks whether a list contains a value, with explicit -/// [`ListContainsOptions`]. -pub fn list_contains_opts( - list: Expression, - value: Expression, - options: ListContainsOptions, -) -> Expression { - ListContains.new_expr(options, [list, value]) -} - -/// Creates `needle IN (list)` with SQL null semantics: a null element is an unknown value, so a -/// needle that matches no element yields `null` rather than `false` when the list holds one, and -/// `not(in_list(..))` never admits such a row. -/// -/// ```rust -/// # use vortex_array::expr::{in_list, lit, root}; -/// let expr = in_list(root(), lit(vec![1, 2, 3])); -/// ``` -pub fn in_list(needle: Expression, list: Expression) -> Expression { - list_contains_opts( - list, - needle, - ListContainsOptions { - sql_null_semantics: true, - }, - ) -} - -/// Creates `needle IN (elements)` with SQL null semantics for a list of constant elements of -/// `element_dtype`, ignoring nullability. An element can also be an untyped null. -/// -/// The nullability of `element_dtype` is not used: the list elements are nullable only when a -/// null element is present. -/// -/// Non-null integers that hold every value between their smallest and largest become a range -/// check, which is cheaper to evaluate and prunes with tighter bounds than a set lookup, as for -/// the keys of a small dimension table that a join pushes down. -/// -/// # Errors -/// -/// Returns an error if an element is not a null and not of `element_dtype`. -/// -/// ```rust -/// # use vortex_array::dtype::{DType, Nullability, PType}; -/// # use vortex_array::expr::{in_list_of, root}; -/// let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); -/// let expr = in_list_of(root(), dtype, vec![3i32.into(), 1i32.into(), 2i32.into()])?; -/// assert_eq!(expr.to_string(), "(1i32 <= $ <= 3i32)"); -/// # Ok::<(), vortex_error::VortexError>(()) -/// ``` -pub fn in_list_of( - needle: Expression, - element_dtype: DType, - elements: Vec, -) -> VortexResult { - let element_dtype = if elements.iter().any(Scalar::is_null) { - element_dtype.as_nullable() - } else { - element_dtype.as_nonnullable() - }; - let mut elements = elements - .iter() - .map(|element| { - vortex_ensure!( - element.dtype() == &DType::Null - || element.dtype().eq_ignore_nullability(&element_dtype), - "IN list element {element} is not of type {element_dtype}" - ); - element.cast(&element_dtype) - }) - .collect::>>()?; - - if let Some((low, high)) = contiguous_integer_bounds(&elements) { - let options = BetweenOptions { - lower_strict: StrictComparison::NonStrict, - upper_strict: StrictComparison::NonStrict, - }; - return Ok(between( - needle, - lit(low.clone()), - lit(high.clone()), - options, - )); - } - - // Ordered elements make equal lists from the same set, so that equal filters deduplicate. - if elements - .windows(2) - .all(|pair| pair[0].partial_cmp(&pair[1]).is_some()) - { - elements.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)); - } - let list = Scalar::list(Arc::new(element_dtype), elements, Nullability::NonNullable); - Ok(in_list(needle, lit(list))) -} - -/// The smallest and largest of non-null integer `scalars` when they hold every integer between -/// those two. -fn contiguous_integer_bounds(scalars: &[Scalar]) -> Option<(&Scalar, &Scalar)> { - let mut values = scalars - .iter() - .map(|scalar| { - let primitive = scalar.as_primitive_opt()?; - if !primitive.ptype().is_int() { - return None; - } - let value = if primitive.ptype().is_unsigned_int() { - primitive.as_opt::()??.into() - } else { - i128::from(primitive.as_opt::()??) - }; - Some((value, scalar)) - }) - .collect::>>()?; - values.sort_unstable_by_key(|(value, _)| *value); - values.dedup_by_key(|(value, _)| *value); - let (low, low_scalar) = values.first()?; - let (high, high_scalar) = values.last()?; - (high - low + 1 == i128::try_from(values.len()).ok()?).then_some((*low_scalar, *high_scalar)) -} - -/// Creates a bound expression that checks whether a list contains a value. -pub fn bound_list_contains(list: BoundExpression, value: BoundExpression) -> BoundExpression { - ListContains - .try_new_bound_expr(ListContainsOptions::default(), [list, value]) - .vortex_expect("list-contains expressions require a list child") -} - -/// Creates a bound `needle IN (list)` with SQL null semantics. See [`in_list`]. -pub fn bound_in_list(needle: BoundExpression, list: BoundExpression) -> BoundExpression { - ListContains - .try_new_bound_expr( - ListContainsOptions { - sql_null_semantics: true, - }, - [list, needle], - ) - .vortex_expect("list-contains expressions require a list child") -} - -// ---- ByteLength ---- - -/// Creates an expression that computes the byte length of each element. -/// This is akin to ANSI SQL OCTET_LENGTH(), or DuckDB's strlen(). -/// -/// ```rust -/// # use vortex_array::expr::{byte_length, root}; -/// let expr = byte_length(root()); -/// ``` -pub fn byte_length(input: Expression) -> Expression { - ByteLength.new_expr(EmptyOptions, [input]) -} - -/// Creates a bound expression that computes each element's byte length. -pub fn bound_byte_length(input: BoundExpression) -> BoundExpression { - ByteLength - .try_new_bound_expr(EmptyOptions, [input]) - .vortex_expect("byte-length expressions require a variable-length binary child") -} - -// ---- ExtStorage ---- - -/// Creates an expression that extracts the storage values from an extension array. -/// -/// ```rust -/// # use vortex_array::expr::{ext_storage, root}; -/// let expr = ext_storage(root()); -/// ``` -pub fn ext_storage(input: Expression) -> Expression { - ExtStorage.new_expr(EmptyOptions, [input]) -} - -/// Creates a bound expression that extracts an extension array's storage values. -pub fn bound_ext_storage(input: BoundExpression) -> BoundExpression { - ExtStorage - .try_new_bound_expr(EmptyOptions, [input]) - .vortex_expect("extension-storage expressions require an extension child") -} - -// ---- ListLength ---- - -/// Creates an expression that computes the number of elements in each list -/// for `List` and `FixedSizeList` inputs. This is akin to ANSI SQL `CARDINALITY()`, -/// or DuckDB's `len()`/`array_length()`. -/// -/// ```rust -/// # use vortex_array::expr::{list_length, root}; -/// let expr = list_length(root()); -/// ``` -pub fn list_length(input: Expression) -> Expression { - ListLength.new_expr(EmptyOptions, [input]) -} - -/// Creates a bound expression that computes the number of elements in each list. -pub fn bound_list_length(input: BoundExpression) -> BoundExpression { - ListLength - .try_new_bound_expr(EmptyOptions, [input]) - .vortex_expect("list-length expressions require a list child") -} - -// ---- ListSum ---- - -/// Creates an expression that sums the elements of each list for `List` and -/// `FixedSizeList` inputs, akin to DuckDB's `list_sum()`. -/// -/// Follows SQL `SUM` semantics per list: null lists, empty lists, and lists whose elements are -/// all null yield null; null elements are skipped; integer and decimal overflow yields a null -/// value. The result dtype follows `sum`'s widening rules and is always nullable. NaN float -/// elements are skipped by default; see [`list_sum_opts`] for the NaN-including variant. -/// -/// ```rust -/// # use vortex_array::expr::{list_sum, root}; -/// let expr = list_sum(root()); -/// ``` -pub fn list_sum(input: Expression) -> Expression { - ListSum.new_expr(NumericalAggregateOpts::default(), [input]) -} - -/// Creates a bound expression that sums the elements of each list. -pub fn bound_list_sum(input: BoundExpression) -> BoundExpression { - ListSum - .try_new_bound_expr(NumericalAggregateOpts::default(), [input]) - .vortex_expect("list-sum expressions require a numeric list child") -} - -/// Creates a [`list_sum`] expression with explicit [`NumericalAggregateOpts`], controlling -/// whether NaN float elements are skipped (the default) or poison the list's sum to NaN. -pub fn list_sum_opts(input: Expression, options: NumericalAggregateOpts) -> Expression { - ListSum.new_expr(options, [input]) -} - -/// Creates a bound list-sum expression with explicit aggregate options. -pub fn bound_list_sum_opts( - input: BoundExpression, - options: NumericalAggregateOpts, -) -> BoundExpression { - ListSum - .try_new_bound_expr(options, [input]) - .vortex_expect("list-sum expressions require a numeric list child") -} - -/// Reinterpret timestamp wall times in `options.time_zone`, preserving the time unit. -/// -/// `ambiguous` is a UTF-8 expression containing `raise`, `earliest`, `latest`, or `null`. -/// Removing the timezone preserves wall time in a timezone-naive timestamp. -pub fn replace_time_zone( - input: Expression, - ambiguous: Expression, - options: ReplaceTimeZoneOptions, -) -> Expression { - ReplaceTimeZone.new_expr(options, [input, ambiguous]) -} - -/// Create a bound timezone replacement expression. -pub fn bound_replace_time_zone( - input: BoundExpression, - ambiguous: BoundExpression, - options: ReplaceTimeZoneOptions, -) -> BoundExpression { - ReplaceTimeZone - .try_new_bound_expr(options, [input, ambiguous]) - .vortex_expect("timezone replacement requires a timestamp and a UTF-8 ambiguity policy") -} - /// Constructors for expressions whose children have already been bound and type-checked. /// /// These mirror the constructors in [`crate::expr`] and panic when the supplied children do not @@ -1422,50 +653,25 @@ pub fn bound_replace_time_zone( pub mod bound { pub use super::bound_and as and; pub use super::bound_and_collect as and_collect; - pub use super::bound_between as between; pub use super::bound_binary as binary; - pub use super::bound_byte_length as byte_length; - pub use super::bound_case_when as case_when; - pub use super::bound_case_when_no_else as case_when_no_else; pub use super::bound_cast as cast; pub use super::bound_checked_add as checked_add; pub use super::bound_col as col; - pub use super::bound_dynamic as dynamic; - pub use super::bound_dynamic_with_options as dynamic_with_options; pub use super::bound_eq as eq; - pub use super::bound_ext_storage as ext_storage; pub use super::bound_fill_null as fill_null; pub use super::bound_get_item as get_item; pub use super::bound_gt as gt; pub use super::bound_gt_eq as gt_eq; - pub use super::bound_ilike as ilike; - pub use super::bound_in_list as in_list; - pub use super::bound_is_nan as is_nan; pub use super::bound_is_not_null as is_not_null; pub use super::bound_is_null as is_null; - pub use super::bound_like as like; - pub use super::bound_list_contains as list_contains; - pub use super::bound_list_length as list_length; - pub use super::bound_list_sum as list_sum; - pub use super::bound_list_sum_opts as list_sum_opts; pub use super::bound_lit as lit; pub use super::bound_lt as lt; pub use super::bound_lt_eq as lt_eq; pub use super::bound_mask as mask; - pub use super::bound_merge as merge; - pub use super::bound_merge_opts as merge_opts; - pub use super::bound_nested_case_when as nested_case_when; pub use super::bound_not as not; pub use super::bound_not_eq as not_eq; - pub use super::bound_not_ilike as not_ilike; - pub use super::bound_not_like as not_like; pub use super::bound_or as or; pub use super::bound_or_collect as or_collect; pub use super::bound_pack as pack; - pub use super::bound_replace_time_zone as replace_time_zone; pub use super::bound_root as root; - pub use super::bound_select as select; - pub use super::bound_select_exclude as select_exclude; - pub use super::bound_variant_get as variant_get; - pub use super::bound_zip_expr as zip_expr; } diff --git a/vortex-array/src/expr/mod.rs b/vortex-array/src/expr/mod.rs index bbdbbfdf01a..beff6e6e6a9 100644 --- a/vortex-array/src/expr/mod.rs +++ b/vortex-array/src/expr/mod.rs @@ -46,7 +46,6 @@ pub mod bound_expression; pub mod display; pub(crate) mod expression; mod exprs; -pub(crate) mod field; pub mod lambda; mod optimize; pub mod proto; @@ -58,57 +57,30 @@ pub use bound_expression::*; pub use expression::*; pub use exprs::and; pub use exprs::and_collect; -pub use exprs::between; pub use exprs::binary; pub use exprs::bound; -pub use exprs::byte_length; -pub use exprs::case_when; -pub use exprs::case_when_no_else; pub use exprs::cast; pub use exprs::checked_add; pub use exprs::col; -pub use exprs::dynamic; -pub use exprs::dynamic_with_options; pub use exprs::eq; -pub use exprs::ext_storage; pub use exprs::fill_null; pub use exprs::get_item; pub use exprs::gt; pub use exprs::gt_eq; -pub use exprs::ilike; -pub use exprs::in_list; -pub use exprs::in_list_of; -pub use exprs::is_nan; pub use exprs::is_not_null; pub use exprs::is_null; pub use exprs::is_root; -pub use exprs::like; -pub use exprs::list_contains; -pub use exprs::list_contains_opts; -pub use exprs::list_length; -pub use exprs::list_sum; -pub use exprs::list_sum_opts; pub use exprs::lit; pub use exprs::lt; pub use exprs::lt_eq; pub use exprs::mask; -pub use exprs::merge; -pub use exprs::merge_opts; -pub use exprs::nested_case_when; pub use exprs::not; pub use exprs::not_eq; -pub use exprs::not_ilike; -pub use exprs::not_like; pub use exprs::or; pub use exprs::or_collect; pub use exprs::pack; -pub use exprs::replace_time_zone; pub use exprs::root; -pub use exprs::select; -pub use exprs::select_exclude; pub use exprs::var; -pub use exprs::variant_get; -pub use exprs::zip_expr; pub use lambda::*; pub use scope::*; @@ -142,167 +114,7 @@ mod tests { use vortex_array::expr::lit; use vortex_array::expr::root; - use super::*; - use crate::dtype::DType; - use crate::dtype::FieldName; - use crate::dtype::FieldNames; - use crate::dtype::Nullability; - use crate::dtype::PType; - use crate::dtype::StructFields; - use crate::expr::and; - use crate::expr::bound; - use crate::expr::case_when; - use crate::expr::col; - use crate::expr::gt; - use crate::expr::gt_eq; - use crate::expr::lt; - use crate::expr::lt_eq; - use crate::expr::not; - use crate::expr::not_eq; - use crate::expr::or; - use crate::expr::select; - use crate::expr::select_exclude; - use crate::scalar::Scalar; - use crate::scalar_fn::fns::literal::Literal; - - #[test] - fn bound_constructors_preserve_order_and_types() -> vortex_error::VortexResult<()> { - let value_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let scope = DType::Struct( - StructFields::from_iter([("value", value_dtype.clone())]), - Nullability::NonNullable, - ); - - let root = bound::root(scope.clone()); - let value = bound::get_item("value", root); - let literal = bound::lit(5i32); - let condition = bound::gt(value.clone(), literal.clone()); - assert_eq!(condition.dtype(), &DType::Bool(Nullability::NonNullable)); - assert_eq!(condition.children(), &[value.clone(), literal.clone()]); - - let case = bound::case_when(condition.clone(), value.clone(), literal.clone()); - assert_eq!(case.dtype(), &value_dtype); - assert_eq!(case.children(), &[condition.clone(), value, literal]); - - let packed = bound::pack( - [("condition", condition.clone()), ("value", case.clone())], - Nullability::NonNullable, - ); - assert_eq!(packed.children(), &[condition, case.clone()]); - assert_eq!( - packed.dtype(), - &DType::Struct( - StructFields::from_iter([ - ("condition", DType::Bool(Nullability::NonNullable)), - ("value", value_dtype), - ]), - Nullability::NonNullable, - ) - ); - - let unbound = case_when(gt(col("value"), lit(5i32)), col("value"), lit(5i32)); - assert_eq!(unbound.bind(&scope)?, case); - Ok(()) - } - - #[test] - fn expr_display() { - assert_eq!(col("a").to_string(), "$.a"); - assert_eq!(root().to_string(), "$"); - - let col1: Expression = col("col1"); - let col2: Expression = col("col2"); - assert_eq!( - and(col1.clone(), col2.clone()).to_string(), - "($.col1 and $.col2)" - ); - assert_eq!( - or(col1.clone(), col2.clone()).to_string(), - "($.col1 or $.col2)" - ); - assert_eq!( - eq(col1.clone(), col2.clone()).to_string(), - "($.col1 = $.col2)" - ); - assert_eq!( - not_eq(col1.clone(), col2.clone()).to_string(), - "($.col1 != $.col2)" - ); - assert_eq!( - gt(col1.clone(), col2.clone()).to_string(), - "($.col1 > $.col2)" - ); - assert_eq!( - gt_eq(col1.clone(), col2.clone()).to_string(), - "($.col1 >= $.col2)" - ); - assert_eq!( - lt(col1.clone(), col2.clone()).to_string(), - "($.col1 < $.col2)" - ); - assert_eq!( - lt_eq(col1.clone(), col2.clone()).to_string(), - "($.col1 <= $.col2)" - ); - - assert_eq!( - or(lt(col1.clone(), col2.clone()), not_eq(col1.clone(), col2),).to_string(), - "(($.col1 < $.col2) or ($.col1 != $.col2))" - ); - - assert_eq!(not(col1).to_string(), "vortex.not($.col1)"); - - assert_eq!( - select(vec![FieldName::from("col1")], root()).to_string(), - "${col1}" - ); - assert_eq!( - select( - vec![FieldName::from("col1"), FieldName::from("col2")], - root() - ) - .to_string(), - "${col1, col2}" - ); - assert_eq!( - select_exclude( - vec![FieldName::from("col1"), FieldName::from("col2")], - root() - ) - .to_string(), - "${~ col1, col2}" - ); - - assert_eq!(lit(Scalar::from(0u8)).to_string(), "0u8"); - assert_eq!(lit(Scalar::from(0.0f32)).to_string(), "0f32"); - assert_eq!( - lit(Scalar::from(i64::MAX)).to_string(), - "9223372036854775807i64" - ); - assert_eq!(lit(Scalar::from(true)).to_string(), "true"); - assert_eq!( - lit(Scalar::null(DType::Bool(Nullability::Nullable))).to_string(), - "null" - ); - - assert_eq!( - lit(Scalar::struct_( - DType::Struct( - StructFields::new( - FieldNames::from(["dog", "cat"]), - vec![ - DType::Primitive(PType::U32, Nullability::NonNullable), - DType::Utf8(Nullability::NonNullable) - ], - ), - Nullability::NonNullable - ), - vec![Scalar::from(32_u32), Scalar::from("rufus".to_string())] - )) - .to_string(), - "{dog: 32u32, cat: \"rufus\"}" - ); - } + use crate::scalar_fn::builtin::literal::Literal; #[test] fn expr_contains() { diff --git a/vortex-array/src/expr/optimize.rs b/vortex-array/src/expr/optimize.rs index ac4356892b6..a96b3e7af32 100644 --- a/vortex-array/src/expr/optimize.rs +++ b/vortex-array/src/expr/optimize.rs @@ -136,7 +136,6 @@ mod tests { use crate::arrays::StructArray; use crate::assert_arrays_eq; use crate::dtype::DType; - use crate::dtype::FieldName; use crate::dtype::Nullability; use crate::dtype::PType; use crate::dtype::StructFields; @@ -147,14 +146,10 @@ mod tests { use crate::expr::get_item; use crate::expr::lit; use crate::expr::lt_eq; - use crate::expr::merge; use crate::expr::or; - use crate::expr::pack; use crate::expr::root; - use crate::expr::select; - use crate::expr::zip_expr; use crate::scalar::Scalar; - use crate::scalar_fn::fns::literal::Literal; + use crate::scalar_fn::builtin::literal::Literal; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); @@ -223,14 +218,6 @@ mod tests { or(root(), lit(Scalar::from(Some(false)))), DType::Bool(Nullability::NonNullable) )] - #[case::zip_true( - zip_expr(lit(true), lit(1i32), root()), - DType::Primitive(PType::I32, Nullability::Nullable) - )] - #[case::zip_false( - zip_expr(lit(false), root(), lit(1i32)), - DType::Primitive(PType::I32, Nullability::Nullable) - )] fn literal_simplification_preserves_result_dtype( #[case] expr: Expression, #[case] scope: DType, @@ -241,88 +228,4 @@ mod tests { Ok(()) } - - /// `Select` over a `pack` rewrites to a `pack` of the selected child fields. It must read - /// those fields out of the child `pack` itself, because the optimizer visits children before - /// their parent and so never revisits the children a rule introduces. - #[test] - fn optimize_select_of_pack_reads_pack_fields() -> VortexResult<()> { - let field_names = (0..4) - .map(|idx| FieldName::from(format!("field_{idx}"))) - .collect::>(); - let fields = StructFields::new( - field_names.clone().into(), - vec![DType::Primitive(PType::U64, Nullability::NonNullable); field_names.len()], - ); - let scope = DType::Struct(fields, Nullability::NonNullable); - - let root_pack = pack( - field_names - .iter() - .map(|name| (name.clone(), get_item(name.clone(), root()))) - .collect::>(), - Nullability::NonNullable, - ); - let expr = select( - field_names - .iter() - .cloned() - .chain(["input_row_idx".into(), "input_file_idx".into()]) - .collect::>(), - merge([ - root_pack, - pack( - [ - ("input_row_idx", lit(0_u64)), - ("input_file_idx", lit(1_u64)), - ], - Nullability::NonNullable, - ), - ]), - ); - - let expected = pack( - field_names - .into_iter() - .map(|name| (name.clone(), get_item(name, root()))) - .chain([ - ("input_row_idx".into(), lit(0_u64)), - ("input_file_idx".into(), lit(1_u64)), - ]) - .collect::>(), - Nullability::NonNullable, - ); - - let expected = expected.bind(&scope)?; - - let optimized = expr.bind(&scope)?.optimize_recursive()?; - assert_eq!(optimized, expected, "got {optimized}"); - assert_eq!(optimized.optimize_recursive()?, expected); - - Ok(()) - } - - /// `Merge` reduces to a `pack` that reads each field out of its children, so it must also - /// read straight out of a `pack` child. - #[test] - fn optimize_merge_of_packs_reads_pack_fields() -> VortexResult<()> { - let scope = DType::Struct(StructFields::empty(), Nullability::NonNullable); - let expr = merge([ - pack([("a", lit(0_u64))], Nullability::NonNullable), - pack([("b", lit(1_u64))], Nullability::NonNullable), - ]); - - let expected = pack( - [("a", lit(0_u64)), ("b", lit(1_u64))], - Nullability::NonNullable, - ); - - let expected = expected.bind(&scope)?; - - let optimized = expr.bind(&scope)?.optimize_recursive()?; - assert_eq!(optimized, expected, "got {optimized}"); - assert_eq!(optimized.optimize_recursive()?, expected); - - Ok(()) - } } diff --git a/vortex-array/src/expr/proto.rs b/vortex-array/src/expr/proto.rs index 273bb0d1166..ba49bde94ff 100644 --- a/vortex-array/src/expr/proto.rs +++ b/vortex-array/src/expr/proto.rs @@ -109,15 +109,13 @@ mod tests { use super::ExprSerializeProtoExt; use crate::expr::Expression; use crate::expr::and; - use crate::expr::between; use crate::expr::eq; use crate::expr::get_item; use crate::expr::lit; + use crate::expr::lt; use crate::expr::or; use crate::expr::root; use crate::proto::expr as pb; - use crate::scalar_fn::fns::between::BetweenOptions; - use crate::scalar_fn::fns::between::StrictComparison; use crate::scalar_fn::session::ScalarFnSession; static SESSION: LazyLock = LazyLock::new(|| { @@ -129,18 +127,7 @@ mod tests { #[test] fn expression_serde() { let expr: Expression = or( - and( - between( - lit(1), - root(), - get_item("a", root()), - BetweenOptions { - lower_strict: StrictComparison::Strict, - upper_strict: StrictComparison::Strict, - }, - ), - lit(1), - ), + and(lt(lit(1), get_item("a", root())), lit(1)), eq(lit(1), root()), ); diff --git a/vortex-array/src/expr/traversal/fold.rs b/vortex-array/src/expr/traversal/fold.rs index 5b6c7c94944..d7c2bcddeba 100644 --- a/vortex-array/src/expr/traversal/fold.rs +++ b/vortex-array/src/expr/traversal/fold.rs @@ -81,9 +81,9 @@ mod tests { use crate::expr::gt; use crate::expr::lit; use crate::expr::traversal::NodeExt; - use crate::scalar_fn::fns::binary::Binary; - use crate::scalar_fn::fns::literal::Literal; - use crate::scalar_fn::fns::operators::Operator; + use crate::scalar_fn::builtin::binary::Binary; + use crate::scalar_fn::builtin::literal::Literal; + use crate::scalar_fn::builtin::operators::Operator; struct AddFold; impl NodeFolderContext for AddFold { diff --git a/vortex-array/src/expr/traversal/mod.rs b/vortex-array/src/expr/traversal/mod.rs index 5b80816c30a..8129703216b 100644 --- a/vortex-array/src/expr/traversal/mod.rs +++ b/vortex-array/src/expr/traversal/mod.rs @@ -566,10 +566,10 @@ mod tests { use crate::expr::lit; use crate::expr::not_eq; use crate::expr::root; - use crate::scalar_fn::fns::binary::Binary; - use crate::scalar_fn::fns::get_item::GetItem; - use crate::scalar_fn::fns::literal::Literal; - use crate::scalar_fn::fns::operators::Operator; + use crate::scalar_fn::builtin::binary::Binary; + use crate::scalar_fn::builtin::get_item::GetItem; + use crate::scalar_fn::builtin::literal::Literal; + use crate::scalar_fn::builtin::operators::Operator; #[derive(Default)] pub struct ExprLitCollector<'a>(pub Vec<&'a Expression>); diff --git a/vortex-array/src/expression.rs b/vortex-array/src/expression.rs index 4f0628a1166..968420d60c1 100644 --- a/vortex-array/src/expression.rs +++ b/vortex-array/src/expression.rs @@ -12,8 +12,7 @@ use crate::arrays::ScalarFnArray; use crate::expr::BoundExpression; use crate::expr::Expression; use crate::expr::expression::unbound_variable; -use crate::optimizer::ArrayOptimizer; -use crate::scalar_fn::fns::literal::Literal; +use crate::scalar_fn::builtin::literal::Literal; impl ArrayRef { /// Apply a bound expression to this array, producing a new array in constant time. @@ -37,9 +36,7 @@ impl ArrayRef { .map(|child| self.clone().apply_bound(child)) .collect::>>()?; - let array = ScalarFnArray::try_new_with_len(scalar_fn.clone(), children, len)?.into_array(); - - array.optimize() + Ok(ScalarFnArray::try_new_with_len(scalar_fn.clone(), children, len)?.into_array()) } /// Apply the expression to this array, producing a new array in constant time. @@ -69,9 +66,6 @@ impl ArrayRef { let scalar_fn = expr .as_scalar() .vortex_expect("leaves and literals were handled above, so this is a scalar node"); - let array = ScalarFnArray::try_new_with_len(scalar_fn.clone(), children, len)?.into_array(); - - // Optimize the resulting array's root. - array.optimize() + Ok(ScalarFnArray::try_new_with_len(scalar_fn.clone(), children, len)?.into_array()) } } diff --git a/vortex-array/src/extension/datetime/timestamp.rs b/vortex-array/src/extension/datetime/timestamp.rs index 7c0cdaeb1bb..f43f8f23cef 100644 --- a/vortex-array/src/extension/datetime/timestamp.rs +++ b/vortex-array/src/extension/datetime/timestamp.rs @@ -78,7 +78,7 @@ impl fmt::Display for TimestampOptions { } /// Resolve an Arrow timestamp timezone, accepting IANA names and fixed offsets. -pub(crate) fn resolve_time_zone(name: Option<&str>) -> VortexResult { +pub fn resolve_time_zone(name: Option<&str>) -> VortexResult { Ok(match name { Some(name) if name.starts_with('+') || name.starts_with('-') => { DateTimeParser::new().parse_time_zone(name)? diff --git a/vortex-array/src/optimizer/kernels.rs b/vortex-array/src/optimizer/kernels.rs index 6a7eca5d6d6..3fb5936d89c 100644 --- a/vortex-array/src/optimizer/kernels.rs +++ b/vortex-array/src/optimizer/kernels.rs @@ -48,8 +48,9 @@ use crate::arrays::Struct; use crate::arrays::struct_::compute::rules::struct_cast_reduce_parent; use crate::kernel::ExecuteParentKernel; use crate::matcher::Matcher; +use crate::optimizer::rules::ArrayParentReduceRule; use crate::scalar_fn::ScalarFnVTable; -use crate::scalar_fn::fns::cast::Cast; +use crate::scalar_fn::builtin::cast::Cast; /// Shared hasher used to combine `(outer, child)` tuples into registry keys. static FN_HASHER: LazyLock = LazyLock::new(DefaultHashBuilder::default); @@ -65,6 +66,60 @@ static FN_HASHER: LazyLock = LazyLock::new(DefaultHashBuilde pub type ReduceParentFn = fn(child: &ArrayRef, parent: &ArrayRef, child_idx: usize) -> VortexResult>; +/// Type-erased parent-reduce rule stored in the session registry. +pub trait DynReduceParentRule: Debug + Send + Sync + 'static { + /// Attempt to rewrite `parent` given its `child` at slot `child_idx`. + fn reduce_parent( + &self, + child: &ArrayRef, + parent: &ArrayRef, + child_idx: usize, + ) -> VortexResult>; +} + +pub(crate) type ReduceParentRuleRef = Arc; + +#[derive(Debug)] +struct ReduceParentFnRule(ReduceParentFn); + +impl DynReduceParentRule for ReduceParentFnRule { + fn reduce_parent( + &self, + child: &ArrayRef, + parent: &ArrayRef, + child_idx: usize, + ) -> VortexResult> { + self.0(child, parent, child_idx) + } +} + +#[derive(Debug)] +struct RegisteredReduceParentRule { + _child: V, + rule: R, +} + +impl DynReduceParentRule for RegisteredReduceParentRule +where + V: VTable, + R: ArrayParentReduceRule, +{ + fn reduce_parent( + &self, + child: &ArrayRef, + parent: &ArrayRef, + child_idx: usize, + ) -> VortexResult> { + let Some(child) = child.as_opt::() else { + return Ok(None); + }; + let Some(parent) = R::Parent::try_match(parent) else { + return Ok(None); + }; + self.rule.reduce_parent(child, parent, child_idx) + } +} + #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Ord, PartialOrd)] #[repr(transparent)] struct ReduceParentFnId(u64); @@ -172,8 +227,8 @@ impl Borrow for ExecuteParentFnId { } } -/// Registry of [`ReduceParentFn`]s, keyed by hashed `(parent, child)` pair. -type ReduceParentRegistry = ArcSwapMap>; +/// Registry of [`ReduceParentRuleRef`]s, keyed by hashed `(parent, child)` pair. +type ReduceParentRegistry = ArcSwapMap>; /// Registry of [`ExecuteParentKernelRef`]s, keyed by hashed `(parent, child)` pair. type ExecuteParentRegistry = ArcSwapMap>; @@ -223,15 +278,46 @@ impl ArrayKernels { /// If functions have already been registered for the same pair, these functions are appended /// after them. pub fn register_reduce_parent(&self, parent: Id, child: Id, fns: &[ReduceParentFn]) { + let rules: Vec = fns + .iter() + .map(|f| Arc::new(ReduceParentFnRule(*f)) as ReduceParentRuleRef) + .collect(); self.reduce_parent - .extend(hash_fn_id(parent, child).into(), fns); + .extend(hash_fn_id(parent, child).into(), rules.as_slice()); } - /// Look up the [`ReduceParentFn`]s registered for `(parent, child)`. + /// Register a typed [`ArrayParentReduceRule`] for `(parent, child.id())`, so a crate that + /// defines a scalar function can add parent-reduce rules for encodings it does not own. + /// + /// The rule applies when the child is a `V` array and the parent matches the rule's + /// [`Parent`](ArrayParentReduceRule::Parent) matcher. `parent` is usually the parent array's + /// encoding id. For `ScalarFnArray`, it is the scalar function id, for example `Cast.id()`. + /// + /// If rules have already been registered for the same pair, this rule is appended after them. + pub fn register_reduce_parent_rule(&self, parent: Id, child: V, rule: R) + where + V: VTable, + R: ArrayParentReduceRule, + { + let child_id = child.id(); + self.reduce_parent.push( + hash_fn_id(parent, child_id).into(), + Arc::new(RegisteredReduceParentRule { + _child: child, + rule, + }) as ReduceParentRuleRef, + ); + } + + /// Look up the parent-reduce rules registered for `(parent, child)`. /// /// Returns an owned [`Arc`] so the session-variable borrow can be dropped before invoking the - /// functions. - pub fn find_reduce_parent(&self, parent: Id, child: Id) -> Option> { + /// rules. + pub(crate) fn find_reduce_parent( + &self, + parent: Id, + child: Id, + ) -> Option> { self.reduce_parent.get(&hash_fn_id(parent, child)) } @@ -386,7 +472,7 @@ mod tests { use crate::ArrayVTable; use crate::arrays::Bool; use crate::scalar_fn::ScalarFnVTable; - use crate::scalar_fn::fns::binary::Binary; + use crate::scalar_fn::builtin::binary::Binary; static SESSION: LazyLock = LazyLock::new(|| VortexSession::empty().with::()); diff --git a/vortex-array/src/optimizer/mod.rs b/vortex-array/src/optimizer/mod.rs index 1d9c2c87da0..26dc59364ad 100644 --- a/vortex-array/src/optimizer/mod.rs +++ b/vortex-array/src/optimizer/mod.rs @@ -7,19 +7,22 @@ //! Optimization runs between execution steps, which is what enables cross-step optimizations: //! after a child is decoded, new `reduce_parent` rules may match that were previously blocked. //! -//! There are three public entry points on [`ArrayOptimizer`]: +//! There are two public entry points on [`ArrayOptimizer`], and both take the session whose +//! [`kernels::ArrayKernels`] registry is consulted before the static parent-reduce rules: //! -//! - [`ArrayOptimizer::optimize`] uses only static rules registered on encoding vtables. -//! - [`ArrayOptimizer::optimize_ctx`] also consults the session's active -//! [`kernels::ArrayKernels`] registry before static parent-reduce rules, so this is the entry -//! point used by execution. -//! - [`ArrayOptimizer::optimize_recursive`] applies the session-aware optimizer to the root and -//! every descendant. +//! - [`ArrayOptimizer::optimize`] optimizes the root array node. +//! - [`ArrayOptimizer::optimize_recursive`] optimizes the root and every descendant. +//! +//! Array constructors such as [`ArrayRef::slice`](crate::ArrayRef::slice) do not optimize: they +//! return the wrapped array. Execution optimizes the whole tree with its session before it starts, +//! and serialization optimizes any node that has no serialized form. Whenever a rule rewrites a +//! node, the nodes the rule built are optimized too, as constructing them used to do. use smallvec::SmallVec; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_session::VortexSession; +use vortex_utils::aliases::hash_set::HashSet; use crate::ArrayRef; use crate::optimizer::kernels::ArrayKernels; @@ -37,33 +40,21 @@ const MAX_OPTIMIZER_REWRITE_PASS: usize = 100; pub trait ArrayOptimizer { /// Optimize the root array node by running reduce and reduce_parent rules to fixpoint. /// - /// This uses only static rules registered on encoding vtables. Use [`Self::optimize_ctx`] - /// when a session-scoped [`kernels::ArrayKernels`] registry should participate. - fn optimize(&self) -> VortexResult; - - /// Optimize the root array node using static rules and the active - /// [`kernels::ArrayKernels`] registry on `session`, if any. - /// /// Session kernels are checked for each `(parent_encoding_id, child_encoding_id)` pair before /// the child's static `PARENT_RULES`. The registry comes from the [`kernels::KernelSession`] on - /// `session`, if any. If `session` does not contain a [`kernels::KernelSession`], this behaves - /// like [`Self::optimize`]. - fn optimize_ctx(&self, session: &VortexSession) -> VortexResult; + /// `session`. + fn optimize(&self, session: &VortexSession) -> VortexResult; /// Optimize the entire array tree recursively (root and all descendants). /// - /// This uses the same session-aware rule ordering as [`Self::optimize_ctx`] for every node in + /// This uses the same session-aware rule ordering as [`Self::optimize`] for every node in /// the tree. fn optimize_recursive(&self, session: &VortexSession) -> VortexResult; } impl ArrayOptimizer for ArrayRef { - fn optimize(&self) -> VortexResult { - Ok(optimize_owned(self.clone(), None)?.0) - } - - fn optimize_ctx(&self, session: &VortexSession) -> VortexResult { - Ok(optimize_owned(self.clone(), Some(session))?.0) + fn optimize(&self, session: &VortexSession) -> VortexResult { + Ok(optimize_owned(self.clone(), session)?.0) } fn optimize_recursive(&self, session: &VortexSession) -> VortexResult { @@ -73,18 +64,18 @@ impl ArrayOptimizer for ArrayRef { fn optimize_owned( mut current_array: ArrayRef, - session: Option<&VortexSession>, + session: &VortexSession, ) -> VortexResult<(ArrayRef, bool)> { let mut any_optimizations = false; - let session_kernels = session.map(|session| session.kernels()); + let session_kernels = session.kernels(); - trace_op!(record_optimize_start(¤t_array, session.is_some())); + trace_op!(record_optimize_start(¤t_array)); for _ in 0..=MAX_OPTIMIZER_REWRITE_PASS { trace_op!(record_optimize_loop_start(¤t_array)); if let Some(new_array) = current_array.reduce()? { - current_array = new_array; + current_array = optimize_new_nodes(new_array, ¤t_array, session)?; any_optimizations = true; trace_op!(record_optimize_loop_end()); continue; @@ -100,9 +91,8 @@ fn optimize_owned( }; // Session kernels take precedence over the child's static parent-reduce rules. - if let Some(session_kernels) = &session_kernels - && let Some(new_array) = - try_session_parent_reduce(session_kernels, ¤t_array, child, slot_idx)? + if let Some(new_array) = + try_session_parent_reduce(&session_kernels, ¤t_array, child, slot_idx)? { reduced_parent = Some(new_array); break; @@ -115,7 +105,7 @@ fn optimize_owned( } if let Some(new_array) = reduced_parent { - current_array = new_array; + current_array = optimize_new_nodes(new_array, ¤t_array, session)?; any_optimizations = true; trace_op!(record_optimize_loop_end()); continue; @@ -132,6 +122,54 @@ fn optimize_owned( vortex_bail!("Exceeded maximum optimization iterations (possible infinite loop)"); } +/// Optimizes the nodes a rule built below its result `array`, which replaces `old`. +/// +/// Rules construct their result from lazy array constructors, so the nodes they build are not yet +/// optimized, at any depth. Optimizing them as each rewrite happens keeps the optimizer from +/// undoing one rule with its inverse: `slice(get_item(s))` and `get_item(slice(s))` rewrite into +/// each other until the inner `slice(s)` collapses into `s`'s own encoding. Subtrees the rule +/// carried over from `old` are left as they are. +fn optimize_new_nodes( + array: ArrayRef, + old: &ArrayRef, + session: &VortexSession, +) -> VortexResult { + let mut existing = HashSet::new(); + let mut stack = vec![old]; + while let Some(node) = stack.pop() { + if existing.insert(node.addr()) { + stack.extend(node.slots().iter().flatten()); + } + } + optimize_new_slots(array, &existing, session) +} + +fn optimize_new_slots( + array: ArrayRef, + existing: &HashSet, + session: &VortexSession, +) -> VortexResult { + let mut new_slots = SmallVec::with_capacity(array.slots().len()); + let mut any_slot_optimized = false; + for slot in array.slots() { + match slot { + Some(child) if !existing.contains(&child.addr()) => { + let (optimized, _) = optimize_owned(child.clone(), session)?; + let optimized = optimize_new_slots(optimized, existing, session)?; + any_slot_optimized |= !ArrayRef::ptr_eq(&optimized, child); + new_slots.push(Some(optimized)); + } + slot => new_slots.push(slot.clone()), + } + } + if !any_slot_optimized { + return Ok(array); + } + // SAFETY: optimizer rules only replace child slots with logically equivalent arrays, so the + // parent's logical values and statistics remain valid. + unsafe { array.with_slots(new_slots) } +} + fn try_session_parent_reduce( kernels: &ArrayKernels, parent: &ArrayRef, @@ -146,7 +184,7 @@ fn try_session_parent_reduce( #[allow(clippy::unused_enumerate_index)] for (_kernel_idx, reduce_parent) in reduce_parent_fns.iter().enumerate() { - if let Some(new_array) = reduce_parent(child, parent, slot_idx)? { + if let Some(new_array) = reduce_parent.reduce_parent(child, parent, slot_idx)? { trace_op!(record_session_parent_reduce_applied( parent, child, @@ -175,7 +213,7 @@ fn try_optimize_recursive( ) -> VortexResult<(ArrayRef, bool)> { trace_op!(record_optimize_recursive_start(¤t_array)); - let (mut current_array, mut any_optimizations) = optimize_owned(current_array, Some(session))?; + let (mut current_array, mut any_optimizations) = optimize_owned(current_array, session)?; let mut new_slots = SmallVec::with_capacity(current_array.slots().len()); let mut any_slot_optimized = false; diff --git a/vortex-array/src/patches.rs b/vortex-array/src/patches.rs index 3dab7b37c3d..832f3239c9b 100644 --- a/vortex-array/src/patches.rs +++ b/vortex-array/src/patches.rs @@ -662,7 +662,7 @@ impl Patches { /// Unlike filter, this preserves the patch indices. /// Unlike mask on a single array, this does not set masked values to null. /// - /// Masking an array always yields a nullable result (see [`crate::scalar_fn::fns::mask`]), + /// Masking an array always yields a nullable result (see [`crate::scalar_fn::builtin::mask`]), /// so the surviving patch values are widened to nullable to match the masked parent. Callers /// therefore receive patches with the correct dtype and do not need to re-cast them. // TODO(joe): make this lazy and remove the ctx. diff --git a/vortex-array/src/scalar/typed_view/list.rs b/vortex-array/src/scalar/typed_view/list.rs index 01e91800432..35b24012771 100644 --- a/vortex-array/src/scalar/typed_view/list.rs +++ b/vortex-array/src/scalar/typed_view/list.rs @@ -164,7 +164,7 @@ impl<'a> ListScalar<'a> { /// scalar for each. /// /// Returns None if the list is null. - pub(crate) fn element_values(&self) -> Option<&'a [Option]> { + pub fn element_values(&self) -> Option<&'a [Option]> { self.elements } diff --git a/vortex-array/src/scalar/typed_view/primitive/numeric_operator.rs b/vortex-array/src/scalar/typed_view/primitive/numeric_operator.rs index d7bcc8d0b4c..bdb33005ec0 100644 --- a/vortex-array/src/scalar/typed_view/primitive/numeric_operator.rs +++ b/vortex-array/src/scalar/typed_view/primitive/numeric_operator.rs @@ -8,7 +8,7 @@ use std::fmt; use vortex_error::VortexError; use vortex_error::vortex_err; -use crate::scalar_fn::fns::operators::Operator; +use crate::scalar_fn::builtin::operators::Operator; #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] /// Binary element-wise operations. diff --git a/vortex-array/src/scalar_fn/fns/binary/boolean.rs b/vortex-array/src/scalar_fn/builtin/binary/boolean.rs similarity index 99% rename from vortex-array/src/scalar_fn/fns/binary/boolean.rs rename to vortex-array/src/scalar_fn/builtin/binary/boolean.rs index df23e5a87cb..20b6dd7204a 100644 --- a/vortex-array/src/scalar_fn/fns/binary/boolean.rs +++ b/vortex-array/src/scalar_fn/builtin/binary/boolean.rs @@ -33,8 +33,8 @@ use crate::dtype::Nullability; use crate::kernel::ExecuteParentKernel; use crate::scalar::BoolScalar; use crate::scalar::Scalar; -use crate::scalar_fn::fns::binary::Binary; -use crate::scalar_fn::fns::operators::Operator; +use crate::scalar_fn::builtin::binary::Binary; +use crate::scalar_fn::builtin::operators::Operator; use crate::validity::Validity; /// Trait for encoding-specific boolean kernels that operate in encoded space. @@ -764,7 +764,7 @@ mod tests { use crate::dtype::DType; use crate::dtype::Nullability; use crate::scalar::Scalar; - use crate::scalar_fn::fns::operators::Operator; + use crate::scalar_fn::builtin::operators::Operator; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/boolean.rs b/vortex-array/src/scalar_fn/builtin/binary/compare/boolean.rs similarity index 97% rename from vortex-array/src/scalar_fn/fns/binary/compare/boolean.rs rename to vortex-array/src/scalar_fn/builtin/binary/compare/boolean.rs index b569e791f32..816b471e2fa 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/boolean.rs +++ b/vortex-array/src/scalar_fn/builtin/binary/compare/boolean.rs @@ -14,8 +14,8 @@ use crate::IntoArray; use crate::arrays::BoolArray; use crate::arrays::Constant; use crate::dtype::Nullability; -use crate::scalar_fn::fns::binary::compare::compare_validity; -use crate::scalar_fn::fns::operators::CompareOperator; +use crate::scalar_fn::builtin::binary::compare::compare_validity; +use crate::scalar_fn::builtin::operators::CompareOperator; use crate::validity::Validity; enum BoolOperand { diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/bytes.rs b/vortex-array/src/scalar_fn/builtin/binary/compare/bytes.rs similarity index 98% rename from vortex-array/src/scalar_fn/fns/binary/compare/bytes.rs rename to vortex-array/src/scalar_fn/builtin/binary/compare/bytes.rs index ef833adf6c4..0287cc2882b 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/bytes.rs +++ b/vortex-array/src/scalar_fn/builtin/binary/compare/bytes.rs @@ -27,9 +27,9 @@ use crate::arrays::varbinview::ResolvedViews; use crate::dtype::DType; use crate::dtype::Nullability; use crate::scalar::Scalar; -use crate::scalar_fn::fns::binary::compare::compare_validity; -use crate::scalar_fn::fns::binary::compare::ordering_predicate; -use crate::scalar_fn::fns::operators::CompareOperator; +use crate::scalar_fn::builtin::binary::compare::compare_validity; +use crate::scalar_fn::builtin::binary::compare::ordering_predicate; +use crate::scalar_fn::builtin::operators::CompareOperator; use crate::validity::Validity; enum BytesOperand { diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/decimal.rs b/vortex-array/src/scalar_fn/builtin/binary/compare/decimal.rs similarity index 98% rename from vortex-array/src/scalar_fn/fns/binary/compare/decimal.rs rename to vortex-array/src/scalar_fn/builtin/binary/compare/decimal.rs index e88e12f11d0..bac46c35d70 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/decimal.rs +++ b/vortex-array/src/scalar_fn/builtin/binary/compare/decimal.rs @@ -28,8 +28,8 @@ use crate::dtype::Nullability; use crate::dtype::i256; use crate::match_each_decimal_value_type; use crate::scalar::DecimalValue; -use crate::scalar_fn::fns::binary::compare::compare_validity; -use crate::scalar_fn::fns::operators::CompareOperator; +use crate::scalar_fn::builtin::binary::compare::compare_validity; +use crate::scalar_fn::builtin::operators::CompareOperator; use crate::validity::Validity; enum DecimalOperand { diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs b/vortex-array/src/scalar_fn/builtin/binary/compare/mod.rs similarity index 98% rename from vortex-array/src/scalar_fn/fns/binary/compare/mod.rs rename to vortex-array/src/scalar_fn/builtin/binary/compare/mod.rs index 6593af28c6e..36041f1f7d3 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs +++ b/vortex-array/src/scalar_fn/builtin/binary/compare/mod.rs @@ -35,15 +35,16 @@ use crate::dtype::DType; use crate::dtype::Nullability; use crate::kernel::ExecuteParentKernel; use crate::scalar::Scalar; -use crate::scalar_fn::fns::binary::Binary; -use crate::scalar_fn::fns::operators::CompareOperator; +use crate::scalar_fn::builtin::binary::Binary; +use crate::scalar_fn::builtin::operators::CompareOperator; use crate::validity::Validity; mod boolean; mod bytes; mod decimal; mod nested; -pub(crate) use nested::build_comparator as build_row_comparator; +pub use nested::RowComparator; +pub use nested::build_comparator as build_row_comparator; mod primitive; #[cfg(test)] mod tests; diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/nested.rs b/vortex-array/src/scalar_fn/builtin/binary/compare/nested.rs similarity index 97% rename from vortex-array/src/scalar_fn/fns/binary/compare/nested.rs rename to vortex-array/src/scalar_fn/builtin/binary/compare/nested.rs index ccccc3ca955..b95f8b55bde 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/nested.rs +++ b/vortex-array/src/scalar_fn/builtin/binary/compare/nested.rs @@ -45,11 +45,12 @@ use crate::dtype::NativePType; use crate::dtype::Nullability; use crate::match_each_integer_ptype; use crate::match_each_native_ptype; -use crate::scalar_fn::fns::binary::compare::compare_validity; -use crate::scalar_fn::fns::operators::CompareOperator; +use crate::scalar_fn::builtin::binary::compare::compare_validity; +use crate::scalar_fn::builtin::operators::CompareOperator; /// A row comparator: compares row `i` of the left operand against row `j` of the right operand. -pub(crate) type RowComparator = Box Ordering>; +/// Compares row `i` of the left array with row `j` of the right array. +pub type RowComparator = Box Ordering>; /// Compare two nested arrays row by row. pub(super) fn compare_nested( @@ -116,7 +117,7 @@ fn validity_mask(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult /// Build a row comparator over two recursively canonical arrays of the same logical dtype /// (ignoring nullability). Null values order before all non-null values at every level. -pub(crate) fn build_comparator( +pub fn build_comparator( lhs: &ArrayRef, rhs: &ArrayRef, ctx: &mut ExecutionCtx, diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs b/vortex-array/src/scalar_fn/builtin/binary/compare/primitive.rs similarity index 97% rename from vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs rename to vortex-array/src/scalar_fn/builtin/binary/compare/primitive.rs index a1da5a84097..64e836387b7 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs +++ b/vortex-array/src/scalar_fn/builtin/binary/compare/primitive.rs @@ -24,9 +24,9 @@ use crate::dtype::Nullability; use crate::dtype::PType; use crate::match_each_native_ptype; use crate::scalar::Scalar; -use crate::scalar_fn::fns::binary::compare::compare_validity; -use crate::scalar_fn::fns::binary::primitive_operand::PrimitiveOperand; -use crate::scalar_fn::fns::operators::CompareOperator; +use crate::scalar_fn::builtin::binary::compare::compare_validity; +use crate::scalar_fn::builtin::binary::primitive_operand::PrimitiveOperand; +use crate::scalar_fn::builtin::operators::CompareOperator; /// Compare two primitive arrays of the same [`PType`]. /// diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs b/vortex-array/src/scalar_fn/builtin/binary/compare/tests.rs similarity index 99% rename from vortex-array/src/scalar_fn/fns/binary/compare/tests.rs rename to vortex-array/src/scalar_fn/builtin/binary/compare/tests.rs index 83d01e2d831..cc9ad9a6348 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs +++ b/vortex-array/src/scalar_fn/builtin/binary/compare/tests.rs @@ -43,9 +43,9 @@ use crate::memory::MemorySessionExt; use crate::memory::test_allocator::counting_allocator; use crate::scalar::DecimalValue; use crate::scalar::Scalar; -use crate::scalar_fn::fns::binary::scalar_cmp; -use crate::scalar_fn::fns::operators::CompareOperator; -use crate::scalar_fn::fns::operators::Operator; +use crate::scalar_fn::builtin::binary::scalar_cmp; +use crate::scalar_fn::builtin::operators::CompareOperator; +use crate::scalar_fn::builtin::operators::Operator; use crate::validity::Validity; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); diff --git a/vortex-array/src/scalar_fn/fns/binary/mod.rs b/vortex-array/src/scalar_fn/builtin/binary/mod.rs similarity index 97% rename from vortex-array/src/scalar_fn/fns/binary/mod.rs rename to vortex-array/src/scalar_fn/builtin/binary/mod.rs index 42cebc8401c..2b90f2106b2 100644 --- a/vortex-array/src/scalar_fn/fns/binary/mod.rs +++ b/vortex-array/src/scalar_fn/builtin/binary/mod.rs @@ -30,9 +30,9 @@ use crate::scalar_fn::ReduceNode; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::ScalarFnVTableExt; -use crate::scalar_fn::fns::literal::Literal; -use crate::scalar_fn::fns::operators::CompareOperator; -use crate::scalar_fn::fns::operators::Operator; +use crate::scalar_fn::builtin::literal::Literal; +use crate::scalar_fn::builtin::operators::CompareOperator; +use crate::scalar_fn::builtin::operators::Operator; pub mod boolean; pub use boolean::BooleanExecuteAdaptor; @@ -597,40 +597,40 @@ mod tests { let array = Binary::try_new(left.clone(), right.clone(), Operator::And)? .into_array() - .optimize()?; // calls reduce() + .optimize(&SESSION)?; // calls reduce() assert!(array.is::()); // and(left, const) -> left assert_arrays_eq!(array, left, &mut ctx); let array = Binary::try_new(right.clone(), left.clone(), Operator::And)? .into_array() - .optimize()?; + .optimize(&SESSION)?; assert!(array.is::()); // and(const, left) -> left assert_arrays_eq!(array, left, &mut ctx); let array = Binary::try_new(left.clone(), right.clone(), Operator::Or)? .into_array() - .optimize()?; + .optimize(&SESSION)?; assert_eq!(array.as_constant(), Some(true.into())); // or(left, const) -> true let array = Binary::try_new(right.clone(), left, Operator::Or)? .into_array() - .optimize()?; + .optimize(&SESSION)?; assert_eq!(array.as_constant(), Some(true.into())); // or(const, left) -> true let array = Binary::try_new(right.clone(), right.clone(), Operator::Or)? .into_array() - .optimize()?; + .optimize(&SESSION)?; assert_eq!(array.as_constant(), Some(true.into())); // or(const, const) -> const let left_false = ConstantArray::new(false, 3).into_array(); let array = Binary::try_new(right.clone(), left_false.clone(), Operator::Or)? .into_array() - .optimize()?; + .optimize(&SESSION)?; assert_eq!(array.as_constant(), Some(true.into())); let array = Binary::try_new(left_false, right, Operator::And)? .into_array() - .optimize()?; + .optimize(&SESSION)?; assert_eq!(array.as_constant(), Some(false.into())); Ok(()) @@ -645,7 +645,7 @@ mod tests { let array = Binary::try_new(nullable_true, right.clone(), Operator::And)? .into_array() - .optimize()?; + .optimize(&SESSION)?; assert_arrays_eq!(array, right, &mut ctx); Ok(()) } @@ -669,7 +669,7 @@ mod tests { let array = lhs .binary(rhs.clone(), Operator::And)? .is_null()? - .optimize()?; + .optimize(&SESSION)?; assert_eq!(*array.as_::>().options, Operator::And); assert_arrays_eq!( array.as_::().get_child(0), // not(left.validity()) @@ -694,7 +694,7 @@ mod tests { let array = lhs .binary(rhs.clone(), Operator::Or)? .is_not_null()? - .optimize()?; + .optimize(&SESSION)?; assert_eq!(*array.as_::>().options, Operator::Or); assert_arrays_eq!( array.as_::().get_child(0), // left.validity() diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/checked.rs b/vortex-array/src/scalar_fn/builtin/binary/numeric/checked.rs similarity index 100% rename from vortex-array/src/scalar_fn/fns/binary/numeric/checked.rs rename to vortex-array/src/scalar_fn/builtin/binary/numeric/checked.rs diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/decimal.rs b/vortex-array/src/scalar_fn/builtin/binary/numeric/decimal.rs similarity index 100% rename from vortex-array/src/scalar_fn/fns/binary/numeric/decimal.rs rename to vortex-array/src/scalar_fn/builtin/binary/numeric/decimal.rs diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs b/vortex-array/src/scalar_fn/builtin/binary/numeric/mod.rs similarity index 100% rename from vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs rename to vortex-array/src/scalar_fn/builtin/binary/numeric/mod.rs diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs b/vortex-array/src/scalar_fn/builtin/binary/numeric/primitive.rs similarity index 100% rename from vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs rename to vortex-array/src/scalar_fn/builtin/binary/numeric/primitive.rs diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs b/vortex-array/src/scalar_fn/builtin/binary/numeric/row.rs similarity index 98% rename from vortex-array/src/scalar_fn/fns/binary/numeric/row.rs rename to vortex-array/src/scalar_fn/builtin/binary/numeric/row.rs index ff7fbe7ab81..8ca4f49edaa 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs +++ b/vortex-array/src/scalar_fn/builtin/binary/numeric/row.rs @@ -26,7 +26,7 @@ use crate::scalar::NumericOperator; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::VecExecutionArgs; -use crate::scalar_fn::fns::binary::Binary; +use crate::scalar_fn::builtin::binary::Binary; use crate::scalar_fn::unstable::row::InitializedElement; use crate::scalar_fn::unstable::row::RowFn; use crate::scalar_fn::unstable::row::RowVisitor; diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs b/vortex-array/src/scalar_fn/builtin/binary/numeric/tests.rs similarity index 99% rename from vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs rename to vortex-array/src/scalar_fn/builtin/binary/numeric/tests.rs index 325e55665f8..f6eb344908a 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs +++ b/vortex-array/src/scalar_fn/builtin/binary/numeric/tests.rs @@ -31,7 +31,7 @@ use crate::dtype::i256; use crate::scalar::DecimalValue; use crate::scalar::NumericOperator; use crate::scalar::Scalar; -use crate::scalar_fn::fns::operators::Operator; +use crate::scalar_fn::builtin::operators::Operator; use crate::validity::Validity; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); diff --git a/vortex-array/src/scalar_fn/fns/binary/primitive_operand.rs b/vortex-array/src/scalar_fn/builtin/binary/primitive_operand.rs similarity index 100% rename from vortex-array/src/scalar_fn/fns/binary/primitive_operand.rs rename to vortex-array/src/scalar_fn/builtin/binary/primitive_operand.rs diff --git a/vortex-array/src/scalar_fn/fns/cast/kernel.rs b/vortex-array/src/scalar_fn/builtin/cast/kernel.rs similarity index 98% rename from vortex-array/src/scalar_fn/fns/cast/kernel.rs rename to vortex-array/src/scalar_fn/builtin/cast/kernel.rs index ca9783cc22e..3f9d6f45eb4 100644 --- a/vortex-array/src/scalar_fn/fns/cast/kernel.rs +++ b/vortex-array/src/scalar_fn/builtin/cast/kernel.rs @@ -13,7 +13,7 @@ use crate::dtype::DType; use crate::kernel::ExecuteParentKernel; use crate::matcher::Matcher; use crate::optimizer::rules::ArrayParentReduceRule; -use crate::scalar_fn::fns::cast::Cast; +use crate::scalar_fn::builtin::cast::Cast; /// Reduce rule for cast: restructure the array without reading buffers. /// diff --git a/vortex-array/src/scalar_fn/fns/cast/mod.rs b/vortex-array/src/scalar_fn/builtin/cast/mod.rs similarity index 99% rename from vortex-array/src/scalar_fn/fns/cast/mod.rs rename to vortex-array/src/scalar_fn/builtin/cast/mod.rs index f3d90cf05fc..cb06ca41c66 100644 --- a/vortex-array/src/scalar_fn/fns/cast/mod.rs +++ b/vortex-array/src/scalar_fn/builtin/cast/mod.rs @@ -49,7 +49,7 @@ use crate::scalar_fn::ReduceNodeValidity; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::ScalarFnVTableExt; -use crate::scalar_fn::fns::literal::Literal; +use crate::scalar_fn::builtin::literal::Literal; use crate::scalar_fn::is_not_null_node; /// A cast expression that converts values to a target data type. @@ -265,7 +265,7 @@ mod tests { use crate::expr::test_harness; use crate::scalar::DecimalValue; use crate::scalar::Scalar; - use crate::scalar_fn::fns::literal::Literal; + use crate::scalar_fn::builtin::literal::Literal; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); diff --git a/vortex-array/src/scalar_fn/fns/fill_null/kernel.rs b/vortex-array/src/scalar_fn/builtin/fill_null/kernel.rs similarity index 98% rename from vortex-array/src/scalar_fn/fns/fill_null/kernel.rs rename to vortex-array/src/scalar_fn/builtin/fill_null/kernel.rs index 1fdcf1c2d76..a37ce62a70c 100644 --- a/vortex-array/src/scalar_fn/fns/fill_null/kernel.rs +++ b/vortex-array/src/scalar_fn/builtin/fill_null/kernel.rs @@ -20,7 +20,7 @@ use crate::builtins::ArrayBuiltins; use crate::kernel::ExecuteParentKernel; use crate::optimizer::rules::ArrayParentReduceRule; use crate::scalar::Scalar; -use crate::scalar_fn::fns::fill_null::FillNull as FillNullExpr; +use crate::scalar_fn::builtin::fill_null::FillNull as FillNullExpr; use crate::validity::Validity; /// Fill nulls in an array with a scalar value without reading buffers. diff --git a/vortex-array/src/scalar_fn/fns/fill_null/mod.rs b/vortex-array/src/scalar_fn/builtin/fill_null/mod.rs similarity index 98% rename from vortex-array/src/scalar_fn/fns/fill_null/mod.rs rename to vortex-array/src/scalar_fn/builtin/fill_null/mod.rs index 5d17aaa68b8..74d0cfa49d0 100644 --- a/vortex-array/src/scalar_fn/fns/fill_null/mod.rs +++ b/vortex-array/src/scalar_fn/builtin/fill_null/mod.rs @@ -58,7 +58,7 @@ impl FillNull { /// These are the dtypes whose canonical arrays have a [`FillNullKernel`], so update this list /// whenever `fill_null_canonical` gains or loses one. Rewrites into fill_null, such as the /// CASE WHEN simplification, check it first because other dtypes fail at execution time. - pub(crate) fn supports_dtype(dtype: &DType) -> bool { + pub fn supports_dtype(dtype: &DType) -> bool { matches!( dtype, DType::Bool(_) @@ -156,7 +156,7 @@ impl ScalarFnVTable for FillNull { // kernels get a chance before we decode the input. return FillNull::try_new(input, fill_value)? .into_array() - .optimize_ctx(ctx.session()); + .optimize(ctx.session()); } let input = input.execute::(ctx)?; return Ok(FillNull::try_new(input, fill_value)?.into_array()); @@ -275,8 +275,8 @@ mod tests { use crate::matcher::Matcher; use crate::optimizer::rules::ArrayParentReduceRule; use crate::scalar::Scalar; - use crate::scalar_fn::fns::binary::Binary; - use crate::scalar_fn::fns::operators::Operator; + use crate::scalar_fn::builtin::binary::Binary; + use crate::scalar_fn::builtin::operators::Operator; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); diff --git a/vortex-array/src/scalar_fn/fns/get_item.rs b/vortex-array/src/scalar_fn/builtin/get_item.rs similarity index 97% rename from vortex-array/src/scalar_fn/fns/get_item.rs rename to vortex-array/src/scalar_fn/builtin/get_item.rs index 663e3eebdbf..c3a24588180 100644 --- a/vortex-array/src/scalar_fn/fns/get_item.rs +++ b/vortex-array/src/scalar_fn/builtin/get_item.rs @@ -34,8 +34,8 @@ use crate::scalar_fn::ReduceNode; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::ScalarFnVTableExt; -use crate::scalar_fn::fns::mask::Mask; -use crate::scalar_fn::fns::pack::Pack; +use crate::scalar_fn::builtin::mask::Mask; +use crate::scalar_fn::builtin::pack::Pack; #[derive(Clone)] pub struct GetItem; @@ -59,10 +59,7 @@ impl GetItem { /// `get_item` node themselves. The optimizer visits children before their parent, so a /// `get_item` that a rule introduces is never simplified again. Removing it here keeps the /// output of the rule a fixed point. - pub(crate) fn apply_now( - struct_node: &T, - field_name: &FieldName, - ) -> VortexResult { + pub fn apply_now(struct_node: &T, field_name: &FieldName) -> VortexResult { if let Some(field) = Self::maybe_apply_now(struct_node, field_name)? { return Ok(field); } @@ -256,8 +253,9 @@ mod tests { use crate::expr::lit; use crate::expr::pack; use crate::expr::root; + use crate::optimizer::ArrayOptimizer; use crate::scalar::Scalar; - use crate::scalar_fn::fns::get_item::StructArray; + use crate::scalar_fn::builtin::get_item::StructArray; use crate::validity::Validity; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); @@ -380,7 +378,7 @@ mod tests { .into_array() .filter(Mask::from_iter([true, false, true]))?; - let item = filtered.get_item("a")?; + let item = filtered.get_item("a")?.optimize(&SESSION)?; assert!(item.is::()); assert!(!item.is::()); @@ -396,7 +394,7 @@ mod tests { fn get_item_pushes_through_slice() -> VortexResult<()> { let sliced = test_array().into_array().slice(1..3)?; - let item = sliced.get_item("a")?; + let item = sliced.get_item("a")?.optimize(&SESSION)?; assert!(item.is::()); assert!(!item.is::()); @@ -413,9 +411,10 @@ mod tests { fn get_item_pushes_through_masked() -> VortexResult<()> { let masked = test_array() .into_array() - .mask(Validity::from_iter([true, false, true]).to_array(3))?; + .mask(Validity::from_iter([true, false, true]).to_array(3))? + .optimize(&SESSION)?; - let item = masked.get_item("a")?; + let item = masked.get_item("a")?.optimize(&SESSION)?; assert!(item.is::()); assert_eq!( diff --git a/vortex-array/src/scalar_fn/fns/is_not_null.rs b/vortex-array/src/scalar_fn/builtin/is_not_null.rs similarity index 98% rename from vortex-array/src/scalar_fn/fns/is_not_null.rs rename to vortex-array/src/scalar_fn/builtin/is_not_null.rs index 253aaa01dcd..26caf1a9afb 100644 --- a/vortex-array/src/scalar_fn/fns/is_not_null.rs +++ b/vortex-array/src/scalar_fn/builtin/is_not_null.rs @@ -24,10 +24,10 @@ use crate::scalar_fn::ReduceNodeValidity; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::ScalarFnVTableExt; -use crate::scalar_fn::fns::binary::Binary; -use crate::scalar_fn::fns::is_null::IsNull; -use crate::scalar_fn::fns::not::Not; -use crate::scalar_fn::fns::operators::Operator; +use crate::scalar_fn::builtin::binary::Binary; +use crate::scalar_fn::builtin::is_null::IsNull; +use crate::scalar_fn::builtin::not::Not; +use crate::scalar_fn::builtin::operators::Operator; use crate::scalar_fn::is_not_null_node; /// In array context, reduce diff --git a/vortex-array/src/scalar_fn/fns/is_null.rs b/vortex-array/src/scalar_fn/builtin/is_null.rs similarity index 99% rename from vortex-array/src/scalar_fn/fns/is_null.rs rename to vortex-array/src/scalar_fn/builtin/is_null.rs index 510ef6d6323..05eceba24a9 100644 --- a/vortex-array/src/scalar_fn/fns/is_null.rs +++ b/vortex-array/src/scalar_fn/builtin/is_null.rs @@ -22,7 +22,7 @@ use crate::scalar_fn::ReduceNode; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::ScalarFnVTableExt; -use crate::scalar_fn::fns::is_not_null::reduce_null; +use crate::scalar_fn::builtin::is_not_null::reduce_null; use crate::validity::Validity; /// Expression that checks for null values. diff --git a/vortex-array/src/scalar_fn/fns/literal.rs b/vortex-array/src/scalar_fn/builtin/literal.rs similarity index 100% rename from vortex-array/src/scalar_fn/fns/literal.rs rename to vortex-array/src/scalar_fn/builtin/literal.rs diff --git a/vortex-array/src/scalar_fn/fns/mask/kernel.rs b/vortex-array/src/scalar_fn/builtin/mask/kernel.rs similarity index 95% rename from vortex-array/src/scalar_fn/fns/mask/kernel.rs rename to vortex-array/src/scalar_fn/builtin/mask/kernel.rs index b03c2dd5509..faf6b19bf88 100644 --- a/vortex-array/src/scalar_fn/fns/mask/kernel.rs +++ b/vortex-array/src/scalar_fn/builtin/mask/kernel.rs @@ -16,7 +16,7 @@ use crate::arrays::scalar_fn::ExactScalarFn; use crate::arrays::scalar_fn::ScalarFnArrayView; use crate::kernel::ExecuteParentKernel; use crate::optimizer::rules::ArrayParentReduceRule; -use crate::scalar_fn::fns::mask::Mask as MaskExpr; +use crate::scalar_fn::builtin::mask::Mask as MaskExpr; /// Mask an array without reading buffers. /// @@ -150,9 +150,9 @@ mod tests { use crate::executor::VortexSessionExecute; use crate::optimizer::ArrayOptimizer; use crate::scalar::Scalar; - use crate::scalar_fn::fns::binary::Binary; - use crate::scalar_fn::fns::mask::Mask as MaskExpr; - use crate::scalar_fn::fns::operators::Operator; + use crate::scalar_fn::builtin::binary::Binary; + use crate::scalar_fn::builtin::mask::Mask as MaskExpr; + use crate::scalar_fn::builtin::operators::Operator; use crate::validity::Validity; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); @@ -172,7 +172,11 @@ mod tests { fn all_valid_input_attaches_lazy_mask(#[case] validity: Validity) -> VortexResult<()> { let input = PrimitiveArray::new(buffer![1i32, 2, 3], validity); let mask = lazy_mask()?; - let output = input.clone().into_array().mask(mask.clone())?; + let output = input + .clone() + .into_array() + .mask(mask.clone())? + .optimize(&SESSION)?; let primitive = output.as_::(); assert!( @@ -236,7 +240,7 @@ mod tests { "expected an un-optimized ScalarFn wrapper before optimization" ); - let optimized = masked.optimize()?; + let optimized = masked.optimize(&SESSION)?; assert!( !optimized.is::(), "constant mask should not fall through to execution, got {}", diff --git a/vortex-array/src/scalar_fn/fns/mask/mod.rs b/vortex-array/src/scalar_fn/builtin/mask/mod.rs similarity index 97% rename from vortex-array/src/scalar_fn/fns/mask/mod.rs rename to vortex-array/src/scalar_fn/builtin/mask/mod.rs index 994c409adc3..6836d33a205 100644 --- a/vortex-array/src/scalar_fn/fns/mask/mod.rs +++ b/vortex-array/src/scalar_fn/builtin/mask/mod.rs @@ -34,9 +34,9 @@ use crate::scalar_fn::ReduceNodeValidity; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::ScalarFnVTableExt; -use crate::scalar_fn::fns::binary::Binary; -use crate::scalar_fn::fns::literal::Literal; -use crate::scalar_fn::fns::operators::Operator; +use crate::scalar_fn::builtin::binary::Binary; +use crate::scalar_fn::builtin::literal::Literal; +use crate::scalar_fn::builtin::operators::Operator; use crate::scalar_fn::is_not_null_node; /// An expression that masks an input based on a boolean mask. diff --git a/vortex-array/src/scalar_fn/builtin/mod.rs b/vortex-array/src/scalar_fn/builtin/mod.rs new file mode 100644 index 00000000000..c70e64273ae --- /dev/null +++ b/vortex-array/src/scalar_fn/builtin/mod.rs @@ -0,0 +1,28 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The scalar functions that vortex-array itself is built on. +//! +//! These are the functions that the core crate's own arrays, validity and expression machinery +//! construct: literals and struct projection and packing for expressions, `is_null`, +//! `is_not_null`, `not` and the boolean operators for validity reduction, and `cast`, `mask`, +//! `fill_null` and the comparison and arithmetic operators, which canonical arrays apply lazily to +//! their own children (offsets, indices, codes and validity). [`ArrayBuiltins`] exposes them as +//! methods. +//! +//! Every other scalar function lives in the `vortex-functions` crate, which registers them, and +//! their kernels for vortex-array's encodings, with a session. +//! +//! [`ArrayBuiltins`]: crate::builtins::ArrayBuiltins + +pub mod binary; +pub mod cast; +pub mod fill_null; +pub mod get_item; +pub mod is_not_null; +pub mod is_null; +pub mod literal; +pub mod mask; +pub mod not; +pub mod operators; +pub mod pack; diff --git a/vortex-array/src/scalar_fn/fns/not/kernel.rs b/vortex-array/src/scalar_fn/builtin/not/kernel.rs similarity index 96% rename from vortex-array/src/scalar_fn/fns/not/kernel.rs rename to vortex-array/src/scalar_fn/builtin/not/kernel.rs index 3da8f66e17f..9e601eb8008 100644 --- a/vortex-array/src/scalar_fn/fns/not/kernel.rs +++ b/vortex-array/src/scalar_fn/builtin/not/kernel.rs @@ -9,7 +9,7 @@ use crate::array::VTable; use crate::arrays::scalar_fn::ExactScalarFn; use crate::arrays::scalar_fn::ScalarFnArrayView; use crate::optimizer::rules::ArrayParentReduceRule; -use crate::scalar_fn::fns::not::Not as NotExpr; +use crate::scalar_fn::builtin::not::Not as NotExpr; /// Invert a boolean array without reading buffers. /// diff --git a/vortex-array/src/scalar_fn/fns/not/mod.rs b/vortex-array/src/scalar_fn/builtin/not/mod.rs similarity index 99% rename from vortex-array/src/scalar_fn/fns/not/mod.rs rename to vortex-array/src/scalar_fn/builtin/not/mod.rs index b4c0c1b3bb5..b09eed78096 100644 --- a/vortex-array/src/scalar_fn/fns/not/mod.rs +++ b/vortex-array/src/scalar_fn/builtin/not/mod.rs @@ -137,7 +137,7 @@ mod tests { use crate::expr::not; use crate::expr::root; use crate::expr::test_harness; - use crate::scalar_fn::fns::not::BoolArray; + use crate::scalar_fn::builtin::not::BoolArray; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); diff --git a/vortex-array/src/scalar_fn/fns/operators.rs b/vortex-array/src/scalar_fn/builtin/operators.rs similarity index 100% rename from vortex-array/src/scalar_fn/fns/operators.rs rename to vortex-array/src/scalar_fn/builtin/operators.rs diff --git a/vortex-array/src/scalar_fn/fns/pack.rs b/vortex-array/src/scalar_fn/builtin/pack.rs similarity index 99% rename from vortex-array/src/scalar_fn/fns/pack.rs rename to vortex-array/src/scalar_fn/builtin/pack.rs index c9e5a60ec3a..47e41a5054d 100644 --- a/vortex-array/src/scalar_fn/fns/pack.rs +++ b/vortex-array/src/scalar_fn/builtin/pack.rs @@ -182,7 +182,7 @@ mod tests { use crate::expr::col; use crate::expr::pack; use crate::scalar_fn::ScalarFnVTableExt; - use crate::scalar_fn::fns::pack::StructArray; + use crate::scalar_fn::builtin::pack::StructArray; use crate::validity::Validity; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); diff --git a/vortex-array/src/scalar_fn/fns/mod.rs b/vortex-array/src/scalar_fn/fns/mod.rs deleted file mode 100644 index 3cb30209dd9..00000000000 --- a/vortex-array/src/scalar_fn/fns/mod.rs +++ /dev/null @@ -1,30 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -pub mod between; -pub mod binary; -pub mod byte_length; -pub mod case_when; -pub mod cast; -pub mod dynamic; -pub mod ext_storage; -pub mod fill_null; -pub mod get_item; -pub mod is_nan; -pub mod is_not_null; -pub mod is_null; -pub mod like; -pub mod list_contains; -pub mod list_length; -pub mod list_sum; -pub mod literal; -pub mod mask; -pub mod merge; -pub mod not; -pub mod operators; -pub mod pack; -pub mod replace_time_zone; -pub mod select; -pub mod stat; -pub mod variant_get; -pub mod zip; diff --git a/vortex-array/src/scalar_fn/mod.rs b/vortex-array/src/scalar_fn/mod.rs index 6be34ce1f34..8046b7bb2d0 100644 --- a/vortex-array/src/scalar_fn/mod.rs +++ b/vortex-array/src/scalar_fn/mod.rs @@ -3,9 +3,10 @@ //! Scalar function vtable machinery. //! -//! This module contains the [`ScalarFnVTable`] trait and all built-in scalar function -//! implementations. Expressions ([`crate::expr::Expression`]) reference scalar functions -//! at each node. +//! This module contains the [`ScalarFnVTable`] trait, the machinery to bind, reduce and execute +//! scalar functions, and in [`builtin`] the few functions vortex-array itself is built on. Other scalar +//! functions live in the `vortex-functions` crate. Expressions ([`crate::expr::Expression`]) +//! reference scalar functions at each node. //! //! Strict functions with row-at-a-time kernels can implement `unstable::row::RowFn`. It handles //! decoding, constants, null propagation, output construction, and validity. This API requires the @@ -15,11 +16,6 @@ use vortex_session::registry::Id; -use crate::scalar_fn::fns::byte_length::ByteLength; -use crate::scalar_fn::fns::ext_storage::ExtStorage; -use crate::scalar_fn::fns::get_item::GetItem; -use crate::scalar_fn::fns::literal::Literal; - mod vtable; pub use vtable::*; @@ -47,7 +43,7 @@ pub mod unstable; #[allow(dead_code, unused_imports)] pub(crate) mod unstable; -pub mod fns; +pub mod builtin; pub mod internal; pub mod session; @@ -65,18 +61,3 @@ mod sealed { /// This can be the **only** implementor for [`super::typed::DynScalarFn`]. impl Sealed for TypedScalarFnInstance {} } - -/// A scalar function has a negative cost if applying it to an array and -/// canonicalizing is cheaper than canonicalizing an array and applying it. -/// -/// Example of negative cost expressions are byte_length(), ext_storage(), and get_item() since -/// they don't depend on input size. -/// -/// Example of non-negative cost expression is like() as it's linear over -/// individual input. -pub fn is_negative_cost(id: ScalarFnId) -> bool { - id == ScalarFnVTable::id(&ByteLength) - || id == ScalarFnVTable::id(&ExtStorage) - || id == ScalarFnVTable::id(&GetItem) - || id == ScalarFnVTable::id(&Literal) -} diff --git a/vortex-array/src/scalar_fn/session.rs b/vortex-array/src/scalar_fn/session.rs index 09ce5d4370e..fe395a46dca 100644 --- a/vortex-array/src/scalar_fn/session.rs +++ b/vortex-array/src/scalar_fn/session.rs @@ -12,32 +12,16 @@ use vortex_session::registry::Id; use crate::scalar_fn::ScalarFnPluginRef; use crate::scalar_fn::ScalarFnVTable; -use crate::scalar_fn::fns::between::Between; -use crate::scalar_fn::fns::binary::Binary; -use crate::scalar_fn::fns::byte_length::ByteLength; -use crate::scalar_fn::fns::case_when::CaseWhen; -use crate::scalar_fn::fns::cast::Cast; -use crate::scalar_fn::fns::ext_storage::ExtStorage; -use crate::scalar_fn::fns::fill_null::FillNull; -use crate::scalar_fn::fns::get_item::GetItem; -use crate::scalar_fn::fns::is_nan::IsNan; -use crate::scalar_fn::fns::is_not_null::IsNotNull; -use crate::scalar_fn::fns::is_null::IsNull; -use crate::scalar_fn::fns::like::Like; -use crate::scalar_fn::fns::list_contains::ListContains; -use crate::scalar_fn::fns::list_contains::PreparedSet; -use crate::scalar_fn::fns::list_length::ListLength; -use crate::scalar_fn::fns::list_sum::ListSum; -use crate::scalar_fn::fns::literal::Literal; -use crate::scalar_fn::fns::mask::Mask; -use crate::scalar_fn::fns::merge::Merge; -use crate::scalar_fn::fns::not::Not; -use crate::scalar_fn::fns::pack::Pack; -use crate::scalar_fn::fns::replace_time_zone::ReplaceTimeZone; -use crate::scalar_fn::fns::select::Select; -use crate::scalar_fn::fns::stat::StatFn; -use crate::scalar_fn::fns::variant_get::VariantGet; -use crate::scalar_fn::fns::zip::Zip; +use crate::scalar_fn::builtin::binary::Binary; +use crate::scalar_fn::builtin::cast::Cast; +use crate::scalar_fn::builtin::fill_null::FillNull; +use crate::scalar_fn::builtin::get_item::GetItem; +use crate::scalar_fn::builtin::is_not_null::IsNotNull; +use crate::scalar_fn::builtin::is_null::IsNull; +use crate::scalar_fn::builtin::literal::Literal; +use crate::scalar_fn::builtin::mask::Mask; +use crate::scalar_fn::builtin::not::Not; +use crate::scalar_fn::builtin::pack::Pack; /// Registry of scalar function vtables. pub type ScalarFnRegistry = ArcSwapMap; @@ -66,33 +50,18 @@ impl Default for ScalarFnSession { registry: ScalarFnRegistry::default(), }; - // Register built-in expressions. - this.register(Between); + // Register the functions vortex-array is built on. `vortex_functions::initialize` registers + // the rest. this.register(Binary); - this.register(ByteLength); - this.register(CaseWhen); this.register(Cast); - this.register(ExtStorage); this.register(FillNull); this.register(GetItem); - this.register(IsNan); this.register(IsNotNull); this.register(IsNull); - this.register(Like); - this.register(ListContains); - this.register(PreparedSet); - this.register(ListLength); - this.register(ListSum); this.register(Literal); this.register(Mask); - this.register(Merge); this.register(Not); this.register(Pack); - this.register(ReplaceTimeZone); - this.register(Select); - this.register(StatFn); - this.register(VariantGet); - this.register(Zip); this } diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/tests.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/tests.rs index 16f639c4254..c265cb4bee0 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/tests.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/tests.rs @@ -26,6 +26,7 @@ use crate::dtype::DType; use crate::dtype::Nullability; use crate::extension::datetime::TimeUnit; use crate::extension::datetime::Timestamp; +use crate::optimizer::ArrayOptimizer; use crate::scalar::Scalar; use crate::scalar_fn::VecExecutionArgs; use crate::scalar_fn::unstable::row::InputElement; @@ -152,7 +153,9 @@ fn test_batch_const_unwraps_filtered_masked_constant() -> VortexResult<()> { let child = ConstantArray::new(7_i64, 3).into_array(); let masked = MaskedArray::try_new(child, Validity::from_iter([true, false, true]), &mut ctx)? .into_array(); - let filtered = masked.filter(Mask::from_iter([true, true, false]))?; + let filtered = masked + .filter(Mask::from_iter([true, true, false]))? + .optimize(&SESSION)?; let Some(const_array) = batch_const(&filtered) else { vortex_bail!("filtered masked constant must remain batch-constant"); @@ -168,7 +171,9 @@ fn test_batch_const_preserves_filtered_extension() -> VortexResult<()> { let ext_dtype = Timestamp::new(TimeUnit::Milliseconds, Nullability::NonNullable).erased(); let extension = ExtensionArray::new(ext_dtype, ConstantArray::new(7_i64, 3).into_array()).into_array(); - let filtered = extension.filter(Mask::from_iter([true, false, true]))?; + let filtered = extension + .filter(Mask::from_iter([true, false, true]))? + .optimize(&SESSION)?; let Some(const_array) = batch_const(&filtered) else { vortex_bail!("filtered extension storage must remain batch-constant"); diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/utf8.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/utf8.rs index 17b034dca94..10c4696769d 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/utf8.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/utf8.rs @@ -344,6 +344,7 @@ mod tests { use crate::buffer::BufferHandle; use crate::dtype::DType; use crate::dtype::Nullability; + use crate::optimizer::ArrayOptimizer; use crate::scalar::Scalar; use crate::scalar_fn::EmptyOptions; use crate::scalar_fn::ScalarFnId; @@ -468,7 +469,8 @@ mod tests { DType::Utf8(Nullability::NonNullable), ) .into_array() - .slice(1..3)?; + .slice(1..3)? + .optimize(&SESSION)?; let mut ctx = SESSION.create_execution_ctx(); let column = Utf8Column::decode(array, &mut ctx)?; diff --git a/vortex-array/src/scalar_fn/vtable.rs b/vortex-array/src/scalar_fn/vtable.rs index 4cd658d4774..ad3b9002e74 100644 --- a/vortex-array/src/scalar_fn/vtable.rs +++ b/vortex-array/src/scalar_fn/vtable.rs @@ -31,10 +31,10 @@ use crate::scalar::Scalar; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnRef; use crate::scalar_fn::TypedScalarFnInstance; -use crate::scalar_fn::fns::binary::Binary; -use crate::scalar_fn::fns::is_not_null::IsNotNull; -use crate::scalar_fn::fns::literal::Literal; -use crate::scalar_fn::fns::operators::Operator; +use crate::scalar_fn::builtin::binary::Binary; +use crate::scalar_fn::builtin::is_not_null::IsNotNull; +use crate::scalar_fn::builtin::literal::Literal; +use crate::scalar_fn::builtin::operators::Operator; use crate::validity::Validity; // Here and beyond we use validity(x) interchangeably with is_not_null(x). @@ -57,7 +57,7 @@ pub enum ReduceNodeValidity { } /// IsNotNull(child) as a reducible node -pub(crate) fn is_not_null_node(child: &T) -> VortexResult { +pub fn is_not_null_node(child: &T) -> VortexResult { child.new_node(IsNotNull.bind(EmptyOptions), std::slice::from_ref(child)) } @@ -263,7 +263,7 @@ pub trait ScalarFnVTable: 'static + Sized + Clone + Send + Sync { /// is not because `false AND null` returns `false`. `is_null` is also not strict. /// /// Strictness does not require valid inputs to produce a valid output. For example, - /// [`crate::expr::list_sum`] returns null for a valid empty list. Implement + /// `list_sum` returns null for a valid empty list. Implement /// [`ScalarFnVTable::validity`] only when the output validity can be derived without /// evaluation. /// @@ -281,7 +281,7 @@ pub trait ScalarFnVTable: 'static + Sized + Clone + Send + Sync { /// /// Return `true` only when a well-typed call cannot error because of its values. `checked_add` /// is fallible on integer overflow, and integer division is fallible when its divisor is zero. - /// A null result is not an error: [`crate::expr::list_sum`] is infallible for an empty list. + /// A null result is not an error: `list_sum` is infallible for an empty list. /// /// Ignore incidental execution errors, such as canonicalization failures, allocation errors, /// and encoding mismatches. They are not part of the function's semantics. diff --git a/vortex-array/src/serde.rs b/vortex-array/src/serde.rs index 3653d051b6a..5bdfc632c27 100644 --- a/vortex-array/src/serde.rs +++ b/vortex-array/src/serde.rs @@ -38,6 +38,7 @@ use crate::flatbuffers::FlatBuffer; use crate::flatbuffers::WriteFlatBuffer; use crate::flatbuffers::array as fba; use crate::flatbuffers::array::Compression; +use crate::optimizer::ArrayOptimizer; use crate::session::ArraySessionExt; use crate::stats::StatsSet; @@ -170,6 +171,16 @@ struct ArraySerializationTree { impl ArraySerializationTree { fn try_new(session: &VortexSession, source: &ArrayRef) -> VortexResult { + // Constructors such as `slice` build lazy wrappers that have no serialized form and that + // optimization rewrites away, so optimize only those nodes. Every other node is written + // as given. + let optimized; + let source = if session.arrays().serializer(&source.encoding_id()).is_none() { + optimized = source.optimize(session)?; + &optimized + } else { + source + }; let Some(serialization) = session.array_serialize(source)? else { vortex_bail!( "Array {} does not support serialization", diff --git a/vortex-array/src/session/mod.rs b/vortex-array/src/session/mod.rs index e9d0d36a55f..ac6075d5b87 100644 --- a/vortex-array/src/session/mod.rs +++ b/vortex-array/src/session/mod.rs @@ -73,7 +73,7 @@ impl ArraySession { } } - fn serializer(&self, id: &ArrayId) -> Option { + pub(crate) fn serializer(&self, id: &ArrayId) -> Option { self.serializers.get(id) } } diff --git a/vortex-array/src/test_harness/trace/mod.rs b/vortex-array/src/test_harness/trace/mod.rs index 1e053cf5a7a..0b4f1b068ce 100644 --- a/vortex-array/src/test_harness/trace/mod.rs +++ b/vortex-array/src/test_harness/trace/mod.rs @@ -53,10 +53,10 @@ //! ```ignore //! use vortex_array::test_harness::trace::trace_op; //! -//! let traced = trace_op(|| filter_array.optimize())?; +//! let traced = trace_op(|| filter_array.optimize(&session))?; //! assert!(traced.output.is::()); //! insta::assert_snapshot!(traced.trace.to_string(), @r" -//! optimize root=vortex.filter(i32, len=4) session=false +//! optimize root=vortex.filter(i32, len=4) //! reduce TrivialFilterRule: vortex.filter(i32, len=4) -> vortex.primitive(i32, len=4) //! done output=vortex.primitive(i32, len=4) //! "); @@ -202,10 +202,10 @@ fn write_indent(f: &mut fmt::Formatter<'_>, depth: usize) -> fmt::Result { /// # Examples /// /// ```ignore -/// let traced = trace_op(|| filter_array.optimize())?; +/// let traced = trace_op(|| filter_array.optimize(&session))?; /// assert!(traced.output.is::()); /// insta::assert_snapshot!(traced.trace.to_string(), @r" -/// optimize root=vortex.filter(i32, len=4) session=false +/// optimize root=vortex.filter(i32, len=4) /// reduce TrivialFilterRule: vortex.filter(i32, len=4) -> vortex.primitive(i32, len=4) /// done output=vortex.primitive(i32, len=4) /// "); @@ -355,10 +355,9 @@ impl Display for ArraySummary { } } -pub(crate) fn record_optimize_start(root: &ArrayRef, session: bool) { +pub(crate) fn record_optimize_start(root: &ArrayRef) { record(TraceEvent::OptimizeStart { root: ArraySummary::new(root), - session, }); } @@ -877,7 +876,6 @@ impl TraceRecorder { enum TraceEvent { OptimizeStart { root: ArraySummary, - session: bool, }, OptimizeLoopStart { array: ArraySummary, @@ -1078,8 +1076,8 @@ impl TraceEvent { fn fmt_line(&self, f: &mut fmt::Formatter<'_>, resolution: TraceResolution) -> fmt::Result { match self { - TraceEvent::OptimizeStart { root, session } => { - write!(f, "optimize root={root} session={session}") + TraceEvent::OptimizeStart { root } => { + write!(f, "optimize root={root}") } TraceEvent::OptimizeLoopStart { array } => { write!(f, "loop input={array}") @@ -1215,9 +1213,9 @@ impl TraceEvent { } => match resolution { TraceResolution::Attempts => write!( f, - "optimize_ctx input={input} output={output} changed={changed}" + "optimize input={input} output={output} changed={changed}" ), - TraceResolution::ExecutedOnly => write!(f, "optimize_ctx {input} -> {output}"), + TraceResolution::ExecutedOnly => write!(f, "optimize {input} -> {output}"), }, TraceEvent::ExecuteEncoding { array } => { write!(f, "execute encoding={array}") diff --git a/vortex-array/src/test_harness/trace/tests.rs b/vortex-array/src/test_harness/trace/tests.rs index 353fc669a54..8ded5408565 100644 --- a/vortex-array/src/test_harness/trace/tests.rs +++ b/vortex-array/src/test_harness/trace/tests.rs @@ -4,6 +4,7 @@ use std::fmt::Display; use std::fmt::Formatter; use std::hash::Hasher; +use std::sync::LazyLock; use rstest::fixture; use rstest::rstest; @@ -56,10 +57,8 @@ use crate::matcher::Matcher; use crate::optimizer::ArrayOptimizer; use crate::optimizer::kernels::ArrayKernelsExt; use crate::scalar::Scalar; -use crate::scalar_fn::fns::binary::Binary; -use crate::scalar_fn::fns::like::Like; -use crate::scalar_fn::fns::like::LikeOptions; -use crate::scalar_fn::fns::operators::Operator; +use crate::scalar_fn::builtin::binary::Binary; +use crate::scalar_fn::builtin::operators::Operator; use crate::serde::ArrayChildren; use crate::session::ArraySession; use crate::test_harness::trace::TraceOptions; @@ -68,6 +67,8 @@ use crate::test_harness::trace::trace_op; use crate::test_harness::trace::trace_op_with; use crate::validity::Validity; +static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + #[fixture] fn stack_parent_fixture() -> VortexResult { stack_parent(stack_child()?) @@ -376,15 +377,15 @@ fn trace_optimize_reduce_fixpoint() -> VortexResult<()> { let values = PrimitiveArray::from_iter([0i32, 1, 2, 3]).into_array(); let filter = FilterArray::try_new(values.clone(), Mask::new_true(values.len()))?.into_array(); - let traced = trace_op(|| filter.optimize())?; + let traced = trace_op(|| filter.optimize(&SESSION))?; assert!(traced.output.is::()); assert_arrays_eq!(traced.output, values, &mut execution_ctx()); - insta::assert_snapshot!(traced.trace.to_string(), @r" -optimize root=vortex.filter(i32, len=4) session=false - reduce TrivialFilterRule: vortex.filter(i32, len=4) -> vortex.primitive(i32, len=4) - done output=vortex.primitive(i32, len=4) -"); + insta::assert_snapshot!(traced.trace.to_string(), @" + optimize root=vortex.filter(i32, len=4) + reduce TrivialFilterRule: vortex.filter(i32, len=4) -> vortex.primitive(i32, len=4) + done output=vortex.primitive(i32, len=4) + "); Ok(()) } @@ -404,7 +405,7 @@ fn trace_optimize_contiguous_filter() -> VortexResult<()> { TraceOptions { resolution: TraceResolution::ExecutedOnly, }, - || outer.optimize(), + || outer.optimize(&SESSION), )?; let optimized_slice = traced.output.as_::(); @@ -414,8 +415,8 @@ fn trace_optimize_contiguous_filter() -> VortexResult<()> { PrimitiveArray::from_iter([2i32, 3]), &mut execution_ctx() ); - insta::assert_snapshot!(traced.trace.to_string(), @r" - optimize root=vortex.filter(i32, len=2) session=false + insta::assert_snapshot!(traced.trace.to_string(), @" + optimize root=vortex.filter(i32, len=2) reduce TrivialFilterRule: vortex.filter(i32, len=2) -> vortex.slice(i32, len=2) done output=vortex.slice(i32, len=2) "); @@ -439,21 +440,26 @@ fn trace_optimize_contiguous_filter() -> VortexResult<()> { &mut execution_ctx() ); insta::assert_snapshot!(traced.trace.to_string(), @" - execute_until target=AnyCanonical root=vortex.filter(i32, len=2) - iter 0 current=vortex.filter(i32, len=2) builder_active=false - Done array=vortex.slice(i32, len=2) - iter 1 current=vortex.slice(i32, len=2) builder_active=false + optimize_recursive root=vortex.filter(i32, len=2) + optimize root=vortex.filter(i32, len=2) + reduce TrivialFilterRule: vortex.filter(i32, len=2) -> vortex.slice(i32, len=2) + done output=vortex.slice(i32, len=2) + optimize_recursive root=vortex.filter(i32, len=4) + optimize_recursive root=vortex.primitive(i32, len=6) + execute_until target=AnyCanonical root=vortex.slice(i32, len=2) + iter 0 current=vortex.slice(i32, len=2) builder_active=false ExecuteSlot slot=0 parent=vortex.slice(i32, len=2) child=vortex.filter(i32, len=4) - iter 2 current=vortex.filter(i32, len=4) stack_parent=vortex.slice(i32, len=2) slot=0 builder_active=false + iter 1 current=vortex.filter(i32, len=4) stack_parent=vortex.slice(i32, len=2) slot=0 builder_active=false Done array=vortex.primitive(i32, len=4) - iter 3 current=vortex.primitive(i32, len=4) stack_parent=vortex.slice(i32, len=2) slot=0 builder_active=false + iter 2 current=vortex.primitive(i32, len=4) stack_parent=vortex.slice(i32, len=2) slot=0 builder_active=false pop_frame slot=0 output=vortex.slice(i32, len=2) - iter 4 current=vortex.slice(i32, len=2) builder_active=false - optimize root=vortex.slice(i32, len=2) session=false + iter 3 current=vortex.slice(i32, len=2) builder_active=false + optimize root=vortex.slice(i32, len=2) reduce_parent static:SliceReduceAdaptor(Primitive) slot=0 parent=vortex.slice(i32, len=2) child=vortex.primitive(i32, len=4) -> vortex.primitive(i32, len=2) done output=vortex.primitive(i32, len=2) Done array=vortex.primitive(i32, len=2) - iter 5 current=vortex.primitive(i32, len=2) builder_active=false + iter 4 current=vortex.primitive(i32, len=2) builder_active=false + optimize_recursive root=vortex.primitive(i32, len=2) return output=vortex.primitive(i32, len=2) "); @@ -492,16 +498,16 @@ fn trace_execute_filter_with_scattered_mask() -> VortexResult<()> { &mut execution_ctx() ); insta::assert_snapshot!(traced.trace.to_string(), @" + optimize_recursive root=vortex.filter(i32, len=2) + optimize root=vortex.filter(i32, len=2) + reduce_parent static:FilterReduceAdaptor(Filter) slot=0 parent=vortex.filter(i32, len=2) child=vortex.filter(i32, len=4) -> vortex.filter(i32, len=2) + done output=vortex.filter(i32, len=2) + optimize_recursive root=vortex.primitive(i32, len=6) execute_until target=AnyCanonical root=vortex.filter(i32, len=2) iter 0 current=vortex.filter(i32, len=2) builder_active=false - ExecuteSlot slot=0 parent=vortex.filter(i32, len=2) child=vortex.filter(i32, len=4) - iter 1 current=vortex.filter(i32, len=4) stack_parent=vortex.filter(i32, len=2) slot=0 builder_active=false - Done array=vortex.primitive(i32, len=4) - iter 2 current=vortex.primitive(i32, len=4) stack_parent=vortex.filter(i32, len=2) slot=0 builder_active=false - pop_frame slot=0 output=vortex.filter(i32, len=2) - iter 3 current=vortex.filter(i32, len=2) builder_active=false Done array=vortex.primitive(i32, len=2) - iter 4 current=vortex.primitive(i32, len=2) builder_active=false + iter 1 current=vortex.primitive(i32, len=2) builder_active=false + optimize_recursive root=vortex.primitive(i32, len=2) return output=vortex.primitive(i32, len=2) "); @@ -529,6 +535,18 @@ fn trace_execution_stack_parent_kernel_attempts( &mut ctx ); insta::assert_snapshot!(traced.trace.to_string(), @" + optimize_recursive root=vortex.test.stack-parent(i32, len=3) + optimize root=vortex.test.stack-parent(i32, len=3) + loop input=vortex.test.stack-parent(i32, len=3) + reduce none array=vortex.test.stack-parent(i32, len=3) + reduce_parent none array=vortex.test.stack-parent(i32, len=3) + done output=vortex.test.stack-parent(i32, len=3) changed=false + optimize_recursive root=vortex.test.stack-child(i32, len=3) + optimize root=vortex.test.stack-child(i32, len=3) + loop input=vortex.test.stack-child(i32, len=3) + reduce none array=vortex.test.stack-child(i32, len=3) + reduce_parent none array=vortex.test.stack-child(i32, len=3) + done output=vortex.test.stack-child(i32, len=3) changed=false execute_until target=AnyCanonical root=vortex.test.stack-parent(i32, len=3) iter 0 current=vortex.test.stack-parent(i32, len=3) builder_active=false done_check target=false canonical=false @@ -541,14 +559,20 @@ fn trace_execution_stack_parent_kernel_attempts( done_check target=false canonical=false stack_execute_parent attempt slot=0 parent=vortex.test.stack-parent(i32, len=3) child=vortex.test.stack-child(i32, len=3) source=session[0] kernel=execute_parent_fn outcome=declined stack_execute_parent applied slot=0 parent=vortex.test.stack-parent(i32, len=3) child=vortex.test.stack-child(i32, len=3) source=session[1] kernel=execute_parent_fn output=vortex.primitive(i32, len=3) - optimize root=vortex.primitive(i32, len=3) session=true + optimize root=vortex.primitive(i32, len=3) loop input=vortex.primitive(i32, len=3) reduce none array=vortex.primitive(i32, len=3) reduce_parent none array=vortex.primitive(i32, len=3) done output=vortex.primitive(i32, len=3) changed=false - optimize_ctx input=vortex.primitive(i32, len=3) output=vortex.primitive(i32, len=3) changed=false + optimize input=vortex.primitive(i32, len=3) output=vortex.primitive(i32, len=3) changed=false iter 2 current=vortex.primitive(i32, len=3) builder_active=false done_check target=true canonical=true + optimize_recursive root=vortex.primitive(i32, len=3) + optimize root=vortex.primitive(i32, len=3) + loop input=vortex.primitive(i32, len=3) + reduce none array=vortex.primitive(i32, len=3) + reduce_parent none array=vortex.primitive(i32, len=3) + done output=vortex.primitive(i32, len=3) changed=false return output=vortex.primitive(i32, len=3) "); @@ -578,6 +602,11 @@ fn trace_execution_chunked_append_child_flow() -> VortexResult<()> { &mut ctx ); insta::assert_snapshot!(traced.trace.to_string(), @" + optimize_recursive root=vortex.chunked(i32, len=5) + optimize_recursive root=vortex.primitive(u64, len=4) + optimize_recursive root=vortex.primitive(i32, len=2) + optimize_recursive root=vortex.primitive(i32, len=1) + optimize_recursive root=vortex.primitive(i32, len=2) execute_until target=AnyCanonical root=vortex.chunked(i32, len=5) iter 0 current=vortex.chunked(i32, len=5) builder_active=false builder start array=vortex.chunked(i32, len=5) @@ -593,22 +622,13 @@ fn trace_execution_chunked_append_child_flow() -> VortexResult<()> { Done array=vortex.primitive(i32, len=0) builder finish output=vortex.primitive(i32, len=5) iter 4 current=vortex.primitive(i32, len=5) builder_active=false + optimize_recursive root=vortex.primitive(i32, len=5) return output=vortex.primitive(i32, len=5) "); Ok(()) } -/// A dictionary of strings: codes `[0, 1, 2, 1, 0, 2]` over values -/// `["alpha", "beta", "gamma"]`. -fn dict_of_strings() -> VortexResult { - Ok(DictArray::try_new( - PrimitiveArray::from_iter([0u32, 1, 2, 1, 0, 2]).into_array(), - VarBinViewArray::from_iter_str(["alpha", "beta", "gamma"]).into_array(), - )? - .into_array()) -} - fn execution_ctx() -> ExecutionCtx { ExecutionCtx::new(VortexSession::empty().with::()) } @@ -630,7 +650,7 @@ fn trace_take_on_chunked() -> VortexResult<()> { // No reduce rule rewrites a take over a chunked array: the work all happens at execution // time, so the optimizer trace is empty. - let traced = trace_op(|| take.optimize())?; + let traced = trace_op(|| take.optimize(&SESSION))?; insta::assert_snapshot!(traced.trace.to_string(), @""); let optimized = traced.output; @@ -646,12 +666,19 @@ fn trace_take_on_chunked() -> VortexResult<()> { &mut execution_ctx() ); insta::assert_snapshot!(traced.trace.to_string(), @" + optimize_recursive root=vortex.dict(i32, len=4) + optimize_recursive root=vortex.primitive(u64, len=4) + optimize_recursive root=vortex.chunked(i32, len=5) + optimize_recursive root=vortex.primitive(u64, len=3) + optimize_recursive root=vortex.primitive(i32, len=2) + optimize_recursive root=vortex.primitive(i32, len=3) execute_until target=AnyCanonical root=vortex.dict(i32, len=4) iter 0 current=vortex.dict(i32, len=4) builder_active=false child_execute_parent session[0]:execute_parent_fn slot=1 parent=vortex.dict(i32, len=4) child=vortex.chunked(i32, len=5) -> vortex.dict(i32, len=4) iter 1 current=vortex.dict(i32, len=4) builder_active=false child_execute_parent session[0]:execute_parent_fn slot=1 parent=vortex.dict(i32, len=4) child=vortex.primitive(i32, len=4) -> vortex.primitive(i32, len=4) iter 2 current=vortex.primitive(i32, len=4) builder_active=false + optimize_recursive root=vortex.primitive(i32, len=4) return output=vortex.primitive(i32, len=4) "); @@ -679,13 +706,13 @@ fn trace_filter_on_struct_with_complex_children() -> VortexResult<()> { FilterArray::try_new(struct_, Mask::from_iter([true, false, true, true, false]))? .into_array(); - let traced = trace_op(|| filtered.optimize())?; + let traced = trace_op(|| filtered.optimize(&SESSION))?; insta::assert_snapshot!(traced.trace.to_string(), @" - optimize root=vortex.filter({name=utf8, score=i64}, len=3) session=false - optimize root=vortex.filter(utf8, len=3) session=false + optimize root=vortex.filter({name=utf8, score=i64}, len=3) + reduce FilterStructRule: vortex.filter({name=utf8, score=i64}, len=3) -> vortex.struct({name=utf8, score=i64}, len=3) + optimize root=vortex.filter(utf8, len=3) reduce_parent static:FilterReduceAdaptor(Dict) slot=0 parent=vortex.filter(utf8, len=3) child=vortex.dict(utf8, len=5) -> vortex.dict(utf8, len=3) done output=vortex.dict(utf8, len=3) - reduce FilterStructRule: vortex.filter({name=utf8, score=i64}, len=3) -> vortex.struct({name=utf8, score=i64}, len=3) done output=vortex.struct({name=utf8, score=i64}, len=3) "); @@ -709,6 +736,16 @@ fn trace_filter_on_struct_with_complex_children() -> VortexResult<()> { .into_array(); assert_arrays_eq!(traced.output, expected, &mut execution_ctx()); insta::assert_snapshot!(traced.trace.to_string(), @" + optimize_recursive root=vortex.struct({name=utf8, score=i64}, len=3) + optimize_recursive root=vortex.dict(utf8, len=3) + optimize_recursive root=vortex.filter(u32, len=3) + optimize_recursive root=vortex.primitive(u32, len=5) + optimize_recursive root=vortex.varbinview(utf8, len=3) + optimize_recursive root=vortex.filter(i64, len=3) + optimize_recursive root=vortex.chunked(i64, len=5) + optimize_recursive root=vortex.primitive(u64, len=3) + optimize_recursive root=vortex.primitive(i64, len=2) + optimize_recursive root=vortex.primitive(i64, len=3) execute_until target=AnyCanonical root=vortex.struct({name=utf8, score=i64}, len=3) iter 0 current=vortex.struct({name=utf8, score=i64}, len=3) builder_active=false return output=vortex.struct({name=utf8, score=i64}, len=3) @@ -728,9 +765,9 @@ fn trace_compare_on_dict() -> VortexResult<()> { let compared = Binary::try_new(dict, rhs, Operator::Eq)?.into_array(); - let traced = trace_op(|| compared.optimize())?; + let traced = trace_op(|| compared.optimize(&SESSION))?; insta::assert_snapshot!(traced.trace.to_string(), @" - optimize root=vortex.binary(bool, len=5) session=false + optimize root=vortex.binary(bool, len=5) reduce_parent static:DictionaryScalarFnValuesPushDownRule slot=0 parent=vortex.binary(bool, len=5) child=vortex.dict(i32, len=5) -> vortex.dict(bool, len=5) done output=vortex.dict(bool, len=5) "); @@ -748,6 +785,11 @@ fn trace_compare_on_dict() -> VortexResult<()> { &mut execution_ctx() ); insta::assert_snapshot!(traced.trace.to_string(), @" + optimize_recursive root=vortex.dict(bool, len=5) + optimize_recursive root=vortex.primitive(u32, len=5) + optimize_recursive root=vortex.binary(bool, len=3) + optimize_recursive root=vortex.primitive(i32, len=3) + optimize_recursive root=vortex.constant(i32, len=3) execute_until target=AnyCanonical root=vortex.dict(bool, len=5) iter 0 current=vortex.dict(bool, len=5) builder_active=false ExecuteSlot slot=1 parent=vortex.dict(bool, len=5) child=vortex.binary(bool, len=3) @@ -758,59 +800,9 @@ fn trace_compare_on_dict() -> VortexResult<()> { iter 3 current=vortex.dict(bool, len=5) builder_active=false child_execute_parent session[0]:execute_parent_fn slot=1 parent=vortex.dict(bool, len=5) child=vortex.bool(bool, len=3) -> vortex.bool(bool, len=5) iter 4 current=vortex.bool(bool, len=5) builder_active=false + optimize_recursive root=vortex.bool(bool, len=5) return output=vortex.bool(bool, len=5) "); Ok(()) } - -#[test] -fn trace_like_on_dict() -> VortexResult<()> { - let strings = dict_of_strings()?; - let pattern = ConstantArray::new(Scalar::from("b%"), strings.len()).into_array(); - - let like = Like::try_new( - strings, - pattern, - LikeOptions { - negated: false, - case_insensitive: false, - }, - )? - .into_array(); - - let traced = trace_op(|| like.optimize())?; - insta::assert_snapshot!(traced.trace.to_string(), @" - optimize root=vortex.like(bool, len=6) session=false - reduce_parent static:LikeReduceAdaptor(Dict) slot=0 parent=vortex.like(bool, len=6) child=vortex.dict(utf8, len=6) -> vortex.dict(bool, len=6) - done output=vortex.dict(bool, len=6) - "); - - let optimized = traced.output; - let traced = trace_op(|| { - optimized - .execute::(&mut execution_ctx()) - .map(IntoArray::into_array) - })?; - - assert_arrays_eq!( - traced.output, - BoolArray::from_iter([false, true, false, true, false, false]), - &mut execution_ctx() - ); - insta::assert_snapshot!(traced.trace.to_string(), @" - execute_until target=AnyCanonical root=vortex.dict(bool, len=6) - iter 0 current=vortex.dict(bool, len=6) builder_active=false - ExecuteSlot slot=1 parent=vortex.dict(bool, len=6) child=vortex.like(bool, len=3) - iter 1 current=vortex.like(bool, len=3) stack_parent=vortex.dict(bool, len=6) slot=1 builder_active=false - Done array=vortex.bool(bool, len=3) - iter 2 current=vortex.bool(bool, len=3) stack_parent=vortex.dict(bool, len=6) slot=1 builder_active=false - pop_frame slot=1 output=vortex.dict(bool, len=6) - iter 3 current=vortex.dict(bool, len=6) builder_active=false - child_execute_parent session[0]:execute_parent_fn slot=1 parent=vortex.dict(bool, len=6) child=vortex.bool(bool, len=3) -> vortex.bool(bool, len=6) - iter 4 current=vortex.bool(bool, len=6) builder_active=false - return output=vortex.bool(bool, len=6) - "); - - Ok(()) -} diff --git a/vortex-array/src/validity.rs b/vortex-array/src/validity.rs index 449f66b07d9..4093793d7aa 100644 --- a/vortex-array/src/validity.rs +++ b/vortex-array/src/validity.rs @@ -30,11 +30,10 @@ use crate::arrays::ConstantArray; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::Nullability; -use crate::optimizer::ArrayOptimizer; use crate::patches::Patches; use crate::scalar::Scalar; -use crate::scalar_fn::fns::binary::Binary; -use crate::scalar_fn::fns::operators::Operator; +use crate::scalar_fn::builtin::binary::Binary; +use crate::scalar_fn::builtin::operators::Operator; /// Validity information for an array. #[derive(Clone)] @@ -309,11 +308,9 @@ impl Validity { | (Validity::AllValid, Validity::NonNullable) | (Validity::AllValid, Validity::AllValid) => Validity::AllValid, // Here we actually have to do some work - (Validity::Array(lhs), Validity::Array(rhs)) => Validity::Array( - Binary::try_new(lhs, rhs, Operator::And)? - .into_array() - .optimize()?, - ), + (Validity::Array(lhs), Validity::Array(rhs)) => { + Validity::Array(Binary::try_new(lhs, rhs, Operator::And)?.into_array()) + } }) } diff --git a/vortex-arrow/src/executor/byte.rs b/vortex-arrow/src/executor/byte.rs index c593bd24074..8f7dc512e9f 100644 --- a/vortex-arrow/src/executor/byte.rs +++ b/vortex-arrow/src/executor/byte.rs @@ -207,7 +207,7 @@ mod tests { use vortex_array::arrays::VarBinViewArray; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; - use vortex_array::scalar_fn::fns::mask::Mask as MaskFn; + use vortex_array::scalar_fn::builtin::mask::Mask as MaskFn; use vortex_array::validity::Validity; use vortex_buffer::ByteBuffer; use vortex_buffer::buffer; diff --git a/vortex-arrow/src/executor/dictionary.rs b/vortex-arrow/src/executor/dictionary.rs index 343c08357af..7186696fde6 100644 --- a/vortex-arrow/src/executor/dictionary.rs +++ b/vortex-arrow/src/executor/dictionary.rs @@ -173,7 +173,7 @@ mod tests { use vortex_array::dtype::DType; use vortex_array::dtype::Nullability::Nullable; use vortex_array::scalar::Scalar; - use vortex_array::scalar_fn::fns::mask::Mask; + use vortex_array::scalar_fn::builtin::mask::Mask; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; diff --git a/vortex-arrow/src/executor/list.rs b/vortex-arrow/src/executor/list.rs index 8246f36155f..831dd1110d3 100644 --- a/vortex-arrow/src/executor/list.rs +++ b/vortex-arrow/src/executor/list.rs @@ -236,7 +236,7 @@ mod tests { use vortex_array::arrays::SliceArray; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability::NonNullable; - use vortex_array::scalar_fn::fns::mask::Mask; + use vortex_array::scalar_fn::builtin::mask::Mask; use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; diff --git a/vortex-arrow/src/executor/struct_.rs b/vortex-arrow/src/executor/struct_.rs index 76946d8f01a..f321e4ccf82 100644 --- a/vortex-arrow/src/executor/struct_.rs +++ b/vortex-arrow/src/executor/struct_.rs @@ -22,7 +22,7 @@ use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::FieldNames; use vortex_array::matcher::Matcher; -use vortex_array::scalar_fn::fns::pack::Pack; +use vortex_array::scalar_fn::builtin::pack::Pack; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_ensure_eq; @@ -234,7 +234,7 @@ mod tests { use vortex_array::arrays::StructArray; use vortex_array::assert_arrays_eq; use vortex_array::dtype::FieldNames; - use vortex_array::scalar_fn::fns::mask::Mask; + use vortex_array::scalar_fn::builtin::mask::Mask; use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; diff --git a/vortex-cuda/src/dynamic_dispatch/plan_builder.rs b/vortex-cuda/src/dynamic_dispatch/plan_builder.rs index 2374267f75b..b2abc1d5de0 100644 --- a/vortex-cuda/src/dynamic_dispatch/plan_builder.rs +++ b/vortex-cuda/src/dynamic_dispatch/plan_builder.rs @@ -45,7 +45,7 @@ use vortex::error::VortexResult; use vortex::error::vortex_bail; use vortex::error::vortex_err; use vortex::scalar_fn::ScalarFnVTable; -use vortex::scalar_fn::fns::cast::Cast; +use vortex::scalar_fn::builtin::cast::Cast; use super::CudaDispatchPlan; use super::MaterializedStage; diff --git a/vortex-datafusion/src/convert/exprs/mod.rs b/vortex-datafusion/src/convert/exprs/mod.rs index b3b4563db09..678b71b3d88 100644 --- a/vortex-datafusion/src/convert/exprs/mod.rs +++ b/vortex-datafusion/src/convert/exprs/mod.rs @@ -39,12 +39,12 @@ use vortex::expr::nested_case_when; use vortex::expr::not; use vortex::expr::pack; use vortex::expr::root; +use vortex::functions::like::Like; +use vortex::functions::like::LikeOptions; use vortex::scalar_fn::ScalarFnVTableExt; -use vortex::scalar_fn::fns::binary::Binary; -use vortex::scalar_fn::fns::like::Like; -use vortex::scalar_fn::fns::like::LikeOptions; -use vortex::scalar_fn::fns::literal::Literal; -use vortex::scalar_fn::fns::operators::Operator; +use vortex::scalar_fn::builtin::binary::Binary; +use vortex::scalar_fn::builtin::literal::Literal; +use vortex::scalar_fn::builtin::operators::Operator; use vortex::session::VortexSession; use vortex_arrow::ArrowSessionExt; diff --git a/vortex-datafusion/src/convert/exprs/tests.rs b/vortex-datafusion/src/convert/exprs/tests.rs index 6e5a31b7a30..81bd0f2d64f 100644 --- a/vortex-datafusion/src/convert/exprs/tests.rs +++ b/vortex-datafusion/src/convert/exprs/tests.rs @@ -49,12 +49,12 @@ use vortex::array::builtins::ArrayBuiltins; use vortex::dtype::PType; use vortex::encodings::uuid::Uuid; use vortex::expr::between; +use vortex::functions::between::BetweenOptions; +use vortex::functions::between::StrictComparison; +use vortex::functions::list_contains::ListContains; use vortex::scalar::Scalar; -use vortex::scalar_fn::fns::between::BetweenOptions; -use vortex::scalar_fn::fns::between::StrictComparison; -use vortex::scalar_fn::fns::cast::Cast; -use vortex::scalar_fn::fns::list_contains::ListContains; -use vortex::scalar_fn::fns::literal::Literal; +use vortex::scalar_fn::builtin::cast::Cast; +use vortex::scalar_fn::builtin::literal::Literal; use vortex_array::assert_arrays_eq; use super::*; diff --git a/vortex-duckdb/src/convert/expr.rs b/vortex-duckdb/src/convert/expr.rs index 7d7cf0e6aa8..2431b24c323 100644 --- a/vortex-duckdb/src/convert/expr.rs +++ b/vortex-duckdb/src/convert/expr.rs @@ -39,18 +39,18 @@ use vortex::expr::lit; use vortex::expr::not; use vortex::expr::or_collect; use vortex::expr::root; +use vortex::functions::between::Between; +use vortex::functions::between::BetweenOptions; +use vortex::functions::between::StrictComparison; +use vortex::functions::like::Like; +use vortex::functions::like::LikeOptions; use vortex::layout::layouts::row_idx::row_idx; use vortex::scalar::Scalar; use vortex::scalar_fn::EmptyOptions as ScalarEmptyOptions; use vortex::scalar_fn::ScalarFnVTableExt; -use vortex::scalar_fn::fns::between::Between; -use vortex::scalar_fn::fns::between::BetweenOptions; -use vortex::scalar_fn::fns::between::StrictComparison; -use vortex::scalar_fn::fns::binary::Binary; -use vortex::scalar_fn::fns::like::Like; -use vortex::scalar_fn::fns::like::LikeOptions; -use vortex::scalar_fn::fns::literal::Literal; -use vortex::scalar_fn::fns::operators::Operator; +use vortex::scalar_fn::builtin::binary::Binary; +use vortex::scalar_fn::builtin::literal::Literal; +use vortex::scalar_fn::builtin::operators::Operator; use vortex_spatial::extension::LineString; use vortex_spatial::extension::MultiLineString; use vortex_spatial::extension::MultiPoint; diff --git a/vortex-duckdb/src/convert/table_filter.rs b/vortex-duckdb/src/convert/table_filter.rs index 00d3decaf5e..869192d8224 100644 --- a/vortex-duckdb/src/convert/table_filter.rs +++ b/vortex-duckdb/src/convert/table_filter.rs @@ -21,8 +21,8 @@ use vortex::expr::lit; use vortex::expr::or_collect; use vortex::scalar::Scalar; use vortex::scalar_fn::ScalarFnVTableExt; -use vortex::scalar_fn::fns::binary::Binary; -use vortex::scalar_fn::fns::operators::CompareOperator; +use vortex::scalar_fn::builtin::binary::Binary; +use vortex::scalar_fn::builtin::operators::CompareOperator; use vortex::scan::selection::Selection; use vortex::scan::strict_sorted_buffer::StrictSortedBuffer; diff --git a/vortex-expr/Cargo.toml b/vortex-expr/Cargo.toml index cc030485545..1aa7e66bf2b 100644 --- a/vortex-expr/Cargo.toml +++ b/vortex-expr/Cargo.toml @@ -18,6 +18,7 @@ parking_lot = { workspace = true } tracing = { workspace = true } vortex-array = { workspace = true } vortex-error = { workspace = true } +vortex-functions = { workspace = true } vortex-session = { workspace = true } vortex-utils = { workspace = true } diff --git a/vortex-expr/src/analysis/immediate_access.rs b/vortex-expr/src/analysis/immediate_access.rs index d0e31577bd6..151bf10144e 100644 --- a/vortex-expr/src/analysis/immediate_access.rs +++ b/vortex-expr/src/analysis/immediate_access.rs @@ -4,9 +4,9 @@ use vortex_array::dtype::FieldName; use vortex_array::dtype::StructFields; use vortex_array::expr::BoundExpression; -use vortex_array::scalar_fn::fns::get_item::GetItem; -use vortex_array::scalar_fn::fns::select::Select; +use vortex_array::scalar_fn::builtin::get_item::GetItem; use vortex_error::VortexExpect; +use vortex_functions::select::Select; use crate::analysis::AnnotationFn; diff --git a/vortex-expr/src/analysis/referenced_field_paths.rs b/vortex-expr/src/analysis/referenced_field_paths.rs index 55cec7afc5e..5ee15691271 100644 --- a/vortex-expr/src/analysis/referenced_field_paths.rs +++ b/vortex-expr/src/analysis/referenced_field_paths.rs @@ -9,10 +9,10 @@ use vortex_array::expr::traversal::FoldDownContext; use vortex_array::expr::traversal::FoldUp; use vortex_array::expr::traversal::NodeExt; use vortex_array::expr::traversal::NodeFolderContext; -use vortex_array::scalar_fn::fns::get_item::GetItem; -use vortex_array::scalar_fn::fns::select::Select; +use vortex_array::scalar_fn::builtin::get_item::GetItem; use vortex_error::VortexResult; use vortex_error::vortex_err; +use vortex_functions::select::Select; /// Returns the rooted field paths referenced by an expression. /// @@ -132,8 +132,8 @@ mod tests { use vortex_array::expr::get_item; use vortex_array::expr::pack; use vortex_array::expr::root; - use vortex_array::expr::select; - use vortex_array::expr::select_exclude; + use vortex_functions::expr::select; + use vortex_functions::expr::select_exclude; use vortex_utils::aliases::hash_set::HashSet; use super::*; diff --git a/vortex-expr/src/forms/extract_conjuncts.rs b/vortex-expr/src/forms/extract_conjuncts.rs index 86e34b97964..3718e21b273 100644 --- a/vortex-expr/src/forms/extract_conjuncts.rs +++ b/vortex-expr/src/forms/extract_conjuncts.rs @@ -2,8 +2,8 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_array::expr::Expression; -use vortex_array::scalar_fn::fns::binary::Binary; -use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::scalar_fn::builtin::binary::Binary; +use vortex_array::scalar_fn::builtin::operators::Operator; /// Converting an expression to a conjunctive normal form can lead to a large number of expression /// nodes. diff --git a/vortex-expr/src/stats/bind.rs b/vortex-expr/src/stats/bind.rs index 337752d7a15..107d69d8db7 100644 --- a/vortex-expr/src/stats/bind.rs +++ b/vortex-expr/src/stats/bind.rs @@ -20,8 +20,8 @@ use vortex_array::expr::bound::lit; use vortex_array::expr::traversal::NodeExt; use vortex_array::expr::traversal::Transformed; use vortex_array::scalar::Scalar; -use vortex_array::scalar_fn::fns::stat::StatFn; use vortex_error::VortexResult; +use vortex_functions::stat::StatFn; /// A target that can bind abstract statistics to concrete expressions. /// diff --git a/vortex-expr/src/stats/expr.rs b/vortex-expr/src/stats/expr.rs index 2952ccc761c..553f49cb9dd 100644 --- a/vortex-expr/src/stats/expr.rs +++ b/vortex-expr/src/stats/expr.rs @@ -18,9 +18,9 @@ use vortex_array::aggregate_fn::fns::sum::Sum; use vortex_array::expr::BoundExpression; use vortex_array::expr::Expression; use vortex_array::scalar_fn::ScalarFnVTableExt; -pub use vortex_array::scalar_fn::fns::stat::StatFn; -pub use vortex_array::scalar_fn::fns::stat::StatOptions; use vortex_error::VortexExpect; +pub use vortex_functions::stat::StatFn; +pub use vortex_functions::stat::StatOptions; /// Creates an expression that reads a stored aggregate statistic for `expr`. /// @@ -184,6 +184,7 @@ mod tests { use vortex_array::expr::root; use vortex_array::expr::stats::Precision; use vortex_array::expr::stats::Stat; + use vortex_array::optimizer::ArrayOptimizer; use vortex_array::scalar::Scalar; use vortex_array::scalar::ScalarValue; use vortex_array::validity::Validity; @@ -267,7 +268,7 @@ mod tests { )? .into_array(); - let result = chunked.apply(&sum(root()))?; + let result = chunked.apply(&sum(root()))?.optimize(&SESSION)?; let chunked_result = result .as_opt::() diff --git a/vortex-expr/src/stats/rewrite.rs b/vortex-expr/src/stats/rewrite.rs index 968fbfad9ba..a2115884e65 100644 --- a/vortex-expr/src/stats/rewrite.rs +++ b/vortex-expr/src/stats/rewrite.rs @@ -10,8 +10,8 @@ use vortex_array::dtype::DType; use vortex_array::expr::BoundExpression; use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::ScalarFnVTableExt; -use vortex_array::scalar_fn::fns::binary::Binary; -use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::scalar_fn::builtin::binary::Binary; +use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_session::VortexSession; @@ -176,7 +176,7 @@ mod tests { use vortex_array::expr::or; use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::ScalarFnVTable; - use vortex_array::scalar_fn::fns::not::Not; + use vortex_array::scalar_fn::builtin::not::Not; use vortex_error::VortexResult; use vortex_session::VortexSession; diff --git a/vortex-expr/src/stats/rewrite/builtins.rs b/vortex-expr/src/stats/rewrite/builtins.rs index d7c65947195..860adc3b60e 100644 --- a/vortex-expr/src/stats/rewrite/builtins.rs +++ b/vortex-expr/src/stats/rewrite/builtins.rs @@ -16,7 +16,6 @@ use vortex_array::expr::bound::and; use vortex_array::expr::bound::and_collect; use vortex_array::expr::bound::binary; use vortex_array::expr::bound::cast; -use vortex_array::expr::bound::dynamic_with_options; use vortex_array::expr::bound::eq; use vortex_array::expr::bound::gt; use vortex_array::expr::bound::gt_eq; @@ -32,21 +31,22 @@ use vortex_array::scalar_fn::EmptyOptions; use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::ScalarFnVTableExt; -use vortex_array::scalar_fn::fns::between::Between; -use vortex_array::scalar_fn::fns::binary::Binary; -use vortex_array::scalar_fn::fns::cast::Cast; -use vortex_array::scalar_fn::fns::dynamic::DynamicComparison; -use vortex_array::scalar_fn::fns::is_nan::IsNan; -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::like::Like; -use vortex_array::scalar_fn::fns::list_contains::ListContains; -use vortex_array::scalar_fn::fns::literal::Literal; -use vortex_array::scalar_fn::fns::operators::CompareOperator; -use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::scalar_fn::builtin::binary::Binary; +use vortex_array::scalar_fn::builtin::cast::Cast; +use vortex_array::scalar_fn::builtin::is_not_null::IsNotNull; +use vortex_array::scalar_fn::builtin::is_null::IsNull; +use vortex_array::scalar_fn::builtin::literal::Literal; +use vortex_array::scalar_fn::builtin::operators::CompareOperator; +use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_array::scalar_fn::internal::row_count::RowCount; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_functions::between::Between; +use vortex_functions::dynamic::DynamicComparison; +use vortex_functions::expr::bound::dynamic_with_options; +use vortex_functions::is_nan::IsNan; +use vortex_functions::like::Like; +use vortex_functions::list_contains::ListContains; use vortex_session::VortexSession; use crate::stats::bound::stat; @@ -852,19 +852,14 @@ mod tests { use vortex_array::expr::BoundExpression; use vortex_array::expr::Expression; use vortex_array::expr::and; - use vortex_array::expr::between; use vortex_array::expr::cast; use vortex_array::expr::col; - use vortex_array::expr::dynamic; 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::is_nan; use vortex_array::expr::is_not_null; use vortex_array::expr::is_null; - use vortex_array::expr::like; - use vortex_array::expr::list_contains; use vortex_array::expr::lit; use vortex_array::expr::lt; use vortex_array::expr::lt_eq; @@ -875,13 +870,18 @@ mod tests { use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::ScalarFnVTableExt; - use vortex_array::scalar_fn::fns::between::BetweenOptions; - use vortex_array::scalar_fn::fns::between::StrictComparison; - use vortex_array::scalar_fn::fns::binary::Binary; - use vortex_array::scalar_fn::fns::dynamic::DynamicComparison; - use vortex_array::scalar_fn::fns::operators::CompareOperator; + use vortex_array::scalar_fn::builtin::binary::Binary; + use vortex_array::scalar_fn::builtin::operators::CompareOperator; use vortex_array::scalar_fn::internal::row_count::RowCount; use vortex_error::VortexResult; + use vortex_functions::between::BetweenOptions; + use vortex_functions::between::StrictComparison; + use vortex_functions::dynamic::DynamicComparison; + use vortex_functions::expr::between; + use vortex_functions::expr::dynamic; + use vortex_functions::expr::is_nan; + use vortex_functions::expr::like; + use vortex_functions::expr::list_contains; use vortex_session::VortexSession; use crate::stats::expr::StatFn; diff --git a/vortex-expr/src/stats/rewrite/fn_tests.rs b/vortex-expr/src/stats/rewrite/fn_tests.rs index a83a0f20ae7..f14a8fada81 100644 --- a/vortex-expr/src/stats/rewrite/fn_tests.rs +++ b/vortex-expr/src/stats/rewrite/fn_tests.rs @@ -17,7 +17,6 @@ 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_contains; use vortex_array::expr::lit; use vortex_array::expr::lt; use vortex_array::expr::or; @@ -28,6 +27,7 @@ use vortex_array::scalar_fn::EmptyOptions; use vortex_array::scalar_fn::ScalarFnVTableExt; use vortex_array::scalar_fn::internal::row_count::RowCount; use vortex_error::VortexResult; +use vortex_functions::expr::list_contains; use vortex_session::VortexSession; use crate::stats::StatsSession; diff --git a/vortex-expr/src/transform/bound_partition.rs b/vortex-expr/src/transform/bound_partition.rs index 88ef5287062..3f99596c337 100644 --- a/vortex-expr/src/transform/bound_partition.rs +++ b/vortex-expr/src/transform/bound_partition.rs @@ -299,9 +299,9 @@ mod tests { use vortex_array::expr::col; use vortex_array::expr::get_item; use vortex_array::expr::lit; - use vortex_array::expr::merge; use vortex_array::expr::pack; use vortex_array::expr::root; + use vortex_functions::expr::merge; use super::*; use crate::analysis::make_bound_free_field_annotator; diff --git a/vortex-expr/src/transform/match_between.rs b/vortex-expr/src/transform/match_between.rs index 71a5f2b6c30..0a80cdb24c0 100644 --- a/vortex-expr/src/transform/match_between.rs +++ b/vortex-expr/src/transform/match_between.rs @@ -7,15 +7,15 @@ use std::ptr; use vortex_array::expr::BoundExpression; use vortex_array::scalar_fn::ScalarFnVTableExt; -use vortex_array::scalar_fn::fns::between::Between; -use vortex_array::scalar_fn::fns::between::BetweenOptions; -use vortex_array::scalar_fn::fns::between::StrictComparison; -use vortex_array::scalar_fn::fns::binary::Binary; -use vortex_array::scalar_fn::fns::get_item::GetItem; -use vortex_array::scalar_fn::fns::literal::Literal; -use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::scalar_fn::builtin::binary::Binary; +use vortex_array::scalar_fn::builtin::get_item::GetItem; +use vortex_array::scalar_fn::builtin::literal::Literal; +use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_functions::between::Between; +use vortex_functions::between::BetweenOptions; +use vortex_functions::between::StrictComparison; /// Merge pairs of range comparisons on the same column within any conjunction of `expr` into a /// single [`Between`]. @@ -228,7 +228,6 @@ mod tests { use vortex_array::dtype::StructFields; use vortex_array::expr::Expression; use vortex_array::expr::and; - use vortex_array::expr::between; use vortex_array::expr::col; use vortex_array::expr::gt; use vortex_array::expr::gt_eq; @@ -237,10 +236,11 @@ mod tests { use vortex_array::expr::lt_eq; use vortex_array::expr::or; use vortex_array::scalar::Scalar; - use vortex_array::scalar_fn::fns::between::BetweenOptions; - use vortex_array::scalar_fn::fns::between::StrictComparison; use vortex_buffer::buffer; use vortex_error::VortexResult; + use vortex_functions::between::BetweenOptions; + use vortex_functions::between::StrictComparison; + use vortex_functions::expr::between; use vortex_session::VortexSession; use super::find_between; diff --git a/vortex-ffi/src/expression.rs b/vortex-ffi/src/expression.rs index 970e81fa126..4b83b5d3d42 100644 --- a/vortex-ffi/src/expression.rs +++ b/vortex-ffi/src/expression.rs @@ -19,8 +19,8 @@ use vortex::expr::or_collect; use vortex::expr::root; use vortex::expr::select; use vortex::scalar_fn::ScalarFnVTableExt; -use vortex::scalar_fn::fns::binary::Binary; -use vortex::scalar_fn::fns::operators::Operator; +use vortex::scalar_fn::builtin::binary::Binary; +use vortex::scalar_fn::builtin::operators::Operator; use crate::box_wrapper; use crate::error::try_or; diff --git a/vortex-file/Cargo.toml b/vortex-file/Cargo.toml index a35b6c888f1..87b62831fff 100644 --- a/vortex-file/Cargo.toml +++ b/vortex-file/Cargo.toml @@ -42,6 +42,7 @@ vortex-error = { workspace = true } vortex-expr = { workspace = true } vortex-fastlanes = { workspace = true } vortex-fsst = { workspace = true } +vortex-functions = { workspace = true } vortex-io = { workspace = true } vortex-layout = { workspace = true } vortex-mask = { workspace = true } diff --git a/vortex-file/src/lib.rs b/vortex-file/src/lib.rs index 1c309112078..f0662e85452 100644 --- a/vortex-file/src/lib.rs +++ b/vortex-file/src/lib.rs @@ -169,6 +169,7 @@ mod forever_constant { /// Registration covers reading: a session can decode every encoding registered here. The /// writer is gated separately by the editions enabled on its session. pub fn register_default_encodings(session: &VortexSession) { + vortex_functions::initialize(session); vortex_bytebool::initialize(session); vortex_fsst::initialize(session); vortex_onpair::initialize(session); diff --git a/vortex-file/src/pruning.rs b/vortex-file/src/pruning.rs index 4ef9a8b75a1..66c10569d9c 100644 --- a/vortex-file/src/pruning.rs +++ b/vortex-file/src/pruning.rs @@ -11,9 +11,9 @@ use vortex_array::dtype::DType; use vortex_array::dtype::FieldPath; use vortex_array::expr::BoundExpression; use vortex_array::expr::bound::lit; -use vortex_array::scalar_fn::fns::cast::Cast; -use vortex_array::scalar_fn::fns::get_item::GetItem; -use vortex_array::scalar_fn::fns::literal::Literal; +use vortex_array::scalar_fn::builtin::cast::Cast; +use vortex_array::scalar_fn::builtin::get_item::GetItem; +use vortex_array::scalar_fn::builtin::literal::Literal; use vortex_array::scalar_fn::internal::row_count::substitute_row_count; use vortex_error::VortexResult; use vortex_expr::stats::bind::StatBinder; diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index 7e15c104226..baba9630a23 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -59,7 +59,6 @@ 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::expr::stats::Precision; use vortex_array::expr::stats::Stat; use vortex_array::extension::datetime::TimeUnit; @@ -70,8 +69,8 @@ use vortex_array::flatbuffers::WriteFlatBufferExt; use vortex_array::scalar::Scalar; use vortex_array::scalar::ScalarValue; 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::builtin::pack::Pack; +use vortex_array::scalar_fn::builtin::pack::PackOptions; use vortex_array::stream::ArrayStreamAdapter; use vortex_array::stream::ArrayStreamExt; use vortex_array::validity::Validity; @@ -90,6 +89,7 @@ use vortex_edition::EditionSessionExt; use vortex_edition::declarations::core::CORE_2026_08_3; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_functions::expr::select; use vortex_io::VortexWrite; use vortex_io::runtime::tokio::TokioRuntime; use vortex_io::session::RuntimeSession; diff --git a/vortex-functions/Cargo.toml b/vortex-functions/Cargo.toml new file mode 100644 index 00000000000..db5aed6004f --- /dev/null +++ b/vortex-functions/Cargo.toml @@ -0,0 +1,79 @@ +[package] +name = "vortex-functions" +authors = { workspace = true } +categories = { workspace = true } +description = "Vortex scalar functions beyond those the core array crate relies on" +edition = { workspace = true } +include = { workspace = true } +keywords = { workspace = true } +license = { workspace = true } +readme = { workspace = true } +repository = { workspace = true } +rust-version = { workspace = true } +version = { workspace = true } + +[dependencies] +arrow-buffer = { workspace = true } +itertools = { workspace = true } +jiff = { workspace = true } +memchr = { workspace = true } +num-traits = { workspace = true } +parking_lot = { workspace = true } +prost = { workspace = true } +regex = { workspace = true } +regex-syntax = { workspace = true } +tracing = { workspace = true } +vortex-array = { workspace = true, features = ["unstable_row_fns"] } +vortex-buffer = { workspace = true } +vortex-compute = { workspace = true } +vortex-error = { workspace = true } +vortex-mask = { workspace = true } +vortex-session = { workspace = true } +vortex-utils = { workspace = true } + +[lints] +workspace = true + +[dev-dependencies] +divan = { workspace = true } +insta = { workspace = true } +mimalloc = { workspace = true } +rand = { workspace = true } +rstest = { workspace = true } +vortex-array = { workspace = true, features = ["_test-harness", "table-display"] } + +[[bench]] +name = "bool_zip" +harness = false + +[[bench]] +name = "varbinview_zip" +harness = false + +[[bench]] +name = "case_when" +harness = false + +[[bench]] +name = "like" +harness = false + +[[bench]] +name = "list_contains_set" +harness = false + +[[bench]] +name = "list_length" +harness = false + +[[bench]] +name = "list_sum" +harness = false + +[[bench]] +name = "listview_zip" +harness = false + +[[bench]] +name = "primitive_zip" +harness = false diff --git a/vortex-array/benches/bool_zip.rs b/vortex-functions/benches/bool_zip.rs similarity index 89% rename from vortex-array/benches/bool_zip.rs rename to vortex-functions/benches/bool_zip.rs index 40994da33d7..c6423b0d0d1 100644 --- a/vortex-array/benches/bool_zip.rs +++ b/vortex-functions/benches/bool_zip.rs @@ -12,8 +12,8 @@ use vortex_array::IntoArray; use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; -use vortex_array::builtins::ArrayBuiltins; use vortex_array::optimizer::kernels::KernelSession; +use vortex_functions::builtins::ArrayFunctions; use vortex_mask::Mask; use vortex_session::VortexSession; @@ -25,8 +25,11 @@ fn main() { divan::main(); } -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty().with::(); + vortex_functions::initialize(&session); + session +}); const LEN: usize = 65_536; diff --git a/vortex-array/benches/expr/case_when_bench.rs b/vortex-functions/benches/case_when.rs similarity index 93% rename from vortex-array/benches/expr/case_when_bench.rs rename to vortex-functions/benches/case_when.rs index 01f721dffc3..ad502a3b459 100644 --- a/vortex-array/benches/expr/case_when_bench.rs +++ b/vortex-functions/benches/case_when.rs @@ -12,6 +12,7 @@ use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::aggregate_fn::session::AggregateFnSession; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::DecimalArray; @@ -20,33 +21,42 @@ use vortex_array::arrays::StructArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DecimalDType; -use vortex_array::expr::case_when; -use vortex_array::expr::case_when_no_else; +use vortex_array::dtype::session::DTypeSession; use vortex_array::expr::eq; use vortex_array::expr::get_item; use vortex_array::expr::gt; use vortex_array::expr::lit; use vortex_array::expr::lt; -use vortex_array::expr::nested_case_when; use vortex_array::expr::root; use vortex_array::memory::MemorySession; use vortex_array::optimizer::kernels::KernelSession; use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::ScalarFnVTableExt; -use vortex_array::scalar_fn::fns::case_when::CaseWhen; -use vortex_array::scalar_fn::fns::case_when::CaseWhenOptions; -use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::scalar_fn::builtin::operators::Operator; +use vortex_array::scalar_fn::session::ScalarFnSession; +use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::Buffer; +use vortex_functions::case_when::CaseWhen; +use vortex_functions::case_when::CaseWhenOptions; +use vortex_functions::expr::case_when; +use vortex_functions::expr::case_when_no_else; +use vortex_functions::expr::nested_case_when; use vortex_session::VortexSession; #[global_allocator] static GLOBAL: MiMalloc = MiMalloc; static SESSION: LazyLock = LazyLock::new(|| { - VortexSession::empty() + let session = VortexSession::empty() + .with::() .with::() - .with::() + .with::() + .with::() + .with::() + .with::(); + vortex_functions::initialize(&session); + session }); fn main() { diff --git a/vortex-array/benches/like.rs b/vortex-functions/benches/like.rs similarity index 94% rename from vortex-array/benches/like.rs rename to vortex-functions/benches/like.rs index bce9e8fa5fb..b14a04edfbe 100644 --- a/vortex-array/benches/like.rs +++ b/vortex-functions/benches/like.rs @@ -19,14 +19,16 @@ use vortex_array::arrays::ConstantArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::memory::MemorySession; use vortex_array::optimizer::kernels::KernelSession; -use vortex_array::scalar_fn::fns::like::Like; -use vortex_array::scalar_fn::fns::like::LikeOptions; +use vortex_functions::like::Like; +use vortex_functions::like::LikeOptions; use vortex_session::VortexSession; static SESSION: LazyLock = LazyLock::new(|| { - VortexSession::empty() + let session = VortexSession::empty() .with::() - .with::() + .with::(); + vortex_functions::initialize(&session); + session }); #[global_allocator] diff --git a/vortex-array/benches/list_contains_set.rs b/vortex-functions/benches/list_contains_set.rs similarity index 96% rename from vortex-array/benches/list_contains_set.rs rename to vortex-functions/benches/list_contains_set.rs index 56b486dd404..a33e9f3d41b 100644 --- a/vortex-array/benches/list_contains_set.rs +++ b/vortex-functions/benches/list_contains_set.rs @@ -28,18 +28,20 @@ use vortex_array::arrays::VarBinViewArray; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; -use vortex_array::expr::list_contains; use vortex_array::expr::lit; use vortex_array::expr::root; use vortex_array::memory::MemorySession; use vortex_array::optimizer::kernels::KernelSession; use vortex_array::scalar::Scalar; +use vortex_functions::expr::list_contains; use vortex_session::VortexSession; static SESSION: LazyLock = LazyLock::new(|| { - VortexSession::empty() + let session = VortexSession::empty() .with::() - .with::() + .with::(); + vortex_functions::initialize(&session); + session }); #[global_allocator] diff --git a/vortex-array/benches/list_length.rs b/vortex-functions/benches/list_length.rs similarity index 92% rename from vortex-array/benches/list_length.rs rename to vortex-functions/benches/list_length.rs index 238f8ab3be2..e487a9d8db1 100644 --- a/vortex-array/benches/list_length.rs +++ b/vortex-functions/benches/list_length.rs @@ -26,12 +26,15 @@ use vortex_array::arrays::BoolArray; use vortex_array::arrays::ListArray; use vortex_array::arrays::ListViewArray; use vortex_array::arrays::PrimitiveArray; -use vortex_array::expr::list_length; +use vortex_array::dtype::session::DTypeSession; use vortex_array::expr::root; use vortex_array::memory::MemorySession; use vortex_array::optimizer::kernels::KernelSession; +use vortex_array::scalar_fn::session::ScalarFnSession; +use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::Buffer; +use vortex_functions::expr::list_length; use vortex_session::VortexSession; #[global_allocator] @@ -43,10 +46,15 @@ fn main() { } static SESSION: LazyLock = LazyLock::new(|| { - VortexSession::empty() + let session = VortexSession::empty() + .with::() .with::() + .with::() + .with::() .with::() - .with::() + .with::(); + vortex_functions::initialize(&session); + session }); const BASE_LIST_SIZE: usize = 8; diff --git a/vortex-array/benches/list_sum.rs b/vortex-functions/benches/list_sum.rs similarity index 93% rename from vortex-array/benches/list_sum.rs rename to vortex-functions/benches/list_sum.rs index fe366820140..33b778a83c4 100644 --- a/vortex-array/benches/list_sum.rs +++ b/vortex-functions/benches/list_sum.rs @@ -25,12 +25,15 @@ use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::ListArray; use vortex_array::arrays::ListViewArray; use vortex_array::arrays::PrimitiveArray; -use vortex_array::expr::list_sum; +use vortex_array::dtype::session::DTypeSession; use vortex_array::expr::root; use vortex_array::memory::MemorySession; use vortex_array::optimizer::kernels::KernelSession; +use vortex_array::scalar_fn::session::ScalarFnSession; +use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::Buffer; +use vortex_functions::expr::list_sum; use vortex_session::VortexSession; #[global_allocator] @@ -42,10 +45,15 @@ fn main() { } static SESSION: LazyLock = LazyLock::new(|| { - VortexSession::empty() + let session = VortexSession::empty() + .with::() .with::() + .with::() + .with::() .with::() - .with::() + .with::(); + vortex_functions::initialize(&session); + session }); const BASE_LIST_SIZE: usize = 8; diff --git a/vortex-array/benches/listview_zip.rs b/vortex-functions/benches/listview_zip.rs similarity index 94% rename from vortex-array/benches/listview_zip.rs rename to vortex-functions/benches/listview_zip.rs index eafde33ea8e..8f8d4039bb2 100644 --- a/vortex-array/benches/listview_zip.rs +++ b/vortex-functions/benches/listview_zip.rs @@ -13,11 +13,11 @@ use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ListViewArray; -use vortex_array::builtins::ArrayBuiltins; use vortex_array::memory::MemorySession; use vortex_array::optimizer::kernels::KernelSession; use vortex_array::validity::Validity; use vortex_buffer::BufferMut; +use vortex_functions::builtins::ArrayFunctions; use vortex_mask::Mask; use vortex_session::VortexSession; @@ -30,9 +30,11 @@ fn main() { } static SESSION: LazyLock = LazyLock::new(|| { - VortexSession::empty() + let session = VortexSession::empty() .with::() - .with::() + .with::(); + vortex_functions::initialize(&session); + session }); // Smaller than the value-path benches: listview zip cost is dominated by element concatenation and diff --git a/vortex-array/benches/primitive_zip.rs b/vortex-functions/benches/primitive_zip.rs similarity index 91% rename from vortex-array/benches/primitive_zip.rs rename to vortex-functions/benches/primitive_zip.rs index 0df9a15c6b5..6629955b86e 100644 --- a/vortex-array/benches/primitive_zip.rs +++ b/vortex-functions/benches/primitive_zip.rs @@ -16,10 +16,10 @@ use vortex_array::IntoArray; use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; -use vortex_array::builtins::ArrayBuiltins; use vortex_array::optimizer::kernels::KernelSession; use vortex_array::validity::Validity; use vortex_buffer::BufferMut; +use vortex_functions::builtins::ArrayFunctions; use vortex_mask::Mask; use vortex_session::VortexSession; @@ -31,8 +31,11 @@ fn main() { divan::main(); } -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty().with::(); + vortex_functions::initialize(&session); + session +}); // Sized so the bench stays well under a few hundred microseconds under CodSpeed's instruction-count // simulation, which runs ~10x the local walltime; the branchless value blend is still exercised. diff --git a/vortex-array/benches/varbinview_zip.rs b/vortex-functions/benches/varbinview_zip.rs similarity index 94% rename from vortex-array/benches/varbinview_zip.rs rename to vortex-functions/benches/varbinview_zip.rs index 1a39b83ddfc..44e52723359 100644 --- a/vortex-array/benches/varbinview_zip.rs +++ b/vortex-functions/benches/varbinview_zip.rs @@ -11,11 +11,11 @@ use vortex_array::IntoArray; use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; use vortex_array::arrays::VarBinViewArray; -use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::memory::MemorySession; use vortex_array::optimizer::kernels::KernelSession; +use vortex_functions::builtins::ArrayFunctions; use vortex_mask::Mask; use vortex_session::VortexSession; @@ -28,9 +28,11 @@ fn main() { } static SESSION: LazyLock = LazyLock::new(|| { - VortexSession::empty() + let session = VortexSession::empty() .with::() - .with::() + .with::(); + vortex_functions::initialize(&session); + session }); /// Benchmarks zip on VarBinView arrays with a highly fragmented mask (worst case for per-slice lookup paths). diff --git a/vortex-array/src/arrays/constant/compute/between.rs b/vortex-functions/src/between/constant.rs similarity index 84% rename from vortex-array/src/arrays/constant/compute/between.rs rename to vortex-functions/src/between/constant.rs index d79dff2fc17..1e53d4d1389 100644 --- a/vortex-array/src/arrays/constant/compute/between.rs +++ b/vortex-functions/src/between/constant.rs @@ -1,16 +1,16 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::IntoArray; +use vortex_array::arrays::Constant; +use vortex_array::arrays::ConstantArray; +use vortex_array::scalar::Scalar; use vortex_error::VortexResult; -use crate::ArrayRef; -use crate::IntoArray; -use crate::array::ArrayView; -use crate::arrays::Constant; -use crate::arrays::ConstantArray; -use crate::scalar::Scalar; -use crate::scalar_fn::fns::between::BetweenOptions; -use crate::scalar_fn::fns::between::BetweenReduce; +use crate::between::BetweenOptions; +use crate::between::BetweenReduce; impl BetweenReduce for Constant { fn between( diff --git a/vortex-array/src/arrays/decimal/compute/between.rs b/vortex-functions/src/between/decimal.rs similarity index 91% rename from vortex-array/src/arrays/decimal/compute/between.rs rename to vortex-functions/src/between/decimal.rs index c3c647cc96e..5e29323c1bc 100644 --- a/vortex-array/src/arrays/decimal/compute/between.rs +++ b/vortex-functions/src/between/decimal.rs @@ -1,24 +1,24 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::Decimal; +use vortex_array::dtype::NativeDecimalType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::i256; +use vortex_array::match_each_decimal_value_type; +use vortex_array::scalar::Scalar; use vortex_buffer::BitBuffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::array::ArrayView; -use crate::arrays::BoolArray; -use crate::arrays::Decimal; -use crate::dtype::NativeDecimalType; -use crate::dtype::Nullability; -use crate::dtype::i256; -use crate::match_each_decimal_value_type; -use crate::scalar::Scalar; -use crate::scalar_fn::fns::between::BetweenKernel; -use crate::scalar_fn::fns::between::BetweenOptions; -use crate::scalar_fn::fns::between::StrictComparison; +use crate::between::BetweenKernel; +use crate::between::BetweenOptions; +use crate::between::StrictComparison; impl BetweenKernel for Decimal { fn between( diff --git a/vortex-array/src/scalar_fn/fns/between/kernel.rs b/vortex-functions/src/between/kernel.rs similarity index 89% rename from vortex-array/src/scalar_fn/fns/between/kernel.rs rename to vortex-functions/src/between/kernel.rs index f7a42252164..9027030d35a 100644 --- a/vortex-array/src/scalar_fn/fns/between/kernel.rs +++ b/vortex-functions/src/between/kernel.rs @@ -1,22 +1,22 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::VTable; +use vortex_array::arrays::ScalarFn; +use vortex_array::arrays::scalar_fn::ExactScalarFn; +use vortex_array::arrays::scalar_fn::ScalarFnArrayExt; +use vortex_array::arrays::scalar_fn::ScalarFnArrayView; +use vortex_array::kernel::ExecuteParentKernel; +use vortex_array::optimizer::rules::ArrayParentReduceRule; use vortex_error::VortexExpect; use vortex_error::VortexResult; use super::Between; use super::BetweenOptions; use super::short_circuit; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::array::ArrayView; -use crate::array::VTable; -use crate::arrays::ScalarFn; -use crate::arrays::scalar_fn::ExactScalarFn; -use crate::arrays::scalar_fn::ScalarFnArrayExt; -use crate::arrays::scalar_fn::ScalarFnArrayView; -use crate::kernel::ExecuteParentKernel; -use crate::optimizer::rules::ArrayParentReduceRule; /// Reduce rule for between: restructure the array without reading buffers. /// diff --git a/vortex-array/src/scalar_fn/fns/between/mod.rs b/vortex-functions/src/between/mod.rs similarity index 85% rename from vortex-array/src/scalar_fn/fns/between/mod.rs rename to vortex-functions/src/between/mod.rs index 9ba0173975d..b7eb6c2d5c8 100644 --- a/vortex-array/src/scalar_fn/fns/between/mod.rs +++ b/vortex-functions/src/between/mod.rs @@ -1,39 +1,59 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +mod constant; +mod decimal; mod kernel; +mod primitive; use std::fmt::Display; use std::fmt::Formatter; pub use kernel::*; use prost::Message; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::Constant; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::Decimal; +use vortex_array::arrays::Primitive; +use vortex_array::arrays::ScalarFnArray; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::DType::Bool; +use vortex_array::expr::display::ExprDisplay; +use vortex_array::optimizer::kernels::ArrayKernelsExt; +use vortex_array::proto::expr as pb; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::Arity; +use vortex_array::scalar_fn::ChildName; +use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::ScalarFnVTableExt; +use vortex_array::scalar_fn::builtin::operators::Operator; +use vortex_array::scalar_fn::session::ScalarFnSessionExt; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::ArrayRef; -use crate::Canonical; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::arrays::ConstantArray; -use crate::arrays::Decimal; -use crate::arrays::Primitive; -use crate::arrays::ScalarFnArray; -use crate::builtins::ArrayBuiltins; -use crate::dtype::DType; -use crate::dtype::DType::Bool; -use crate::expr::display::ExprDisplay; -use crate::proto::expr as pb; -use crate::scalar::Scalar; -use crate::scalar_fn::Arity; -use crate::scalar_fn::ChildName; -use crate::scalar_fn::ExecutionArgs; -use crate::scalar_fn::ScalarFnId; -use crate::scalar_fn::ScalarFnVTable; -use crate::scalar_fn::ScalarFnVTableExt; -use crate::scalar_fn::fns::operators::Operator; +use crate::builtins::ArrayFunctions; + +/// Registers [`Between`] and its kernels for vortex-array's canonical encodings. +pub(crate) fn initialize(session: &VortexSession) { + session.scalar_fns().register(Between); + let kernels = session.kernels(); + kernels.register_execute_parent_kernel( + Between.id(), + Primitive, + BetweenExecuteAdaptor(Primitive), + ); + kernels.register_execute_parent_kernel(Between.id(), Decimal, BetweenExecuteAdaptor(Decimal)); + kernels.register_reduce_parent_rule(Between.id(), Constant, BetweenReduceAdaptor(Constant)); +} #[derive(Debug, Clone, PartialEq, Eq, Hash)] pub struct BetweenOptions { @@ -88,7 +108,7 @@ impl StrictComparison { /// The result can be a lazy [`ScalarFn`] array, so a caller that needs a computed array /// **must** execute it. /// -/// [`ScalarFn`]: crate::arrays::ScalarFn +/// [`ScalarFn`]: vortex_array::arrays::ScalarFn pub(super) fn short_circuit( arr: &ArrayRef, lower: &ArrayRef, @@ -340,31 +360,39 @@ mod tests { use std::sync::LazyLock; use rstest::rstest; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::BoolArray; + use vortex_array::arrays::DecimalArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::StructArray; + use vortex_array::assert_arrays_eq; + use vortex_array::dtype::DType; + use vortex_array::dtype::DecimalDType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::expr::Expression; + use vortex_array::expr::col; + use vortex_array::expr::get_item; + use vortex_array::expr::lit; + use vortex_array::expr::proto::ExprSerializeProtoExt; + use vortex_array::expr::root; + use vortex_array::scalar::DecimalValue; + use vortex_array::scalar::Scalar; + use vortex_array::test_harness::to_int_indices; + use vortex_array::validity::Validity; use vortex_buffer::buffer; use super::*; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::arrays::BoolArray; - use crate::arrays::DecimalArray; - use crate::arrays::PrimitiveArray; - use crate::arrays::StructArray; - use crate::assert_arrays_eq; - use crate::dtype::DType; - use crate::dtype::DecimalDType; - use crate::dtype::Nullability; - use crate::dtype::PType; use crate::expr::between; - use crate::expr::col; - use crate::expr::get_item; - use crate::expr::lit; - use crate::expr::root; - use crate::scalar::DecimalValue; - use crate::scalar::Scalar; - use crate::test_harness::to_int_indices; - use crate::validity::Validity; - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + + crate::initialize(&session); + + session + }); const NON_STRICT: BetweenOptions = BetweenOptions { lower_strict: StrictComparison::NonStrict, @@ -738,4 +766,41 @@ mod tests { assert_eq!(to_int_indices(result, ctx).unwrap(), expected_indices); } + + #[test] + fn test_display_tree_between() { + use insta::assert_snapshot; + let between_expr = between( + get_item("score", root()), + lit(0), + lit(100), + BetweenOptions { + lower_strict: StrictComparison::NonStrict, + upper_strict: StrictComparison::NonStrict, + }, + ); + assert_snapshot!(between_expr.display_tree().to_string(), @r" + vortex.between(lower_strict: <=, upper_strict: <=) + ├── array: vortex.get_item(score) + │ └── input: vortex.root() + ├── lower: vortex.literal(0i32) + └── upper: vortex.literal(100i32) + "); + } + + #[test] + fn between_serde() -> VortexResult<()> { + let expr = between( + lit(1), + root(), + get_item("a", root()), + BetweenOptions { + lower_strict: StrictComparison::Strict, + upper_strict: StrictComparison::Strict, + }, + ); + let proto = expr.serialize_proto()?; + assert_eq!(Expression::from_proto(&proto, &SESSION)?, expr); + Ok(()) + } } diff --git a/vortex-array/src/arrays/primitive/compute/between.rs b/vortex-functions/src/between/primitive.rs similarity index 88% rename from vortex-array/src/arrays/primitive/compute/between.rs rename to vortex-functions/src/between/primitive.rs index 7012cb1be5f..adc5488eb1d 100644 --- a/vortex-array/src/arrays/primitive/compute/between.rs +++ b/vortex-functions/src/between/primitive.rs @@ -1,22 +1,22 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::Primitive; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::Nullability; +use vortex_array::match_each_native_ptype; use vortex_buffer::BitBuffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::array::ArrayView; -use crate::arrays::BoolArray; -use crate::arrays::Primitive; -use crate::dtype::NativePType; -use crate::dtype::Nullability; -use crate::match_each_native_ptype; -use crate::scalar_fn::fns::between::BetweenKernel; -use crate::scalar_fn::fns::between::BetweenOptions; -use crate::scalar_fn::fns::between::StrictComparison; +use crate::between::BetweenKernel; +use crate::between::BetweenOptions; +use crate::between::StrictComparison; impl BetweenKernel for Primitive { fn between( diff --git a/vortex-functions/src/builtins.rs b/vortex-functions/src/builtins.rs new file mode 100644 index 00000000000..6c261a418dc --- /dev/null +++ b/vortex-functions/src/builtins.rs @@ -0,0 +1,74 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Builder methods for this crate's scalar functions, complementing +//! [`ExprBuiltins`](vortex_array::builtins::ExprBuiltins) and +//! [`ArrayBuiltins`](vortex_array::builtins::ArrayBuiltins). + +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::expr::Expression; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ScalarFnVTableExt; +use vortex_error::VortexResult; + +use crate::between::Between; +use crate::between::BetweenOptions; +use crate::list_contains::ListContains; +use crate::list_contains::ListContainsOptions; +use crate::zip::Zip; + +/// Builder methods on [`Expression`] for this crate's scalar functions. +pub trait ExprFunctions: Sized { + /// Conditional selection: `result[i] = if mask[i] then if_true[i] else if_false[i]`. + fn zip(&self, if_true: Expression, if_false: Expression) -> VortexResult; + + /// Check if a list contains a value. + fn list_contains(&self, value: Expression) -> VortexResult; +} + +impl ExprFunctions for Expression { + fn zip(&self, if_true: Expression, if_false: Expression) -> VortexResult { + Zip.try_new_expr(EmptyOptions, [if_true, if_false, self.clone()]) + } + + fn list_contains(&self, value: Expression) -> VortexResult { + ListContains.try_new_expr(ListContainsOptions::default(), [self.clone(), value]) + } +} + +/// Builder methods on [`ArrayRef`] for this crate's scalar functions. +pub trait ArrayFunctions: Sized { + /// Conditional selection: `result[i] = if mask[i] then if_true[i] else if_false[i]`. + fn zip(&self, if_true: ArrayRef, if_false: ArrayRef) -> VortexResult; + + /// Check if a list contains a value. + fn list_contains(&self, value: ArrayRef) -> VortexResult; + + /// Compare a values between lower VortexResult; +} + +impl ArrayFunctions for ArrayRef { + fn zip(&self, if_true: ArrayRef, if_false: ArrayRef) -> VortexResult { + Ok(Zip::try_new(if_true, if_false, self.clone())?.into_array()) + } + + fn list_contains(&self, value: ArrayRef) -> VortexResult { + Ok(ListContains::try_new(self.clone(), value)?.into_array()) + } + + fn between( + self, + lower: ArrayRef, + upper: ArrayRef, + options: BetweenOptions, + ) -> VortexResult { + Ok(Between::try_new(self, lower, upper, options)?.into_array()) + } +} diff --git a/vortex-array/src/scalar_fn/fns/byte_length.rs b/vortex-functions/src/byte_length.rs similarity index 83% rename from vortex-array/src/scalar_fn/fns/byte_length.rs rename to vortex-functions/src/byte_length.rs index afdf77d118d..552b7dc82d6 100644 --- a/vortex-array/src/scalar_fn/fns/byte_length.rs +++ b/vortex-functions/src/byte_length.rs @@ -2,6 +2,28 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use num_traits::AsPrimitive; +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VTable; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::VarBinViewArray; +use vortex_array::arrays::scalar_fn::ExactScalarFn; +use vortex_array::arrays::scalar_fn::ScalarFnArrayView; +use vortex_array::arrays::varbinview::VarBinViewArrayExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::kernel::ExecuteParentKernel; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::Arity; +use vortex_array::scalar_fn::ChildName; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::ScalarFnVTable; use vortex_buffer::Buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -10,29 +32,6 @@ use vortex_error::vortex_ensure_eq; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::array::ArrayView; -use crate::array::VTable; -use crate::arrays::ConstantArray; -use crate::arrays::PrimitiveArray; -use crate::arrays::VarBinViewArray; -use crate::arrays::scalar_fn::ExactScalarFn; -use crate::arrays::scalar_fn::ScalarFnArrayView; -use crate::arrays::varbinview::VarBinViewArrayExt; -use crate::dtype::DType; -use crate::dtype::Nullability; -use crate::dtype::PType; -use crate::kernel::ExecuteParentKernel; -use crate::scalar::Scalar; -use crate::scalar_fn::Arity; -use crate::scalar_fn::ChildName; -use crate::scalar_fn::EmptyOptions; -use crate::scalar_fn::ExecutionArgs; -use crate::scalar_fn::ScalarFnId; -use crate::scalar_fn::ScalarFnVTable; - pub trait ByteLengthKernel: VTable { fn byte_length( array: ArrayView<'_, Self>, @@ -169,24 +168,30 @@ mod tests { use std::sync::LazyLock; use rstest::rstest; + use vortex_array::ArrayRef; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::ConstantArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::VarBinArray; + use vortex_array::arrays::VarBinViewArray; + use vortex_array::assert_arrays_eq; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::expr::root; + use vortex_array::scalar::Scalar; use vortex_error::VortexResult; use vortex_session::VortexSession; - use crate::ArrayRef; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::arrays::ConstantArray; - use crate::arrays::PrimitiveArray; - use crate::arrays::VarBinArray; - use crate::arrays::VarBinViewArray; - use crate::assert_arrays_eq; - use crate::dtype::DType; - use crate::dtype::Nullability; use crate::expr::byte_length; - use crate::expr::root; - use crate::scalar::Scalar; - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + + crate::initialize(&session); + + session + }); #[rstest] #[case(VarBinArray::from_strs(vec!["hello", "world", ""]).into_array(), vec![5u64, 5, 0])] diff --git a/vortex-array/src/scalar_fn/fns/case_when.rs b/vortex-functions/src/case_when.rs similarity index 96% rename from vortex-array/src/scalar_fn/fns/case_when.rs rename to vortex-functions/src/case_when.rs index 1a2e7f6abb5..d031d55793e 100644 --- a/vortex-array/src/scalar_fn/fns/case_when.rs +++ b/vortex-functions/src/case_when.rs @@ -19,6 +19,30 @@ use std::hash::Hash; use std::sync::Arc; use prost::Message; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::bool::BoolArrayExt; +use vortex_array::builders::builder_with_capacity_in; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::expr::BoundExpression; +use vortex_array::expr::display::ExprDisplay; +use vortex_array::proto::expr as pb; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::Arity; +use vortex_array::scalar_fn::ChildName; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::ScalarFnVTableExt; +use vortex_array::scalar_fn::builtin::fill_null::FillNull; +use vortex_array::scalar_fn::builtin::is_not_null::IsNotNull; +use vortex_array::scalar_fn::builtin::is_null::IsNull; +use vortex_array::scalar_fn::builtin::literal::Literal; use vortex_buffer::BufferMut; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -28,31 +52,7 @@ use vortex_mask::Mask; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::arrays::BoolArray; -use crate::arrays::ConstantArray; -use crate::arrays::bool::BoolArrayExt; -use crate::builders::builder_with_capacity_in; -use crate::builtins::ArrayBuiltins; -use crate::dtype::DType; -use crate::expr::BoundExpression; -use crate::expr::display::ExprDisplay; -use crate::proto::expr as pb; -use crate::scalar::Scalar; -use crate::scalar_fn::Arity; -use crate::scalar_fn::ChildName; -use crate::scalar_fn::EmptyOptions; -use crate::scalar_fn::ExecutionArgs; -use crate::scalar_fn::ScalarFnId; -use crate::scalar_fn::ScalarFnVTable; -use crate::scalar_fn::ScalarFnVTableExt; -use crate::scalar_fn::fns::fill_null::FillNull; -use crate::scalar_fn::fns::is_not_null::IsNotNull; -use crate::scalar_fn::fns::is_null::IsNull; -use crate::scalar_fn::fns::literal::Literal; -use crate::scalar_fn::fns::zip::zip_nullability_union; +use crate::zip::zip_nullability_union; /// Options for the n-ary CaseWhen expression. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] @@ -477,48 +477,52 @@ mod tests { use std::sync::LazyLock; use rstest::rstest; + use vortex_array::Canonical; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::BoolArray; + use vortex_array::arrays::DecimalArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::ScalarFnArray; + use vortex_array::arrays::StructArray; + use vortex_array::arrays::VarBinViewArray; + use vortex_array::assert_arrays_eq; + use vortex_array::dtype::DType; + use vortex_array::dtype::DecimalDType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::dtype::StructFields; + use vortex_array::expr::Expression; + 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::is_not_null; + use vortex_array::expr::is_null; + use vortex_array::expr::lit; + use vortex_array::expr::proto::ExprSerializeProtoExt; + use vortex_array::expr::root; + use vortex_array::expr::test_harness; + use vortex_array::scalar::Scalar; + use vortex_array::scalar_fn::ScalarFnVTableExt; + use vortex_array::scalar_fn::VecExecutionArgs; + use vortex_array::scalar_fn::builtin::operators::Operator; + use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::buffer; use vortex_error::VortexExpect as _; use vortex_session::VortexSession; use super::*; - use crate::Canonical; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::arrays::BoolArray; - use crate::arrays::DecimalArray; - use crate::arrays::PrimitiveArray; - use crate::arrays::ScalarFnArray; - use crate::arrays::StructArray; - use crate::arrays::VarBinViewArray; - use crate::assert_arrays_eq; - use crate::dtype::DType; - use crate::dtype::DecimalDType; - use crate::dtype::Nullability; - use crate::dtype::PType; - use crate::dtype::StructFields; - use crate::expr::Expression; use crate::expr::case_when; use crate::expr::case_when_no_else; - use crate::expr::col; - use crate::expr::eq; - use crate::expr::get_item; - use crate::expr::gt; - use crate::expr::is_not_null; - use crate::expr::is_null; - use crate::expr::lit; use crate::expr::nested_case_when; - use crate::expr::proto::ExprSerializeProtoExt; - use crate::expr::root; - use crate::expr::test_harness; - use crate::scalar::Scalar; - use crate::scalar_fn::ScalarFnVTableExt; - use crate::scalar_fn::VecExecutionArgs; - use crate::scalar_fn::fns::operators::Operator; - use crate::validity::Validity; - - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + crate::initialize(&session); + session + }); #[test] fn compact_and_run_merges_agree_for_nested_and_nullable_values() -> VortexResult<()> { diff --git a/vortex-array/src/scalar_fn/fns/dynamic.rs b/vortex-functions/src/dynamic.rs similarity index 89% rename from vortex-array/src/scalar_fn/fns/dynamic.rs rename to vortex-functions/src/dynamic.rs index 7ce5a17dd3b..ee47b87a9bb 100644 --- a/vortex-array/src/scalar_fn/fns/dynamic.rs +++ b/vortex-functions/src/dynamic.rs @@ -9,34 +9,33 @@ use std::hash::Hasher; use std::sync::Arc; use parking_lot::Mutex; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::dtype::DType; +use vortex_array::expr::BoundExpression; +use vortex_array::expr::display::ExprDisplay; +use vortex_array::expr::traversal::NodeExt; +use vortex_array::expr::traversal::NodeVisitor; +use vortex_array::expr::traversal::TraversalOrder; +use vortex_array::scalar::Scalar; +use vortex_array::scalar::ScalarValue; +use vortex_array::scalar_fn::Arity; +use vortex_array::scalar_fn::ChildName; +use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::ScalarFnVTableExt; +use vortex_array::scalar_fn::VecExecutionArgs; +use vortex_array::scalar_fn::builtin::binary::Binary; +use vortex_array::scalar_fn::builtin::operators::CompareOperator; +use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_session::registry::CachedId; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::arrays::ConstantArray; -use crate::dtype::DType; -use crate::expr::BoundExpression; -use crate::expr::display::ExprDisplay; -use crate::expr::traversal::NodeExt; -use crate::expr::traversal::NodeVisitor; -use crate::expr::traversal::TraversalOrder; -use crate::scalar::Scalar; -use crate::scalar::ScalarValue; -use crate::scalar_fn::Arity; -use crate::scalar_fn::ChildName; -use crate::scalar_fn::ExecutionArgs; -use crate::scalar_fn::ScalarFnId; -use crate::scalar_fn::ScalarFnVTable; -use crate::scalar_fn::ScalarFnVTableExt; -use crate::scalar_fn::VecExecutionArgs; -use crate::scalar_fn::fns::binary::Binary; -use crate::scalar_fn::fns::operators::CompareOperator; -use crate::scalar_fn::fns::operators::Operator; - /// A dynamic comparison expression can be used to capture a comparison to a value that can change /// during the execution of a query, such as when a compute engine pushes down an ORDER BY + LIMIT /// operation and is able to progressively tighten the bounds of the filter. @@ -302,22 +301,28 @@ mod tests { use std::sync::atomic::AtomicI32; use std::sync::atomic::Ordering; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::BoolArray; + use vortex_array::assert_arrays_eq; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::expr::root; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; use super::*; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::arrays::BoolArray; - use crate::assert_arrays_eq; - use crate::dtype::DType; - use crate::dtype::Nullability; - use crate::dtype::PType; use crate::expr::dynamic; - use crate::expr::root; - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + + crate::initialize(&session); + + session + }); #[test] fn is_not_strict() { diff --git a/vortex-functions/src/expr.rs b/vortex-functions/src/expr.rs new file mode 100644 index 00000000000..3b9e1b6bde8 --- /dev/null +++ b/vortex-functions/src/expr.rs @@ -0,0 +1,815 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Expression constructors for this crate's scalar functions. +//! +//! These complement the constructors in `vortex_array::expr`. + +use std::sync::Arc; + +use vortex_array::aggregate_fn::NumericalAggregateOpts; +use vortex_array::dtype::DType; +use vortex_array::dtype::FieldNames; +use vortex_array::dtype::Nullability; +use vortex_array::expr::BoundExpression; +use vortex_array::expr::Expression; +use vortex_array::expr::lit; +use vortex_array::scalar::Scalar; +use vortex_array::scalar::ScalarValue; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ScalarFnVTableExt; +use vortex_array::scalar_fn::builtin::operators::CompareOperator; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_panic; + +use crate::between::Between; +use crate::between::BetweenOptions; +use crate::between::StrictComparison; +use crate::byte_length::ByteLength; +use crate::case_when::CaseWhen; +use crate::case_when::CaseWhenOptions; +use crate::dynamic::DynamicComparison; +use crate::dynamic::DynamicComparisonExpr; +use crate::dynamic::Rhs; +use crate::ext_storage::ExtStorage; +use crate::is_nan::IsNan; +use crate::like::Like; +use crate::like::LikeOptions; +use crate::list_contains::ListContains; +use crate::list_contains::ListContainsOptions; +use crate::list_length::ListLength; +use crate::list_sum::ListSum; +use crate::merge::DuplicateHandling; +use crate::merge::Merge; +use crate::replace_time_zone::ReplaceTimeZone; +use crate::replace_time_zone::ReplaceTimeZoneOptions; +use crate::select::FieldSelection; +use crate::select::Select; +use crate::variant_get::VariantGet; +use crate::variant_get::VariantGetOptions; +use crate::variant_get::VariantPath; +use crate::zip::Zip; + +/// Creates a CASE WHEN expression with one WHEN/THEN pair and an ELSE value. +pub fn case_when( + condition: Expression, + then_value: Expression, + else_value: Expression, +) -> Expression { + let options = CaseWhenOptions { + num_when_then_pairs: 1, + has_else: true, + }; + CaseWhen.new_expr(options, [condition, then_value, else_value]) +} + +/// Creates a bound CASE WHEN expression with one WHEN/THEN pair and an ELSE value. +pub fn bound_case_when( + condition: BoundExpression, + then_value: BoundExpression, + else_value: BoundExpression, +) -> BoundExpression { + let options = CaseWhenOptions { + num_when_then_pairs: 1, + has_else: true, + }; + CaseWhen + .try_new_bound_expr(options, [condition, then_value, else_value]) + .vortex_expect("case expressions must have boolean conditions and matching branch dtypes") +} + +/// Creates a CASE WHEN expression with one WHEN/THEN pair and no ELSE value. +pub fn case_when_no_else(condition: Expression, then_value: Expression) -> Expression { + let options = CaseWhenOptions { + num_when_then_pairs: 1, + has_else: false, + }; + CaseWhen.new_expr(options, [condition, then_value]) +} + +/// Creates a bound CASE WHEN expression with one WHEN/THEN pair and no ELSE value. +pub fn bound_case_when_no_else( + condition: BoundExpression, + then_value: BoundExpression, +) -> BoundExpression { + let options = CaseWhenOptions { + num_when_then_pairs: 1, + has_else: false, + }; + CaseWhen + .try_new_bound_expr(options, [condition, then_value]) + .vortex_expect("case expressions must have boolean conditions") +} + +/// Creates an n-ary CASE WHEN expression from WHEN/THEN pairs and an optional ELSE value. +pub fn nested_case_when( + when_then_pairs: Vec<(Expression, Expression)>, + else_value: Option, +) -> Expression { + assert!( + !when_then_pairs.is_empty(), + "nested_case_when requires at least one when/then pair" + ); + + let has_else = else_value.is_some(); + let mut children = Vec::with_capacity(when_then_pairs.len() * 2 + usize::from(has_else)); + for (condition, then_value) in &when_then_pairs { + children.push(condition.clone()); + children.push(then_value.clone()); + } + if let Some(else_expr) = else_value { + children.push(else_expr); + } + + let Ok(num_when_then_pairs) = u32::try_from(when_then_pairs.len()) else { + vortex_panic!("nested_case_when has too many when/then pairs"); + }; + let options = CaseWhenOptions { + num_when_then_pairs, + has_else, + }; + CaseWhen.new_expr(options, children) +} + +/// Creates a bound n-ary CASE WHEN expression from WHEN/THEN pairs and an optional ELSE value. +pub fn bound_nested_case_when( + when_then_pairs: Vec<(BoundExpression, BoundExpression)>, + else_value: Option, +) -> BoundExpression { + assert!( + !when_then_pairs.is_empty(), + "nested_case_when requires at least one when/then pair" + ); + + let Ok(num_when_then_pairs) = u32::try_from(when_then_pairs.len()) else { + vortex_panic!("nested_case_when has too many when/then pairs"); + }; + let has_else = else_value.is_some(); + let mut children = Vec::with_capacity(when_then_pairs.len() * 2 + usize::from(has_else)); + for (condition, then_value) in when_then_pairs { + children.push(condition); + children.push(then_value); + } + if let Some(else_expr) = else_value { + children.push(else_expr); + } + + let options = CaseWhenOptions { + num_when_then_pairs, + has_else, + }; + CaseWhen + .try_new_bound_expr(options, children) + .vortex_expect("case expressions must have boolean conditions and matching branch dtypes") +} + +/// Creates an expression that selects (includes) specific fields from an array. +/// +/// Projects only the specified fields from the child expression, which must be of DType struct. +/// ```rust +/// # use vortex_functions::expr::select; +/// # use vortex_array::expr::root; +/// let expr = select(["name", "age"], root()); +/// ``` +pub fn select(field_names: impl Into, child: Expression) -> Expression { + Select + .try_new_expr(FieldSelection::Include(field_names.into()), [child]) + .vortex_expect("Failed to create Select expression") +} + +/// Creates a bound expression that selects specific fields from a struct expression. +pub fn bound_select(field_names: impl Into, child: BoundExpression) -> BoundExpression { + Select + .try_new_bound_expr(FieldSelection::Include(field_names.into()), [child]) + .vortex_expect("select expressions require fields from a struct child") +} + +/// Creates an expression that excludes specific fields from an array. +/// +/// Projects all fields except the specified ones from the input struct expression. +/// +/// ```rust +/// # use vortex_functions::expr::select_exclude; +/// # use vortex_array::expr::root; +/// let expr = select_exclude(["internal_id", "metadata"], root()); +/// ``` +pub fn select_exclude(fields: impl Into, child: Expression) -> Expression { + Select + .try_new_expr(FieldSelection::Exclude(fields.into()), [child]) + .vortex_expect("Failed to create Select expression") +} + +/// Creates a bound expression that excludes specific fields from a struct expression. +pub fn bound_select_exclude( + fields: impl Into, + child: BoundExpression, +) -> BoundExpression { + Select + .try_new_bound_expr(FieldSelection::Exclude(fields.into()), [child]) + .vortex_expect("select expressions require fields from a struct child") +} + +/// Creates an expression that merges struct expressions into a single struct. +/// +/// Combines fields from all input expressions. If field names are duplicated, +/// later expressions win. Fields are not recursively merged. +/// +/// ```rust +/// # use vortex_array::dtype::Nullability; +/// # use vortex_functions::expr::merge; +/// # use vortex_array::expr::{get_item, root}; +/// let expr = merge([get_item("a", root()), get_item("b", root())]); +/// ``` +pub fn merge(elements: impl IntoIterator>) -> Expression { + use itertools::Itertools as _; + let values = elements.into_iter().map(|value| value.into()).collect_vec(); + Merge.new_expr(DuplicateHandling::default(), values) +} + +/// Creates a bound expression that merges struct expressions. +pub fn bound_merge(elements: impl IntoIterator) -> BoundExpression { + bound_merge_opts(elements, DuplicateHandling::default()) +} + +/// Creates a merge expression with explicit duplicate handling. +pub fn merge_opts( + elements: impl IntoIterator>, + duplicate_handling: DuplicateHandling, +) -> Expression { + use itertools::Itertools as _; + let values = elements.into_iter().map(|value| value.into()).collect_vec(); + Merge.new_expr(duplicate_handling, values) +} + +/// Creates a bound merge expression with explicit duplicate handling. +pub fn bound_merge_opts( + elements: impl IntoIterator, + duplicate_handling: DuplicateHandling, +) -> BoundExpression { + Merge + .try_new_bound_expr(duplicate_handling, elements) + .vortex_expect("merge expressions require non-nullable struct children") +} + +/// Creates an expression that computes the byte length of each element. +/// This is akin to ANSI SQL OCTET_LENGTH(), or DuckDB's strlen(). +/// +/// ```rust +/// # use vortex_functions::expr::byte_length; +/// # use vortex_array::expr::root; +/// let expr = byte_length(root()); +/// ``` +pub fn byte_length(input: Expression) -> Expression { + ByteLength.new_expr(EmptyOptions, [input]) +} + +/// Creates a bound expression that computes each element's byte length. +pub fn bound_byte_length(input: BoundExpression) -> BoundExpression { + ByteLength + .try_new_bound_expr(EmptyOptions, [input]) + .vortex_expect("byte-length expressions require a variable-length binary child") +} + +/// Creates an expression that extracts the storage values from an extension array. +/// +/// ```rust +/// # use vortex_functions::expr::ext_storage; +/// # use vortex_array::expr::root; +/// let expr = ext_storage(root()); +/// ``` +pub fn ext_storage(input: Expression) -> Expression { + ExtStorage.new_expr(EmptyOptions, [input]) +} + +/// Creates a bound expression that extracts an extension array's storage values. +pub fn bound_ext_storage(input: BoundExpression) -> BoundExpression { + ExtStorage + .try_new_bound_expr(EmptyOptions, [input]) + .vortex_expect("extension-storage expressions require an extension child") +} + +/// Creates an expression that computes the number of elements in each list +/// for `List` and `FixedSizeList` inputs. This is akin to ANSI SQL `CARDINALITY()`, +/// or DuckDB's `len()`/`array_length()`. +/// +/// ```rust +/// # use vortex_functions::expr::list_length; +/// # use vortex_array::expr::root; +/// let expr = list_length(root()); +/// ``` +pub fn list_length(input: Expression) -> Expression { + ListLength.new_expr(EmptyOptions, [input]) +} + +/// Creates a bound expression that computes the number of elements in each list. +pub fn bound_list_length(input: BoundExpression) -> BoundExpression { + ListLength + .try_new_bound_expr(EmptyOptions, [input]) + .vortex_expect("list-length expressions require a list child") +} + +/// Creates an expression that sums the elements of each list for `List` and +/// `FixedSizeList` inputs, akin to DuckDB's `list_sum()`. +/// +/// Follows SQL `SUM` semantics per list: null lists, empty lists, and lists whose elements are +/// all null yield null; null elements are skipped; integer and decimal overflow yields a null +/// value. The result dtype follows `sum`'s widening rules and is always nullable. NaN float +/// elements are skipped by default; see [`list_sum_opts`] for the NaN-including variant. +/// +/// ```rust +/// # use vortex_functions::expr::list_sum; +/// # use vortex_array::expr::root; +/// let expr = list_sum(root()); +/// ``` +pub fn list_sum(input: Expression) -> Expression { + ListSum.new_expr(NumericalAggregateOpts::default(), [input]) +} + +/// Creates a bound expression that sums the elements of each list. +pub fn bound_list_sum(input: BoundExpression) -> BoundExpression { + ListSum + .try_new_bound_expr(NumericalAggregateOpts::default(), [input]) + .vortex_expect("list-sum expressions require a numeric list child") +} + +/// Creates a [`list_sum`] expression with explicit [`NumericalAggregateOpts`], controlling +/// whether NaN float elements are skipped (the default) or poison the list's sum to NaN. +pub fn list_sum_opts(input: Expression, options: NumericalAggregateOpts) -> Expression { + ListSum.new_expr(options, [input]) +} + +/// Creates a bound list-sum expression with explicit aggregate options. +pub fn bound_list_sum_opts( + input: BoundExpression, + options: NumericalAggregateOpts, +) -> BoundExpression { + ListSum + .try_new_bound_expr(options, [input]) + .vortex_expect("list-sum expressions require a numeric list child") +} + +/// Reinterpret timestamp wall times in `options.time_zone`, preserving the time unit. +/// +/// `ambiguous` is a UTF-8 expression containing `raise`, `earliest`, `latest`, or `null`. +/// Removing the timezone preserves wall time in a timezone-naive timestamp. +pub fn replace_time_zone( + input: Expression, + ambiguous: Expression, + options: ReplaceTimeZoneOptions, +) -> Expression { + ReplaceTimeZone.new_expr(options, [input, ambiguous]) +} + +/// Create a bound timezone replacement expression. +pub fn bound_replace_time_zone( + input: BoundExpression, + ambiguous: BoundExpression, + options: ReplaceTimeZoneOptions, +) -> BoundExpression { + ReplaceTimeZone + .try_new_bound_expr(options, [input, ambiguous]) + .vortex_expect("timezone replacement requires a timestamp and a UTF-8 ambiguity policy") +} + +/// Constructors for expressions whose children have already been bound and type-checked. +/// +/// These mirror the constructors in [`crate::expr`] and panic when the supplied children do not +/// form a well-typed expression. +/// Creates a dynamic comparison expression from its complete options. +pub fn dynamic_with_options(options: DynamicComparisonExpr, lhs: Expression) -> Expression { + DynamicComparison.new_expr(options, [lhs]) +} + +/// Creates a bound dynamic comparison expression from its complete options. +pub fn bound_dynamic_with_options( + options: DynamicComparisonExpr, + lhs: BoundExpression, +) -> BoundExpression { + DynamicComparison + .try_new_bound_expr(options, [lhs]) + .vortex_expect("dynamic comparisons require a compatible left-hand dtype") +} + +/// Creates a dynamic comparison expression. +pub fn dynamic( + operator: CompareOperator, + rhs_value: impl Fn() -> Option + Send + Sync + 'static, + rhs_dtype: DType, + default: bool, + lhs: Expression, +) -> Expression { + dynamic_with_options( + DynamicComparisonExpr { + operator, + rhs: Arc::new(Rhs { + value: Arc::new(rhs_value), + dtype: rhs_dtype, + }), + default, + }, + lhs, + ) +} + +/// Creates a bound dynamic comparison expression. +pub fn bound_dynamic( + operator: CompareOperator, + rhs_value: impl Fn() -> Option + Send + Sync + 'static, + rhs_dtype: DType, + default: bool, + lhs: BoundExpression, +) -> BoundExpression { + bound_dynamic_with_options( + DynamicComparisonExpr { + operator, + rhs: Arc::new(Rhs { + value: Arc::new(rhs_value), + dtype: rhs_dtype, + }), + default, + }, + lhs, + ) +} + +// ---- Between ---- + +/// Creates an expression that checks if values are between two bounds. +/// +/// Returns a boolean array indicating which values fall within the specified range. +/// The comparison strictness is controlled by the options parameter. +/// +/// ```rust +/// # use vortex_functions::between::BetweenOptions; +/// # use vortex_functions::between::StrictComparison; +/// # use vortex_array::expr::{lit, root}; +/// # use vortex_functions::expr::between; +/// let opts = BetweenOptions { +/// lower_strict: StrictComparison::NonStrict, +/// upper_strict: StrictComparison::NonStrict, +/// }; +/// let expr = between(root(), lit(10), lit(20), opts); +/// ``` +pub fn between( + arr: Expression, + lower: Expression, + upper: Expression, + options: BetweenOptions, +) -> Expression { + Between + .try_new_expr(options, [arr, lower, upper]) + .vortex_expect("Failed to create Between expression") +} + +/// Creates a bound expression that checks if values are between two bounds. +pub fn bound_between( + arr: BoundExpression, + lower: BoundExpression, + upper: BoundExpression, + options: BetweenOptions, +) -> BoundExpression { + Between + .try_new_bound_expr(options, [arr, lower, upper]) + .vortex_expect("between expressions require compatible operand dtypes") +} + +// ---- ListContains ---- + +/// Creates an expression that checks whether a list contains a value. +/// +/// A null element never matches anything: a needle that matches no element yields `false`. For +/// SQL's three-valued `IN`, where a null element makes that `null`, use [`in_list`]. +/// +/// ```rust +/// # use vortex_array::expr::{lit, root}; +/// # use vortex_functions::expr::list_contains; +/// let expr = list_contains(root(), lit(42)); +/// ``` +pub fn list_contains(list: Expression, value: Expression) -> Expression { + list_contains_opts(list, value, ListContainsOptions::default()) +} + +/// Creates an expression that checks whether a list contains a value, with explicit +/// [`ListContainsOptions`]. +pub fn list_contains_opts( + list: Expression, + value: Expression, + options: ListContainsOptions, +) -> Expression { + ListContains.new_expr(options, [list, value]) +} + +/// Creates `needle IN (list)` with SQL null semantics: a null element is an unknown value, so a +/// needle that matches no element yields `null` rather than `false` when the list holds one, and +/// `not(in_list(..))` never admits such a row. +/// +/// ```rust +/// # use vortex_array::expr::{lit, root}; +/// # use vortex_functions::expr::in_list; +/// let expr = in_list(root(), lit(vec![1, 2, 3])); +/// ``` +pub fn in_list(needle: Expression, list: Expression) -> Expression { + list_contains_opts( + list, + needle, + ListContainsOptions { + sql_null_semantics: true, + }, + ) +} + +/// Creates `needle IN (elements)` with SQL null semantics for a list of constant elements of +/// `element_dtype`, ignoring nullability. An element can also be an untyped null. +/// +/// The nullability of `element_dtype` is not used: the list elements are nullable only when a +/// null element is present. +/// +/// Non-null integers that hold every value between their smallest and largest become a range +/// check, which is cheaper to evaluate and prunes with tighter bounds than a set lookup, as for +/// the keys of a small dimension table that a join pushes down. +/// +/// # Errors +/// +/// Returns an error if an element is not a null and not of `element_dtype`. +/// +/// ```rust +/// # use vortex_array::dtype::{DType, Nullability, PType}; +/// # use vortex_array::expr::root; +/// # use vortex_functions::expr::in_list_of; +/// let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); +/// let expr = in_list_of(root(), dtype, vec![3i32.into(), 1i32.into(), 2i32.into()])?; +/// assert_eq!(expr.to_string(), "(1i32 <= $ <= 3i32)"); +/// # Ok::<(), vortex_error::VortexError>(()) +/// ``` +pub fn in_list_of( + needle: Expression, + element_dtype: DType, + elements: Vec, +) -> VortexResult { + let element_dtype = if elements.iter().any(Scalar::is_null) { + element_dtype.as_nullable() + } else { + element_dtype.as_nonnullable() + }; + let mut elements = elements + .iter() + .map(|element| { + vortex_ensure!( + element.dtype() == &DType::Null + || element.dtype().eq_ignore_nullability(&element_dtype), + "IN list element {element} is not of type {element_dtype}" + ); + element.cast(&element_dtype) + }) + .collect::>>()?; + + if let Some((low, high)) = contiguous_integer_bounds(&elements) { + let options = BetweenOptions { + lower_strict: StrictComparison::NonStrict, + upper_strict: StrictComparison::NonStrict, + }; + return Ok(between( + needle, + lit(low.clone()), + lit(high.clone()), + options, + )); + } + + // Ordered elements make equal lists from the same set, so that equal filters deduplicate. + if elements + .windows(2) + .all(|pair| pair[0].partial_cmp(&pair[1]).is_some()) + { + elements.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)); + } + let list = Scalar::list(Arc::new(element_dtype), elements, Nullability::NonNullable); + Ok(in_list(needle, lit(list))) +} + +/// The smallest and largest of non-null integer `scalars` when they hold every integer between +/// those two. +fn contiguous_integer_bounds(scalars: &[Scalar]) -> Option<(&Scalar, &Scalar)> { + let mut values = scalars + .iter() + .map(|scalar| { + let primitive = scalar.as_primitive_opt()?; + if !primitive.ptype().is_int() { + return None; + } + let value = if primitive.ptype().is_unsigned_int() { + primitive.as_opt::()??.into() + } else { + i128::from(primitive.as_opt::()??) + }; + Some((value, scalar)) + }) + .collect::>>()?; + values.sort_unstable_by_key(|(value, _)| *value); + values.dedup_by_key(|(value, _)| *value); + let (low, low_scalar) = values.first()?; + let (high, high_scalar) = values.last()?; + (high - low + 1 == i128::try_from(values.len()).ok()?).then_some((*low_scalar, *high_scalar)) +} + +/// Creates a bound expression that checks whether a list contains a value. +pub fn bound_list_contains(list: BoundExpression, value: BoundExpression) -> BoundExpression { + ListContains + .try_new_bound_expr(ListContainsOptions::default(), [list, value]) + .vortex_expect("list-contains expressions require a list child") +} + +/// Creates a bound `needle IN (list)` with SQL null semantics. See [`in_list`]. +pub fn bound_in_list(needle: BoundExpression, list: BoundExpression) -> BoundExpression { + ListContains + .try_new_bound_expr( + ListContainsOptions { + sql_null_semantics: true, + }, + [list, needle], + ) + .vortex_expect("list-contains expressions require a list child") +} + +// ---- VariantGet ---- + +/// Creates an expression that extracts a path from a Variant expression. +/// +/// Missing paths, traversal mismatches, and failed casts return null. When `dtype` is `None`, +/// results are nullable Variant values; otherwise results are nullable values of `dtype`. +pub fn variant_get( + child: Expression, + path: impl Into, + dtype: Option, +) -> Expression { + VariantGet.new_expr(VariantGetOptions::new(path.into(), dtype), vec![child]) +} + +/// Creates a bound expression that extracts a path from a Variant expression. +pub fn bound_variant_get( + child: BoundExpression, + path: impl Into, + dtype: Option, +) -> BoundExpression { + VariantGet + .try_new_bound_expr(VariantGetOptions::new(path.into(), dtype), [child]) + .vortex_expect("variant-get expressions require a Variant child") +} + +// ---- Like ---- + +/// Creates a SQL LIKE expression. +pub fn like(child: Expression, pattern: Expression) -> Expression { + Like.new_expr( + LikeOptions { + negated: false, + case_insensitive: false, + }, + [child, pattern], + ) +} + +/// Creates a bound SQL LIKE expression. +pub fn bound_like(child: BoundExpression, pattern: BoundExpression) -> BoundExpression { + bound_like_with_options(child, pattern, false, false) +} + +/// Creates a case-insensitive SQL ILIKE expression. +pub fn ilike(child: Expression, pattern: Expression) -> Expression { + Like.new_expr( + LikeOptions { + negated: false, + case_insensitive: true, + }, + [child, pattern], + ) +} + +/// Creates a bound case-insensitive SQL ILIKE expression. +pub fn bound_ilike(child: BoundExpression, pattern: BoundExpression) -> BoundExpression { + bound_like_with_options(child, pattern, false, true) +} + +/// Creates a negated SQL NOT LIKE expression. +pub fn not_like(child: Expression, pattern: Expression) -> Expression { + Like.new_expr( + LikeOptions { + negated: true, + case_insensitive: false, + }, + [child, pattern], + ) +} + +/// Creates a bound negated SQL NOT LIKE expression. +pub fn bound_not_like(child: BoundExpression, pattern: BoundExpression) -> BoundExpression { + bound_like_with_options(child, pattern, true, false) +} + +/// Creates a negated case-insensitive SQL NOT ILIKE expression. +pub fn not_ilike(child: Expression, pattern: Expression) -> Expression { + Like.new_expr( + LikeOptions { + negated: true, + case_insensitive: true, + }, + [child, pattern], + ) +} + +/// Creates a bound negated case-insensitive SQL NOT ILIKE expression. +pub fn bound_not_ilike(child: BoundExpression, pattern: BoundExpression) -> BoundExpression { + bound_like_with_options(child, pattern, true, true) +} + +fn bound_like_with_options( + child: BoundExpression, + pattern: BoundExpression, + negated: bool, + case_insensitive: bool, +) -> BoundExpression { + Like.try_new_bound_expr( + LikeOptions { + negated, + case_insensitive, + }, + [child, pattern], + ) + .vortex_expect("like expressions require UTF-8 or binary operands") +} + +// ---- IsNan ---- + +/// Creates an expression that checks for NaN values. +/// +/// The expression is strict: null inputs produce null outputs, so the output nullability +/// follows the input. Only primitive float inputs are supported. +/// +/// ```rust +/// # use vortex_array::expr::root; +/// # use vortex_functions::expr::is_nan; +/// let expr = is_nan(root()); +/// ``` +pub fn is_nan(child: Expression) -> Expression { + IsNan.new_expr(EmptyOptions, vec![child]) +} + +/// Creates a bound expression that checks for NaN values. +pub fn bound_is_nan(child: BoundExpression) -> BoundExpression { + IsNan + .try_new_bound_expr(EmptyOptions, [child]) + .vortex_expect("is-nan expressions are always well-typed") +} + +// ---- Zip ---- + +/// Creates a zip expression that conditionally selects between two arrays. +/// +/// ```rust +/// # use vortex_array::expr::{lit, root}; +/// # use vortex_functions::expr::zip_expr; +/// let expr = zip_expr(lit(true), root(), lit(0i32)); +/// ``` +pub fn zip_expr(mask: Expression, if_true: Expression, if_false: Expression) -> Expression { + Zip.new_expr(EmptyOptions, [if_true, if_false, mask]) +} + +/// Creates a bound zip expression that conditionally selects between two arrays. +pub fn bound_zip_expr( + mask: BoundExpression, + if_true: BoundExpression, + if_false: BoundExpression, +) -> BoundExpression { + Zip.try_new_bound_expr(EmptyOptions, [if_true, if_false, mask]) + .vortex_expect("zip expressions require a boolean mask and compatible value dtypes") +} + +pub mod bound { + pub use super::bound_between as between; + pub use super::bound_byte_length as byte_length; + pub use super::bound_case_when as case_when; + pub use super::bound_case_when_no_else as case_when_no_else; + pub use super::bound_dynamic as dynamic; + pub use super::bound_dynamic_with_options as dynamic_with_options; + pub use super::bound_ext_storage as ext_storage; + pub use super::bound_ilike as ilike; + pub use super::bound_in_list as in_list; + pub use super::bound_is_nan as is_nan; + pub use super::bound_like as like; + pub use super::bound_list_contains as list_contains; + pub use super::bound_list_length as list_length; + pub use super::bound_list_sum as list_sum; + pub use super::bound_list_sum_opts as list_sum_opts; + pub use super::bound_merge as merge; + pub use super::bound_merge_opts as merge_opts; + pub use super::bound_nested_case_when as nested_case_when; + pub use super::bound_not_ilike as not_ilike; + pub use super::bound_not_like as not_like; + pub use super::bound_replace_time_zone as replace_time_zone; + pub use super::bound_select as select; + pub use super::bound_select_exclude as select_exclude; + pub use super::bound_variant_get as variant_get; + pub use super::bound_zip_expr as zip_expr; +} diff --git a/vortex-array/src/scalar_fn/fns/ext_storage.rs b/vortex-functions/src/ext_storage.rs similarity index 79% rename from vortex-array/src/scalar_fn/fns/ext_storage.rs rename to vortex-functions/src/ext_storage.rs index 1e708cf2414..6edfede078c 100644 --- a/vortex-array/src/scalar_fn/fns/ext_storage.rs +++ b/vortex-functions/src/ext_storage.rs @@ -1,25 +1,24 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::dtype::DType; +use vortex_array::scalar_fn::Arity; +use vortex_array::scalar_fn::ChildName; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::ScalarFnVTable; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::arrays::ConstantArray; -use crate::arrays::ExtensionArray; -use crate::arrays::extension::ExtensionArrayExt; -use crate::dtype::DType; -use crate::scalar_fn::Arity; -use crate::scalar_fn::ChildName; -use crate::scalar_fn::EmptyOptions; -use crate::scalar_fn::ExecutionArgs; -use crate::scalar_fn::ScalarFnId; -use crate::scalar_fn::ScalarFnVTable; - /// Extract the storage values from an extension array. #[derive(Clone)] pub struct ExtStorage; @@ -97,27 +96,33 @@ impl ScalarFnVTable for ExtStorage { mod tests { use std::sync::LazyLock; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::ConstantArray; + use vortex_array::arrays::ExtensionArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::assert_arrays_eq; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::dtype::extension::ExtDTypeRef; + use vortex_array::expr::root; + use vortex_array::extension::datetime::TimeUnit; + use vortex_array::extension::datetime::Timestamp; + use vortex_array::scalar::Scalar; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::arrays::ConstantArray; - use crate::arrays::ExtensionArray; - use crate::arrays::PrimitiveArray; - use crate::assert_arrays_eq; - use crate::dtype::DType; - use crate::dtype::Nullability; - use crate::dtype::PType; - use crate::dtype::extension::ExtDTypeRef; use crate::expr::ext_storage; - use crate::expr::root; - use crate::extension::datetime::TimeUnit; - use crate::extension::datetime::Timestamp; - use crate::scalar::Scalar; - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + + crate::initialize(&session); + + session + }); fn ext_dtype(nullability: Nullability) -> ExtDTypeRef { Timestamp::new(TimeUnit::Nanoseconds, nullability).erased() diff --git a/vortex-array/src/expr/field.rs b/vortex-functions/src/field.rs similarity index 91% rename from vortex-array/src/expr/field.rs rename to vortex-functions/src/field.rs index 7a96f949ad5..e0125adfc8f 100644 --- a/vortex-array/src/expr/field.rs +++ b/vortex-functions/src/field.rs @@ -4,8 +4,7 @@ use std::fmt::Display; use itertools::Itertools; - -use crate::dtype::FieldNames; +use vortex_array::dtype::FieldNames; pub struct DisplayFieldNames<'a>(pub &'a FieldNames); diff --git a/vortex-array/src/scalar_fn/fns/is_nan.rs b/vortex-functions/src/is_nan.rs old mode 100755 new mode 100644 similarity index 89% rename from vortex-array/src/scalar_fn/fns/is_nan.rs rename to vortex-functions/src/is_nan.rs index c8059a4e6ab..e799f198b48 --- a/vortex-array/src/scalar_fn/fns/is_nan.rs +++ b/vortex-functions/src/is_nan.rs @@ -1,17 +1,16 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::dtype::DType; +use vortex_array::match_each_float_ptype; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::unstable::row::RowFn; +use vortex_array::scalar_fn::unstable::row::RowVisitor; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_session::registry::CachedId; -use crate::dtype::DType; -use crate::match_each_float_ptype; -use crate::scalar_fn::EmptyOptions; -use crate::scalar_fn::ScalarFnId; -use crate::scalar_fn::unstable::row::RowFn; -use crate::scalar_fn::unstable::row::RowVisitor; - /// Expression that checks for NaN values. /// /// The function is strict: null inputs produce null outputs. Only primitive float inputs @@ -50,23 +49,27 @@ impl RowFn for IsNan { mod tests { use std::sync::LazyLock; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::expr::get_item; + use vortex_array::expr::lit; + use vortex_array::expr::root; + use vortex_array::scalar::Scalar; + use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_session::VortexSession; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::arrays::PrimitiveArray; - use crate::dtype::DType; - use crate::dtype::Nullability; - use crate::dtype::PType; - use crate::expr::get_item; use crate::expr::is_nan; - use crate::expr::lit; - use crate::expr::root; - use crate::scalar::Scalar; - use crate::validity::Validity; - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + crate::initialize(&session); + session + }); #[test] fn dtype() { diff --git a/vortex-functions/src/lib.rs b/vortex-functions/src/lib.rs new file mode 100644 index 00000000000..0f3657d4f50 --- /dev/null +++ b/vortex-functions/src/lib.rs @@ -0,0 +1,80 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Scalar functions built on the vortex-array scalar function framework that the core array +//! crate does not itself depend on. +//! +//! The scalar functions that canonical arrays, validity and the optimizer rely on (cast, +//! fill_null, mask, binary operators, get_item, pack, literal, zip, ...) live in +//! `vortex_array::scalar_fn::builtin`. This crate holds struct reshaping ([`merge`], [`select`]), +//! conditionals ([`case_when`]), list reductions ([`list_length`], [`list_sum`]), and +//! [`byte_length`], [`ext_storage`] and [`replace_time_zone`], and the scan-time functions +//! [`dynamic`] (comparisons against a value that changes during a scan) and [`stat`] (reads of +//! a statistic), along with their expression constructors in [`expr`]. +//! +//! Call [`initialize`] to register them with a session. + +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::builtin::get_item::GetItem; +use vortex_array::scalar_fn::builtin::literal::Literal; +use vortex_array::scalar_fn::session::ScalarFnSessionExt; +use vortex_session::VortexSession; + +pub mod between; +pub mod builtins; +pub mod byte_length; +pub mod case_when; +pub mod dynamic; +pub mod expr; +pub mod ext_storage; +mod field; +pub mod is_nan; +pub mod like; +pub mod list_contains; +pub mod list_length; +pub mod list_sum; +pub mod merge; +pub mod replace_time_zone; +pub mod select; +pub mod stat; +pub mod variant_get; +pub mod zip; + +#[cfg(test)] +mod tests; + +/// Register this crate's scalar functions with the given session. +pub fn initialize(session: &VortexSession) { + between::initialize(session); + like::initialize(session); + list_contains::initialize(session); + variant_get::initialize(session); + zip::initialize(session); + let scalar_fns = session.scalar_fns(); + scalar_fns.register(byte_length::ByteLength); + scalar_fns.register(case_when::CaseWhen); + scalar_fns.register(ext_storage::ExtStorage); + scalar_fns.register(is_nan::IsNan); + scalar_fns.register(list_length::ListLength); + scalar_fns.register(list_sum::ListSum); + scalar_fns.register(merge::Merge); + scalar_fns.register(replace_time_zone::ReplaceTimeZone); + scalar_fns.register(select::Select); + scalar_fns.register(stat::StatFn); +} + +/// A scalar function has a negative cost if applying it to an array and +/// canonicalizing is cheaper than canonicalizing an array and applying it. +/// +/// Example of negative cost expressions are byte_length(), ext_storage(), and get_item() since +/// they don't depend on input size. +/// +/// Example of non-negative cost expression is like() as it's linear over +/// individual input. +pub fn is_negative_cost(id: ScalarFnId) -> bool { + id == ScalarFnVTable::id(&byte_length::ByteLength) + || id == ScalarFnVTable::id(&ext_storage::ExtStorage) + || id == ScalarFnVTable::id(&GetItem) + || id == ScalarFnVTable::id(&Literal) +} diff --git a/vortex-array/src/arrays/dict/compute/like.rs b/vortex-functions/src/like/dict.rs similarity index 66% rename from vortex-array/src/arrays/dict/compute/like.rs rename to vortex-functions/src/like/dict.rs index ba4bf2f9f8c..606688b61a7 100644 --- a/vortex-array/src/arrays/dict/compute/like.rs +++ b/vortex-functions/src/like/dict.rs @@ -1,20 +1,19 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::IntoArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::DictArray; +use vortex_array::arrays::dict::DictArrayExt; +use vortex_array::arrays::dict::DictArraySlotsExt; use vortex_error::VortexResult; use super::Dict; -use super::DictArray; -use crate::ArrayRef; -use crate::IntoArray; -use crate::array::ArrayView; -use crate::arrays::ConstantArray; -use crate::arrays::dict::DictArrayExt; -use crate::arrays::dict::DictArraySlotsExt; -use crate::optimizer::ArrayOptimizer; -use crate::scalar_fn::fns::like::Like; -use crate::scalar_fn::fns::like::LikeOptions; -use crate::scalar_fn::fns::like::LikeReduce; +use crate::like::Like; +use crate::like::LikeOptions; +use crate::like::LikeReduce; impl LikeReduce for Dict { fn like( @@ -29,9 +28,7 @@ impl LikeReduce for Dict { if let Some(pattern) = pattern.as_constant() { let pattern = ConstantArray::new(pattern, array.values().len()).into_array(); - let values = Like::try_new(array.values().clone(), pattern, options)? - .into_array() - .optimize()?; + let values = Like::try_new(array.values().clone(), pattern, options)?.into_array(); // SAFETY: LIKE preserves the len of the values, so codes are still pointing at // valid positions. @@ -53,22 +50,27 @@ impl LikeReduce for Dict { mod tests { use std::sync::LazyLock; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::BoolArray; + use vortex_array::arrays::ConstantArray; + use vortex_array::arrays::DictArray; + use vortex_array::arrays::VarBinArray; + use vortex_array::assert_arrays_eq; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::arrays::BoolArray; - use crate::arrays::DictArray; - use crate::arrays::VarBinArray; - use crate::arrays::dict::compute::like::ConstantArray; - use crate::assert_arrays_eq; - use crate::optimizer::ArrayOptimizer; - use crate::scalar_fn::fns::like::Like; - use crate::scalar_fn::fns::like::LikeOptions; + use crate::like::Like; + use crate::like::LikeOptions; - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + + crate::initialize(&session); + + session + }); #[test] fn like_reduce_dict() -> VortexResult<()> { @@ -80,9 +82,7 @@ mod tests { .into_array(); let pattern = ConstantArray::new("hello%", 4).into_array(); - let result = Like::try_new(dict, pattern, LikeOptions::default())? - .into_array() - .optimize()?; + let result = Like::try_new(dict, pattern, LikeOptions::default())?.into_array(); assert_arrays_eq!( result, diff --git a/vortex-array/src/scalar_fn/fns/like/kernel.rs b/vortex-functions/src/like/kernel.rs similarity index 86% rename from vortex-array/src/scalar_fn/fns/like/kernel.rs rename to vortex-functions/src/like/kernel.rs index b3b683212ff..d586590da0f 100644 --- a/vortex-array/src/scalar_fn/fns/like/kernel.rs +++ b/vortex-functions/src/like/kernel.rs @@ -1,21 +1,21 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::VTable; +use vortex_array::arrays::ScalarFn; +use vortex_array::arrays::scalar_fn::ExactScalarFn; +use vortex_array::arrays::scalar_fn::ScalarFnArrayExt; +use vortex_array::arrays::scalar_fn::ScalarFnArrayView; +use vortex_array::kernel::ExecuteParentKernel; +use vortex_array::optimizer::rules::ArrayParentReduceRule; use vortex_error::VortexExpect; use vortex_error::VortexResult; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::array::ArrayView; -use crate::array::VTable; -use crate::arrays::ScalarFn; -use crate::arrays::scalar_fn::ExactScalarFn; -use crate::arrays::scalar_fn::ScalarFnArrayExt; -use crate::arrays::scalar_fn::ScalarFnArrayView; -use crate::kernel::ExecuteParentKernel; -use crate::optimizer::rules::ArrayParentReduceRule; -use crate::scalar_fn::fns::like::Like as LikeExpr; -use crate::scalar_fn::fns::like::LikeOptions; +use crate::like::Like as LikeExpr; +use crate::like::LikeOptions; /// Like pattern matching on an array without reading buffers. /// diff --git a/vortex-array/src/scalar_fn/fns/like/mod.rs b/vortex-functions/src/like/mod.rs similarity index 91% rename from vortex-array/src/scalar_fn/fns/like/mod.rs rename to vortex-functions/src/like/mod.rs index da0d094c201..2ffbb23f5ec 100644 --- a/vortex-array/src/scalar_fn/fns/like/mod.rs +++ b/vortex-functions/src/like/mod.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +mod dict; mod kernel; mod pattern; @@ -10,6 +11,30 @@ use std::fmt::Formatter; pub use kernel::*; use pattern::LikePattern; use prost::Message; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::Dict; +use vortex_array::arrays::ScalarFnArray; +use vortex_array::arrays::VarBinViewArray; +use vortex_array::arrays::varbinview::BinaryView; +use vortex_array::arrays::varbinview::ResolvedViews; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::expr::display::ExprDisplay; +use vortex_array::optimizer::kernels::ArrayKernelsExt; +use vortex_array::proto::expr as pb; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::Arity; +use vortex_array::scalar_fn::ChildName; +use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::ScalarFnVTableExt; +use vortex_array::scalar_fn::session::ScalarFnSessionExt; use vortex_buffer::BitBuffer; use vortex_buffer::BitBufferMut; use vortex_buffer::BufferAllocatorRef; @@ -19,27 +44,13 @@ use vortex_error::vortex_err; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::ArrayRef; -use crate::Canonical; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::arrays::BoolArray; -use crate::arrays::ConstantArray; -use crate::arrays::ScalarFnArray; -use crate::arrays::VarBinViewArray; -use crate::arrays::varbinview::BinaryView; -use crate::arrays::varbinview::ResolvedViews; -use crate::dtype::DType; -use crate::dtype::Nullability; -use crate::expr::display::ExprDisplay; -use crate::proto::expr as pb; -use crate::scalar::Scalar; -use crate::scalar_fn::Arity; -use crate::scalar_fn::ChildName; -use crate::scalar_fn::ExecutionArgs; -use crate::scalar_fn::ScalarFnId; -use crate::scalar_fn::ScalarFnVTable; -use crate::scalar_fn::ScalarFnVTableExt; +/// Registers [`Like`] and its rule for dictionary-encoded arrays. +pub(crate) fn initialize(session: &VortexSession) { + session.scalar_fns().register(Like); + session + .kernels() + .register_reduce_parent_rule(Like.id(), Dict, LikeReduceAdaptor(Dict)); +} /// Options for SQL LIKE function #[derive(Default, Debug, Clone, Copy, PartialEq, Eq, Hash)] @@ -423,37 +434,43 @@ mod tests { use std::sync::LazyLock; use rstest::rstest; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::BoolArray; + use vortex_array::arrays::ConstantArray; + use vortex_array::arrays::VarBinArray; + use vortex_array::arrays::VarBinViewArray; + use vortex_array::assert_arrays_eq; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::expr::get_item; + use vortex_array::expr::lit; + use vortex_array::expr::root; + use vortex_array::scalar::Scalar; use vortex_session::VortexSession; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::arrays::BoolArray; - use crate::arrays::ConstantArray; - use crate::arrays::VarBinArray; - use crate::arrays::VarBinViewArray; - use crate::assert_arrays_eq; - use crate::dtype::DType; - use crate::dtype::Nullability; - use crate::expr::get_item; use crate::expr::like; - use crate::expr::lit; use crate::expr::not_ilike; - use crate::expr::root; - use crate::scalar::Scalar; - use crate::scalar_fn::fns::like::Like; - use crate::scalar_fn::fns::like::LikeOptions; - use crate::scalar_fn::fns::like::is_ascii; - use crate::scalar_fn::fns::like::suffix_bytes_unchecked; + use crate::like::Like; + use crate::like::LikeOptions; + use crate::like::is_ascii; + use crate::like::suffix_bytes_unchecked; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + + crate::initialize(&session); - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + session + }); const LONG: &str = "a value far too long to live inside its own view"; fn run_like( - array: crate::ArrayRef, - pattern: crate::ArrayRef, + array: vortex_array::ArrayRef, + pattern: vortex_array::ArrayRef, options: LikeOptions, - ) -> crate::ArrayRef { + ) -> vortex_array::ArrayRef { Like::try_new(array, pattern, options).unwrap().into_array() } diff --git a/vortex-array/src/scalar_fn/fns/like/pattern.rs b/vortex-functions/src/like/pattern.rs similarity index 100% rename from vortex-array/src/scalar_fn/fns/like/pattern.rs rename to vortex-functions/src/like/pattern.rs diff --git a/vortex-array/src/arrays/chunked/compute/list_contains.rs b/vortex-functions/src/list_contains/chunked.rs similarity index 67% rename from vortex-array/src/arrays/chunked/compute/list_contains.rs rename to vortex-functions/src/list_contains/chunked.rs index 2c59d5c0dae..6bfce251f00 100644 --- a/vortex-array/src/arrays/chunked/compute/list_contains.rs +++ b/vortex-functions/src/list_contains/chunked.rs @@ -1,20 +1,19 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::IntoArray; +use vortex_array::arrays::Chunked; +use vortex_array::arrays::ChunkedArray; +use vortex_array::arrays::chunked::ChunkedArrayExt; +use vortex_array::dtype::DType; use vortex_error::VortexResult; -use crate::ArrayRef; -use crate::IntoArray; -use crate::array::ArrayView; -use crate::arrays::Chunked; -use crate::arrays::ChunkedArray; -use crate::arrays::chunked::ChunkedArrayExt; -use crate::dtype::DType; -use crate::optimizer::ArrayOptimizer; -use crate::scalar_fn::fns::list_contains::ListContains; -use crate::scalar_fn::fns::list_contains::ListContainsElementReduce; -use crate::scalar_fn::fns::list_contains::ListContainsOptions; -use crate::scalar_fn::fns::list_contains::PreparedSet; +use crate::list_contains::ListContains; +use crate::list_contains::ListContainsElementReduce; +use crate::list_contains::ListContainsOptions; +use crate::list_contains::PreparedSet; /// Pushes `list_contains` of a prepared set into each chunk of the needles. /// @@ -36,9 +35,7 @@ impl ListContainsElementReduce for Chunked { .map(|chunk| { let set = list.slice(offset..offset + chunk.len())?; offset += chunk.len(); - ListContains::try_new_opts(set, chunk.clone(), *options)? - .into_array() - .optimize() + Ok(ListContains::try_new_opts(set, chunk.clone(), *options)?.into_array()) }) .collect::>>()?; diff --git a/vortex-array/src/array/haystack/mod.rs b/vortex-functions/src/list_contains/haystack.rs similarity index 94% rename from vortex-array/src/array/haystack/mod.rs rename to vortex-functions/src/list_contains/haystack.rs index f1e04f9c686..89dbca6bbfe 100644 --- a/vortex-array/src/array/haystack/mod.rs +++ b/vortex-functions/src/list_contains/haystack.rs @@ -7,6 +7,26 @@ use std::hash::BuildHasher; use num_traits::ToPrimitive; use num_traits::WrappingSub; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::RecursiveCanonical; +use vortex_array::arrays::DecimalArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::VarBinViewArray; +use vortex_array::arrays::decimal::converted_buffer; +use vortex_array::arrays::primitive::PrimitiveArrayExt; +use vortex_array::arrays::varbinview::BinaryView; +use vortex_array::dtype::DType; +use vortex_array::dtype::DecimalType; +use vortex_array::dtype::NativeDecimalType; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::PType; +use vortex_array::dtype::i256; +use vortex_array::match_each_decimal_value_type; +use vortex_array::match_each_integer_ptype; +use vortex_array::scalar_fn::builtin::binary::build_row_comparator; +use vortex_array::validity::Validity; use vortex_buffer::BitBuffer; use vortex_buffer::BitBufferMut; use vortex_buffer::Buffer; @@ -20,27 +40,6 @@ use vortex_utils::aliases::hash_map::HashTable; use vortex_utils::aliases::hash_map::HashTableEntry; use vortex_utils::aliases::hash_map::RandomState; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::RecursiveCanonical; -use crate::arrays::DecimalArray; -use crate::arrays::PrimitiveArray; -use crate::arrays::VarBinViewArray; -use crate::arrays::decimal::converted_buffer; -use crate::arrays::primitive::PrimitiveArrayExt; -use crate::arrays::varbinview::BinaryView; -use crate::dtype::DType; -use crate::dtype::DecimalType; -use crate::dtype::NativeDecimalType; -use crate::dtype::NativePType; -use crate::dtype::PType; -use crate::dtype::i256; -use crate::match_each_decimal_value_type; -use crate::match_each_integer_ptype; -use crate::scalar_fn::fns::binary::build_row_comparator; -use crate::validity::Validity; - /// A set whose span of values needs at most this many bits per element is looked up through a /// bitmap over the span, bounding the bitmap to a few words per element. const BITMAP_BITS_PER_ELEMENT: u128 = 64; @@ -49,7 +48,7 @@ const BITMAP_MIN_BITS: u128 = 1 << 12; /// Looks up needles among the elements of an array, for many batches of needles. /// -/// The haystack counterpart of [`RepeatedArrayProbe`](crate::RepeatedArrayProbe): where a probe +/// The haystack counterpart of [`RepeatedArrayProbe`](vortex_array::RepeatedArrayProbe): where a probe /// reads the rows of an array by index, a haystack answers whether each row of a batch of /// needles is an element of the array. A haystack pays off only over repeated lookups, so it /// has no one-off form. It is built eagerly. A holder that must build it lazily, or share it diff --git a/vortex-array/src/scalar_fn/fns/list_contains/kernel.rs b/vortex-functions/src/list_contains/kernel.rs similarity index 78% rename from vortex-array/src/scalar_fn/fns/list_contains/kernel.rs rename to vortex-functions/src/list_contains/kernel.rs index 28830b71dd0..8eaf59c604e 100644 --- a/vortex-array/src/scalar_fn/fns/list_contains/kernel.rs +++ b/vortex-functions/src/list_contains/kernel.rs @@ -1,21 +1,21 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::VTable; +use vortex_array::arrays::ScalarFn; +use vortex_array::arrays::scalar_fn::ExactScalarFn; +use vortex_array::arrays::scalar_fn::ScalarFnArrayExt; +use vortex_array::arrays::scalar_fn::ScalarFnArrayView; +use vortex_array::kernel::ExecuteParentKernel; +use vortex_array::optimizer::rules::ArrayParentReduceRule; use vortex_error::VortexExpect; use vortex_error::VortexResult; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::array::ArrayView; -use crate::array::VTable; -use crate::arrays::ScalarFn; -use crate::arrays::scalar_fn::ExactScalarFn; -use crate::arrays::scalar_fn::ScalarFnArrayExt; -use crate::arrays::scalar_fn::ScalarFnArrayView; -use crate::kernel::ExecuteParentKernel; -use crate::optimizer::rules::ArrayParentReduceRule; -use crate::scalar_fn::fns::list_contains::ListContains as ListContainsExpr; -use crate::scalar_fn::fns::list_contains::ListContainsOptions; +use crate::list_contains::ListContains as ListContainsExpr; +use crate::list_contains::ListContainsOptions; /// Check list-contains without reading buffers (metadata-only). /// @@ -43,11 +43,11 @@ pub trait ListContainsElementReduce: VTable { /// [`PreparedSetData::contains`] or [`PreparedSetData::contains_scalar`], or make its result from /// matches it found another way with [`PreparedSetData::result_from_bits`]. /// -/// [`PreparedSet`]: crate::scalar_fn::fns::list_contains::PreparedSet -/// [`PreparedSet::set_of`]: crate::scalar_fn::fns::list_contains::PreparedSet::set_of -/// [`PreparedSetData::contains`]: crate::scalar_fn::fns::list_contains::PreparedSetData::contains -/// [`PreparedSetData::contains_scalar`]: crate::scalar_fn::fns::list_contains::PreparedSetData::contains_scalar -/// [`PreparedSetData::result_from_bits`]: crate::scalar_fn::fns::list_contains::PreparedSetData::result_from_bits +/// [`PreparedSet`]: crate::list_contains::PreparedSet +/// [`PreparedSet::set_of`]: crate::list_contains::PreparedSet::set_of +/// [`PreparedSetData::contains`]: crate::list_contains::PreparedSetData::contains +/// [`PreparedSetData::contains_scalar`]: crate::list_contains::PreparedSetData::contains_scalar +/// [`PreparedSetData::result_from_bits`]: crate::list_contains::PreparedSetData::result_from_bits pub trait ListContainsElementKernel: VTable { fn list_contains( list: &ArrayRef, diff --git a/vortex-array/src/scalar_fn/fns/list_contains/mod.rs b/vortex-functions/src/list_contains/mod.rs similarity index 95% rename from vortex-array/src/scalar_fn/fns/list_contains/mod.rs rename to vortex-functions/src/list_contains/mod.rs index 5d92baac6f6..9035fe4ce83 100644 --- a/vortex-array/src/scalar_fn/fns/list_contains/mod.rs +++ b/vortex-functions/src/list_contains/mod.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +mod chunked; +mod haystack; mod kernel; mod prepared; @@ -16,6 +18,42 @@ use num_traits::Zero; pub use prepared::PreparedSet; pub use prepared::PreparedSetData; use prost::Message; +use vortex_array::ArrayRef; +use vortex_array::Columnar; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::Chunked; +use vortex_array::arrays::Constant; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::ListViewArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::ScalarFnArray; +use vortex_array::arrays::bool::BoolArrayExt; +use vortex_array::arrays::listview::ListViewArraySlotsExt; +use vortex_array::arrays::primitive::PrimitiveArrayExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::IntegerPType; +use vortex_array::dtype::Nullability; +use vortex_array::match_each_integer_ptype; +use vortex_array::match_each_unsigned_integer_ptype; +use vortex_array::optimizer::kernels::ArrayKernelsExt; +use vortex_array::proto::expr as pb; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::Arity; +use vortex_array::scalar_fn::ChildName; +use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::ReduceNode; +use vortex_array::scalar_fn::ReduceNodeValidity; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::ScalarFnVTableExt; +use vortex_array::scalar_fn::builtin::binary::Binary; +use vortex_array::scalar_fn::builtin::literal::Literal; +use vortex_array::scalar_fn::builtin::operators::Operator; +use vortex_array::scalar_fn::is_not_null_node; +use vortex_array::scalar_fn::session::ScalarFnSessionExt; +use vortex_array::validity::Validity; use vortex_buffer::BitBuffer; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; @@ -26,39 +64,18 @@ use vortex_mask::Mask; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::ArrayRef; -use crate::Columnar; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::arrays::BoolArray; -use crate::arrays::Constant; -use crate::arrays::ConstantArray; -use crate::arrays::ListViewArray; -use crate::arrays::PrimitiveArray; -use crate::arrays::ScalarFnArray; -use crate::arrays::bool::BoolArrayExt; -use crate::arrays::listview::ListViewArraySlotsExt; -use crate::arrays::primitive::PrimitiveArrayExt; -use crate::dtype::DType; -use crate::dtype::IntegerPType; -use crate::dtype::Nullability; -use crate::match_each_integer_ptype; -use crate::match_each_unsigned_integer_ptype; -use crate::proto::expr as pb; -use crate::scalar::Scalar; -use crate::scalar_fn::Arity; -use crate::scalar_fn::ChildName; -use crate::scalar_fn::ExecutionArgs; -use crate::scalar_fn::ReduceNode; -use crate::scalar_fn::ReduceNodeValidity; -use crate::scalar_fn::ScalarFnId; -use crate::scalar_fn::ScalarFnVTable; -use crate::scalar_fn::ScalarFnVTableExt; -use crate::scalar_fn::fns::binary::Binary; -use crate::scalar_fn::fns::literal::Literal; -use crate::scalar_fn::fns::operators::Operator; -use crate::scalar_fn::is_not_null_node; -use crate::validity::Validity; +/// Registers [`ListContains`], [`PreparedSet`] and the chunked-element rule for +/// [`ListContains`]. +pub(crate) fn initialize(session: &VortexSession) { + let scalar_fns = session.scalar_fns(); + scalar_fns.register(ListContains); + scalar_fns.register(PreparedSet); + session.kernels().register_reduce_parent_rule( + ListContains.id(), + Chunked, + ListContainsElementReduceAdaptor(Chunked), + ); +} #[derive(Clone)] pub struct ListContains; @@ -818,6 +835,36 @@ mod tests { use itertools::Itertools; use num_traits::PrimInt; use rstest::rstest; + use vortex_array::ArrayRef; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::BoolArray; + use vortex_array::arrays::ChunkedArray; + use vortex_array::arrays::ConstantArray; + use vortex_array::arrays::DictArray; + use vortex_array::arrays::ListArray; + use vortex_array::arrays::ListViewArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::VarBinArray; + use vortex_array::arrays::VarBinViewArray; + use vortex_array::assert_arrays_eq; + use vortex_array::dtype::DType; + use vortex_array::dtype::NativePType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::dtype::PType::I32; + use vortex_array::dtype::StructFields; + use vortex_array::expr::Expression; + use vortex_array::expr::bound; + use vortex_array::expr::col; + use vortex_array::expr::get_item; + use vortex_array::expr::lit; + use vortex_array::expr::root; + use vortex_array::optimizer::ArrayOptimizer; + use vortex_array::scalar::PValue; + use vortex_array::scalar::Scalar; + use vortex_array::scalar_fn::builtin::literal::Literal; + use vortex_array::validity::Validity; use vortex_buffer::BitBuffer; use vortex_buffer::Buffer; use vortex_buffer::buffer; @@ -826,45 +873,21 @@ mod tests { use vortex_session::VortexSession; use super::process_matches; - use crate::ArrayRef; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::arrays::BoolArray; - use crate::arrays::ChunkedArray; - use crate::arrays::ConstantArray; - use crate::arrays::DictArray; - use crate::arrays::ListArray; - use crate::arrays::ListViewArray; - use crate::arrays::PrimitiveArray; - use crate::arrays::VarBinArray; - use crate::arrays::VarBinViewArray; - use crate::assert_arrays_eq; - use crate::dtype::DType; - use crate::dtype::NativePType; - use crate::dtype::Nullability; - use crate::dtype::PType; - use crate::dtype::PType::I32; - use crate::dtype::StructFields; - use crate::expr::Expression; - use crate::expr::bound; - use crate::expr::col; - use crate::expr::get_item; use crate::expr::in_list; use crate::expr::list_contains; use crate::expr::list_contains_opts; - use crate::expr::lit; - use crate::expr::root; - use crate::optimizer::ArrayOptimizer; - use crate::scalar::PValue; - use crate::scalar::Scalar; - use crate::scalar_fn::fns::list_contains::ListContains; - use crate::scalar_fn::fns::list_contains::ListContainsOptions; - use crate::scalar_fn::fns::list_contains::PreparedSet; - use crate::scalar_fn::fns::list_contains::PreparedSetData; - use crate::scalar_fn::fns::literal::Literal; - use crate::validity::Validity; - - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + use crate::list_contains::ListContains; + use crate::list_contains::ListContainsOptions; + use crate::list_contains::PreparedSet; + use crate::list_contains::PreparedSetData; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + + crate::initialize(&session); + + session + }); fn test_array() -> ArrayRef { let mut ctx = SESSION.create_execution_ctx(); @@ -1541,7 +1564,7 @@ mod tests { [None, Some(true), None, Some(true)], )?; assert_result( - needles.apply(&crate::expr::not(in_list(root(), lit(set)))), + needles.apply(&vortex_array::expr::not(in_list(root(), lit(set)))), [None, Some(false), None, Some(false)], ) } @@ -1972,7 +1995,7 @@ mod tests { #[test] fn options_round_trip_through_serde() -> VortexResult<()> { - use crate::scalar_fn::ScalarFnVTable; + use vortex_array::scalar_fn::ScalarFnVTable; for options in [ListContainsOptions::default(), SQL] { let bytes = ListContains .serialize(&options)? @@ -2013,7 +2036,7 @@ mod tests { options, )? .into_array() - .optimize()?; + .optimize(&SESSION)?; assert_eq!(array.as_constant(), Some(expected)); Ok(()) } diff --git a/vortex-array/src/scalar_fn/fns/list_contains/prepared/data.rs b/vortex-functions/src/list_contains/prepared/data.rs similarity index 93% rename from vortex-array/src/scalar_fn/fns/list_contains/prepared/data.rs rename to vortex-functions/src/list_contains/prepared/data.rs index aaaf047c1d9..2c9913d83fb 100644 --- a/vortex-array/src/scalar_fn/fns/list_contains/prepared/data.rs +++ b/vortex-functions/src/list_contains/prepared/data.rs @@ -9,6 +9,19 @@ use std::hash::Hasher; use std::sync::Arc; use std::sync::OnceLock; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::ListViewArray; +use vortex_array::arrays::listview::ListViewArraySlotsExt; +use vortex_array::dtype::DType; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::ScalarFnRef; +use vortex_array::scalar_fn::ScalarFnVTableExt; +use vortex_array::scalar_fn::builtin::literal::Literal; +use vortex_array::validity::Validity; use vortex_buffer::BitBuffer; use vortex_error::SharedVortexResult; use vortex_error::VortexResult; @@ -16,21 +29,8 @@ use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_panic; -use crate::ArrayHaystack; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::arrays::BoolArray; -use crate::arrays::ConstantArray; -use crate::arrays::ListViewArray; -use crate::arrays::listview::ListViewArraySlotsExt; -use crate::dtype::DType; -use crate::scalar::Scalar; -use crate::scalar_fn::ScalarFnRef; -use crate::scalar_fn::ScalarFnVTableExt; -use crate::scalar_fn::fns::list_contains::ListContainsOptions; -use crate::scalar_fn::fns::literal::Literal; -use crate::validity::Validity; +use crate::list_contains::ListContainsOptions; +use crate::list_contains::haystack::ArrayHaystack; /// The options of a [`PreparedSet`]: the list that every row holds, and its set. /// @@ -38,7 +38,7 @@ use crate::validity::Validity; /// slice or filter of the array it gives, looks up needles in the one set. The haystack needs /// an [`ExecutionCtx`] to materialize the elements, so it is built on the first lookup. /// -/// [`PreparedSet`]: crate::scalar_fn::fns::list_contains::PreparedSet +/// [`PreparedSet`]: crate::list_contains::PreparedSet #[derive(Clone)] pub struct PreparedSetData(Arc); diff --git a/vortex-array/src/scalar_fn/fns/list_contains/prepared/function.rs b/vortex-functions/src/list_contains/prepared/function.rs similarity index 80% rename from vortex-array/src/scalar_fn/fns/list_contains/prepared/function.rs rename to vortex-functions/src/list_contains/prepared/function.rs index eb18ed35372..42efbf01d78 100644 --- a/vortex-array/src/scalar_fn/fns/list_contains/prepared/function.rs +++ b/vortex-functions/src/list_contains/prepared/function.rs @@ -4,30 +4,30 @@ use std::fmt::Formatter; use prost::Message; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::ScalarFn; +use vortex_array::arrays::ScalarFnArray; +use vortex_array::dtype::DType; +use vortex_array::expr::display::ExprDisplay; +use vortex_array::proto::scalar as pb; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::Arity; +use vortex_array::scalar_fn::ChildName; +use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::ReduceNode; +use vortex_array::scalar_fn::ReduceNodeValidity; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::ScalarFnVTableExt; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_session::VortexSession; use vortex_session::registry::CachedId; use super::PreparedSetData; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::arrays::ConstantArray; -use crate::arrays::ScalarFn; -use crate::arrays::ScalarFnArray; -use crate::dtype::DType; -use crate::expr::display::ExprDisplay; -use crate::proto::scalar as pb; -use crate::scalar::Scalar; -use crate::scalar_fn::Arity; -use crate::scalar_fn::ChildName; -use crate::scalar_fn::ExecutionArgs; -use crate::scalar_fn::ReduceNode; -use crate::scalar_fn::ReduceNodeValidity; -use crate::scalar_fn::ScalarFnId; -use crate::scalar_fn::ScalarFnVTable; -use crate::scalar_fn::ScalarFnVTableExt; /// A constant, non-null list whose elements are prepared as a set for membership lookups. /// @@ -45,9 +45,9 @@ use crate::scalar_fn::ScalarFnVTableExt; /// /// It serializes as the list, and deserializes into a set prepared again. /// -/// [`ListContains`]: crate::scalar_fn::fns::list_contains::ListContains -/// [`ListContainsElementKernel`]: crate::scalar_fn::fns::list_contains::ListContainsElementKernel -/// [`ListContainsElementReduce`]: crate::scalar_fn::fns::list_contains::ListContainsElementReduce +/// [`ListContains`]: crate::list_contains::ListContains +/// [`ListContainsElementKernel`]: crate::list_contains::ListContainsElementKernel +/// [`ListContainsElementReduce`]: crate::list_contains::ListContainsElementReduce #[derive(Clone, Debug)] pub struct PreparedSet; @@ -57,8 +57,8 @@ impl PreparedSet { /// [`ListContains`] prepares a constant list itself. A kernel crate can use this to test its /// [`ListContainsElementKernel`] against a prepared set. /// - /// [`ListContains`]: crate::scalar_fn::fns::list_contains::ListContains - /// [`ListContainsElementKernel`]: crate::scalar_fn::fns::list_contains::ListContainsElementKernel + /// [`ListContains`]: crate::list_contains::ListContains + /// [`ListContainsElementKernel`]: crate::list_contains::ListContainsElementKernel /// /// # Errors /// diff --git a/vortex-array/src/scalar_fn/fns/list_contains/prepared/mod.rs b/vortex-functions/src/list_contains/prepared/mod.rs similarity index 74% rename from vortex-array/src/scalar_fn/fns/list_contains/prepared/mod.rs rename to vortex-functions/src/list_contains/prepared/mod.rs index b077ab2e4db..5d02735f49d 100644 --- a/vortex-array/src/scalar_fn/fns/list_contains/prepared/mod.rs +++ b/vortex-functions/src/list_contains/prepared/mod.rs @@ -2,7 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors //! Membership in a constant set. [`PreparedSet`] holds a constant list together with the -//! [`ArrayHaystack`](crate::ArrayHaystack) of its elements, so that the kernels of a needle +//! [`ArrayHaystack`](crate::list_contains::haystack::ArrayHaystack) of its elements, so that the kernels of a needle //! encoding can look up their own values. mod data; diff --git a/vortex-array/src/scalar_fn/fns/list_contains/prepared/tests.rs b/vortex-functions/src/list_contains/prepared/tests.rs similarity index 83% rename from vortex-array/src/scalar_fn/fns/list_contains/prepared/tests.rs rename to vortex-functions/src/list_contains/prepared/tests.rs index fe3c6b46da6..2443212f489 100644 --- a/vortex-array/src/scalar_fn/fns/list_contains/prepared/tests.rs +++ b/vortex-functions/src/list_contains/prepared/tests.rs @@ -4,6 +4,42 @@ use std::sync::LazyLock; use rstest::rstest; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::Bool; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::Chunked; +use vortex_array::arrays::ChunkedArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::DecimalArray; +use vortex_array::arrays::FilterArray; +use vortex_array::arrays::FixedSizeListArray; +use vortex_array::arrays::ListArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::ScalarFn; +use vortex_array::arrays::ScalarFnArray; +use vortex_array::arrays::StructArray; +use vortex_array::arrays::VarBinViewArray; +use vortex_array::arrays::chunked::ChunkedArrayExt; +use vortex_array::arrays::scalar_fn::ScalarFnArrayExt; +use vortex_array::assert_arrays_eq; +use vortex_array::builders::builder_with_capacity_in; +use vortex_array::dtype::DType; +use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::DecimalType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::dtype::i256; +use vortex_array::expr::lit; +use vortex_array::expr::root; +use vortex_array::match_each_decimal_value_type; +use vortex_array::optimizer::ArrayOptimizer; +use vortex_array::scalar::DecimalValue; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::TypedScalarFnInstance; +use vortex_array::validity::Validity; use vortex_buffer::BitBuffer; use vortex_buffer::buffer; use vortex_error::VortexResult; @@ -11,45 +47,18 @@ use vortex_mask::Mask; use vortex_session::VortexSession; use super::PreparedSet; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::VortexSessionExecute; -use crate::arrays::Bool; -use crate::arrays::BoolArray; -use crate::arrays::Chunked; -use crate::arrays::ChunkedArray; -use crate::arrays::ConstantArray; -use crate::arrays::DecimalArray; -use crate::arrays::FixedSizeListArray; -use crate::arrays::ListArray; -use crate::arrays::PrimitiveArray; -use crate::arrays::ScalarFn; -use crate::arrays::StructArray; -use crate::arrays::VarBinViewArray; -use crate::arrays::chunked::ChunkedArrayExt; -use crate::arrays::scalar_fn::ScalarFnArrayExt; -use crate::assert_arrays_eq; -use crate::builders::builder_with_capacity_in; -use crate::dtype::DType; -use crate::dtype::DecimalDType; -use crate::dtype::DecimalType; -use crate::dtype::Nullability; -use crate::dtype::PType; -use crate::dtype::i256; use crate::expr::list_contains; -use crate::expr::lit; -use crate::expr::root; -use crate::match_each_decimal_value_type; -use crate::optimizer::ArrayOptimizer; -use crate::scalar::DecimalValue; -use crate::scalar::Scalar; -use crate::scalar_fn::fns::list_contains::ListContains; -use crate::scalar_fn::fns::list_contains::ListContainsOptions; -use crate::scalar_fn::fns::list_contains::PreparedSetData; -use crate::validity::Validity; - -static SESSION: LazyLock = LazyLock::new(VortexSession::empty); +use crate::list_contains::ListContains; +use crate::list_contains::ListContainsOptions; +use crate::list_contains::PreparedSetData; + +static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + + crate::initialize(&session); + + session +}); /// Prepares the non-null list scalar `list`. fn prepare(list: &Scalar) -> VortexResult { @@ -321,12 +330,12 @@ fn test_prepared_set_rows_are_the_constant_list() -> VortexResult<()> { // A slice and a filter keep the haystack instead of building it again. let prepared = PreparedSet::set_of(&set).expect("prepared set"); - let sliced = set.slice(1..3)?.optimize()?; + let sliced = set.slice(1..3)?.optimize(&SESSION)?; assert_eq!(sliced.len(), 2); assert!(PreparedSet::set_of(&sliced).is_some_and(|sliced| sliced.shares_set_with(prepared))); let filtered = set .filter(Mask::from_iter([true, false, true, true]))? - .optimize()?; + .optimize(&SESSION)?; assert_eq!(filtered.len(), 3); assert!( PreparedSet::set_of(&filtered).is_some_and(|filtered| filtered.shares_set_with(prepared)) @@ -535,7 +544,7 @@ fn test_constant_list_over_chunked_needles( // The list is prepared before the push-down, so every chunk gets a slice of the one set. let array = ListContains::try_new_opts(list, needles, options)? .into_array() - .optimize()?; + .optimize(&SESSION)?; assert_chunks_share_one_set(&array); assert_arrays_eq!(array, BoolArray::from_iter(expected), &mut ctx); @@ -550,7 +559,7 @@ fn test_optimized_expression_over_chunked_needles() -> VortexResult<()> { .bind(needles.dtype())? .optimize()?; - let array = needles.apply_bound(&expr)?; + let array = needles.apply_bound(&expr)?.optimize(&SESSION)?; assert_chunks_share_one_set(&array); assert_arrays_eq!( @@ -560,3 +569,56 @@ fn test_optimized_expression_over_chunked_needles() -> VortexResult<()> { ); Ok(()) } + +/// A prepared set has no children, so its `ScalarFnArray` takes an explicit length. +#[test] +fn test_scalar_fn_without_children_requires_explicit_len() -> VortexResult<()> { + let list = Scalar::list( + DType::Primitive(PType::I32, Nullability::NonNullable), + vec![Scalar::from(1i32)], + Nullability::NonNullable, + ); + let scalar_fn = + TypedScalarFnInstance::new(PreparedSet, PreparedSetData::try_new(list.clone())?).erased(); + + let Err(err) = ScalarFnArray::try_new(scalar_fn.clone(), vec![]) else { + panic!("ScalarFnArray::try_new should reject zero children"); + }; + assert!( + err.to_string() + .contains("ScalarFnArray length cannot be inferred without children") + ); + + let scalar_fn_array = ScalarFnArray::try_new_with_len(scalar_fn, vec![], 3)?; + assert_eq!(scalar_fn_array.len(), 3); + assert_eq!(scalar_fn_array.child_count(), 0); + assert_arrays_eq!( + scalar_fn_array, + ConstantArray::new(list, 3), + &mut SESSION.create_execution_ctx() + ); + Ok(()) +} + +/// Filtering a prepared set, which has no children, keeps a prepared set of the filtered length. +#[test] +fn test_filter_pushdown_without_children() -> VortexResult<()> { + let list = Scalar::list( + DType::Primitive(PType::I32, Nullability::NonNullable), + vec![Scalar::from(7i32)], + Nullability::NonNullable, + ); + let array = PreparedSet::try_new_array(list.clone(), 3)?; + let result = FilterArray::try_new(array, Mask::from_iter([true, false, true]))? + .into_array() + .optimize(&SESSION)?; + + assert_eq!(result.as_::().nchildren(), 0); + assert_eq!(result.len(), 2); + assert_arrays_eq!( + result, + ConstantArray::new(list, 2), + &mut SESSION.create_execution_ctx() + ); + Ok(()) +} diff --git a/vortex-array/src/scalar_fn/fns/list_length.rs b/vortex-functions/src/list_length.rs similarity index 86% rename from vortex-array/src/scalar_fn/fns/list_length.rs rename to vortex-functions/src/list_length.rs index f88e2baf80d..0d1226a7a5a 100644 --- a/vortex-array/src/scalar_fn/fns/list_length.rs +++ b/vortex-functions/src/list_length.rs @@ -2,38 +2,37 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use num_traits::AsPrimitive; +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::FixedSizeList; +use vortex_array::arrays::List; +use vortex_array::arrays::ListView; +use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; +use vortex_array::arrays::list::ListArrayExt; +use vortex_array::arrays::list::ListArraySlotsExt; +use vortex_array::arrays::listview::ListViewArrayExt; +use vortex_array::arrays::listview::ListViewArraySlotsExt; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::matcher::Matcher; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::Arity; +use vortex_array::scalar_fn::ChildName; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::array::ArrayView; -use crate::arrays::ConstantArray; -use crate::arrays::FixedSizeList; -use crate::arrays::List; -use crate::arrays::ListView; -use crate::arrays::fixed_size_list::FixedSizeListArrayExt; -use crate::arrays::list::ListArrayExt; -use crate::arrays::list::ListArraySlotsExt; -use crate::arrays::listview::ListViewArrayExt; -use crate::arrays::listview::ListViewArraySlotsExt; -use crate::builtins::ArrayBuiltins; -use crate::dtype::DType; -use crate::dtype::Nullability; -use crate::dtype::PType; -use crate::matcher::Matcher; -use crate::scalar::Scalar; -use crate::scalar_fn::Arity; -use crate::scalar_fn::ChildName; -use crate::scalar_fn::EmptyOptions; -use crate::scalar_fn::ExecutionArgs; -use crate::scalar_fn::ScalarFnId; -use crate::scalar_fn::ScalarFnVTable; -use crate::scalar_fn::fns::operators::Operator; - /// Number of elements in each list of a `List` or `FixedSizeList` typed array. /// /// This is computed purely from the list's offsets (`ListArray`), sizes (`ListViewArray`), or @@ -187,30 +186,36 @@ mod tests { use std::sync::LazyLock; use rstest::rstest; + use vortex_array::ArrayRef; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::BoolArray; + use vortex_array::arrays::ConstantArray; + use vortex_array::arrays::FixedSizeListArray; + use vortex_array::arrays::ListArray; + use vortex_array::arrays::ListViewArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::assert_arrays_eq; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::expr::cast; + use vortex_array::expr::root; + use vortex_array::scalar::Scalar; + use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; - use crate::ArrayRef; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::arrays::BoolArray; - use crate::arrays::ConstantArray; - use crate::arrays::FixedSizeListArray; - use crate::arrays::ListArray; - use crate::arrays::ListViewArray; - use crate::arrays::PrimitiveArray; - use crate::assert_arrays_eq; - use crate::dtype::DType; - use crate::dtype::Nullability; - use crate::dtype::PType; - use crate::expr::cast; use crate::expr::list_length; - use crate::expr::root; - use crate::scalar::Scalar; - use crate::validity::Validity; - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + + crate::initialize(&session); + + session + }); fn create_list_elements() -> ArrayRef { PrimitiveArray::from_option_iter::([ diff --git a/vortex-array/src/scalar_fn/fns/list_sum.rs b/vortex-functions/src/list_sum.rs similarity index 90% rename from vortex-array/src/scalar_fn/fns/list_sum.rs rename to vortex-functions/src/list_sum.rs index 0e83a22ed3b..fb25a96e4ba 100644 --- a/vortex-array/src/scalar_fn/fns/list_sum.rs +++ b/vortex-functions/src/list_sum.rs @@ -1,30 +1,29 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::Columnar; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::aggregate_fn::AggregateFnVTable; +use vortex_array::aggregate_fn::DynGroupedAccumulator; +use vortex_array::aggregate_fn::GroupedAccumulator; +use vortex_array::aggregate_fn::NumericalAggregateOpts; +use vortex_array::aggregate_fn::fns::sum_v2::SumV2; +use vortex_array::arrays::ConstantArray; +use vortex_array::dtype::DType; +use vortex_array::scalar_fn::Arity; +use vortex_array::scalar_fn::ChildName; +use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::ScalarFnVTable; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::ArrayRef; -use crate::Canonical; -use crate::Columnar; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::aggregate_fn::AggregateFnVTable; -use crate::aggregate_fn::DynGroupedAccumulator; -use crate::aggregate_fn::GroupedAccumulator; -use crate::aggregate_fn::NumericalAggregateOpts; -use crate::aggregate_fn::fns::sum_v2::SumV2; -use crate::arrays::ConstantArray; -use crate::dtype::DType; -use crate::scalar_fn::Arity; -use crate::scalar_fn::ChildName; -use crate::scalar_fn::ExecutionArgs; -use crate::scalar_fn::ScalarFnId; -use crate::scalar_fn::ScalarFnVTable; - /// Sum of the elements in each list of a `List` or `FixedSizeList` typed array. /// /// Follows SQL `SUM` semantics per list, matching DuckDB's `list_sum`: null lists, empty @@ -142,36 +141,42 @@ mod tests { use prost::Message; use rstest::rstest; + use vortex_array::ArrayRef; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::NumericalAggregateOpts; + use vortex_array::arrays::BoolArray; + use vortex_array::arrays::ConstantArray; + use vortex_array::arrays::FixedSizeListArray; + use vortex_array::arrays::ListArray; + use vortex_array::arrays::ListViewArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::assert_arrays_eq; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::expr::Expression; + use vortex_array::expr::proto::ExprSerializeProtoExt; + use vortex_array::expr::root; + use vortex_array::proto::expr as pb; + use vortex_array::scalar::Scalar; + use vortex_array::scalar_fn::ScalarFnVTable; + use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; - use crate::ArrayRef; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::aggregate_fn::NumericalAggregateOpts; - use crate::arrays::BoolArray; - use crate::arrays::ConstantArray; - use crate::arrays::FixedSizeListArray; - use crate::arrays::ListArray; - use crate::arrays::ListViewArray; - use crate::arrays::PrimitiveArray; - use crate::assert_arrays_eq; - use crate::dtype::DType; - use crate::dtype::Nullability; - use crate::dtype::PType; - use crate::expr::Expression; use crate::expr::list_sum; use crate::expr::list_sum_opts; - use crate::expr::proto::ExprSerializeProtoExt; - use crate::expr::root; - use crate::proto::expr as pb; - use crate::scalar::Scalar; - use crate::scalar_fn::ScalarFnVTable; - use crate::scalar_fn::fns::list_sum::ListSum; - use crate::validity::Validity; + use crate::list_sum::ListSum; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + + crate::initialize(&session); - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + session + }); fn create_list_elements() -> ArrayRef { PrimitiveArray::from_option_iter::([ @@ -564,7 +569,9 @@ mod tests { let proto = expr.serialize_proto()?; let buf = proto.encode_to_vec(); let decoded = pb::Expr::decode(buf.as_slice())?; - let deser = Expression::from_proto(&decoded, &SESSION)?; + let session = VortexSession::empty(); + crate::initialize(&session); + let deser = Expression::from_proto(&decoded, &session)?; assert_eq!(expr, deser); } Ok(()) diff --git a/vortex-array/src/scalar_fn/fns/merge.rs b/vortex-functions/src/merge.rs similarity index 90% rename from vortex-array/src/scalar_fn/fns/merge.rs rename to vortex-functions/src/merge.rs index 4913b120000..0481e439ea7 100644 --- a/vortex-array/src/scalar_fn/fns/merge.rs +++ b/vortex-functions/src/merge.rs @@ -7,6 +7,26 @@ use std::hash::Hash; use std::sync::Arc; use itertools::Itertools as _; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray as _; +use vortex_array::arrays::StructArray; +use vortex_array::arrays::struct_::StructArrayExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::FieldNames; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::StructFields; +use vortex_array::scalar_fn::Arity; +use vortex_array::scalar_fn::ChildName; +use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::ReduceNode; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::ScalarFnVTableExt; +use vortex_array::scalar_fn::builtin::get_item::GetItem; +use vortex_array::scalar_fn::builtin::pack::Pack; +use vortex_array::scalar_fn::builtin::pack::PackOptions; +use vortex_array::validity::Validity; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -14,27 +34,6 @@ use vortex_session::VortexSession; use vortex_session::registry::CachedId; use vortex_utils::aliases::hash_set::HashSet; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::IntoArray as _; -use crate::arrays::StructArray; -use crate::arrays::struct_::StructArrayExt; -use crate::dtype::DType; -use crate::dtype::FieldNames; -use crate::dtype::Nullability; -use crate::dtype::StructFields; -use crate::scalar_fn::Arity; -use crate::scalar_fn::ChildName; -use crate::scalar_fn::ExecutionArgs; -use crate::scalar_fn::ReduceNode; -use crate::scalar_fn::ScalarFnId; -use crate::scalar_fn::ScalarFnVTable; -use crate::scalar_fn::ScalarFnVTableExt; -use crate::scalar_fn::fns::get_item::GetItem; -use crate::scalar_fn::fns::pack::Pack; -use crate::scalar_fn::fns::pack::PackOptions; -use crate::validity::Validity; - /// Merge zero or more expressions that ALL return structs. /// /// If any field names are duplicated, the field from later expressions wins. @@ -256,33 +255,39 @@ impl Display for DuplicateHandling { mod tests { use std::sync::LazyLock; + use vortex_array::ArrayRef; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::struct_::StructArrayExt; + use vortex_array::assert_arrays_eq; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability::NonNullable; + use vortex_array::dtype::PType::I32; + use vortex_array::dtype::PType::I64; + use vortex_array::dtype::PType::U32; + use vortex_array::dtype::PType::U64; + use vortex_array::expr::Expression; + use vortex_array::expr::get_item; + use vortex_array::expr::root; + use vortex_array::scalar_fn::builtin::pack::Pack; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_session::VortexSession; - use crate::ArrayRef; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::arrays::PrimitiveArray; - use crate::arrays::struct_::StructArrayExt; - use crate::assert_arrays_eq; - use crate::dtype::DType; - use crate::dtype::Nullability::NonNullable; - use crate::dtype::PType::I32; - use crate::dtype::PType::I64; - use crate::dtype::PType::U32; - use crate::dtype::PType::U64; - use crate::expr::Expression; - use crate::expr::get_item; use crate::expr::merge; use crate::expr::merge_opts; - use crate::expr::root; - use crate::scalar_fn::fns::merge::DuplicateHandling; - use crate::scalar_fn::fns::merge::StructArray; - use crate::scalar_fn::fns::pack::Pack; + use crate::merge::DuplicateHandling; + use crate::merge::StructArray; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + + crate::initialize(&session); - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + session + }); fn field_array(array: &ArrayRef, field_path: &[&str]) -> VortexResult { let mut ctx = SESSION.create_execution_ctx(); @@ -444,7 +449,7 @@ mod tests { Default::default(), vec![], 3, - crate::validity::Validity::NonNullable, + vortex_array::validity::Validity::NonNullable, ) .unwrap(); assert_arrays_eq!(actual_array, expected, &mut SESSION.create_execution_ctx()); diff --git a/vortex-array/src/scalar_fn/fns/replace_time_zone.rs b/vortex-functions/src/replace_time_zone.rs similarity index 89% rename from vortex-array/src/scalar_fn/fns/replace_time_zone.rs rename to vortex-functions/src/replace_time_zone.rs index a07029cfe70..4ba367f2222 100644 --- a/vortex-array/src/scalar_fn/fns/replace_time_zone.rs +++ b/vortex-functions/src/replace_time_zone.rs @@ -10,6 +10,32 @@ use jiff::Timestamp as JiffTimestamp; use jiff::tz::AmbiguousOffset; use jiff::tz::TimeZone; use prost::Message; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::VarBinViewArray; +use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::expr::BoundExpression; +use vortex_array::expr::bound; +use vortex_array::extension::datetime::TimeUnit; +use vortex_array::extension::datetime::Timestamp; +use vortex_array::extension::datetime::TimestampOptions; +use vortex_array::extension::datetime::resolve_time_zone as resolve_zone; +use vortex_array::proto::expr as pb; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::Arity; +use vortex_array::scalar_fn::ChildName; +use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::builtin::literal::Literal; +use vortex_array::validity::Validity; use vortex_buffer::BitBufferMut; use vortex_buffer::BufferMut; use vortex_error::VortexResult; @@ -18,33 +44,6 @@ use vortex_error::vortex_err; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::arrays::ConstantArray; -use crate::arrays::ExtensionArray; -use crate::arrays::PrimitiveArray; -use crate::arrays::VarBinViewArray; -use crate::arrays::extension::ExtensionArrayExt; -use crate::dtype::DType; -use crate::dtype::Nullability; -use crate::dtype::PType; -use crate::expr::BoundExpression; -use crate::expr::bound; -use crate::extension::datetime::TimeUnit; -use crate::extension::datetime::Timestamp; -use crate::extension::datetime::TimestampOptions; -use crate::extension::datetime::resolve_time_zone as resolve_zone; -use crate::proto::expr as pb; -use crate::scalar::Scalar; -use crate::scalar_fn::Arity; -use crate::scalar_fn::ChildName; -use crate::scalar_fn::ExecutionArgs; -use crate::scalar_fn::ScalarFnId; -use crate::scalar_fn::ScalarFnVTable; -use crate::scalar_fn::fns::literal::Literal; -use crate::validity::Validity; - /// Timezone replacement options. The ambiguity policy is a UTF-8 expression child. #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub struct ReplaceTimeZoneOptions { @@ -316,6 +315,24 @@ mod tests { use std::sync::LazyLock; use rstest::rstest; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::ExtensionArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::StructArray; + use vortex_array::arrays::VarBinViewArray; + use vortex_array::assert_arrays_eq; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::StructFields; + use vortex_array::expr::get_item; + use vortex_array::expr::lit; + use vortex_array::expr::root; + use vortex_array::extension::datetime::TimeUnit; + use vortex_array::extension::datetime::Timestamp; + use vortex_array::scalar::Scalar; + use vortex_array::scalar_fn::ScalarFnVTable; + use vortex_array::scalar_fn::builtin::literal::Literal; use vortex_error::VortexResult; use vortex_error::vortex_err; use vortex_session::VortexSession; @@ -324,27 +341,15 @@ mod tests { use super::ReplaceTimeZoneOptions; use super::replace; use super::resolve_zone; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::arrays::ExtensionArray; - use crate::arrays::PrimitiveArray; - use crate::arrays::StructArray; - use crate::arrays::VarBinViewArray; - use crate::assert_arrays_eq; - use crate::dtype::DType; - use crate::dtype::Nullability; - use crate::dtype::StructFields; - use crate::expr::get_item; - use crate::expr::lit; use crate::expr::replace_time_zone; - use crate::expr::root; - use crate::extension::datetime::TimeUnit; - use crate::extension::datetime::Timestamp; - use crate::scalar::Scalar; - use crate::scalar_fn::ScalarFnVTable; - use crate::scalar_fn::fns::literal::Literal; - - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + + crate::initialize(&session); + + session + }); #[rstest] #[case(TimeUnit::Seconds, 1)] diff --git a/vortex-array/src/scalar_fn/fns/select.rs b/vortex-functions/src/select.rs similarity index 89% rename from vortex-array/src/scalar_fn/fns/select.rs rename to vortex-functions/src/select.rs index e410105d1b4..bd0003984fe 100644 --- a/vortex-array/src/scalar_fn/fns/select.rs +++ b/vortex-functions/src/select.rs @@ -6,6 +6,29 @@ use std::fmt::Formatter; use itertools::Itertools; use prost::Message; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::Constant; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::StructArray; +use vortex_array::arrays::struct_::StructArrayExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::FieldName; +use vortex_array::dtype::FieldNames; +use vortex_array::expr::BoundExpression; +use vortex_array::expr::display::ExprDisplay; +use vortex_array::proto::expr::FieldNames as ProtoFieldNames; +use vortex_array::proto::expr::SelectOpts; +use vortex_array::proto::expr::select_opts::Opts; +use vortex_array::scalar_fn::Arity; +use vortex_array::scalar_fn::ChildName; +use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::ScalarFnVTableExt; +use vortex_array::scalar_fn::builtin::pack::Pack; +use vortex_array::scalar_fn::builtin::pack::PackOptions; use vortex_error::VortexError; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -15,30 +38,7 @@ use vortex_session::VortexSession; use vortex_session::registry::CachedId; use vortex_utils::aliases::hash_set::HashSet; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::arrays::Constant; -use crate::arrays::ConstantArray; -use crate::arrays::StructArray; -use crate::arrays::struct_::StructArrayExt; -use crate::dtype::DType; -use crate::dtype::FieldName; -use crate::dtype::FieldNames; -use crate::expr::BoundExpression; -use crate::expr::display::ExprDisplay; -use crate::expr::field::DisplayFieldNames; -use crate::proto::expr::FieldNames as ProtoFieldNames; -use crate::proto::expr::SelectOpts; -use crate::proto::expr::select_opts::Opts; -use crate::scalar_fn::Arity; -use crate::scalar_fn::ChildName; -use crate::scalar_fn::ExecutionArgs; -use crate::scalar_fn::ScalarFnId; -use crate::scalar_fn::ScalarFnVTable; -use crate::scalar_fn::ScalarFnVTableExt; -use crate::scalar_fn::fns::pack::Pack; -use crate::scalar_fn::fns::pack::PackOptions; +use crate::field::DisplayFieldNames; #[derive(Debug, Clone, PartialEq, Eq, Hash)] pub enum FieldSelection { @@ -348,35 +348,41 @@ impl Display for FieldSelection { mod tests { use std::sync::LazyLock; + use vortex_array::ArrayRef; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::Constant; + use vortex_array::arrays::ConstantArray; + use vortex_array::arrays::ScalarFnArray; + use vortex_array::assert_arrays_eq; + use vortex_array::dtype::DType; + use vortex_array::dtype::FieldName; + use vortex_array::dtype::FieldNames; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::Nullability::Nullable; + use vortex_array::dtype::PType::I32; + use vortex_array::dtype::StructFields; + use vortex_array::expr::root; + use vortex_array::expr::test_harness; + use vortex_array::scalar::Scalar; + use vortex_array::scalar_fn::ScalarFnVTableExt; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; - use crate::ArrayRef; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::arrays::Constant; - use crate::arrays::ConstantArray; - use crate::arrays::ScalarFnArray; - use crate::assert_arrays_eq; - use crate::dtype::DType; - use crate::dtype::FieldName; - use crate::dtype::FieldNames; - use crate::dtype::Nullability; - use crate::dtype::Nullability::Nullable; - use crate::dtype::PType::I32; - use crate::dtype::StructFields; - use crate::expr::root; use crate::expr::select; use crate::expr::select_exclude; - use crate::expr::test_harness; - use crate::scalar::Scalar; - use crate::scalar_fn::ScalarFnVTableExt; - use crate::scalar_fn::fns::select::FieldSelection; - use crate::scalar_fn::fns::select::Select; - use crate::scalar_fn::fns::select::StructArray; - - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + use crate::select::FieldSelection; + use crate::select::Select; + use crate::select::StructArray; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + + crate::initialize(&session); + + session + }); fn test_array() -> StructArray { StructArray::from_fields(&[ diff --git a/vortex-array/src/scalar_fn/fns/stat.rs b/vortex-functions/src/stat.rs similarity index 86% rename from vortex-array/src/scalar_fn/fns/stat.rs rename to vortex-functions/src/stat.rs index c961ac10d53..e88d3a652d9 100644 --- a/vortex-array/src/scalar_fn/fns/stat.rs +++ b/vortex-functions/src/stat.rs @@ -6,33 +6,32 @@ use std::fmt::Display; use std::fmt::Formatter; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::aggregate_fn::AggregateFnRef; +use vortex_array::aggregate_fn::fns::all_nan::AllNan; +use vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan; +use vortex_array::aggregate_fn::fns::all_non_null::AllNonNull; +use vortex_array::aggregate_fn::fns::all_null::AllNull; +use vortex_array::arrays::ConstantArray; +use vortex_array::dtype::DType; +use vortex_array::expr::display::ExprDisplay; +use vortex_array::expr::stats::Precision; +use vortex_array::expr::stats::Stat; +use vortex_array::expr::stats::StatsProvider; +use vortex_array::expr::stats::StatsProviderExt; +use vortex_array::scalar::Scalar; +use vortex_array::scalar::ScalarValue; +use vortex_array::scalar_fn::Arity; +use vortex_array::scalar_fn::ChildName; +use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::ScalarFnVTable; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_session::registry::CachedId; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::aggregate_fn::AggregateFnRef; -use crate::aggregate_fn::fns::all_nan::AllNan; -use crate::aggregate_fn::fns::all_non_nan::AllNonNan; -use crate::aggregate_fn::fns::all_non_null::AllNonNull; -use crate::aggregate_fn::fns::all_null::AllNull; -use crate::arrays::ConstantArray; -use crate::dtype::DType; -use crate::expr::display::ExprDisplay; -use crate::expr::stats::Precision; -use crate::expr::stats::Stat; -use crate::expr::stats::StatsProvider; -use crate::expr::stats::StatsProviderExt; -use crate::scalar::Scalar; -use crate::scalar::ScalarValue; -use crate::scalar_fn::Arity; -use crate::scalar_fn::ChildName; -use crate::scalar_fn::ExecutionArgs; -use crate::scalar_fn::ScalarFnId; -use crate::scalar_fn::ScalarFnVTable; - /// Options for the `stat` scalar function. #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub struct StatOptions { diff --git a/vortex-functions/src/tests.rs b/vortex-functions/src/tests.rs new file mode 100644 index 00000000000..65ce6b88e0d --- /dev/null +++ b/vortex-functions/src/tests.rs @@ -0,0 +1,337 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Expression display, binding and optimizer tests that involve this crate's functions. + +use vortex_array::dtype::DType; +use vortex_array::dtype::FieldName; +use vortex_array::dtype::FieldNames; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::dtype::StructFields; +use vortex_array::expr::Expression; +use vortex_array::expr::and; +use vortex_array::expr::bound; +use vortex_array::expr::cast; +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::not; +use vortex_array::expr::not_eq; +use vortex_array::expr::or; +use vortex_array::expr::pack; +use vortex_array::expr::root; +use vortex_array::scalar::Scalar; +use vortex_error::VortexResult; + +use crate::between::BetweenOptions; +use crate::between::StrictComparison; +use crate::expr::between; +use crate::expr::bound::case_when as bound_case_when; +use crate::expr::case_when; +use crate::expr::merge; +use crate::expr::select; +use crate::expr::select_exclude; + +#[test] +fn test_display_tree_select() { + use insta::assert_snapshot; + let select_expr = select(["name", "age"], root()); + assert_snapshot!(select_expr.display_tree().to_string(), @r" + vortex.select({name, age}) + └── child: vortex.root() + "); +} + +#[test] +fn test_display_tree_select_exclude() { + use insta::assert_snapshot; + let select_exclude_expr = select_exclude(["internal_id", "metadata"], root()); + assert_snapshot!(select_exclude_expr.display_tree().to_string(), @r" + vortex.select(~{internal_id, metadata}) + └── child: vortex.root() + "); +} + +#[test] +fn test_display_tree_nested() { + use insta::assert_snapshot; + let nested_expr = select( + ["result"], + cast( + between( + get_item("score", root()), + lit(50), + lit(100), + BetweenOptions { + lower_strict: StrictComparison::Strict, + upper_strict: StrictComparison::NonStrict, + }, + ), + DType::Bool(Nullability::NonNullable), + ), + ); + assert_snapshot!(nested_expr.display_tree().to_string(), @r" + vortex.select({result}) + └── child: vortex.cast(bool) + └── input: vortex.between(lower_strict: <, upper_strict: <=) + ├── array: vortex.get_item(score) + │ └── input: vortex.root() + ├── lower: vortex.literal(50i32) + └── upper: vortex.literal(100i32) + "); +} + +#[test] +fn test_display_tree_pack() { + use insta::assert_snapshot; + let select_from_pack_expr = select( + ["fizz", "buzz"], + pack( + [ + ("fizz", root()), + ("bar", lit(5)), + ("buzz", eq(lit(42), get_item("answer", root()))), + ], + Nullability::Nullable, + ), + ); + assert_snapshot!(select_from_pack_expr.display_tree().to_string(), @r" + vortex.select({fizz, buzz}) + └── child: vortex.pack(names: [fizz, bar, buzz], nullability: Nullable) + ├── fizz: vortex.root() + ├── bar: vortex.literal(5i32) + └── buzz: vortex.binary(=) + ├── lhs: vortex.literal(42i32) + └── rhs: vortex.get_item(answer) + └── input: vortex.root() + "); +} + +/// `Select` over a `pack` rewrites to a `pack` of the selected child fields. It must read +/// those fields out of the child `pack` itself, because the optimizer visits children before +/// their parent and so never revisits the children a rule introduces. +#[test] +fn optimize_select_of_pack_reads_pack_fields() -> VortexResult<()> { + let field_names = (0..4) + .map(|idx| FieldName::from(format!("field_{idx}"))) + .collect::>(); + let fields = StructFields::new( + field_names.clone().into(), + vec![DType::Primitive(PType::U64, Nullability::NonNullable); field_names.len()], + ); + let scope = DType::Struct(fields, Nullability::NonNullable); + + let root_pack = pack( + field_names + .iter() + .map(|name| (name.clone(), get_item(name.clone(), root()))) + .collect::>(), + Nullability::NonNullable, + ); + let expr = select( + field_names + .iter() + .cloned() + .chain(["input_row_idx".into(), "input_file_idx".into()]) + .collect::>(), + merge([ + root_pack, + pack( + [ + ("input_row_idx", lit(0_u64)), + ("input_file_idx", lit(1_u64)), + ], + Nullability::NonNullable, + ), + ]), + ); + + let expected = pack( + field_names + .into_iter() + .map(|name| (name.clone(), get_item(name, root()))) + .chain([ + ("input_row_idx".into(), lit(0_u64)), + ("input_file_idx".into(), lit(1_u64)), + ]) + .collect::>(), + Nullability::NonNullable, + ); + + let expected = expected.bind(&scope)?; + + let optimized = expr.bind(&scope)?.optimize_recursive()?; + assert_eq!(optimized, expected, "got {optimized}"); + assert_eq!(optimized.optimize_recursive()?, expected); + + Ok(()) +} + +/// `Merge` reduces to a `pack` that reads each field out of its children, so it must also +/// read straight out of a `pack` child. +#[test] +fn optimize_merge_of_packs_reads_pack_fields() -> VortexResult<()> { + let scope = DType::Struct(StructFields::empty(), Nullability::NonNullable); + let expr = merge([ + pack([("a", lit(0_u64))], Nullability::NonNullable), + pack([("b", lit(1_u64))], Nullability::NonNullable), + ]); + + let expected = pack( + [("a", lit(0_u64)), ("b", lit(1_u64))], + Nullability::NonNullable, + ); + + let expected = expected.bind(&scope)?; + + let optimized = expr.bind(&scope)?.optimize_recursive()?; + assert_eq!(optimized, expected, "got {optimized}"); + assert_eq!(optimized.optimize_recursive()?, expected); + + Ok(()) +} + +#[test] +fn bound_constructors_preserve_order_and_types() -> VortexResult<()> { + let value_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let scope = DType::Struct( + StructFields::from_iter([("value", value_dtype.clone())]), + Nullability::NonNullable, + ); + + let root = bound::root(scope.clone()); + let value = bound::get_item("value", root); + let literal = bound::lit(5i32); + let condition = bound::gt(value.clone(), literal.clone()); + assert_eq!(condition.dtype(), &DType::Bool(Nullability::NonNullable)); + assert_eq!(condition.children(), &[value.clone(), literal.clone()]); + + let case = bound_case_when(condition.clone(), value.clone(), literal.clone()); + assert_eq!(case.dtype(), &value_dtype); + assert_eq!(case.children(), &[condition.clone(), value, literal]); + + let packed = bound::pack( + [("condition", condition.clone()), ("value", case.clone())], + Nullability::NonNullable, + ); + assert_eq!(packed.children(), &[condition, case.clone()]); + assert_eq!( + packed.dtype(), + &DType::Struct( + StructFields::from_iter([ + ("condition", DType::Bool(Nullability::NonNullable)), + ("value", value_dtype), + ]), + Nullability::NonNullable, + ) + ); + + let unbound = case_when(gt(col("value"), lit(5i32)), col("value"), lit(5i32)); + assert_eq!(unbound.bind(&scope)?, case); + Ok(()) +} + +#[test] +fn expr_display() { + assert_eq!(col("a").to_string(), "$.a"); + assert_eq!(root().to_string(), "$"); + + let col1: Expression = col("col1"); + let col2: Expression = col("col2"); + assert_eq!( + and(col1.clone(), col2.clone()).to_string(), + "($.col1 and $.col2)" + ); + assert_eq!( + or(col1.clone(), col2.clone()).to_string(), + "($.col1 or $.col2)" + ); + assert_eq!( + eq(col1.clone(), col2.clone()).to_string(), + "($.col1 = $.col2)" + ); + assert_eq!( + not_eq(col1.clone(), col2.clone()).to_string(), + "($.col1 != $.col2)" + ); + assert_eq!( + gt(col1.clone(), col2.clone()).to_string(), + "($.col1 > $.col2)" + ); + assert_eq!( + gt_eq(col1.clone(), col2.clone()).to_string(), + "($.col1 >= $.col2)" + ); + assert_eq!( + lt(col1.clone(), col2.clone()).to_string(), + "($.col1 < $.col2)" + ); + assert_eq!( + lt_eq(col1.clone(), col2.clone()).to_string(), + "($.col1 <= $.col2)" + ); + + assert_eq!( + or(lt(col1.clone(), col2.clone()), not_eq(col1.clone(), col2),).to_string(), + "(($.col1 < $.col2) or ($.col1 != $.col2))" + ); + + assert_eq!(not(col1).to_string(), "vortex.not($.col1)"); + + assert_eq!( + select(vec![FieldName::from("col1")], root()).to_string(), + "${col1}" + ); + assert_eq!( + select( + vec![FieldName::from("col1"), FieldName::from("col2")], + root() + ) + .to_string(), + "${col1, col2}" + ); + assert_eq!( + select_exclude( + vec![FieldName::from("col1"), FieldName::from("col2")], + root() + ) + .to_string(), + "${~ col1, col2}" + ); + + assert_eq!(lit(Scalar::from(0u8)).to_string(), "0u8"); + assert_eq!(lit(Scalar::from(0.0f32)).to_string(), "0f32"); + assert_eq!( + lit(Scalar::from(i64::MAX)).to_string(), + "9223372036854775807i64" + ); + assert_eq!(lit(Scalar::from(true)).to_string(), "true"); + assert_eq!( + lit(Scalar::null(DType::Bool(Nullability::Nullable))).to_string(), + "null" + ); + + assert_eq!( + lit(Scalar::struct_( + DType::Struct( + StructFields::new( + FieldNames::from(["dog", "cat"]), + vec![ + DType::Primitive(PType::U32, Nullability::NonNullable), + DType::Utf8(Nullability::NonNullable) + ], + ), + Nullability::NonNullable + ), + vec![Scalar::from(32_u32), Scalar::from("rufus".to_string())] + )) + .to_string(), + "{dog: 32u32, cat: \"rufus\"}" + ); +} diff --git a/vortex-array/src/scalar_fn/fns/variant_get/mod.rs b/vortex-functions/src/variant_get/mod.rs similarity index 82% rename from vortex-array/src/scalar_fn/fns/variant_get/mod.rs rename to vortex-functions/src/variant_get/mod.rs index ef493d0146e..92947f16092 100644 --- a/vortex-array/src/scalar_fn/fns/variant_get/mod.rs +++ b/vortex-functions/src/variant_get/mod.rs @@ -1,11 +1,34 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +mod variant; use std::fmt; use std::fmt::Display; use std::fmt::Formatter; use prost::Message; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::ChunkedArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::ScalarFnArray; +use vortex_array::arrays::VariantArray; +use vortex_array::builders::builder_with_capacity_in; +use vortex_array::dtype::DType; +use vortex_array::dtype::FieldName; +use vortex_array::dtype::Nullability; +use vortex_array::expr::display::ExprDisplay; +use vortex_array::proto::expr as pb; +use vortex_array::proto::expr::variant_path_element; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::Arity; +use vortex_array::scalar_fn::ChildName; +use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::ScalarFnVTableExt; +use vortex_array::scalar_fn::session::ScalarFnSessionExt; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_err; @@ -13,27 +36,11 @@ use vortex_session::VortexSession; use vortex_session::registry::CachedId; use vortex_utils::aliases::StringEscape; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::arrays::ChunkedArray; -use crate::arrays::ConstantArray; -use crate::arrays::ScalarFnArray; -use crate::arrays::VariantArray; -use crate::builders::builder_with_capacity_in; -use crate::dtype::DType; -use crate::dtype::FieldName; -use crate::dtype::Nullability; -use crate::expr::display::ExprDisplay; -use crate::proto::expr as pb; -use crate::proto::expr::variant_path_element; -use crate::scalar::Scalar; -use crate::scalar_fn::Arity; -use crate::scalar_fn::ChildName; -use crate::scalar_fn::ExecutionArgs; -use crate::scalar_fn::ScalarFnId; -use crate::scalar_fn::ScalarFnVTable; -use crate::scalar_fn::ScalarFnVTableExt; +/// Registers [`VariantGet`] and its kernel for canonical variant arrays. +pub(crate) fn initialize(session: &VortexSession) { + session.scalar_fns().register(VariantGet); + variant::initialize(session); +} /// Extracts a field/index path from Variant values. /// @@ -411,43 +418,51 @@ impl From for VariantPathElement { mod tests { use std::sync::LazyLock; + use vortex_array::ArrayRef; + use vortex_array::Canonical; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::Chunked; + use vortex_array::arrays::ChunkedArray; + use vortex_array::arrays::ConstantArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::StructArray; + use vortex_array::arrays::VariantArray; + use vortex_array::arrays::variant::VariantArraySlotsExt; + use vortex_array::assert_arrays_eq; + use vortex_array::assert_nth_scalar_is_null; + use vortex_array::dtype::DType; + use vortex_array::dtype::FieldName; + use vortex_array::dtype::FieldNames; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::dtype::StructFields; + use vortex_array::expr::Expression; + use vortex_array::expr::proto::ExprSerializeProtoExt; + use vortex_array::expr::root; + use vortex_array::optimizer::ArrayOptimizer; + use vortex_array::scalar::Scalar; + use vortex_array::scalar::ScalarValue; + use vortex_array::scalar_fn::ScalarFnVTable; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_err; use vortex_session::VortexSession; - use crate::ArrayRef; - use crate::Canonical; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::arrays::Chunked; - use crate::arrays::ChunkedArray; - use crate::arrays::ConstantArray; - use crate::arrays::PrimitiveArray; - use crate::arrays::VariantArray; - use crate::arrays::variant::VariantArraySlotsExt; - use crate::assert_arrays_eq; - use crate::assert_nth_scalar_is_null; - use crate::dtype::DType; - use crate::dtype::FieldName; - use crate::dtype::FieldNames; - use crate::dtype::Nullability; - use crate::dtype::PType; - use crate::dtype::StructFields; - use crate::expr::Expression; - use crate::expr::proto::ExprSerializeProtoExt; - use crate::expr::root; use crate::expr::variant_get; - use crate::scalar::Scalar; - use crate::scalar::ScalarValue; - use crate::scalar_fn::ScalarFnVTable; - use crate::scalar_fn::fns::variant_get::VariantGet; - use crate::scalar_fn::fns::variant_get::VariantGetOptions; - use crate::scalar_fn::fns::variant_get::VariantPath; - use crate::scalar_fn::fns::variant_get::VariantPathElement; + use crate::variant_get::VariantGet; + use crate::variant_get::VariantGetOptions; + use crate::variant_get::VariantPath; + use crate::variant_get::VariantPathElement; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + + crate::initialize(&session); - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + session + }); fn variant_object(fields: impl IntoIterator) -> Scalar { let fields = fields.into_iter().collect::>(); @@ -564,6 +579,7 @@ mod tests { let expr = variant_get(root(), parse_path(path)?, dtype); array .apply(&expr)? + .optimize(&SESSION)? .execute::(&mut SESSION.create_execution_ctx()) } @@ -848,4 +864,85 @@ mod tests { } Ok(()) } + + fn row_storage(values: impl IntoIterator) -> VortexResult { + Ok(ChunkedArray::try_new( + values.into_iter().map(|value| { + ConstantArray::new( + Scalar::variant(Scalar::primitive(value, Nullability::NonNullable)), + 1, + ) + .into_array() + }), + DType::Variant(Nullability::NonNullable), + )? + .into_array()) + } + + #[test] + fn variant_get_keeps_valid_shredded_rows_for_matching_dtype() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let core_storage = row_storage([1, 2, 3])?; + let shredded = StructArray::try_from_iter([( + "a", + PrimitiveArray::from_iter([10i32, 20, 30]).into_array(), + )])?; + let variant = VariantArray::try_new(core_storage, Some(shredded.into_array()))?; + let expr = variant_get( + root(), + VariantPath::field("a"), + Some(DType::Primitive(PType::I32, Nullability::NonNullable)), + ); + + let result = variant.into_array().apply(&expr)?; + + assert_arrays_eq!( + result, + PrimitiveArray::from_option_iter([Some(10i32), Some(20), Some(30)]), + &mut ctx + ); + Ok(()) + } + + #[test] + fn variant_get_treats_value_and_typed_value_as_logical_field_names() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let core_storage = row_storage([1, 2, 3])?; + let shredded = StructArray::try_from_iter([ + ( + "value", + PrimitiveArray::from_iter([10i32, 20, 30]).into_array(), + ), + ( + "typed_value", + PrimitiveArray::from_iter([40i32, 50, 60]).into_array(), + ), + ])?; + let variant = VariantArray::try_new(core_storage, Some(shredded.into_array()))?; + + let value_expr = variant_get( + root(), + VariantPath::field("value"), + Some(DType::Primitive(PType::I32, Nullability::NonNullable)), + ); + let value_result = variant.clone().into_array().apply(&value_expr)?; + assert_arrays_eq!( + value_result, + PrimitiveArray::from_option_iter([Some(10i32), Some(20), Some(30)]), + &mut ctx + ); + + let typed_value_expr = variant_get( + root(), + VariantPath::field("typed_value"), + Some(DType::Primitive(PType::I32, Nullability::NonNullable)), + ); + let typed_value_result = variant.into_array().apply(&typed_value_expr)?; + assert_arrays_eq!( + typed_value_result, + PrimitiveArray::from_option_iter([Some(40i32), Some(50), Some(60)]), + &mut ctx + ); + Ok(()) + } } diff --git a/vortex-array/src/arrays/variant/vtable/kernel.rs b/vortex-functions/src/variant_get/variant.rs similarity index 86% rename from vortex-array/src/arrays/variant/vtable/kernel.rs rename to vortex-functions/src/variant_get/variant.rs index 22d52ecd6a7..8815d57339b 100644 --- a/vortex-array/src/arrays/variant/vtable/kernel.rs +++ b/vortex-functions/src/variant_get/variant.rs @@ -1,35 +1,36 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::Array; +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::ChunkedArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::Struct; +use vortex_array::arrays::Variant; +use vortex_array::arrays::VariantArray; +use vortex_array::arrays::scalar_fn::ExactScalarFn; +use vortex_array::arrays::scalar_fn::ScalarFnArrayView; +use vortex_array::arrays::struct_::StructArrayExt; +use vortex_array::arrays::variant::VariantArraySlotsExt; +use vortex_array::arrays::variant::merge_typed_scalar_as_variant; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::kernel::ExecuteParentKernel; +use vortex_array::optimizer::kernels::ArrayKernelsExt; +use vortex_array::scalar_fn::ScalarFnVTable; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_session::VortexSession; -use super::merge_typed_scalar_as_variant; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::array::Array; -use crate::array::ArrayView; -use crate::arrays::ChunkedArray; -use crate::arrays::ConstantArray; -use crate::arrays::Struct; -use crate::arrays::Variant; -use crate::arrays::VariantArray; -use crate::arrays::scalar_fn::ExactScalarFn; -use crate::arrays::scalar_fn::ScalarFnArrayView; -use crate::arrays::struct_::StructArrayExt; -use crate::arrays::variant::VariantArraySlotsExt; -use crate::builtins::ArrayBuiltins; -use crate::dtype::DType; -use crate::dtype::Nullability; -use crate::kernel::ExecuteParentKernel; -use crate::optimizer::kernels::ArrayKernelsExt; -use crate::scalar_fn::ScalarFnVTable; -use crate::scalar_fn::fns::variant_get::VariantGet; -use crate::scalar_fn::fns::variant_get::VariantGetOptions; -use crate::scalar_fn::fns::variant_get::VariantPath; -use crate::scalar_fn::fns::variant_get::VariantPathElement; +use crate::builtins::ArrayFunctions; +use crate::variant_get::VariantGet; +use crate::variant_get::VariantGetOptions; +use crate::variant_get::VariantPath; +use crate::variant_get::VariantPathElement; pub(crate) fn initialize(session: &VortexSession) { let kernels = session.kernels(); diff --git a/vortex-array/src/arrays/bool/compute/zip.rs b/vortex-functions/src/zip/bool.rs similarity index 88% rename from vortex-array/src/arrays/bool/compute/zip.rs rename to vortex-functions/src/zip/bool.rs index 70e22a27752..70144c2c243 100644 --- a/vortex-array/src/arrays/bool/compute/zip.rs +++ b/vortex-functions/src/zip/bool.rs @@ -1,21 +1,21 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::Bool; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::bool::BoolArrayExt; +use vortex_array::builtins::ArrayBuiltins; use vortex_buffer::BitBuffer; use vortex_buffer::BufferMut; use vortex_error::VortexResult; use vortex_mask::Mask; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::array::ArrayView; -use crate::arrays::Bool; -use crate::arrays::BoolArray; -use crate::arrays::bool::BoolArrayExt; -use crate::builtins::ArrayBuiltins; -use crate::scalar_fn::fns::zip::ZipKernel; -use crate::scalar_fn::fns::zip::zip_validity; +use crate::zip::ZipKernel; +use crate::zip::zip_validity; /// A branchless boolean zip kernel that blends the two value bitmaps with the mask in one pass. /// @@ -86,20 +86,26 @@ fn zip_value_bits(if_true: &BitBuffer, if_false: &BitBuffer, mask: &BitBuffer) - mod tests { use std::sync::LazyLock; + use vortex_array::ArrayRef; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::Bool; + use vortex_array::arrays::BoolArray; + use vortex_array::assert_arrays_eq; use vortex_error::VortexResult; use vortex_mask::Mask; use vortex_session::VortexSession; use super::zip_value_bits; - use crate::ArrayRef; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::arrays::Bool; - use crate::arrays::BoolArray; - use crate::assert_arrays_eq; - use crate::builtins::ArrayBuiltins; - - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + use crate::builtins::ArrayFunctions; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + + crate::initialize(&session); + + session + }); #[test] fn blend_value_bits_boundaries() { diff --git a/vortex-array/src/arrays/chunked/compute/zip.rs b/vortex-functions/src/zip/chunked.rs similarity index 75% rename from vortex-array/src/arrays/chunked/compute/zip.rs rename to vortex-functions/src/zip/chunked.rs index dc43043b6c7..44b0ce36359 100644 --- a/vortex-array/src/arrays/chunked/compute/zip.rs +++ b/vortex-functions/src/zip/chunked.rs @@ -1,18 +1,18 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::Chunked; +use vortex_array::arrays::ChunkedArray; +use vortex_array::arrays::chunked::ChunkedArrayExt; use vortex_error::VortexResult; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::array::ArrayView; -use crate::arrays::Chunked; -use crate::arrays::ChunkedArray; -use crate::arrays::chunked::ChunkedArrayExt; -use crate::arrays::chunked::paired_chunks::PairedChunksExt; -use crate::builtins::ArrayBuiltins; -use crate::scalar_fn::fns::zip::ZipKernel; +use crate::builtins::ArrayFunctions; +use crate::zip::ZipKernel; +use crate::zip::paired_chunks::PairedChunksExt; // Push down the zip call to the chunks. Without this rule // the default implementation canonicalises the chunked array @@ -48,23 +48,29 @@ impl ZipKernel for Chunked { mod tests { use std::sync::LazyLock; + use vortex_array::ArrayRef; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::Chunked; + use vortex_array::arrays::ChunkedArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::chunked::ChunkedArrayExt; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; use vortex_buffer::buffer; use vortex_mask::Mask; use vortex_session::VortexSession; - use crate::ArrayRef; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::arrays::Chunked; - use crate::arrays::ChunkedArray; - use crate::arrays::PrimitiveArray; - use crate::arrays::chunked::ChunkedArrayExt; - use crate::builtins::ArrayBuiltins; - use crate::dtype::DType; - use crate::dtype::Nullability; - use crate::dtype::PType; - - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + use crate::builtins::ArrayFunctions; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + + crate::initialize(&session); + + session + }); #[test] fn test_chunked_zip_aligns_across_boundaries() { diff --git a/vortex-array/src/scalar_fn/fns/zip/kernel.rs b/vortex-functions/src/zip/kernel.rs similarity index 87% rename from vortex-array/src/scalar_fn/fns/zip/kernel.rs rename to vortex-functions/src/zip/kernel.rs index 575c2f4c55a..d0db3cc07aa 100644 --- a/vortex-array/src/scalar_fn/fns/zip/kernel.rs +++ b/vortex-functions/src/zip/kernel.rs @@ -1,20 +1,20 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::VTable; +use vortex_array::arrays::ScalarFn; +use vortex_array::arrays::scalar_fn::ExactScalarFn; +use vortex_array::arrays::scalar_fn::ScalarFnArrayExt; +use vortex_array::arrays::scalar_fn::ScalarFnArrayView; +use vortex_array::kernel::ExecuteParentKernel; +use vortex_array::optimizer::rules::ArrayParentReduceRule; use vortex_error::VortexExpect; use vortex_error::VortexResult; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::array::ArrayView; -use crate::array::VTable; -use crate::arrays::ScalarFn; -use crate::arrays::scalar_fn::ExactScalarFn; -use crate::arrays::scalar_fn::ScalarFnArrayExt; -use crate::arrays::scalar_fn::ScalarFnArrayView; -use crate::kernel::ExecuteParentKernel; -use crate::optimizer::rules::ArrayParentReduceRule; -use crate::scalar_fn::fns::zip::Zip as ZipExpr; +use crate::zip::Zip as ZipExpr; /// Zip two arrays using a mask without reading buffers. /// diff --git a/vortex-array/src/arrays/listview/compute/zip.rs b/vortex-functions/src/zip/listview.rs similarity index 81% rename from vortex-array/src/arrays/listview/compute/zip.rs rename to vortex-functions/src/zip/listview.rs index a50f1909395..2b34455d25f 100644 --- a/vortex-array/src/arrays/listview/compute/zip.rs +++ b/vortex-functions/src/zip/listview.rs @@ -6,33 +6,33 @@ use std::ops::BitAnd; use std::ops::BitOr; use std::ops::Not; +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::Chunked; +use vortex_array::arrays::ChunkedArray; +use vortex_array::arrays::ListView; +use vortex_array::arrays::ListViewArray; +use vortex_array::arrays::chunked::ChunkedArrayExt; +use vortex_array::arrays::listview::ListViewArraySlotsExt; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_mask::Mask; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::array::ArrayView; -use crate::arrays::Chunked; -use crate::arrays::ChunkedArray; -use crate::arrays::ListView; -use crate::arrays::ListViewArray; -use crate::arrays::chunked::ChunkedArrayExt; -use crate::arrays::listview::ListViewArraySlotsExt; -use crate::builtins::ArrayBuiltins; -use crate::dtype::DType; -use crate::dtype::Nullability; -use crate::dtype::PType; -use crate::scalar_fn::fns::zip::ZipKernel; -use crate::validity::Validity; +use crate::zip::ZipKernel; /// Zip two [`ListViewArray`]s by selecting whole list views per row. /// /// A [`ListViewArray`] addresses each list by an `(offset, size)` pair into a shared `elements` -/// array, and unlike [`ListArray`](crate::arrays::ListArray) it does not require lists to be stored +/// array, and unlike [`ListArray`](vortex_array::arrays::ListArray) it does not require lists to be stored /// contiguously or in order. Zipping two list views is therefore a metadata-only operation over the /// `offsets`, `sizes` and `validity` child arrays: we concatenate the two `elements` arrays /// (without rewriting them) and, for each row, select the `(offset, size)` pair from `if_true` or @@ -236,33 +236,40 @@ mod tests { clippy::cast_possible_truncation, reason = "test fixtures use small indices that fit the target widths" )] - + use std::sync::Arc; use std::sync::LazyLock; + use vortex_array::ArrayRef; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::BoolArray; + use vortex_array::arrays::Chunked; + use vortex_array::arrays::ChunkedArray; + use vortex_array::arrays::ListView; + use vortex_array::arrays::ListViewArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::chunked::ChunkedArrayExt; + use vortex_array::arrays::listview::ListViewArraySlotsExt; + use vortex_array::assert_arrays_eq; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_mask::Mask; use vortex_session::VortexSession; - use crate::ArrayRef; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::arrays::BoolArray; - use crate::arrays::Chunked; - use crate::arrays::ChunkedArray; - use crate::arrays::ListView; - use crate::arrays::ListViewArray; - use crate::arrays::chunked::ChunkedArrayExt; - use crate::arrays::listview::ListViewArraySlotsExt; - use crate::assert_arrays_eq; - use crate::builtins::ArrayBuiltins; - use crate::dtype::DType; - use crate::dtype::Nullability; - use crate::dtype::PType; - use crate::validity::Validity; - - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + use crate::builtins::ArrayFunctions; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + + crate::initialize(&session); + + session + }); fn list_view( elements: ArrayRef, @@ -583,4 +590,102 @@ mod tests { assert_arrays_eq!(result, expected, &mut ctx); Ok(()) } + + #[test] + fn test_zip_widens_false_element_nullability() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + // [[1, 2], [3], [4]] + let if_true = ListViewArray::new( + buffer![1i32, 2, 3, 4].into_array(), + buffer![0u32, 2, 3].into_array(), + buffer![2u32, 1, 1].into_array(), + Validity::NonNullable, + &mut ctx, + ) + .into_array(); + // [[10, null], [30], [40]] + let if_false = ListViewArray::new( + PrimitiveArray::from_option_iter([Some(10i32), None, Some(30), Some(40)]).into_array(), + buffer![0u32, 2, 3].into_array(), + buffer![2u32, 1, 1].into_array(), + Validity::NonNullable, + &mut ctx, + ) + .into_array(); + let mask = Mask::from_iter([false, true, false]); + + let result = mask + .into_array() + .zip(if_true, if_false)? + .execute::(&mut SESSION.create_execution_ctx())?; + assert!(result.is::()); + assert_eq!( + result.dtype(), + &DType::List( + Arc::new(DType::Primitive(PType::I32, Nullability::Nullable)), + Nullability::NonNullable, + ) + ); + + // [[10, null], [3], [40]] + let expected = ListViewArray::new( + PrimitiveArray::from_option_iter([Some(10i32), None, Some(3), Some(40)]).into_array(), + buffer![0u32, 2, 3].into_array(), + buffer![2u32, 1, 1].into_array(), + Validity::NonNullable, + &mut ctx, + ) + .into_array(); + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_zip_widens_true_element_nullability() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + // [[1, null], [3], [4]] + let if_true = ListViewArray::new( + PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), Some(4)]).into_array(), + buffer![0u32, 2, 3].into_array(), + buffer![2u32, 1, 1].into_array(), + Validity::NonNullable, + &mut ctx, + ) + .into_array(); + // [[10], [20], [30]] + let if_false = ListViewArray::new( + buffer![10i32, 20, 30].into_array(), + buffer![0u32, 1, 2].into_array(), + buffer![1u32, 1, 1].into_array(), + Validity::NonNullable, + &mut ctx, + ) + .into_array(); + let mask = Mask::from_iter([true, false, true]); + + let result = mask + .into_array() + .zip(if_true, if_false)? + .execute::(&mut SESSION.create_execution_ctx())?; + assert!(result.is::()); + assert_eq!( + result.dtype(), + &DType::List( + Arc::new(DType::Primitive(PType::I32, Nullability::Nullable)), + Nullability::NonNullable, + ) + ); + + // [[1, null], [20], [4]] + let expected = ListViewArray::new( + PrimitiveArray::from_option_iter([Some(1i32), None, Some(20), Some(4)]).into_array(), + buffer![0u32, 2, 3].into_array(), + buffer![2u32, 1, 1].into_array(), + Validity::NonNullable, + &mut ctx, + ) + .into_array(); + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } } diff --git a/vortex-array/src/scalar_fn/fns/zip/mod.rs b/vortex-functions/src/zip/mod.rs similarity index 81% rename from vortex-array/src/scalar_fn/fns/zip/mod.rs rename to vortex-functions/src/zip/mod.rs index fa46512a4f0..c3d86ffa0b1 100644 --- a/vortex-array/src/scalar_fn/fns/zip/mod.rs +++ b/vortex-functions/src/zip/mod.rs @@ -1,13 +1,51 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +mod bool; +mod chunked; mod kernel; +mod listview; +mod paired_chunks; +mod primitive; +mod struct_; +mod varbinview; use std::fmt::Display; use std::fmt::Formatter; use std::sync::Arc; pub use kernel::*; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::Bool; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::Chunked; +use vortex_array::arrays::ListView; +use vortex_array::arrays::Primitive; +use vortex_array::arrays::ScalarFnArray; +use vortex_array::arrays::Struct; +use vortex_array::arrays::VarBinView; +use vortex_array::arrays::bool::BoolArrayExt; +use vortex_array::builders::ArrayBuilder; +use vortex_array::builders::builder_with_capacity_in; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::StructFields; +use vortex_array::expr::BoundExpression; +use vortex_array::expr::bound; +use vortex_array::expr::display::ExprDisplay; +use vortex_array::optimizer::kernels::ArrayKernelsExt; +use vortex_array::scalar_fn::Arity; +use vortex_array::scalar_fn::ChildName; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::ScalarFnVTableExt; +use vortex_array::scalar_fn::builtin::literal::Literal; +use vortex_array::scalar_fn::session::ScalarFnSessionExt; +use vortex_array::validity::Validity; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_err; @@ -16,29 +54,19 @@ use vortex_mask::MaskValues; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::arrays::BoolArray; -use crate::arrays::ScalarFnArray; -use crate::arrays::bool::BoolArrayExt; -use crate::builders::ArrayBuilder; -use crate::builders::builder_with_capacity_in; -use crate::builtins::ArrayBuiltins; -use crate::dtype::DType; -use crate::dtype::StructFields; -use crate::expr::BoundExpression; -use crate::expr::bound; -use crate::expr::display::ExprDisplay; -use crate::scalar_fn::Arity; -use crate::scalar_fn::ChildName; -use crate::scalar_fn::EmptyOptions; -use crate::scalar_fn::ExecutionArgs; -use crate::scalar_fn::ScalarFnId; -use crate::scalar_fn::ScalarFnVTable; -use crate::scalar_fn::ScalarFnVTableExt; -use crate::scalar_fn::fns::literal::Literal; -use crate::validity::Validity; +use crate::builtins::ArrayFunctions; + +/// Registers [`Zip`] and its kernels for vortex-array's canonical encodings. +pub(crate) fn initialize(session: &VortexSession) { + session.scalar_fns().register(Zip); + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Zip.id(), Bool, ZipExecuteAdaptor(Bool)); + kernels.register_execute_parent_kernel(Zip.id(), Chunked, ZipExecuteAdaptor(Chunked)); + kernels.register_execute_parent_kernel(Zip.id(), ListView, ZipExecuteAdaptor(ListView)); + kernels.register_execute_parent_kernel(Zip.id(), Primitive, ZipExecuteAdaptor(Primitive)); + kernels.register_execute_parent_kernel(Zip.id(), Struct, ZipExecuteAdaptor(Struct)); + kernels.register_execute_parent_kernel(Zip.id(), VarBinView, ZipExecuteAdaptor(VarBinView)); +} /// An expression that conditionally selects between two arrays based on a boolean mask. /// @@ -237,7 +265,7 @@ fn zip_return_dtype(if_true: &DType, if_false: &DType) -> VortexResult { /// Unions the nullability of two dtypes at every nesting level. /// /// Returns `None` when the dtypes differ in more than nullability. -pub(crate) fn zip_nullability_union(lhs: &DType, rhs: &DType) -> Option { +pub fn zip_nullability_union(lhs: &DType, rhs: &DType) -> Option { let nullability = lhs.nullability() | rhs.nullability(); match (lhs, rhs) { @@ -344,35 +372,42 @@ pub(crate) fn zip_validity( mod tests { use std::sync::LazyLock; + use vortex_array::ArrayRef; + use vortex_array::Columnar; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::ConstantArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::Struct; + use vortex_array::arrays::StructArray; + use vortex_array::arrays::VarBinView; + use vortex_array::arrays::VarBinViewArray; + use vortex_array::assert_arrays_eq; + use vortex_array::builders::ArrayBuilder; + use vortex_array::builders::BufferGrowthStrategy; + use vortex_array::builders::VarBinViewBuilder; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::expr::Expression; + use vortex_array::expr::lit; + use vortex_array::expr::root; + use vortex_array::scalar::Scalar; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_mask::Mask; use vortex_session::VortexSession; - use crate::ArrayRef; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::arrays::ConstantArray; - use crate::arrays::PrimitiveArray; - use crate::arrays::Struct; - use crate::arrays::StructArray; - use crate::arrays::VarBinView; - use crate::arrays::VarBinViewArray; - use crate::assert_arrays_eq; - use crate::builders::ArrayBuilder; - use crate::builders::BufferGrowthStrategy; - use crate::builders::VarBinViewBuilder; - use crate::builtins::ArrayBuiltins; - use crate::columnar::Columnar; - use crate::dtype::DType; - use crate::dtype::Nullability; - use crate::dtype::PType; - use crate::expr::lit; - use crate::expr::root; + use crate::builtins::ArrayFunctions; use crate::expr::zip_expr; - use crate::scalar::Scalar; - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + + crate::initialize(&session); + + session + }); #[test] fn dtype() { @@ -578,4 +613,23 @@ mod tests { })); assert_arrays_eq!(zipped.array().clone(), expected, &mut ctx); } + + #[rstest::rstest] + #[case::zip_true( + zip_expr(lit(true), lit(1i32), root()), + DType::Primitive(PType::I32, Nullability::Nullable) + )] + #[case::zip_false( + zip_expr(lit(false), root(), lit(1i32)), + DType::Primitive(PType::I32, Nullability::Nullable) + )] + fn literal_simplification_preserves_result_dtype( + #[case] expr: Expression, + #[case] scope: DType, + ) -> VortexResult<()> { + let original = expr.bind(&scope)?; + let optimized = expr.bind(&scope)?.optimize_recursive()?; + assert_eq!(optimized.dtype(), original.dtype()); + Ok(()) + } } diff --git a/vortex-array/src/arrays/chunked/paired_chunks.rs b/vortex-functions/src/zip/paired_chunks.rs similarity index 92% rename from vortex-array/src/arrays/chunked/paired_chunks.rs rename to vortex-functions/src/zip/paired_chunks.rs index 796daee1b37..12f23b806d8 100644 --- a/vortex-array/src/arrays/chunked/paired_chunks.rs +++ b/vortex-functions/src/zip/paired_chunks.rs @@ -3,11 +3,10 @@ use std::ops::Range; +use vortex_array::ArrayRef; +use vortex_array::arrays::chunked::ChunkedArrayExt; use vortex_error::VortexResult; -use crate::ArrayRef; -use crate::arrays::chunked::ChunkedArrayExt; - pub(crate) struct AlignedPair { pub left: ArrayRef, pub right: ArrayRef, @@ -119,20 +118,26 @@ impl Iterator for PairedChunks { mod tests { use std::sync::LazyLock; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::ChunkedArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::arrays::ChunkedArray; - use crate::arrays::PrimitiveArray; - use crate::arrays::chunked::paired_chunks::PairedChunksExt; - use crate::dtype::DType; - use crate::dtype::Nullability; - use crate::dtype::PType; + use crate::zip::paired_chunks::PairedChunksExt; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + + crate::initialize(&session); - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + session + }); fn i32_dtype() -> DType { DType::Primitive(PType::I32, Nullability::NonNullable) diff --git a/vortex-array/src/arrays/primitive/compute/zip.rs b/vortex-functions/src/zip/primitive.rs similarity index 89% rename from vortex-array/src/arrays/primitive/compute/zip.rs rename to vortex-functions/src/zip/primitive.rs index a653f01099e..293af7fe814 100644 --- a/vortex-array/src/arrays/primitive/compute/zip.rs +++ b/vortex-functions/src/zip/primitive.rs @@ -3,23 +3,23 @@ use std::mem::MaybeUninit; +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::Primitive; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::NativePType; +use vortex_array::match_each_native_ptype; use vortex_buffer::BufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_mask::Mask; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::array::ArrayView; -use crate::arrays::Primitive; -use crate::arrays::PrimitiveArray; -use crate::builtins::ArrayBuiltins; -use crate::dtype::NativePType; -use crate::match_each_native_ptype; -use crate::scalar_fn::fns::zip::ZipKernel; -use crate::scalar_fn::fns::zip::zip_validity; +use crate::zip::ZipKernel; +use crate::zip::zip_validity; /// A dedicated primitive zip kernel that selects values branchlessly per row. /// @@ -150,19 +150,25 @@ mod tests { use std::sync::LazyLock; + use vortex_array::ArrayRef; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::Primitive; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::assert_arrays_eq; use vortex_error::VortexResult; use vortex_mask::Mask; use vortex_session::VortexSession; - use crate::ArrayRef; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::arrays::Primitive; - use crate::arrays::PrimitiveArray; - use crate::assert_arrays_eq; - use crate::builtins::ArrayBuiltins; + use crate::builtins::ArrayFunctions; - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + + crate::initialize(&session); + + session + }); /// The branchless kernel must agree with the scalar reference across the chunk boundary (index /// 63/64) and the trailing remainder, for non-nullable inputs. diff --git a/vortex-array/src/arrays/struct_/compute/zip.rs b/vortex-functions/src/zip/struct_.rs similarity index 84% rename from vortex-array/src/arrays/struct_/compute/zip.rs rename to vortex-functions/src/zip/struct_.rs index a6dc944d019..fde3ba4f380 100644 --- a/vortex-array/src/arrays/struct_/compute/zip.rs +++ b/vortex-functions/src/zip/struct_.rs @@ -5,19 +5,20 @@ use std::ops::BitAnd; use std::ops::BitOr; use std::ops::Not; +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::Struct; +use vortex_array::arrays::StructArray; +use vortex_array::arrays::struct_::StructArrayExt; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::validity::Validity; use vortex_error::VortexResult; use vortex_mask::Mask; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::array::ArrayView; -use crate::arrays::Struct; -use crate::arrays::StructArray; -use crate::arrays::struct_::StructArrayExt; -use crate::builtins::ArrayBuiltins; -use crate::scalar_fn::fns::zip::ZipKernel; -use crate::validity::Validity; +use crate::builtins::ArrayFunctions; +use crate::zip::ZipKernel; impl ZipKernel for Struct { fn zip( @@ -72,18 +73,24 @@ impl ZipKernel for Struct { mod tests { use std::sync::LazyLock; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::StructArray; + use vortex_array::dtype::FieldNames; + use vortex_array::validity::Validity; use vortex_mask::Mask; use vortex_session::VortexSession; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::arrays::PrimitiveArray; - use crate::arrays::StructArray; - use crate::builtins::ArrayBuiltins; - use crate::dtype::FieldNames; - use crate::validity::Validity; + use crate::builtins::ArrayFunctions; - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + + crate::initialize(&session); + + session + }); #[test] fn test_validity_zip_both_validity_array() { diff --git a/vortex-array/src/arrays/varbinview/compute/zip.rs b/vortex-functions/src/zip/varbinview.rs similarity index 90% rename from vortex-array/src/arrays/varbinview/compute/zip.rs rename to vortex-functions/src/zip/varbinview.rs index fdcfccbd0d5..6a42731a9bf 100644 --- a/vortex-array/src/arrays/varbinview/compute/zip.rs +++ b/vortex-functions/src/zip/varbinview.rs @@ -3,24 +3,24 @@ use std::ops::Range; +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::VarBinView; +use vortex_array::arrays::VarBinViewArray; +use vortex_array::arrays::varbinview::BinaryView; +use vortex_array::arrays::varbinview::VarBinViewArrayExt; +use vortex_array::builders::DeduplicatedBuffers; +use vortex_array::builders::LazyBitBufferBuilder; +use vortex_array::builtins::ArrayBuiltins; use vortex_buffer::BufferMut; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_mask::AllOr; use vortex_mask::Mask; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::array::ArrayView; -use crate::arrays::VarBinView; -use crate::arrays::VarBinViewArray; -use crate::arrays::varbinview::BinaryView; -use crate::arrays::varbinview::VarBinViewArrayExt; -use crate::builders::DeduplicatedBuffers; -use crate::builders::LazyBitBufferBuilder; -use crate::builtins::ArrayBuiltins; -use crate::scalar_fn::fns::zip::ZipKernel; +use crate::zip::ZipKernel; // A dedicated VarBinView zip kernel that builds the result directly by adjusting views and validity, // instead of routing through the generic builder (which would redo buffer lookups per mask slice). @@ -212,18 +212,24 @@ fn push_view( mod tests { use std::sync::LazyLock; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::VarBinViewArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; use vortex_error::VortexResult; use vortex_mask::Mask; use vortex_session::VortexSession; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::arrays::VarBinViewArray; - use crate::builtins::ArrayBuiltins; - use crate::dtype::DType; - use crate::dtype::Nullability; + use crate::builtins::ArrayFunctions; - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + + crate::initialize(&session); + + session + }); #[test] fn zip_varbinview_kernel_zips() -> VortexResult<()> { diff --git a/vortex-jni/src/expression.rs b/vortex-jni/src/expression.rs index afcefe990c4..02e2d631a68 100644 --- a/vortex-jni/src/expression.rs +++ b/vortex-jni/src/expression.rs @@ -57,21 +57,21 @@ use vortex::extension::datetime::Date; use vortex::extension::datetime::Time; use vortex::extension::datetime::TimeUnit; use vortex::extension::datetime::Timestamp; +use vortex::functions::between::BetweenOptions; +use vortex::functions::between::StrictComparison; +use vortex::functions::like::Like; +use vortex::functions::like::LikeOptions; +use vortex::functions::list_contains::ListContainsOptions; +use vortex::functions::merge::DuplicateHandling; use vortex::layout::layouts::row_idx::row_idx; use vortex::scalar::DecimalValue; use vortex::scalar::Scalar; use vortex::scalar::ScalarValue; use vortex::scalar_fn::EmptyOptions; use vortex::scalar_fn::ScalarFnVTableExt; -use vortex::scalar_fn::fns::between::BetweenOptions; -use vortex::scalar_fn::fns::between::StrictComparison; -use vortex::scalar_fn::fns::binary::Binary; -use vortex::scalar_fn::fns::like::Like; -use vortex::scalar_fn::fns::like::LikeOptions; -use vortex::scalar_fn::fns::list_contains::ListContainsOptions; -use vortex::scalar_fn::fns::literal::Literal; -use vortex::scalar_fn::fns::merge::DuplicateHandling; -use vortex::scalar_fn::fns::operators::Operator; +use vortex::scalar_fn::builtin::binary::Binary; +use vortex::scalar_fn::builtin::literal::Literal; +use vortex::scalar_fn::builtin::operators::Operator; use vortex::session::VortexSession; use vortex_spatial::extension::native_geometry_scalar_from_wkb; use vortex_spatial::scalar_fn::area::SpatialArea; diff --git a/vortex-json/Cargo.toml b/vortex-json/Cargo.toml index 27c8c86fea2..c8742dc41f7 100644 --- a/vortex-json/Cargo.toml +++ b/vortex-json/Cargo.toml @@ -23,6 +23,7 @@ vortex-array = { workspace = true, default-features = false } vortex-arrow = { workspace = true } vortex-edition = { workspace = true } vortex-error = { workspace = true, default-features = false } +vortex-functions = { workspace = true } vortex-session = { workspace = true } [dev-dependencies] diff --git a/vortex-json/src/json_to_variant.rs b/vortex-json/src/json_to_variant.rs index 32538b7ad3d..634c1c333ba 100644 --- a/vortex-json/src/json_to_variant.rs +++ b/vortex-json/src/json_to_variant.rs @@ -23,12 +23,12 @@ use vortex_array::scalar_fn::ExecutionArgs; use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::ScalarFnVTableExt; -use vortex_array::scalar_fn::fns::variant_get::VariantPath; -use vortex_array::scalar_fn::fns::variant_get::VariantPathElement; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_err; +use vortex_functions::variant_get::VariantPath; +use vortex_functions::variant_get::VariantPathElement; use vortex_session::VortexSession; use vortex_session::registry::CachedId; diff --git a/vortex-layout/Cargo.toml b/vortex-layout/Cargo.toml index 81cd54e9a38..a5dc9506093 100644 --- a/vortex-layout/Cargo.toml +++ b/vortex-layout/Cargo.toml @@ -48,6 +48,7 @@ vortex-btrblocks = { workspace = true } vortex-buffer = { workspace = true } vortex-error = { workspace = true } vortex-expr = { workspace = true } +vortex-functions = { workspace = true } vortex-io = { workspace = true } vortex-mask = { workspace = true } vortex-metrics = { workspace = true } diff --git a/vortex-layout/examples/reverse_index/main.rs b/vortex-layout/examples/reverse_index/main.rs index a03a16d2d20..dd5cb7e3bea 100644 --- a/vortex-layout/examples/reverse_index/main.rs +++ b/vortex-layout/examples/reverse_index/main.rs @@ -45,9 +45,9 @@ use vortex_array::expr::lit; use vortex_array::memory::BufferAllocatorRef; use vortex_array::memory::MemorySessionExt; use vortex_array::scalar::Scalar; -use vortex_array::scalar_fn::fns::binary::Binary; -use vortex_array::scalar_fn::fns::literal::Literal; -use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::scalar_fn::builtin::binary::Binary; +use vortex_array::scalar_fn::builtin::literal::Literal; +use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_array::search_sorted::SearchSorted; use vortex_array::search_sorted::SearchSortedArray; use vortex_array::search_sorted::SearchSortedSide; diff --git a/vortex-layout/src/layouts/dict/reader.rs b/vortex-layout/src/layouts/dict/reader.rs index 0beab32f938..447c0794aab 100644 --- a/vortex-layout/src/layouts/dict/reader.rs +++ b/vortex-layout/src/layouts/dict/reader.rs @@ -24,13 +24,13 @@ use vortex_array::expr::ExactBoundExpr; use vortex_array::expr::bound::pack as bound_pack; use vortex_array::expr::root; use vortex_array::optimizer::ArrayOptimizer; -use vortex_array::scalar_fn::is_negative_cost; use vortex_error::VortexError; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_expr::direct_bound_annotations; use vortex_expr::label_bound_tree; use vortex_expr::transform::partition_bound_annotations; +use vortex_functions::is_negative_cost; use vortex_mask::Mask; use vortex_session::VortexSession; use vortex_utils::aliases::dash_map::DashMap; @@ -328,6 +328,7 @@ impl LayoutReader for DictReader { self.values_array() }; let all_values_referenced = self.layout.has_all_values_referenced(); + let session = self.session.clone(); Ok(async move { let (values, codes) = try_join!(values_eval.map_err(VortexError::from), codes_eval)?; @@ -341,7 +342,7 @@ impl LayoutReader for DictReader { .set_all_values_referenced(all_values_referenced) } .into_array() - .optimize()?; + .optimize(&session)?; array.apply_bound(&expr_outer) } @@ -377,18 +378,18 @@ mod tests { use vortex_array::expr::BoundExpression; use vortex_array::expr::Expression; use vortex_array::expr::bound::pack as bound_pack; - use vortex_array::expr::byte_length; use vortex_array::expr::cast; use vortex_array::expr::eq; use vortex_array::expr::get_item; use vortex_array::expr::is_not_null; - use vortex_array::expr::like; use vortex_array::expr::lit; use vortex_array::expr::pack; use vortex_array::expr::root; use vortex_array::validity::Validity; use vortex_btrblocks::BtrBlocksCompressorBuilder; use vortex_error::VortexResult; + use vortex_functions::expr::byte_length; + use vortex_functions::expr::like; use vortex_io::runtime::Handle; use vortex_io::runtime::single::block_on; use vortex_io::session::RuntimeSession; diff --git a/vortex-layout/src/layouts/dict/writer.rs b/vortex-layout/src/layouts/dict/writer.rs index 7808777a8c2..1fae93156fd 100644 --- a/vortex-layout/src/layouts/dict/writer.rs +++ b/vortex-layout/src/layouts/dict/writer.rs @@ -30,6 +30,7 @@ use vortex_array::builders::dict::dict_encoder_in; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; +use vortex_array::optimizer::ArrayOptimizer; use vortex_error::VortexError; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -570,15 +571,24 @@ fn encode_chunk( ctx: &mut ExecutionCtx, ) -> VortexResult { let encoded = encoder.encode(chunk, ctx)?.into_array(); - match remainder(chunk, encoded.len())? { + match remainder(chunk, encoded.len(), ctx)? { None => Ok(EncodingState::Continue((encoder, encoded))), Some(unencoded) => Ok(EncodingState::Done((encoder.reset(), encoded, unencoded))), } } -fn remainder(array: &ArrayRef, encoded_len: usize) -> VortexResult> { +fn remainder( + array: &ArrayRef, + encoded_len: usize, + ctx: &ExecutionCtx, +) -> VortexResult> { if encoded_len < array.len() { - Ok(Some(array.slice(encoded_len..array.len())?)) + // The dictionary encoder takes concrete encodings, so optimize the lazy slice into one. + Ok(Some( + array + .slice(encoded_len..array.len())? + .optimize(ctx.session())?, + )) } else { Ok(None) } diff --git a/vortex-layout/src/layouts/indexed/tests.rs b/vortex-layout/src/layouts/indexed/tests.rs index f2f28b54509..e8fdee1ec73 100644 --- a/vortex-layout/src/layouts/indexed/tests.rs +++ b/vortex-layout/src/layouts/indexed/tests.rs @@ -28,7 +28,6 @@ use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::expr::BoundExpression; use vortex_array::expr::eq; -use vortex_array::expr::like; use vortex_array::expr::lit; use vortex_array::expr::root; use vortex_array::stream::ArrayStreamExt; @@ -36,6 +35,7 @@ use vortex_array::test_harness::check_metadata; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; +use vortex_functions::expr::like; use vortex_mask::Mask; use vortex_session::VortexSession; @@ -1092,9 +1092,9 @@ mod exact_value { use vortex_array::expr::col; use vortex_array::expr::eq; use vortex_array::expr::lit; - use vortex_array::scalar_fn::fns::binary::Binary; - use vortex_array::scalar_fn::fns::literal::Literal; - use vortex_array::scalar_fn::fns::operators::Operator; + use vortex_array::scalar_fn::builtin::binary::Binary; + use vortex_array::scalar_fn::builtin::literal::Literal; + use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_array::validity::Validity; use vortex_error::VortexResult; use vortex_error::vortex_bail; diff --git a/vortex-layout/src/layouts/list/expr.rs b/vortex-layout/src/layouts/list/expr.rs index 8c10e400ad8..d6d795ee80d 100644 --- a/vortex-layout/src/layouts/list/expr.rs +++ b/vortex-layout/src/layouts/list/expr.rs @@ -5,10 +5,10 @@ use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::expr::BoundExpression; use vortex_array::expr::bound::not; -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::builtin::is_not_null::IsNotNull; +use vortex_array::scalar_fn::builtin::is_null::IsNull; use vortex_error::VortexResult; +use vortex_functions::list_length::ListLength; /// The minimal set of list children an expression needs for evaluation. /// diff --git a/vortex-layout/src/layouts/list/reader.rs b/vortex-layout/src/layouts/list/reader.rs index fafa48346bd..0f96f72e845 100644 --- a/vortex-layout/src/layouts/list/reader.rs +++ b/vortex-layout/src/layouts/list/reader.rs @@ -21,7 +21,7 @@ use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::expr::BoundExpression; use vortex_array::expr::root; -use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_array::validity::Validity; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -631,12 +631,11 @@ mod tests { 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_buffer::buffer; + use vortex_functions::expr::list_length; use vortex_io::session::RuntimeSession; use vortex_io::session::RuntimeSessionExt; - use vortex_session::VortexSession; use super::*; use crate::LayoutRef; diff --git a/vortex-layout/src/layouts/list/writer.rs b/vortex-layout/src/layouts/list/writer.rs index 8e3a25cf407..d599b2e6bf3 100644 --- a/vortex-layout/src/layouts/list/writer.rs +++ b/vortex-layout/src/layouts/list/writer.rs @@ -22,7 +22,7 @@ use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::matcher::Matcher; -use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; diff --git a/vortex-layout/src/layouts/row_idx/mod.rs b/vortex-layout/src/layouts/row_idx/mod.rs index f9641cbd6cb..633cc1c771b 100644 --- a/vortex-layout/src/layouts/row_idx/mod.rs +++ b/vortex-layout/src/layouts/row_idx/mod.rs @@ -335,8 +335,9 @@ fn row_idx_array_future( ) -> ArrayFuture { let row_range = row_range.clone(); async move { - // The filter is lazy: trivial masks reduce to the sequence itself or an empty array, and - // value masks are left for the consumer to execute alongside the expression. + // The filter is lazy: when the consumer optimizes or executes it, trivial masks reduce to + // the sequence itself or an empty array, and value masks execute alongside the + // expression. idx_array(row_offset, &row_range) .into_array() .filter(mask.await?)? @@ -363,6 +364,7 @@ mod tests { use vortex_array::expr::lit; use vortex_array::expr::or; use vortex_array::expr::root; + use vortex_array::optimizer::ArrayOptimizer; use vortex_buffer::buffer; use vortex_io::runtime::single::block_on; use vortex_io::session::RuntimeSessionExt; @@ -538,6 +540,8 @@ mod tests { MaskFuture::ready(Mask::new_true(5)), ) .await + .unwrap() + .optimize(&new_session()) .unwrap(); assert!(result.is::()); diff --git a/vortex-layout/src/layouts/struct_/reader.rs b/vortex-layout/src/layouts/struct_/reader.rs index 9d9630473c8..0902458f4c0 100644 --- a/vortex-layout/src/layouts/struct_/reader.rs +++ b/vortex-layout/src/layouts/struct_/reader.rs @@ -29,13 +29,11 @@ use vortex_array::expr::bound::pack; use vortex_array::expr::traversal::NodeExt; use vortex_array::expr::traversal::Transformed; use vortex_array::expr::traversal::TraversalOrder; -use vortex_array::scalar_fn::fns::get_item::GetItem; -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::mask::Mask as MaskFn; -use vortex_array::scalar_fn::fns::pack::Pack; -use vortex_array::scalar_fn::fns::select::FieldSelection; -use vortex_array::scalar_fn::fns::select::Select; +use vortex_array::scalar_fn::builtin::get_item::GetItem; +use vortex_array::scalar_fn::builtin::is_not_null::IsNotNull; +use vortex_array::scalar_fn::builtin::is_null::IsNull; +use vortex_array::scalar_fn::builtin::mask::Mask as MaskFn; +use vortex_array::scalar_fn::builtin::pack::Pack; use vortex_array::validity::Validity; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -45,6 +43,8 @@ use vortex_expr::analysis::AnnotationFn; use vortex_expr::make_bound_free_field_annotator; use vortex_expr::transform::BoundPartitionedExpr; use vortex_expr::transform::partition_bound; +use vortex_functions::select::FieldSelection; +use vortex_functions::select::Select; use vortex_mask::Mask; use vortex_session::VortexSession; use vortex_utils::aliases::dash_map::DashMap; @@ -731,12 +731,12 @@ mod tests { use vortex_array::expr::or; use vortex_array::expr::pack; use vortex_array::expr::root; - use vortex_array::expr::select; - use vortex_array::expr::select_exclude; use vortex_array::scalar::Scalar; use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; + use vortex_functions::expr::select; + use vortex_functions::expr::select_exclude; use vortex_io::runtime::single::block_on; use vortex_io::session::RuntimeSessionExt; use vortex_mask::Mask; diff --git a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/scalar_fn.rs b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/scalar_fn.rs index fc8f91627cb..1de5d1116e5 100644 --- a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/scalar_fn.rs +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/scalar_fn.rs @@ -32,9 +32,9 @@ use vortex_array::scalar_fn::ExecutionArgs; use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::ScalarFnVTableExt; -use vortex_array::scalar_fn::fns::binary::Binary; -use vortex_array::scalar_fn::fns::literal::Literal; -use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::scalar_fn::builtin::binary::Binary; +use vortex_array::scalar_fn::builtin::literal::Literal; +use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_buffer::BitBufferMut; use vortex_buffer::Buffer; use vortex_error::VortexExpect; diff --git a/vortex-layout/src/layouts/zoned/pruning.rs b/vortex-layout/src/layouts/zoned/pruning.rs index 6e39823689e..d508b2f0b49 100644 --- a/vortex-layout/src/layouts/zoned/pruning.rs +++ b/vortex-layout/src/layouts/zoned/pruning.rs @@ -21,11 +21,11 @@ use vortex_array::dtype::DType; use vortex_array::expr::BoundExpression; use vortex_array::expr::ExactBoundExpr; use vortex_array::expr::root; -use vortex_array::scalar_fn::fns::dynamic::DynamicExprUpdates; use vortex_error::SharedVortexResult; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_expr::stats::rewrite::BoundExpressionPruneExt; +use vortex_functions::dynamic::DynamicExprUpdates; use vortex_mask::Mask; use vortex_session::VortexSession; use vortex_utils::aliases::dash_map::DashMap; diff --git a/vortex-layout/src/layouts/zoned/zone_map.rs b/vortex-layout/src/layouts/zoned/zone_map.rs index 2b9011e330c..adffc62b25c 100644 --- a/vortex-layout/src/layouts/zoned/zone_map.rs +++ b/vortex-layout/src/layouts/zoned/zone_map.rs @@ -426,7 +426,6 @@ mod tests { use vortex_array::expr::BoundExpression; use vortex_array::expr::Expression; use vortex_array::expr::cast; - use vortex_array::expr::dynamic; use vortex_array::expr::gt; use vortex_array::expr::gt_eq; use vortex_array::expr::is_not_null; @@ -437,7 +436,7 @@ mod tests { use vortex_array::expr::root; use vortex_array::expr::stats::Stat; use vortex_array::scalar::ScalarValue; - use vortex_array::scalar_fn::fns::operators::CompareOperator; + use vortex_array::scalar_fn::builtin::operators::CompareOperator; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::buffer; @@ -447,6 +446,7 @@ mod tests { use vortex_expr::stats::all_non_null; use vortex_expr::stats::all_null; use vortex_expr::stats::rewrite::BoundExpressionPruneExt; + use vortex_functions::expr::dynamic; use vortex_mask::Mask; use crate::layouts::zoned::zone_map::ZoneMap; diff --git a/vortex-layout/src/plan/physical/plans/eval.rs b/vortex-layout/src/plan/physical/plans/eval.rs index b6a2f1ff620..93390c14860 100644 --- a/vortex-layout/src/plan/physical/plans/eval.rs +++ b/vortex-layout/src/plan/physical/plans/eval.rs @@ -12,7 +12,7 @@ use vortex_array::expr::BoundExpression; use vortex_array::expr::traversal::NodeExt; use vortex_array::expr::traversal::Transformed; use vortex_array::expr::traversal::TraversalOrder; -use vortex_array::scalar_fn::fns::pack::Pack; +use vortex_array::scalar_fn::builtin::pack::Pack; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; diff --git a/vortex-layout/src/plan/physical/plans/pack.rs b/vortex-layout/src/plan/physical/plans/pack.rs index b37f50934e8..43c09ef3f0d 100644 --- a/vortex-layout/src/plan/physical/plans/pack.rs +++ b/vortex-layout/src/plan/physical/plans/pack.rs @@ -16,10 +16,9 @@ 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::pack::Pack as PackFn; -use vortex_array::scalar_fn::fns::pack::PackOptions; -use vortex_array::scalar_fn::fns::select::Select; +use vortex_array::scalar_fn::builtin::get_item::GetItem; +use vortex_array::scalar_fn::builtin::pack::Pack as PackFn; +use vortex_array::scalar_fn::builtin::pack::PackOptions; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -28,6 +27,7 @@ use vortex_error::vortex_err; use vortex_expr::descendent_bound_annotations; use vortex_expr::make_bound_free_field_annotator; use vortex_expr::transform::partition_bound; +use vortex_functions::select::Select; use vortex_mask::Mask; use vortex_session::registry::CachedId; diff --git a/vortex-layout/src/plan/physical/plans/row_idx.rs b/vortex-layout/src/plan/physical/plans/row_idx.rs index 5c45d2cae4f..cc2066ff324 100644 --- a/vortex-layout/src/plan/physical/plans/row_idx.rs +++ b/vortex-layout/src/plan/physical/plans/row_idx.rs @@ -16,7 +16,7 @@ use vortex_array::expr::BoundExpression; use vortex_array::expr::traversal::NodeExt; use vortex_array::expr::traversal::Transformed; use vortex_array::expr::traversal::TraversalOrder; -use vortex_array::scalar_fn::fns::pack::Pack as PackFn; +use vortex_array::scalar_fn::builtin::pack::Pack as PackFn; use vortex_buffer::Buffer; use vortex_error::VortexResult; use vortex_error::vortex_ensure_eq; diff --git a/vortex-layout/src/plan/physical/plans/take.rs b/vortex-layout/src/plan/physical/plans/take.rs index b2a1703ff8a..0fd2faf8c3b 100644 --- a/vortex-layout/src/plan/physical/plans/take.rs +++ b/vortex-layout/src/plan/physical/plans/take.rs @@ -10,11 +10,11 @@ use vortex_array::IntoArray; use vortex_array::dtype::DType; use vortex_array::expr::BoundExpression; use vortex_array::expr::ExactBoundExpr; -use vortex_array::scalar_fn::fns::dynamic::DynamicComparison; -use vortex_array::scalar_fn::is_negative_cost; use vortex_error::VortexResult; use vortex_expr::BoundLabels; use vortex_expr::label_bound_tree; +use vortex_functions::dynamic::DynamicComparison; +use vortex_functions::is_negative_cost; use vortex_mask::Mask; use vortex_session::registry::CachedId; diff --git a/vortex-layout/src/plan/physical/plans/zoned.rs b/vortex-layout/src/plan/physical/plans/zoned.rs index d1a17702ebf..fdfc108540d 100644 --- a/vortex-layout/src/plan/physical/plans/zoned.rs +++ b/vortex-layout/src/plan/physical/plans/zoned.rs @@ -21,14 +21,14 @@ use vortex_array::expr::traversal::NodeExt; use vortex_array::expr::traversal::Transformed; use vortex_array::expr::traversal::TraversalOrder; use vortex_array::scalar::Scalar; -use vortex_array::scalar_fn::fns::dynamic::DynamicExprUpdates; -use vortex_array::scalar_fn::fns::stat::StatFn; use vortex_array::validity::Validity; use vortex_buffer::BitBufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; +use vortex_functions::dynamic::DynamicExprUpdates; +use vortex_functions::stat::StatFn; use vortex_mask::Mask; use vortex_session::VortexSession; use vortex_session::registry::CachedId; diff --git a/vortex-layout/src/plan/physical/tests.rs b/vortex-layout/src/plan/physical/tests.rs index 47bad2dd94f..1abcfa5e2de 100644 --- a/vortex-layout/src/plan/physical/tests.rs +++ b/vortex-layout/src/plan/physical/tests.rs @@ -21,10 +21,8 @@ use vortex_array::dtype::StructFields; use vortex_array::expr::Expression; use vortex_array::expr::and; use vortex_array::expr::bound::and as bound_and; -use vortex_array::expr::byte_length; use vortex_array::expr::cast; use vortex_array::expr::checked_add; -use vortex_array::expr::dynamic; use vortex_array::expr::get_item; use vortex_array::expr::gt; use vortex_array::expr::is_null; @@ -33,11 +31,13 @@ use vortex_array::expr::pack; use vortex_array::expr::root; use vortex_array::expr::stats::Stat; use vortex_array::scalar::ScalarValue; -use vortex_array::scalar_fn::fns::operators::CompareOperator; +use vortex_array::scalar_fn::builtin::operators::CompareOperator; use vortex_array::stats::as_stat_bitset_bytes; use vortex_error::VortexResult; use vortex_error::vortex_err; use vortex_expr::stats::rewrite::BoundExpressionPruneExt; +use vortex_functions::expr::byte_length; +use vortex_functions::expr::dynamic; use vortex_mask::Mask; use vortex_session::VortexSession; use vortex_session::registry::CachedId; diff --git a/vortex-layout/src/scan/filter.rs b/vortex-layout/src/scan/filter.rs index 2cb36068913..1cc2cfee482 100644 --- a/vortex-layout/src/scan/filter.rs +++ b/vortex-layout/src/scan/filter.rs @@ -8,12 +8,12 @@ use itertools::Itertools; use parking_lot::RwLock; use sketches_ddsketch::DDSketch; use vortex_array::expr::BoundExpression; -use vortex_array::scalar_fn::fns::binary::Binary; -use vortex_array::scalar_fn::fns::dynamic::DynamicExprUpdates; -use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::scalar_fn::builtin::binary::Binary; +use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_error::VortexExpect; use vortex_error::vortex_err; use vortex_error::vortex_panic; +use vortex_functions::dynamic::DynamicExprUpdates; /// The selectivity histogram quantile to use for reordering conjuncts. Where 0 == no rows match. const DEFAULT_SELECTIVITY_QUANTILE: f64 = 0.1; diff --git a/vortex-layout/src/scan/scan_builder.rs b/vortex-layout/src/scan/scan_builder.rs index 21dc22c04ca..95479388f0a 100644 --- a/vortex-layout/src/scan/scan_builder.rs +++ b/vortex-layout/src/scan/scan_builder.rs @@ -494,7 +494,6 @@ mod test { use vortex_array::expr::BoundExpression; use vortex_array::expr::ExactBoundExpr; use vortex_array::expr::and; - use vortex_array::expr::between; use vortex_array::expr::col; use vortex_array::expr::eq; use vortex_array::expr::get_item; @@ -503,11 +502,12 @@ mod test { use vortex_array::expr::lit; use vortex_array::expr::lt; use vortex_array::expr::root; - use vortex_array::scalar_fn::fns::between::BetweenOptions; - use vortex_array::scalar_fn::fns::between::StrictComparison; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; + use vortex_functions::between::BetweenOptions; + use vortex_functions::between::StrictComparison; + use vortex_functions::expr::between; use vortex_io::runtime::BlockingRuntime; use vortex_io::runtime::single::SingleThreadRuntime; use vortex_io::runtime::tokio::TokioRuntime; diff --git a/vortex-python/src/arrays/mod.rs b/vortex-python/src/arrays/mod.rs index 842c6c694ae..b84332114ca 100644 --- a/vortex-python/src/arrays/mod.rs +++ b/vortex-python/src/arrays/mod.rs @@ -52,7 +52,7 @@ use vortex::error::VortexResult; use vortex::flatbuffers::WriteFlatBufferExt; use vortex::ipc::messages::EncoderMessage; use vortex::ipc::messages::MessageEncoder; -use vortex::scalar_fn::fns::operators::Operator; +use vortex::scalar_fn::builtin::operators::Operator; use vortex_arrow::ArrowSessionExt; use vortex_python_abi::BUFFER_EXPORT_CAPSULE_NAME; use vortex_python_abi::VORTEX_BUFFER_EXPORT_VERSION; diff --git a/vortex-python/src/expr/mod.rs b/vortex-python/src/expr/mod.rs index b126ba39343..5cb5b305d93 100644 --- a/vortex-python/src/expr/mod.rs +++ b/vortex-python/src/expr/mod.rs @@ -17,17 +17,17 @@ use vortex::expr; use vortex::expr::Expression; use vortex::expr::lit; use vortex::expr::proto::ExprSerializeProtoExt; +use vortex::functions::between::BetweenOptions; +use vortex::functions::between::StrictComparison; +use vortex::functions::list_contains::ListContainsOptions; +use vortex::functions::merge::DuplicateHandling; +use vortex::functions::replace_time_zone::ReplaceTimeZoneOptions; +use vortex::functions::variant_get::VariantPath; +use vortex::functions::variant_get::VariantPathElement; use vortex::proto::expr as pb; use vortex::scalar_fn::ScalarFnVTableExt; -use vortex::scalar_fn::fns::between::BetweenOptions; -use vortex::scalar_fn::fns::between::StrictComparison; -use vortex::scalar_fn::fns::binary::Binary; -use vortex::scalar_fn::fns::list_contains::ListContainsOptions; -use vortex::scalar_fn::fns::merge::DuplicateHandling; -use vortex::scalar_fn::fns::operators::Operator; -use vortex::scalar_fn::fns::replace_time_zone::ReplaceTimeZoneOptions; -use vortex::scalar_fn::fns::variant_get::VariantPath; -use vortex::scalar_fn::fns::variant_get::VariantPathElement; +use vortex::scalar_fn::builtin::binary::Binary; +use vortex::scalar_fn::builtin::operators::Operator; use crate::dtype::PyDType; use crate::error::PyVortexResult; diff --git a/vortex-spatial/Cargo.toml b/vortex-spatial/Cargo.toml index b48831cf879..bd604217afd 100644 --- a/vortex-spatial/Cargo.toml +++ b/vortex-spatial/Cargo.toml @@ -26,6 +26,7 @@ vortex-buffer = { workspace = true } vortex-edition = { workspace = true } vortex-error = { workspace = true } vortex-expr = { workspace = true } +vortex-functions = { workspace = true } vortex-mask = { workspace = true } vortex-session = { workspace = true } wkb = { workspace = true } diff --git a/vortex-spatial/src/prune/distance.rs b/vortex-spatial/src/prune/distance.rs index 7417a23d042..0713adf0434 100644 --- a/vortex-spatial/src/prune/distance.rs +++ b/vortex-spatial/src/prune/distance.rs @@ -7,9 +7,9 @@ use geo::Rect as SpatialRect; use vortex_array::expr::BoundExpression; use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::ScalarFnVTable; -use vortex_array::scalar_fn::fns::binary::Binary; -use vortex_array::scalar_fn::fns::literal::Literal; -use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::scalar_fn::builtin::binary::Binary; +use vortex_array::scalar_fn::builtin::literal::Literal; +use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_error::VortexResult; use vortex_expr::stats::rewrite::StatsRewriteRule; use vortex_session::VortexSession; @@ -138,8 +138,8 @@ mod tests { use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::EmptyOptions; use vortex_array::scalar_fn::ScalarFnVTableExt; - use vortex_array::scalar_fn::fns::binary::Binary; - use vortex_array::scalar_fn::fns::operators::Operator; + use vortex_array::scalar_fn::builtin::binary::Binary; + use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_error::VortexResult; use vortex_expr::stats::rewrite::BoundExpressionPruneExt; use vortex_expr::stats::rewrite::StatsRewriteRule; diff --git a/vortex-spatial/src/prune/mod.rs b/vortex-spatial/src/prune/mod.rs index dd1a6951503..50aea223a33 100644 --- a/vortex-spatial/src/prune/mod.rs +++ b/vortex-spatial/src/prune/mod.rs @@ -24,9 +24,7 @@ use vortex_array::aggregate_fn::AggregateFnVTableExt; use vortex_array::aggregate_fn::EmptyOptions; use vortex_array::expr::BoundExpression; use vortex_array::expr::bound::binary; -use vortex_array::expr::bound::case_when; use vortex_array::expr::bound::checked_add; -use vortex_array::expr::bound::ext_storage; use vortex_array::expr::bound::get_item; use vortex_array::expr::bound::gt; use vortex_array::expr::bound::gt_eq; @@ -34,10 +32,12 @@ use vortex_array::expr::bound::lit; use vortex_array::expr::bound::lt; use vortex_array::expr::bound::lt_eq; use vortex_array::scalar::Scalar; -use vortex_array::scalar_fn::fns::literal::Literal; -use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::scalar_fn::builtin::literal::Literal; +use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_error::VortexResult; use vortex_expr::stats::bound::stat; +use vortex_functions::expr::bound::case_when; +use vortex_functions::expr::bound::ext_storage; use vortex_session::VortexSession; use crate::aggregate_fn::GeometryAabb; diff --git a/vortex-tensor/src/vector_search.rs b/vortex-tensor/src/vector_search.rs index 3603691a1c6..f436e764734 100644 --- a/vortex-tensor/src/vector_search.rs +++ b/vortex-tensor/src/vector_search.rs @@ -40,7 +40,7 @@ use vortex_array::dtype::NativePType; use vortex_array::dtype::Nullability; use vortex_array::scalar::PValue; use vortex_array::scalar::Scalar; -use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::scalar_fn::builtin::operators::Operator; use vortex_error::VortexResult; use crate::scalar_fns::cosine_similarity::CosineSimilarity; diff --git a/vortex/Cargo.toml b/vortex/Cargo.toml index cff73ed08ae..5fd973db292 100644 --- a/vortex/Cargo.toml +++ b/vortex/Cargo.toml @@ -34,6 +34,7 @@ vortex-expr = { workspace = true } vortex-fastlanes = { workspace = true } vortex-file = { workspace = true, optional = true } vortex-fsst = { workspace = true } +vortex-functions = { workspace = true } vortex-io = { workspace = true } vortex-ipc = { workspace = true } vortex-layout = { workspace = true } diff --git a/vortex/src/lib.rs b/vortex/src/lib.rs index 8cc8ee52c1c..6a62e3cde0e 100644 --- a/vortex/src/lib.rs +++ b/vortex/src/lib.rs @@ -111,12 +111,19 @@ pub mod expr { //! Expression trees, their optimizer, and the analyses and rewrites built on them. pub use vortex_array::expr::*; pub use vortex_expr::*; + pub use vortex_functions::expr::*; /// Statistics, stat expressions, and statistics-based pruning rewrites. pub mod stats { pub use vortex_array::expr::stats::*; pub use vortex_expr::stats::*; } + + /// Constructors for expressions whose children have already been bound and type-checked. + pub mod bound { + pub use vortex_array::expr::bound::*; + pub use vortex_functions::expr::bound::*; + } } use vortex_array::memory::MemorySession; use vortex_array::optimizer::kernels::KernelSession; @@ -125,6 +132,7 @@ use vortex_array::scalar_fn::session::ScalarFnSession; use vortex_array::session::ArraySession; use vortex_btrblocks::CompressionSession; use vortex_expr::stats::session::StatsSession; +pub use vortex_functions as functions; use vortex_io::session::RuntimeSession; use vortex_layout::session::LayoutSession; use vortex_session::VortexSession; @@ -346,6 +354,7 @@ impl VortexSessionDefault for VortexSession { .with::() .with::() .with::(); + vortex_functions::initialize(&session); vortex_arrow::initialize(&session); #[cfg(feature = "parquet-variant")] vortex_parquet_variant::initialize(&session); @@ -386,7 +395,6 @@ mod test { use vortex_array::expr::gt; use vortex_array::expr::lit; use vortex_array::expr::root; - use vortex_array::expr::select; use vortex_array::stream::ArrayStreamExt; use vortex_array::validity::Validity; use vortex_btrblocks::BtrBlocksCompressorBuilder; @@ -395,6 +403,7 @@ mod test { use vortex_file::OpenOptionsSessionExt; use vortex_file::WriteOptionsSessionExt; use vortex_file::WriteStrategyBuilder; + use vortex_functions::expr::select; use vortex_session::VortexSession; use crate as vortex;