From a263f024894bf4c89c5db82961e076ac6e6952ae Mon Sep 17 00:00:00 2001 From: Luke Kim <80174+lukekim@users.noreply.github.com> Date: Wed, 14 Jan 2026 16:01:51 -0800 Subject: [PATCH 01/32] fix: Add log dependency to vortex-io --- vortex-io/Cargo.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/vortex-io/Cargo.toml b/vortex-io/Cargo.toml index 7e57683b0b0..6d954acc054 100644 --- a/vortex-io/Cargo.toml +++ b/vortex-io/Cargo.toml @@ -38,6 +38,7 @@ tokio = { workspace = true, features = [ tracing = { workspace = true } vortex-array = { workspace = true } vortex-buffer = { workspace = true } +log = { workspace = true } vortex-error = { workspace = true } vortex-metrics = { workspace = true } vortex-session = { workspace = true } From 9b4bde99dc51a6b9b2f64ac3cf833dfe0a266928 Mon Sep 17 00:00:00 2001 From: Luke Kim <80174+lukekim@users.noreply.github.com> Date: Thu, 26 Feb 2026 15:59:09 -0800 Subject: [PATCH 02/32] feat: WriterStrategyBuilder with Clone, Debug, and VortexSession support Ported from spiceai-51 patches: - Add Clone and Debug derives to WriteStrategyBuilder (#7, #4) - Add Debug bound to CompressorPlugin trait (#7) - Add VortexSession::set() method (#4) - Use session WriteStrategyBuilder in write_options() if available (#4) --- Cargo.lock | 1 + vortex-file/src/strategy.rs | 14 +++++++++++++ vortex-file/src/writer.rs | 27 ++++++++++++++++++++++--- vortex-layout/src/layouts/compressed.rs | 4 ++-- vortex-session/src/lib.rs | 20 ++++++++++++++++++ 5 files changed, 61 insertions(+), 5 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index c22c03572b3..0a768fa56d6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10445,6 +10445,7 @@ dependencies = [ "handle", "itertools 0.14.0", "kanal", + "log", "object_store", "oneshot", "parking_lot", diff --git a/vortex-file/src/strategy.rs b/vortex-file/src/strategy.rs index a9d805d5447..abd273f936d 100644 --- a/vortex-file/src/strategy.rs +++ b/vortex-file/src/strategy.rs @@ -123,6 +123,7 @@ pub static ALLOWED_ENCODINGS: LazyLock = LazyLock::new(|| { /// Vortex provides an out-of-the-box file writer that optimizes the layout of chunks on-disk, /// repartitioning and compressing them to strike a balance between size on-disk, /// bulk decoding performance, and IOPS required to perform an indexed read. +#[derive(Clone)] pub struct WriteStrategyBuilder { compressor: Option>, row_block_size: usize, @@ -131,6 +132,19 @@ pub struct WriteStrategyBuilder { flat_strategy: Option>, } +impl std::fmt::Debug for WriteStrategyBuilder { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("WriteStrategyBuilder") + .field("compressor", &self.compressor) + .field("row_block_size", &self.row_block_size) + .field( + "field_writers", + &format!("<{} entries>", self.field_writers.len()), + ) + .finish() + } +} + impl Default for WriteStrategyBuilder { /// Create a new empty builder. It can be further configured, /// and then finally built yielding the [`LayoutStrategy`]. diff --git a/vortex-file/src/writer.rs b/vortex-file/src/writer.rs index 6bc6aef2dda..0af1b129ee2 100644 --- a/vortex-file/src/writer.rs +++ b/vortex-file/src/writer.rs @@ -69,10 +69,14 @@ pub struct VortexWriteOptions { pub trait WriteOptionsSessionExt: SessionExt { /// Create [`VortexWriteOptions`] for writing to a Vortex file. fn write_options(&self) -> VortexWriteOptions { - let session = self.session(); + let maybe_write_strategy_builder = self.get_opt::(); + let strategy = maybe_write_strategy_builder + .map(|opt| opt.clone().build()) + .unwrap_or_else(|| WriteStrategyBuilder::default().build()); + VortexWriteOptions { - strategy: WriteStrategyBuilder::default().build(), - session, + session: self.session(), + strategy, exclude_dtype: false, file_statistics: PRUNING_STATS.to_vec(), max_variable_length_statistics_size: 64, @@ -470,3 +474,20 @@ impl WriteSummary { self.footer.row_count() } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_write_options_from_session_vars() { + let session = VortexSession::empty(); + let fetched_write_strategy = session.get_opt::(); + assert!(fetched_write_strategy.is_none()); + drop(fetched_write_strategy); + + let session = session.set(WriteStrategyBuilder::default()); + let fetched_write_strategy = session.get_opt::(); + assert!(fetched_write_strategy.is_some()); + } +} diff --git a/vortex-layout/src/layouts/compressed.rs b/vortex-layout/src/layouts/compressed.rs index de672715d9e..91661e92c77 100644 --- a/vortex-layout/src/layouts/compressed.rs +++ b/vortex-layout/src/layouts/compressed.rs @@ -26,7 +26,7 @@ use crate::sequence::SequentialStreamExt; /// A boxed compressor function from arrays into compressed arrays. /// /// API consumers are free to implement this trait to provide new plugin compressors. -pub trait CompressorPlugin: Send + Sync + 'static { +pub trait CompressorPlugin: std::fmt::Debug + Send + Sync + 'static { fn compress_chunk(&self, chunk: &dyn Array) -> VortexResult; } @@ -38,7 +38,7 @@ impl CompressorPlugin for Arc { impl CompressorPlugin for F where - F: Fn(&dyn Array) -> VortexResult + Send + Sync + 'static, + F: Fn(&dyn Array) -> VortexResult + Send + Sync + 'static + std::fmt::Debug, { fn compress_chunk(&self, chunk: &dyn Array) -> VortexResult { self(chunk) diff --git a/vortex-session/src/lib.rs b/vortex-session/src/lib.rs index c53784173e0..d2112caf30c 100644 --- a/vortex-session/src/lib.rs +++ b/vortex-session/src/lib.rs @@ -52,6 +52,26 @@ impl VortexSession { } self } + + /// Inserts a new session variable of type `V` with the supplied value. + /// + /// # Panics + /// + /// If a variable of that type already exists. + pub fn set(self, val: V) -> Self { + match self.0.entry(TypeId::of::()) { + Entry::Occupied(_) => { + vortex_panic!( + "Session variable of type {} already exists", + type_name::() + ); + } + Entry::Vacant(e) => { + e.insert(Box::new(val)); + } + } + self + } } /// Trait for accessing and modifying the state of a Vortex session. From a9ef29deaeeca5f7db03ae28e34dc5358809cf82 Mon Sep 17 00:00:00 2001 From: Luke Kim <80174+lukekim@users.noreply.github.com> Date: Thu, 26 Feb 2026 16:00:26 -0800 Subject: [PATCH 03/32] fix: Handle UncompressedSizeInBytes statistic correctly Ported from spiceai-51 patch #3: - Add UncompressedSizeInBytes to PRUNING_STATS - Propagate UncompressedSizeInBytes and IsConstant through take stats --- vortex-array/src/arrays/dict/take.rs | 9 +++++++-- vortex-array/src/stats/mod.rs | 1 + 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/vortex-array/src/arrays/dict/take.rs b/vortex-array/src/arrays/dict/take.rs index f2517478068..6a6259f8ed3 100644 --- a/vortex-array/src/arrays/dict/take.rs +++ b/vortex-array/src/arrays/dict/take.rs @@ -150,8 +150,13 @@ pub(crate) fn propagate_take_stats( st.set(Stat::IsConstant, Precision::exact(true)); } } - let inexact_min_max = [Stat::Min, Stat::Max] - .into_iter() + let inexact_min_max = [ + Stat::Min, + Stat::Max, + Stat::UncompressedSizeInBytes, + Stat::IsConstant, + ] + .into_iter() .filter_map(|stat| { source .statistics() diff --git a/vortex-array/src/stats/mod.rs b/vortex-array/src/stats/mod.rs index ea8ae6b58a6..5c45970edf5 100644 --- a/vortex-array/src/stats/mod.rs +++ b/vortex-array/src/stats/mod.rs @@ -26,6 +26,7 @@ pub const PRUNING_STATS: &[Stat] = &[ Stat::Sum, Stat::NullCount, Stat::NaNCount, + Stat::UncompressedSizeInBytes, ]; pub fn as_stat_bitset_bytes(stats: &[Stat]) -> Vec { From f46666c8a7ecdb4ba2d36f29f3779920207c9e44 Mon Sep 17 00:00:00 2001 From: Luke Kim <80174+lukekim@users.noreply.github.com> Date: Thu, 26 Feb 2026 16:03:09 -0800 Subject: [PATCH 04/32] fix: Expression handling improvements for DataFusion integration Ported from spiceai-51 patches: - Skip unsupported expressions in make_vortex_predicate instead of erroring (#8) - Handle empty IN lists: x IN () returns false, x NOT IN () returns true (#8) - Allow DynamicFilterPhysicalExpr pushdown (#8) - Remove is_dynamic_physical_expr check that blocked all dynamic filters - Add Debug impl for BtrBlocksCompressor to satisfy CompressorPlugin bounds --- vortex-btrblocks/src/canonical_compressor.rs | 10 +++++++ vortex-datafusion/src/convert/exprs.rs | 29 ++++++++++++-------- 2 files changed, 28 insertions(+), 11 deletions(-) diff --git a/vortex-btrblocks/src/canonical_compressor.rs b/vortex-btrblocks/src/canonical_compressor.rs index 1b4c9250efe..b606d4ba4ae 100644 --- a/vortex-btrblocks/src/canonical_compressor.rs +++ b/vortex-btrblocks/src/canonical_compressor.rs @@ -101,6 +101,16 @@ pub struct BtrBlocksCompressor { pub string_schemes: Vec<&'static dyn StringScheme>, } +impl std::fmt::Debug for BtrBlocksCompressor { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("BtrBlocksCompressor") + .field("int_schemes", &self.int_schemes.len()) + .field("float_schemes", &self.float_schemes.len()) + .field("string_schemes", &self.string_schemes.len()) + .finish() + } +} + impl Default for BtrBlocksCompressor { fn default() -> Self { BtrBlocksCompressorBuilder::default().build() diff --git a/vortex-datafusion/src/convert/exprs.rs b/vortex-datafusion/src/convert/exprs.rs index b0380db5cb3..6091ae4c51d 100644 --- a/vortex-datafusion/src/convert/exprs.rs +++ b/vortex-datafusion/src/convert/exprs.rs @@ -16,7 +16,6 @@ use datafusion_physical_expr::ScalarFunctionExpr; use datafusion_physical_expr::projection::ProjectionExpr; use datafusion_physical_expr::projection::ProjectionExprs; use datafusion_physical_expr::utils::collect_columns; -use datafusion_physical_expr_common::physical_expr::is_dynamic_physical_expr; use datafusion_physical_plan::expressions as df_expr; use itertools::Itertools; use vortex::dtype::DType; @@ -52,10 +51,13 @@ pub(crate) fn make_vortex_predicate( expr_convertor: &dyn ExpressionConvertor, predicate: &[Arc], ) -> DFResult> { - let exprs = predicate + let exprs: Vec<_> = predicate .iter() - .map(|e| expr_convertor.convert(e.as_ref())) - .collect::>>()?; + .filter_map(|e| { + // If conversion fails, skip this expression (equivalent to lit(true) in AND conjunction) + expr_convertor.convert(e.as_ref()).ok() + }) + .collect(); Ok(and_collect(exprs)) } @@ -221,8 +223,13 @@ impl ExpressionConvertor for DefaultExpressionConvertor { }) .try_collect()?; + // Handle empty IN list: `x IN ()` is always false, `x NOT IN ()` is always true + let Some(first_element) = list_elements.first() else { + return Ok(lit(in_list.negated())); + }; + let list = Scalar::list( - list_elements[0].dtype().clone(), + first_element.dtype().clone(), list_elements, Nullability::Nullable, ); @@ -361,12 +368,6 @@ fn try_operator_from_df(value: &DFOperator) -> DFResult { } fn can_be_pushed_down_impl(df_expr: &Arc, schema: &Schema) -> bool { - // We currently do not support pushdown of dynamic expressions in DF. - // See issue: https://github.com/vortex-data/vortex/issues/4034 - if is_dynamic_physical_expr(df_expr) { - return false; - } - let expr = df_expr.as_any(); if let Some(binary) = expr.downcast_ref::() { can_binary_be_pushed_down(binary, schema) @@ -396,6 +397,12 @@ fn can_be_pushed_down_impl(df_expr: &Arc, schema: &Schema) -> .all(|e| can_be_pushed_down_impl(e, schema)) } else if let Some(scalar_fn) = expr.downcast_ref::() { can_scalar_fn_be_pushed_down(scalar_fn) + } else if expr + .downcast_ref::() + .is_some() + { + // Assume dynamic filters can be pushed down - the child won't be specified until execution time + true } else { tracing::debug!(%df_expr, "DataFusion expression can't be pushed down"); false From 2e5af89e850a12daa70515b57ed7f38dd170dbac Mon Sep 17 00:00:00 2001 From: Luke Kim <80174+lukekim@users.noreply.github.com> Date: Thu, 26 Feb 2026 17:24:50 -0800 Subject: [PATCH 05/32] feat(datafusion): add CASE conversion and export convertor Signed-off-by: Luke Kim <80174+lukekim@users.noreply.github.com> --- vortex-datafusion/src/convert/exprs.rs | 99 ++++++++++++++++++++++++++ vortex-datafusion/src/lib.rs | 1 + 2 files changed, 100 insertions(+) diff --git a/vortex-datafusion/src/convert/exprs.rs b/vortex-datafusion/src/convert/exprs.rs index 6091ae4c51d..1c8e61b7bfa 100644 --- a/vortex-datafusion/src/convert/exprs.rs +++ b/vortex-datafusion/src/convert/exprs.rs @@ -31,6 +31,7 @@ use vortex::expr::lit; use vortex::expr::not; use vortex::expr::pack; use vortex::expr::root; +use vortex::expr::zip_expr; use vortex::scalar::Scalar; use vortex::scalar_fn::ScalarFnVTableExt; use vortex::scalar_fn::fns::binary::Binary; @@ -146,6 +147,37 @@ impl DefaultExpressionConvertor { scalar_fn.name() )) } + + fn try_convert_case_expr(&self, case_expr: &df_expr::CaseExpr) -> DFResult { + let mut else_expr = if let Some(else_expr) = case_expr.else_expr() { + self.convert(else_expr.as_ref())? + } else { + return Err(exec_datafusion_err!( + "CASE expression without ELSE is not supported for pushdown" + )); + }; + + if let Some(base_expr) = case_expr.expr() { + let base_expr = self.convert(base_expr.as_ref())?; + for (when_expr, then_expr) in case_expr.when_then_expr().iter().rev() { + let when_expr = self.convert(when_expr.as_ref())?; + let then_expr = self.convert(then_expr.as_ref())?; + else_expr = zip_expr( + then_expr, + else_expr, + Binary.new_expr(Operator::Eq, [base_expr.clone(), when_expr]), + ); + } + } else { + for (when_expr, then_expr) in case_expr.when_then_expr().iter().rev() { + let when_expr = self.convert(when_expr.as_ref())?; + let then_expr = self.convert(then_expr.as_ref())?; + else_expr = zip_expr(then_expr, else_expr, when_expr); + } + } + + Ok(else_expr) + } } impl ExpressionConvertor for DefaultExpressionConvertor { @@ -242,6 +274,10 @@ impl ExpressionConvertor for DefaultExpressionConvertor { return self.try_convert_scalar_function(scalar_fn); } + if let Some(case_expr) = df.as_any().downcast_ref::() { + return self.try_convert_case_expr(case_expr); + } + Err(exec_datafusion_err!( "Couldn't convert DataFusion physical {df} expression to a vortex expression" )) @@ -397,6 +433,8 @@ fn can_be_pushed_down_impl(df_expr: &Arc, schema: &Schema) -> .all(|e| can_be_pushed_down_impl(e, schema)) } else if let Some(scalar_fn) = expr.downcast_ref::() { can_scalar_fn_be_pushed_down(scalar_fn) + } else if let Some(case_expr) = expr.downcast_ref::() { + can_case_be_pushed_down(case_expr, schema) } else if expr .downcast_ref::() .is_some() @@ -416,6 +454,22 @@ fn can_binary_be_pushed_down(binary: &df_expr::BinaryExpr, schema: &Schema) -> b && can_be_pushed_down_impl(binary.right(), schema) } +fn can_case_be_pushed_down(case_expr: &df_expr::CaseExpr, schema: &Schema) -> bool { + case_expr + .expr() + .is_none_or(|base_expr| can_be_pushed_down_impl(base_expr, schema)) + && case_expr + .when_then_expr() + .iter() + .all(|(when_expr, then_expr)| { + can_be_pushed_down_impl(when_expr, schema) + && can_be_pushed_down_impl(then_expr, schema) + }) + && case_expr + .else_expr() + .is_some_and(|else_expr| can_be_pushed_down_impl(else_expr, schema)) +} + fn supported_data_types(dt: &DataType) -> bool { use DataType::*; @@ -632,6 +686,51 @@ mod tests { ); } + #[test] + fn test_expr_from_df_case_when_with_else() { + let when_then_expr = vec![( + Arc::new(df_expr::Column::new("active", 0)) as Arc, + Arc::new(df_expr::Literal::new(ScalarValue::Utf8(Some( + "yes".to_string(), + )))) as Arc, + )]; + let case_expr = df_expr::CaseExpr::try_new( + None, + when_then_expr, + Some(Arc::new(df_expr::Literal::new(ScalarValue::Utf8(Some( + "no".to_string(), + )))) as Arc), + ) + .unwrap(); + + let result = DefaultExpressionConvertor::default() + .convert(&case_expr) + .unwrap(); + + assert_snapshot!(result.display_tree().to_string(), @r#" + vortex.zip() + ├── if_true: vortex.literal("yes") + ├── if_false: vortex.literal("no") + └── mask: vortex.get_item(active) + └── input: vortex.root() + "#); + } + + #[test] + fn test_expr_from_df_case_when_without_else_not_pushable() { + let when_then_expr = vec![( + Arc::new(df_expr::Column::new("active", 0)) as Arc, + Arc::new(df_expr::Literal::new(ScalarValue::Utf8(Some( + "yes".to_string(), + )))) as Arc, + )]; + let case_expr = Arc::new(df_expr::CaseExpr::try_new(None, when_then_expr, None).unwrap()) + as Arc; + + let schema = Schema::new(vec![Field::new("active", DataType::Boolean, false)]); + assert!(!can_be_pushed_down_impl(&case_expr, &schema)); + } + #[rstest] // Supported types #[case::null(DataType::Null, true)] diff --git a/vortex-datafusion/src/lib.rs b/vortex-datafusion/src/lib.rs index 8e7fc57dbf5..0e09886c07c 100644 --- a/vortex-datafusion/src/lib.rs +++ b/vortex-datafusion/src/lib.rs @@ -15,6 +15,7 @@ pub mod v2; #[cfg(test)] mod tests; +pub use convert::exprs::DefaultExpressionConvertor; pub use convert::exprs::ExpressionConvertor; pub use persistent::*; From 315abb8b007585506d0b056b6801f1b55ecaee19 Mon Sep 17 00:00:00 2001 From: Luke Kim <80174+lukekim@users.noreply.github.com> Date: Thu, 26 Feb 2026 17:47:32 -0800 Subject: [PATCH 06/32] fix: streamline filter_map usage in propagate_take_stats function --- vortex-array/src/arrays/dict/take.rs | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/vortex-array/src/arrays/dict/take.rs b/vortex-array/src/arrays/dict/take.rs index 6a6259f8ed3..d38db57c923 100644 --- a/vortex-array/src/arrays/dict/take.rs +++ b/vortex-array/src/arrays/dict/take.rs @@ -157,14 +157,14 @@ pub(crate) fn propagate_take_stats( Stat::IsConstant, ] .into_iter() - .filter_map(|stat| { - source - .statistics() - .get(stat) - .and_then(|v| v.map(|s| s.into_value()).into_inexact().transpose()) - .map(|sv| (stat, sv)) - }) - .collect::>(); + .filter_map(|stat| { + source + .statistics() + .get(stat) + .and_then(|v| v.map(|s| s.into_value()).into_inexact().transpose()) + .map(|sv| (stat, sv)) + }) + .collect::>(); st.combine_sets( &(unsafe { StatsSet::new_unchecked(inexact_min_max) }).as_typed_ref(source.dtype()), ) From 3e251fb007cd2f49f95505a5ac1b56b98c5250d5 Mon Sep 17 00:00:00 2001 From: Luke Kim <80174+lukekim@users.noreply.github.com> Date: Thu, 26 Feb 2026 17:51:37 -0800 Subject: [PATCH 07/32] fix: remove duplicate log dependency in Cargo.toml --- vortex-io/Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vortex-io/Cargo.toml b/vortex-io/Cargo.toml index 6d954acc054..132ca7b55eb 100644 --- a/vortex-io/Cargo.toml +++ b/vortex-io/Cargo.toml @@ -26,6 +26,7 @@ getrandom_v03 = { workspace = true } # Needed to pickup the "wasm_js" feature fo glob = { workspace = true } handle = "1.0.2" kanal = { workspace = true } +log = { workspace = true } object_store = { workspace = true, optional = true, features = ["fs"] } oneshot = { workspace = true } parking_lot = { workspace = true } @@ -38,7 +39,6 @@ tokio = { workspace = true, features = [ tracing = { workspace = true } vortex-array = { workspace = true } vortex-buffer = { workspace = true } -log = { workspace = true } vortex-error = { workspace = true } vortex-metrics = { workspace = true } vortex-session = { workspace = true } From b27e89af53aece9cd112b26d6ca68c9c22fcee33 Mon Sep 17 00:00:00 2001 From: Sergei Grebnov Date: Sun, 11 Jan 2026 20:53:30 +0300 Subject: [PATCH 08/32] Use tokio::sync::oneshot to prevent FuturesUnordered reentrant drop crash (#9) --- vortex-file/src/read/driver.rs | 2 ++ vortex-io/src/runtime/single.rs | 6 +++++- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/vortex-file/src/read/driver.rs b/vortex-file/src/read/driver.rs index d8fd638adc8..752dc754818 100644 --- a/vortex-file/src/read/driver.rs +++ b/vortex-file/src/read/driver.rs @@ -8,6 +8,8 @@ use std::task::Context; use std::task::Poll; use futures::Stream; +#[cfg(all(test, not(feature = "tokio")))] +use oneshot; use pin_project_lite::pin_project; use vortex_buffer::Alignment; use vortex_error::VortexExpect; diff --git a/vortex-io/src/runtime/single.rs b/vortex-io/src/runtime/single.rs index ea6a68f2f24..0a3abf553dc 100644 --- a/vortex-io/src/runtime/single.rs +++ b/vortex-io/src/runtime/single.rs @@ -8,8 +8,13 @@ use futures::Stream; use futures::StreamExt; use futures::future::BoxFuture; use futures::stream::LocalBoxStream; +#[cfg(not(feature = "tokio"))] +use oneshot; use parking_lot::Mutex; use smol::LocalExecutor; +// Prefer tokio::sync::oneshot when tokio feature is enabled +#[cfg(feature = "tokio")] +use tokio::sync::oneshot; use vortex_error::vortex_panic; use crate::runtime::AbortHandle; @@ -18,7 +23,6 @@ use crate::runtime::BlockingRuntime; use crate::runtime::Executor; use crate::runtime::Handle; use crate::runtime::smol::SmolAbortHandle; - /// A runtime that drives all work on the current thread. /// /// This is subtly different from using a current-thread runtime to drive a future since it is From 8044a84704a1c2fe079920e9b290ccf03cdf58b4 Mon Sep 17 00:00:00 2001 From: Luke Kim <80174+lukekim@users.noreply.github.com> Date: Wed, 7 Jan 2026 20:24:48 -0800 Subject: [PATCH 09/32] Don't return an error when we have an unsupported node, bubble up "TRUE" as in keep for that node, up to any `and` or `or` node and handle empty IN list. (#8) --- vortex-datafusion/src/convert/exprs.rs | 303 ++++++++++++++++----- vortex-datafusion/src/persistent/format.rs | 32 +-- 2 files changed, 248 insertions(+), 87 deletions(-) diff --git a/vortex-datafusion/src/convert/exprs.rs b/vortex-datafusion/src/convert/exprs.rs index 1c8e61b7bfa..9173bc7330c 100644 --- a/vortex-datafusion/src/convert/exprs.rs +++ b/vortex-datafusion/src/convert/exprs.rs @@ -39,7 +39,7 @@ use vortex::scalar_fn::fns::like::Like; use vortex::scalar_fn::fns::like::LikeOptions; use vortex::scalar_fn::fns::operators::Operator; -use crate::convert::FromDataFusion; +use crate::convert::{FromDataFusion, TryFromDataFusion}; /// Result of splitting a projection into Vortex expressions and leftover DataFusion projections. pub struct ProcessedProjection { @@ -189,9 +189,9 @@ impl ExpressionConvertor for DefaultExpressionConvertor { // TODO(joe): Don't return an error when we have an unsupported node, bubble up "TRUE" as in keep // for that node, up to any `and` or `or` node. if let Some(binary_expr) = df.as_any().downcast_ref::() { - let left = self.convert(binary_expr.left().as_ref())?; - let right = self.convert(binary_expr.right().as_ref())?; - let operator = try_operator_from_df(binary_expr.op())?; + let left = Expression::try_from_df(binary_expr.left().as_ref())?; + let right = Expression::try_from_df(binary_expr.right().as_ref())?; + let operator = Operator::try_from_df(binary_expr.op())?; return Ok(Binary.new_expr(operator, [left, right])); } @@ -201,8 +201,8 @@ impl ExpressionConvertor for DefaultExpressionConvertor { } if let Some(like) = df.as_any().downcast_ref::() { - let child = self.convert(like.expr().as_ref())?; - let pattern = self.convert(like.pattern().as_ref())?; + let child = Expression::try_from_df(like.expr().as_ref())?; + let pattern = Expression::try_from_df(like.pattern().as_ref())?; return Ok(Like.new_expr( LikeOptions { negated: like.negated(), @@ -219,30 +219,22 @@ impl ExpressionConvertor for DefaultExpressionConvertor { if let Some(cast_expr) = df.as_any().downcast_ref::() { let cast_dtype = DType::from_arrow((cast_expr.cast_type(), Nullability::Nullable)); - let child = self.convert(cast_expr.expr().as_ref())?; + let child = Expression::try_from_df(cast_expr.expr().as_ref())?; return Ok(cast(child, cast_dtype)); } - if let Some(cast_col_expr) = df.as_any().downcast_ref::() { - let target = cast_col_expr.target_field(); - - let target_dtype = DType::from_arrow((target.data_type(), target.is_nullable().into())); - let child = self.convert(cast_col_expr.expr().as_ref())?; - return Ok(cast(child, target_dtype)); - } - if let Some(is_null_expr) = df.as_any().downcast_ref::() { - let arg = self.convert(is_null_expr.arg().as_ref())?; + let arg = Expression::try_from_df(is_null_expr.arg().as_ref())?; return Ok(is_null(arg)); } if let Some(is_not_null_expr) = df.as_any().downcast_ref::() { - let arg = self.convert(is_not_null_expr.arg().as_ref())?; + let arg = Expression::try_from_df(is_not_null_expr.arg().as_ref())?; return Ok(not(is_null(arg))); } if let Some(in_list) = df.as_any().downcast_ref::() { - let value = self.convert(in_list.expr().as_ref())?; + let value = Expression::try_from_df(in_list.expr().as_ref())?; let list_elements: Vec<_> = in_list .list() .iter() @@ -271,7 +263,16 @@ impl ExpressionConvertor for DefaultExpressionConvertor { } if let Some(scalar_fn) = df.as_any().downcast_ref::() { - return self.try_convert_scalar_function(scalar_fn); + return try_convert_scalar_function(scalar_fn); + } + + if let Some(dynamic_expr) = df + .as_any() + .downcast_ref::() + && let Ok(current) = dynamic_expr.current() + { + let returned_expr = Expression::try_from_df(current.as_ref())?; + return Ok(returned_expr); } if let Some(case_expr) = df.as_any().downcast_ref::() { @@ -417,10 +418,8 @@ fn can_be_pushed_down_impl(df_expr: &Arc, schema: &Schema) -> && can_be_pushed_down_impl(like.pattern(), schema) } else if let Some(lit) = expr.downcast_ref::() { supported_data_types(&lit.value().data_type()) - } else if expr.downcast_ref::().is_some() - || expr.downcast_ref::().is_some() - { - true + } else if let Some(cast) = expr.downcast_ref::() { + supported_data_types(cast.cast_type()) && can_be_pushed_down(cast.expr(), schema) } else if let Some(is_null) = expr.downcast_ref::() { can_be_pushed_down_impl(is_null.arg(), schema) } else if let Some(is_not_null) = expr.downcast_ref::() { @@ -448,7 +447,7 @@ fn can_be_pushed_down_impl(df_expr: &Arc, schema: &Schema) -> } fn can_binary_be_pushed_down(binary: &df_expr::BinaryExpr, schema: &Schema) -> bool { - let is_op_supported = try_operator_from_df(binary.op()).is_ok(); + let is_op_supported = Operator::try_from_df(binary.op()).is_ok(); is_op_supported && can_be_pushed_down_impl(binary.left(), schema) && can_be_pushed_down_impl(binary.right(), schema) @@ -497,15 +496,57 @@ fn supported_data_types(dt: &DataType) -> bool { ); if !is_supported { - tracing::debug!("DataFusion data type {dt:?} is not supported"); + tracing::debug!(data_type = ?dt, "DataFusion data type is not supported"); } is_supported } /// Checks if a GetField scalar function can be pushed down. -fn can_scalar_fn_be_pushed_down(scalar_fn: &ScalarFunctionExpr) -> bool { - ScalarFunctionExpr::try_downcast_func::(scalar_fn).is_some() +fn can_scalar_fn_be_pushed_down(scalar_fn: &ScalarFunctionExpr, schema: &Schema) -> bool { + let Some(get_field_fn) = ScalarFunctionExpr::try_downcast_func::(scalar_fn) + else { + // Only get_field pushdown is supported. + return false; + }; + + let args = get_field_fn.args(); + if args.len() != 2 { + tracing::debug!( + "Expected 2 arguments for GetField, not pushing down {} arguments", + args.len() + ); + return false; + } + let source_expr = &args[0]; + let field_name_expr = &args[1]; + let Some(field_name) = field_name_expr + .as_any() + .downcast_ref::() + .and_then(|lit| lit.value().try_as_str().flatten()) + else { + return false; + }; + + let Ok(source_dt) = source_expr.data_type(schema) else { + tracing::debug!( + field_name = field_name, + schema = ?schema, + source_expr = ?source_expr, + "Failed to get source type for GetField, not pushing down" + ); + return false; + }; + let DataType::Struct(fields) = source_dt else { + tracing::debug!( + field_name = field_name, + schema = ?schema, + source_expr = ?source_expr, + "Failed to get source type as struct for GetField, not pushing down" + ); + return false; + }; + fields.find(field_name).is_some() } // TODO(adam): Replace with `DataType::is_decimal` once its released. @@ -523,16 +564,16 @@ fn is_decimal(dt: &DataType) -> bool { mod tests { use std::sync::Arc; - use arrow_schema::DataType; - use arrow_schema::Field; - use arrow_schema::Schema; - use arrow_schema::TimeUnit as ArrowTimeUnit; + use arrow_schema::{DataType, Field, Fields, Schema, TimeUnit as ArrowTimeUnit}; + use datafusion::functions::core::getfield::GetFieldFunc; use datafusion_common::ScalarValue; - use datafusion_expr::Operator as DFOperator; + use datafusion_common::config::ConfigOptions; + use datafusion_expr::{Operator as DFOperator, ScalarUDF}; use datafusion_physical_expr::PhysicalExpr; use datafusion_physical_plan::expressions as df_expr; use insta::assert_snapshot; use rstest::rstest; + use vortex::expr::{Expression, Operator}; use super::*; use crate::common_tests::TestSessionContext; @@ -559,25 +600,22 @@ mod tests { #[test] fn test_make_vortex_predicate_empty() { - let expr_convertor = DefaultExpressionConvertor::default(); - let result = make_vortex_predicate(&expr_convertor, &[]).unwrap(); + let result = make_vortex_predicate(&[]).unwrap(); assert!(result.is_none()); } #[test] fn test_make_vortex_predicate_single() { - let expr_convertor = DefaultExpressionConvertor::default(); let col_expr = Arc::new(df_expr::Column::new("test", 0)) as Arc; - let result = make_vortex_predicate(&expr_convertor, &[col_expr]).unwrap(); + let result = make_vortex_predicate(&[&col_expr]).unwrap(); assert!(result.is_some()); } #[test] fn test_make_vortex_predicate_multiple() { - let expr_convertor = DefaultExpressionConvertor::default(); let col1 = Arc::new(df_expr::Column::new("col1", 0)) as Arc; let col2 = Arc::new(df_expr::Column::new("col2", 1)) as Arc; - let result = make_vortex_predicate(&expr_convertor, &[col1, col2]).unwrap(); + let result = make_vortex_predicate(&[&col1, &col2]).unwrap(); assert!(result.is_some()); // Result should be an AND expression combining the two columns } @@ -599,7 +637,7 @@ mod tests { #[case] df_op: DFOperator, #[case] expected_vortex_op: Operator, ) { - let result = try_operator_from_df(&df_op).unwrap(); + let result = Operator::try_from_df(&df_op).unwrap(); assert_eq!(result, expected_vortex_op); } @@ -609,7 +647,7 @@ mod tests { #[case::regex_match(DFOperator::RegexMatch)] #[case::like_match(DFOperator::LikeMatch)] fn test_operator_conversion_unsupported(#[case] df_op: DFOperator) { - let result = try_operator_from_df(&df_op); + let result = Operator::try_from_df(&df_op); assert!(result.is_err()); assert!( result @@ -622,24 +660,20 @@ mod tests { #[test] fn test_expr_from_df_column() { let col_expr = df_expr::Column::new("test_column", 0); - let result = DefaultExpressionConvertor::default() - .convert(&col_expr) - .unwrap(); + let result = Expression::try_from_df(&col_expr).unwrap(); - assert_snapshot!(result.display_tree().to_string(), @r" - vortex.get_item(test_column) - └── input: vortex.root() - "); + assert_snapshot!(result.display_tree().to_string(), @r#" + vortex.get_item "test_column" + └── input: vortex.root + "#); } #[test] fn test_expr_from_df_literal() { let literal_expr = df_expr::Literal::new(ScalarValue::Int32(Some(42))); - let result = DefaultExpressionConvertor::default() - .convert(&literal_expr) - .unwrap(); + let result = Expression::try_from_df(&literal_expr).unwrap(); - assert_snapshot!(result.display_tree().to_string(), @"vortex.literal(42i32)"); + assert_snapshot!(result.display_tree().to_string(), @"vortex.literal 42i32"); } #[test] @@ -649,16 +683,14 @@ mod tests { Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(42)))) as Arc; let binary_expr = df_expr::BinaryExpr::new(left, DFOperator::Eq, right); - let result = DefaultExpressionConvertor::default() - .convert(&binary_expr) - .unwrap(); + let result = Expression::try_from_df(&binary_expr).unwrap(); - assert_snapshot!(result.display_tree().to_string(), @r" - vortex.binary(=) - ├── lhs: vortex.get_item(left) - │ └── input: vortex.root() - └── rhs: vortex.literal(42i32) - "); + assert_snapshot!(result.display_tree().to_string(), @r#" + vortex.binary = + ├── lhs: vortex.get_item "left" + │ └── input: vortex.root + └── rhs: vortex.literal 42i32 + "#); } #[rstest] @@ -673,12 +705,10 @@ mod tests { )))) as Arc; let like_expr = df_expr::LikeExpr::new(negated, case_insensitive, expr, pattern); - let result = DefaultExpressionConvertor::default() - .convert(&like_expr) - .unwrap(); - let like_opts = result.as_::(); + let result = Expression::try_from_df(&like_expr).unwrap(); + let like_expr = result.as_::(); assert_eq!( - like_opts, + like_expr.data(), &LikeOptions { negated, case_insensitive @@ -763,8 +793,7 @@ mod tests { DataType::List(Arc::new(Field::new("item", DataType::Int32, true))), false )] - #[case::struct_type(DataType::Struct(vec![Field::new("field", DataType::Int32, true)].into() - ), false)] + #[case::struct_type(DataType::Struct(vec![Field::new("field", DataType::Int32, true)].into()), false)] // Dictionary types - should be supported if value type is supported #[case::dict_utf8( DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)), @@ -917,4 +946,144 @@ mod tests { Ok(()) } + + #[test] + fn test_expr_from_df_get_field() { + let struct_col = Arc::new(df_expr::Column::new("my_struct", 0)) as Arc; + let field_name = Arc::new(df_expr::Literal::new(ScalarValue::Utf8(Some( + "field1".to_string(), + )))) as Arc; + let get_field_expr = ScalarFunctionExpr::new( + "get_field", + Arc::new(ScalarUDF::from(GetFieldFunc::new())), + vec![struct_col, field_name], + Arc::new(Field::new("field1", DataType::Utf8, true)), + Arc::new(ConfigOptions::new()), + ); + let result = Expression::try_from_df(&get_field_expr).unwrap(); + assert_snapshot!(result.display_tree().to_string(), @r#" + vortex.get_item "field1" + └── input: vortex.get_item "my_struct" + └── input: vortex.root + "#); + } + + #[rstest] + #[case::valid_field("field1", true)] + #[case::missing_field("nonexistent_field", false)] + fn test_can_be_pushed_down_get_field(#[case] field_name: &str, #[case] expected: bool) { + let struct_fields = Fields::from(vec![ + Field::new("field1", DataType::Utf8, true), + Field::new("field2", DataType::Int32, true), + ]); + let schema = Schema::new(vec![Field::new( + "my_struct", + DataType::Struct(struct_fields), + true, + )]); + + let struct_col = Arc::new(df_expr::Column::new("my_struct", 0)) as Arc; + let field_name_lit = Arc::new(df_expr::Literal::new(ScalarValue::Utf8(Some( + field_name.to_string(), + )))) as Arc; + + let get_field_expr = Arc::new(ScalarFunctionExpr::new( + "get_field", + Arc::new(ScalarUDF::from(GetFieldFunc::new())), + vec![struct_col, field_name_lit], + Arc::new(Field::new(field_name, DataType::Utf8, true)), + Arc::new(ConfigOptions::new()), + )) as Arc; + + assert_eq!(can_be_pushed_down(&get_field_expr, &schema), expected); + } + + /// Create an unsupported scalar function expression (simulating functions like to_timestamp) + fn create_unsupported_scalar_fn() -> Arc { + use datafusion_functions::datetime::to_timestamp::ToTimestampFunc; + + let arg = Arc::new(df_expr::Literal::new(ScalarValue::Utf8(Some( + "2024-01-01".to_string(), + )))) as Arc; + + Arc::new(ScalarFunctionExpr::new( + "to_timestamp", + Arc::new(ScalarUDF::from(ToTimestampFunc::new())), + vec![arg], + Arc::new(Field::new( + "result", + DataType::Timestamp(ArrowTimeUnit::Nanosecond, None), + true, + )), + Arc::new(ConfigOptions::new()), + )) + } + + #[test] + fn test_make_vortex_predicate_skips_unsupported_scalar_function() { + // Unsupported scalar function like to_timestamp should be skipped, not error + let unsupported_fn = create_unsupported_scalar_fn(); + let result = make_vortex_predicate(&[&unsupported_fn]); + + // Should succeed (not error) and return None since the only expression was skipped + assert!(result.is_ok()); + assert!(result.unwrap().is_none()); + } + + #[test] + fn test_make_vortex_predicate_combines_supported_and_skips_unsupported() { + // Mix of supported column expression and unsupported scalar function + let supported_col = Arc::new(df_expr::Column::new("test", 0)) as Arc; + let unsupported_fn = create_unsupported_scalar_fn(); + + let result = make_vortex_predicate(&[&supported_col, &unsupported_fn]); + + // Should succeed and return the supported expression only + assert!(result.is_ok()); + let predicate = result.unwrap(); + assert!(predicate.is_some()); + + // The result should just be the column expression since the unsupported one was skipped + assert_snapshot!(predicate.unwrap().display_tree().to_string(), @r#" + vortex.get_item "test" + └── input: vortex.root + "#); + } + + #[test] + fn test_make_vortex_predicate_multiple_supported_with_unsupported() { + // Two supported columns and one unsupported function + let col1 = Arc::new(df_expr::Column::new("col1", 0)) as Arc; + let col2 = Arc::new(df_expr::Column::new("col2", 1)) as Arc; + let unsupported_fn = create_unsupported_scalar_fn(); + + let result = make_vortex_predicate(&[&col1, &unsupported_fn, &col2]); + + // Should succeed and return AND of the two supported expressions + assert!(result.is_ok()); + let predicate = result.unwrap(); + assert!(predicate.is_some()); + + // The result should be an AND of col1 and col2 + assert_snapshot!(predicate.unwrap().display_tree().to_string(), @r#" + vortex.binary and + ├── lhs: vortex.get_item "col1" + │ └── input: vortex.root + └── rhs: vortex.get_item "col2" + └── input: vortex.root + "#); + } + + #[test] + fn test_make_vortex_predicate_all_unsupported_returns_none() { + // When all expressions are unsupported, should return None (no filter) + let unsupported_fn1 = create_unsupported_scalar_fn(); + let unsupported_fn2 = create_unsupported_scalar_fn(); + + let result = make_vortex_predicate(&[&unsupported_fn1, &unsupported_fn2]); + + // Should succeed and return None + assert!(result.is_ok()); + assert!(result.unwrap().is_none()); + } } diff --git a/vortex-datafusion/src/persistent/format.rs b/vortex-datafusion/src/persistent/format.rs index fb3d5f9db11..99fc6a465d3 100644 --- a/vortex-datafusion/src/persistent/format.rs +++ b/vortex-datafusion/src/persistent/format.rs @@ -2,33 +2,28 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::any::Any; -use std::fmt::Debug; -use std::fmt::Formatter; +use std::fmt::{Debug, Formatter}; use std::sync::Arc; -use arrow_schema::Schema; -use arrow_schema::SchemaRef; +use arrow_schema::{Schema, SchemaRef}; use async_trait::async_trait; use datafusion_catalog::Session; -use datafusion_common::ColumnStatistics; -use datafusion_common::DataFusionError; -use datafusion_common::GetExt; -use datafusion_common::Result as DFResult; -use datafusion_common::Statistics; use datafusion_common::config::ConfigField; use datafusion_common::config_namespace; use datafusion_common::internal_datafusion_err; use datafusion_common::not_impl_err; use datafusion_common::parsers::CompressionTypeVariant; use datafusion_common::stats::Precision; +use datafusion_common::{ + ColumnStatistics, DataFusionError, GetExt, Result as DFResult, Statistics, config_namespace, + not_impl_err, +}; use datafusion_common_runtime::SpawnedTask; use datafusion_datasource::TableSchema; use datafusion_datasource::file::FileSource; use datafusion_datasource::file_compression_type::FileCompressionType; -use datafusion_datasource::file_format::FileFormat; -use datafusion_datasource::file_format::FileFormatFactory; -use datafusion_datasource::file_scan_config::FileScanConfig; -use datafusion_datasource::file_scan_config::FileScanConfigBuilder; +use datafusion_datasource::file_format::{FileFormat, FileFormatFactory}; +use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder}; use datafusion_datasource::file_sink_config::FileSinkConfig; use datafusion_datasource::sink::DataSinkExec; use datafusion_datasource::source::DataSourceExec; @@ -59,6 +54,8 @@ use vortex::io::object_store::ObjectStoreReadAt; use vortex::io::session::RuntimeSessionExt; use vortex::scalar::Scalar; use vortex::session::VortexSession; +use vortex::stats::{Stat, StatsSet}; +use vortex::{VortexSessionDefault, stats}; use super::cache::CachedVortexMetadata; use super::sink::VortexSink; @@ -66,8 +63,6 @@ use super::source::VortexSource; use crate::PrecisionExt as _; use crate::convert::TryToDataFusion; -const DEFAULT_FOOTER_INITIAL_READ_SIZE_BYTES: usize = MAX_POSTSCRIPT_SIZE as usize + EOF_SIZE; - /// Vortex implementation of a DataFusion [`FileFormat`]. pub struct VortexFormat { session: VortexSession, @@ -126,10 +121,7 @@ impl GetExt for VortexFormatFactory { impl VortexFormatFactory { /// Creates a new instance with a default [`VortexSession`] and default options. - #[expect( - clippy::new_without_default, - reason = "FormatFactory defines `default` method, so having `Default` implementation is confusing" - )] + #[allow(clippy::new_without_default)] // FormatFactory defines `default` method, so having `Default` implementation is confusing. pub fn new() -> Self { Self { session: VortexSession::default(), @@ -162,7 +154,7 @@ impl VortexFormatFactory { } impl FileFormatFactory for VortexFormatFactory { - #[expect(clippy::disallowed_types, reason = "required by trait signature")] + #[allow(clippy::disallowed_types)] fn create( &self, _state: &dyn Session, From d8b806b6326e841a25b0b53995c0e2c4819287fb Mon Sep 17 00:00:00 2001 From: William <98815791+peasee@users.noreply.github.com> Date: Wed, 10 Dec 2025 13:24:43 +1000 Subject: [PATCH 10/32] feat: Support retrieving WriterStrategyBuilder from VortexSession (#6) * Fix session get-or-default (#5662) The comments described this get-or-default, but instead it was a panic --------- Signed-off-by: Nicholas Gates * feat: Support retrieving writer strategy builder from vortex session --------- Signed-off-by: Nicholas Gates Co-authored-by: Nicholas Gates --- vortex-session/src/lib.rs | 28 +++++----------------------- 1 file changed, 5 insertions(+), 23 deletions(-) diff --git a/vortex-session/src/lib.rs b/vortex-session/src/lib.rs index d2112caf30c..bfb2a51ba13 100644 --- a/vortex-session/src/lib.rs +++ b/vortex-session/src/lib.rs @@ -3,20 +3,14 @@ pub mod registry; -use std::any::Any; -use std::any::TypeId; -use std::any::type_name; +use std::any::{Any, TypeId, type_name}; use std::fmt::Debug; -use std::hash::BuildHasherDefault; -use std::hash::Hasher; -use std::ops::Deref; -use std::ops::DerefMut; +use std::hash::{BuildHasherDefault, Hasher}; +use std::ops::{Deref, DerefMut}; use std::sync::Arc; -use dashmap::DashMap; -use dashmap::Entry; -use vortex_error::VortexExpect; -use vortex_error::vortex_panic; +use dashmap::{DashMap, Entry}; +use vortex_error::{VortexExpect, vortex_panic}; /// A Vortex session encapsulates the set of extensible arrays, layouts, compute functions, dtypes, /// etc. that are available for use in a given context. @@ -108,18 +102,6 @@ impl SessionExt for VortexSession { /// Returns the scope variable of type `V`, or inserts a default one if it does not exist. fn get(&self) -> Ref<'_, V> { - // NOTE(ngates): we don't use `entry().or_insert_with_key()` here because the DashMap - // would immediately acquire an exclusive write lock. - if let Some(v) = self.0.get(&TypeId::of::()) { - return Ref(v.map(|v| { - (**v) - .as_any() - .downcast_ref::() - .vortex_expect("Type mismatch - this is a bug") - })); - } - - // If we get here, the value was not present, so we insert the default with a write lock. Ref(self .0 .entry(TypeId::of::()) From 6712e9ffa99065c327d3cbc65edcefa10ba60037 Mon Sep 17 00:00:00 2001 From: Luke Kim <80174+lukekim@users.noreply.github.com> Date: Wed, 14 Jan 2026 15:58:31 -0800 Subject: [PATCH 11/32] fix: Handle UncompressedSizeInBytes statistic correctly (#3) --- vortex-datafusion/src/persistent/format.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vortex-datafusion/src/persistent/format.rs b/vortex-datafusion/src/persistent/format.rs index 99fc6a465d3..fbf3a7f0cb7 100644 --- a/vortex-datafusion/src/persistent/format.rs +++ b/vortex-datafusion/src/persistent/format.rs @@ -54,7 +54,7 @@ use vortex::io::object_store::ObjectStoreReadAt; use vortex::io::session::RuntimeSessionExt; use vortex::scalar::Scalar; use vortex::session::VortexSession; -use vortex::stats::{Stat, StatsSet}; +use vortex::stats::{ArrayStats, Stat, StatsSet}; use vortex::{VortexSessionDefault, stats}; use super::cache::CachedVortexMetadata; From 42648c64e3ce2be36fd7fb981eeb1702e22bab2b Mon Sep 17 00:00:00 2001 From: Luke Kim <80174+lukekim@users.noreply.github.com> Date: Wed, 14 Jan 2026 16:01:51 -0800 Subject: [PATCH 12/32] fix: Add log dependency to vortex-io --- vortex-io/Cargo.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/vortex-io/Cargo.toml b/vortex-io/Cargo.toml index 132ca7b55eb..024f0a9144f 100644 --- a/vortex-io/Cargo.toml +++ b/vortex-io/Cargo.toml @@ -39,6 +39,7 @@ tokio = { workspace = true, features = [ tracing = { workspace = true } vortex-array = { workspace = true } vortex-buffer = { workspace = true } +log = { workspace = true } vortex-error = { workspace = true } vortex-metrics = { workspace = true } vortex-session = { workspace = true } From 995c5e39fd6ae992d4fe6a9a1350fee10293bcbd Mon Sep 17 00:00:00 2001 From: Luke Kim <80174+lukekim@users.noreply.github.com> Date: Wed, 14 Jan 2026 16:34:18 -0800 Subject: [PATCH 13/32] fix: Update persistent module to use simplified expression handling from PR #8 --- vortex-datafusion/src/lib.rs | 2 +- vortex-datafusion/src/persistent/opener.rs | 121 ++++++++++----------- vortex-datafusion/src/persistent/source.rs | 14 +-- 3 files changed, 62 insertions(+), 75 deletions(-) diff --git a/vortex-datafusion/src/lib.rs b/vortex-datafusion/src/lib.rs index 0e09886c07c..b976215aa18 100644 --- a/vortex-datafusion/src/lib.rs +++ b/vortex-datafusion/src/lib.rs @@ -6,7 +6,7 @@ use std::fmt::Debug; use datafusion_common::stats::Precision as DFPrecision; -use vortex::expr::stats::Precision; +use vortex::stats::Precision; mod convert; mod persistent; diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs index 5986a06da49..199eb2e4912 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/vortex-datafusion/src/persistent/opener.rs @@ -2,8 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::ops::Range; -use std::sync::Arc; -use std::sync::Weak; +use std::sync::{Arc, Weak}; use arrow_schema::Schema; use datafusion_common::DataFusionError; @@ -43,8 +42,7 @@ use vortex::metrics::Label; use vortex::metrics::MetricsRegistry; use vortex::scan::ScanBuilder; use vortex::session::VortexSession; -use vortex_utils::aliases::dash_map::DashMap; -use vortex_utils::aliases::dash_map::Entry; +use vortex_utils::aliases::dash_map::{DashMap, Entry}; use crate::VortexAccessPlan; use crate::convert::exprs::ExpressionConvertor; @@ -90,6 +88,7 @@ pub(crate) struct VortexOpener { pub layout_readers: Arc>>, /// Whether the query has output ordering specified pub has_output_ordering: bool, +} pub expression_convertor: Arc, pub file_metadata_cache: Option>, @@ -99,7 +98,7 @@ pub(crate) struct VortexOpener { } impl FileOpener for VortexOpener { - fn open(&self, file: PartitionedFile) -> DFResult { + fn open(&self, file_meta: FileMeta, file: PartitionedFile) -> DFResult { let session = self.session.clone(); let metrics_registry = self.metrics_registry.clone(); let labels = vec![ @@ -158,7 +157,7 @@ impl FileOpener for VortexOpener { // - Partition column values (e.g., date=2024-01-01) // - File-level statistics (min/max values per column) let mut file_pruner = file_pruning_predicate - .filter(|p| { + .map(|predicate| { // Only create pruner if we have dynamic expressions or file statistics // to work with. Static predicates without stats won't benefit from pruning. is_dynamic_physical_expr(p) || file.has_statistics() @@ -269,13 +268,13 @@ impl FileOpener for VortexOpener { let projector = leftover_projection.make_projector(&stream_schema)?; // We share our layout readers with others partitions in the scan, so we can only need to read each layout in each file once. - let layout_reader = match layout_reader.entry(file.object_meta.location.clone()) { + let layout_reader = match layout_reader.entry(file_meta.object_meta.location.clone()) { Entry::Occupied(mut occupied_entry) => { if let Some(reader) = occupied_entry.get().upgrade() { - tracing::trace!("reusing layout reader for {}", occupied_entry.key()); + log::trace!("reusing layout reader for {}", occupied_entry.key()); reader } else { - tracing::trace!("creating layout reader for {}", occupied_entry.key()); + log::trace!("creating layout reader for {}", occupied_entry.key()); let reader = vxf.layout_reader().map_err(|e| { DataFusionError::Execution(format!( "Failed to create layout reader: {e}" @@ -286,7 +285,7 @@ impl FileOpener for VortexOpener { } } Entry::Vacant(vacant_entry) => { - tracing::trace!("creating layout reader for {}", vacant_entry.key()); + log::trace!("creating layout reader for {}", vacant_entry.key()); let reader = vxf.layout_reader().map_err(|e| { DataFusionError::Execution(format!("Failed to create layout reader: {e}")) })?; @@ -296,18 +295,11 @@ impl FileOpener for VortexOpener { } }; - let mut scan_builder = ScanBuilder::new(session.clone(), layout_reader); - - if let Some(extensions) = file.extensions - && let Some(vortex_plan) = extensions.downcast_ref::() - { - scan_builder = vortex_plan.apply_to_builder(scan_builder); - } - - if let Some(file_range) = file.range { + let mut scan_builder = ScanBuilder::new(session, layout_reader); + if let Some(file_range) = file_meta.range { scan_builder = apply_byte_range( file_range, - file.object_meta.size, + file_meta.object_meta.size, vxf.row_count(), scan_builder, ); @@ -315,10 +307,10 @@ impl FileOpener for VortexOpener { let filter = filter .and_then(|f| { - // Verify that all filters we've accepted from DataFusion get pushed down. - // This will only fail if the user has not configured a suitable - // PhysicalExprAdapterFactory on the file source to handle rewriting the - // expression to handle missing/reordered columns in the Vortex file. + let exprs = split_conjunction(&f) + .into_iter() + .filter(|expr| can_be_pushed_down(expr, &predicate_file_schema)) + .collect::>(); let (pushed, unpushed): (Vec, Vec) = split_conjunction(&f) @@ -391,7 +383,7 @@ impl FileOpener for VortexOpener { .map_err(move |e: VortexError| { DataFusionError::External(Box::new(e.with_context(format!( "Failed to read Vortex file: {}", - file.object_meta.location + file_meta.object_meta.location )))) }) .try_flatten() @@ -405,9 +397,13 @@ impl FileOpener for VortexOpener { .boxed(); if let Some(file_pruner) = file_pruner { - Ok(PrunableStream::new(file_pruner, stream).boxed()) + Ok(Box::pin(EarlyStoppingStream::new( + stream, + file_pruner, + Count::new(), + ))) } else { - Ok(stream) + Ok(Box::pin(stream)) } } .in_current_span() @@ -415,7 +411,7 @@ impl FileOpener for VortexOpener { } } -/// If the file has a [`FileRange`], we translate it into a row range in the file for the scan. +/// If the file has a [`FileRange`](datafusion::datasource::listing::FileRange), we translate it into a row range in the file for the scan. fn apply_byte_range( file_range: FileRange, total_size: u64, @@ -447,7 +443,6 @@ mod tests { use std::sync::Arc; use std::sync::LazyLock; - use arrow_schema::Field; use arrow_schema::Fields; use arrow_schema::SchemaRef; use datafusion::arrow::array::DictionaryArray; @@ -458,7 +453,6 @@ mod tests { use datafusion::arrow::datatypes::Schema; use datafusion::arrow::datatypes::UInt32Type; use datafusion::arrow::util::display::FormatOptions; - use datafusion::arrow::util::pretty::pretty_format_batches_with_options; use datafusion::common::record_batch; use datafusion::logical_expr::col; use datafusion::logical_expr::lit; @@ -595,7 +589,7 @@ mod tests { // filter matches partition value let filter = col("part").eq(lit(1)); - let filter = logical2physical(&filter, table_schema.table_schema()); + let filter = logical2physical(&filter, table_schema.as_ref()); let opener = make_opener(object_store.clone(), table_schema.clone(), Some(filter)); let stream = opener.open(file.clone()).unwrap().await.unwrap(); @@ -604,11 +598,11 @@ mod tests { let num_batches = data.len(); let num_rows = data.iter().map(|rb| rb.num_rows()).sum::(); - assert_eq!((num_batches, num_rows), (1, 3)); + assert_eq!((num_batches, num_rows), expected_result1); // filter doesn't matches partition value let filter = col("part").eq(lit(2)); - let filter = logical2physical(&filter, table_schema.table_schema()); + let filter = logical2physical(&filter, table_schema.as_ref()); let opener = make_opener(object_store.clone(), table_schema.clone(), Some(filter)); let stream = opener.open(file.clone()).unwrap().await.unwrap(); @@ -625,29 +619,18 @@ mod tests { #[tokio::test] async fn test_open_files_different_table_schema() -> anyhow::Result<()> { let object_store = Arc::new(InMemory::new()) as Arc; + let file1_path = "/path/file1.vortex"; + let batch1 = record_batch!(("a", Int32, vec![Some(1), Some(2), Some(3)])).unwrap(); + let data_size1 = write_arrow_to_vortex(object_store.clone(), file1_path, batch1).await?; + let file1 = PartitionedFile::new(file1_path.to_string(), data_size1); - let file1 = { - let file1_path = "/path/file1.vortex"; - let batch1 = record_batch!(("a", Int32, vec![Some(1), Some(2), Some(3)])).unwrap(); - let data_size1 = - write_arrow_to_vortex(object_store.clone(), file1_path, batch1).await?; - PartitionedFile::new(file1_path.to_string(), data_size1) - }; - - let file2 = { - let file2_path = "/path/file2.vortex"; - let batch2 = record_batch!(("a", Int16, vec![Some(-1), Some(-2), Some(-3)])).unwrap(); - let data_size2 = - write_arrow_to_vortex(object_store.clone(), file2_path, batch2).await?; - PartitionedFile::new(file2_path.to_string(), data_size2) - }; + let file2_path = "/path/file2.vortex"; + let batch2 = record_batch!(("a", Int16, vec![Some(-1), Some(-2), Some(-3)])).unwrap(); + let data_size2 = write_arrow_to_vortex(object_store.clone(), file2_path, batch2).await?; + let file2 = PartitionedFile::new(file1_path.to_string(), data_size1); // Table schema has can accommodate both files - let table_schema = TableSchema::from_file_schema(Arc::new(Schema::new(vec![Field::new( - "a", - DataType::Int32, - true, - )]))); + let table_schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); let make_opener = |filter| VortexOpener { partition: 1, @@ -670,10 +653,12 @@ mod tests { }; let filter = col("a").lt(lit(100_i32)); - let filter = logical2physical(&filter, table_schema.table_schema()); + let filter = logical2physical(&filter, table_schema.as_ref()); let opener1 = make_opener(filter.clone()); - let stream = opener1.open(file1)?.await?; + let stream = opener1 + .open(make_meta(file1_path, data_size1), file1)? + .await?; let format_opts = FormatOptions::new().with_types_info(true); @@ -690,7 +675,9 @@ mod tests { "); let opener2 = make_opener(filter.clone()); - let stream = opener2.open(file2)?.await?; + let stream = opener2 + .open(make_meta(file2_path, data_size2), file2)? + .await?; let data = stream.try_collect::>().await?; assert_snapshot!(pretty_format_batches_with_options(&data, &format_opts)?.to_string(), @r" @@ -806,7 +793,7 @@ mod tests { let data_size = write_arrow_to_vortex(object_store.clone(), file_path, batch).await?; // Table schema has an extra utf8 field. - let table_schema = TableSchema::from_file_schema(Arc::new(Schema::new(vec![Field::new( + let table_schema = Arc::new(Schema::new(vec![Field::new( "my_struct", DataType::Struct(Fields::from(vec![ Field::new( @@ -822,23 +809,25 @@ mod tests { Field::new("field3", DataType::Utf8, true), ])), true, - )]))); + )])); - let opener = make_opener( - object_store.clone(), - table_schema.clone(), - // expression references my_struct column which has different fields in each - // field. - Some(logical2physical( + let opener = VortexOpener { + session: SESSION.clone(), + object_store: object_store.clone(), + projection: None, + filter: Some(logical2physical( &col("my_struct").is_not_null(), - table_schema.table_schema(), + &table_schema, )), ); // The opener should be able to open the file with a filter on the // struct column. let data = opener - .open(PartitionedFile::new(file_path.to_string(), data_size))? + .open( + make_meta(file_path, data_size), + PartitionedFile::new(file_path.to_string(), data_size), + )? .await? .try_collect::>() .await?; diff --git a/vortex-datafusion/src/persistent/source.rs b/vortex-datafusion/src/persistent/source.rs index bac7fe2b39a..180ec18f11e 100644 --- a/vortex-datafusion/src/persistent/source.rs +++ b/vortex-datafusion/src/persistent/source.rs @@ -3,12 +3,11 @@ use std::any::Any; use std::fmt::Formatter; -use std::sync::Arc; -use std::sync::Weak; +use std::sync::{Arc, Weak}; use datafusion_common::Result as DFResult; use datafusion_common::config::ConfigOptions; -use datafusion_datasource::TableSchema; +use datafusion_common::{Result as DFResult, Statistics}; use datafusion_datasource::file::FileSource; use datafusion_datasource::file_scan_config::FileScanConfig; use datafusion_datasource::file_stream::FileOpener; @@ -18,12 +17,11 @@ use datafusion_physical_expr::conjunction; use datafusion_physical_expr::projection::ProjectionExprs; use datafusion_physical_expr_adapter::DefaultPhysicalExprAdapterFactory; use datafusion_physical_expr_common::physical_expr::fmt_sql; -use datafusion_physical_plan::DisplayFormatType; -use datafusion_physical_plan::PhysicalExpr; -use datafusion_physical_plan::filter_pushdown::FilterPushdownPropagation; -use datafusion_physical_plan::filter_pushdown::PushedDown; -use datafusion_physical_plan::filter_pushdown::PushedDownPredicate; +use datafusion_physical_plan::filter_pushdown::{ + FilterPushdownPropagation, PushedDown, PushedDownPredicate, +}; use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; +use datafusion_physical_plan::{DisplayFormatType, PhysicalExpr}; use object_store::ObjectStore; use object_store::path::Path; use vortex::error::VortexExpect; From 8aedcb4de6d868c783d75102d5176e614c7d5bbf Mon Sep 17 00:00:00 2001 From: Luke Kim <80174+lukekim@users.noreply.github.com> Date: Wed, 14 Jan 2026 16:36:37 -0800 Subject: [PATCH 14/32] revert: Restore spiceai-51 versions of persistent module (incompatible with DF51) --- vortex-datafusion/src/convert/exprs.rs | 310 +++++---------------- vortex-datafusion/src/lib.rs | 2 +- vortex-datafusion/src/persistent/opener.rs | 121 ++++---- vortex-datafusion/src/persistent/source.rs | 16 +- 4 files changed, 145 insertions(+), 304 deletions(-) diff --git a/vortex-datafusion/src/convert/exprs.rs b/vortex-datafusion/src/convert/exprs.rs index 9173bc7330c..a047d1e6211 100644 --- a/vortex-datafusion/src/convert/exprs.rs +++ b/vortex-datafusion/src/convert/exprs.rs @@ -39,7 +39,7 @@ use vortex::scalar_fn::fns::like::Like; use vortex::scalar_fn::fns::like::LikeOptions; use vortex::scalar_fn::fns::operators::Operator; -use crate::convert::{FromDataFusion, TryFromDataFusion}; +use crate::convert::FromDataFusion; /// Result of splitting a projection into Vortex expressions and leftover DataFusion projections. pub struct ProcessedProjection { @@ -189,9 +189,9 @@ impl ExpressionConvertor for DefaultExpressionConvertor { // TODO(joe): Don't return an error when we have an unsupported node, bubble up "TRUE" as in keep // for that node, up to any `and` or `or` node. if let Some(binary_expr) = df.as_any().downcast_ref::() { - let left = Expression::try_from_df(binary_expr.left().as_ref())?; - let right = Expression::try_from_df(binary_expr.right().as_ref())?; - let operator = Operator::try_from_df(binary_expr.op())?; + let left = self.convert(binary_expr.left().as_ref())?; + let right = self.convert(binary_expr.right().as_ref())?; + let operator = try_operator_from_df(binary_expr.op())?; return Ok(Binary.new_expr(operator, [left, right])); } @@ -201,8 +201,8 @@ impl ExpressionConvertor for DefaultExpressionConvertor { } if let Some(like) = df.as_any().downcast_ref::() { - let child = Expression::try_from_df(like.expr().as_ref())?; - let pattern = Expression::try_from_df(like.pattern().as_ref())?; + let child = self.convert(like.expr().as_ref())?; + let pattern = self.convert(like.pattern().as_ref())?; return Ok(Like.new_expr( LikeOptions { negated: like.negated(), @@ -219,22 +219,30 @@ impl ExpressionConvertor for DefaultExpressionConvertor { if let Some(cast_expr) = df.as_any().downcast_ref::() { let cast_dtype = DType::from_arrow((cast_expr.cast_type(), Nullability::Nullable)); - let child = Expression::try_from_df(cast_expr.expr().as_ref())?; + let child = self.convert(cast_expr.expr().as_ref())?; return Ok(cast(child, cast_dtype)); } + if let Some(cast_col_expr) = df.as_any().downcast_ref::() { + let target = cast_col_expr.target_field(); + + let target_dtype = DType::from_arrow((target.data_type(), target.is_nullable().into())); + let child = self.convert(cast_col_expr.expr().as_ref())?; + return Ok(cast(child, target_dtype)); + } + if let Some(is_null_expr) = df.as_any().downcast_ref::() { - let arg = Expression::try_from_df(is_null_expr.arg().as_ref())?; + let arg = self.convert(is_null_expr.arg().as_ref())?; return Ok(is_null(arg)); } if let Some(is_not_null_expr) = df.as_any().downcast_ref::() { - let arg = Expression::try_from_df(is_not_null_expr.arg().as_ref())?; + let arg = self.convert(is_not_null_expr.arg().as_ref())?; return Ok(not(is_null(arg))); } if let Some(in_list) = df.as_any().downcast_ref::() { - let value = Expression::try_from_df(in_list.expr().as_ref())?; + let value = self.convert(in_list.expr().as_ref())?; let list_elements: Vec<_> = in_list .list() .iter() @@ -247,13 +255,8 @@ impl ExpressionConvertor for DefaultExpressionConvertor { }) .try_collect()?; - // Handle empty IN list: `x IN ()` is always false, `x NOT IN ()` is always true - let Some(first_element) = list_elements.first() else { - return Ok(lit(in_list.negated())); - }; - let list = Scalar::list( - first_element.dtype().clone(), + list_elements[0].dtype().clone(), list_elements, Nullability::Nullable, ); @@ -263,16 +266,7 @@ impl ExpressionConvertor for DefaultExpressionConvertor { } if let Some(scalar_fn) = df.as_any().downcast_ref::() { - return try_convert_scalar_function(scalar_fn); - } - - if let Some(dynamic_expr) = df - .as_any() - .downcast_ref::() - && let Ok(current) = dynamic_expr.current() - { - let returned_expr = Expression::try_from_df(current.as_ref())?; - return Ok(returned_expr); + return self.try_convert_scalar_function(scalar_fn); } if let Some(case_expr) = df.as_any().downcast_ref::() { @@ -418,8 +412,10 @@ fn can_be_pushed_down_impl(df_expr: &Arc, schema: &Schema) -> && can_be_pushed_down_impl(like.pattern(), schema) } else if let Some(lit) = expr.downcast_ref::() { supported_data_types(&lit.value().data_type()) - } else if let Some(cast) = expr.downcast_ref::() { - supported_data_types(cast.cast_type()) && can_be_pushed_down(cast.expr(), schema) + } else if expr.downcast_ref::().is_some() + || expr.downcast_ref::().is_some() + { + true } else if let Some(is_null) = expr.downcast_ref::() { can_be_pushed_down_impl(is_null.arg(), schema) } else if let Some(is_not_null) = expr.downcast_ref::() { @@ -447,7 +443,7 @@ fn can_be_pushed_down_impl(df_expr: &Arc, schema: &Schema) -> } fn can_binary_be_pushed_down(binary: &df_expr::BinaryExpr, schema: &Schema) -> bool { - let is_op_supported = Operator::try_from_df(binary.op()).is_ok(); + let is_op_supported = try_operator_from_df(binary.op()).is_ok(); is_op_supported && can_be_pushed_down_impl(binary.left(), schema) && can_be_pushed_down_impl(binary.right(), schema) @@ -496,57 +492,15 @@ fn supported_data_types(dt: &DataType) -> bool { ); if !is_supported { - tracing::debug!(data_type = ?dt, "DataFusion data type is not supported"); + tracing::debug!("DataFusion data type {dt:?} is not supported"); } is_supported } /// Checks if a GetField scalar function can be pushed down. -fn can_scalar_fn_be_pushed_down(scalar_fn: &ScalarFunctionExpr, schema: &Schema) -> bool { - let Some(get_field_fn) = ScalarFunctionExpr::try_downcast_func::(scalar_fn) - else { - // Only get_field pushdown is supported. - return false; - }; - - let args = get_field_fn.args(); - if args.len() != 2 { - tracing::debug!( - "Expected 2 arguments for GetField, not pushing down {} arguments", - args.len() - ); - return false; - } - let source_expr = &args[0]; - let field_name_expr = &args[1]; - let Some(field_name) = field_name_expr - .as_any() - .downcast_ref::() - .and_then(|lit| lit.value().try_as_str().flatten()) - else { - return false; - }; - - let Ok(source_dt) = source_expr.data_type(schema) else { - tracing::debug!( - field_name = field_name, - schema = ?schema, - source_expr = ?source_expr, - "Failed to get source type for GetField, not pushing down" - ); - return false; - }; - let DataType::Struct(fields) = source_dt else { - tracing::debug!( - field_name = field_name, - schema = ?schema, - source_expr = ?source_expr, - "Failed to get source type as struct for GetField, not pushing down" - ); - return false; - }; - fields.find(field_name).is_some() +fn can_scalar_fn_be_pushed_down(scalar_fn: &ScalarFunctionExpr) -> bool { + ScalarFunctionExpr::try_downcast_func::(scalar_fn).is_some() } // TODO(adam): Replace with `DataType::is_decimal` once its released. @@ -564,16 +518,16 @@ fn is_decimal(dt: &DataType) -> bool { mod tests { use std::sync::Arc; - use arrow_schema::{DataType, Field, Fields, Schema, TimeUnit as ArrowTimeUnit}; - use datafusion::functions::core::getfield::GetFieldFunc; + use arrow_schema::DataType; + use arrow_schema::Field; + use arrow_schema::Schema; + use arrow_schema::TimeUnit as ArrowTimeUnit; use datafusion_common::ScalarValue; - use datafusion_common::config::ConfigOptions; - use datafusion_expr::{Operator as DFOperator, ScalarUDF}; + use datafusion_expr::Operator as DFOperator; use datafusion_physical_expr::PhysicalExpr; use datafusion_physical_plan::expressions as df_expr; use insta::assert_snapshot; use rstest::rstest; - use vortex::expr::{Expression, Operator}; use super::*; use crate::common_tests::TestSessionContext; @@ -600,22 +554,25 @@ mod tests { #[test] fn test_make_vortex_predicate_empty() { - let result = make_vortex_predicate(&[]).unwrap(); + let expr_convertor = DefaultExpressionConvertor::default(); + let result = make_vortex_predicate(&expr_convertor, &[]).unwrap(); assert!(result.is_none()); } #[test] fn test_make_vortex_predicate_single() { + let expr_convertor = DefaultExpressionConvertor::default(); let col_expr = Arc::new(df_expr::Column::new("test", 0)) as Arc; - let result = make_vortex_predicate(&[&col_expr]).unwrap(); + let result = make_vortex_predicate(&expr_convertor, &[col_expr]).unwrap(); assert!(result.is_some()); } #[test] fn test_make_vortex_predicate_multiple() { + let expr_convertor = DefaultExpressionConvertor::default(); let col1 = Arc::new(df_expr::Column::new("col1", 0)) as Arc; let col2 = Arc::new(df_expr::Column::new("col2", 1)) as Arc; - let result = make_vortex_predicate(&[&col1, &col2]).unwrap(); + let result = make_vortex_predicate(&expr_convertor, &[col1, col2]).unwrap(); assert!(result.is_some()); // Result should be an AND expression combining the two columns } @@ -637,7 +594,7 @@ mod tests { #[case] df_op: DFOperator, #[case] expected_vortex_op: Operator, ) { - let result = Operator::try_from_df(&df_op).unwrap(); + let result = try_operator_from_df(&df_op).unwrap(); assert_eq!(result, expected_vortex_op); } @@ -647,7 +604,7 @@ mod tests { #[case::regex_match(DFOperator::RegexMatch)] #[case::like_match(DFOperator::LikeMatch)] fn test_operator_conversion_unsupported(#[case] df_op: DFOperator) { - let result = Operator::try_from_df(&df_op); + let result = try_operator_from_df(&df_op); assert!(result.is_err()); assert!( result @@ -660,20 +617,24 @@ mod tests { #[test] fn test_expr_from_df_column() { let col_expr = df_expr::Column::new("test_column", 0); - let result = Expression::try_from_df(&col_expr).unwrap(); + let result = DefaultExpressionConvertor::default() + .convert(&col_expr) + .unwrap(); - assert_snapshot!(result.display_tree().to_string(), @r#" - vortex.get_item "test_column" - └── input: vortex.root - "#); + assert_snapshot!(result.display_tree().to_string(), @r" + vortex.get_item(test_column) + └── input: vortex.root() + "); } #[test] fn test_expr_from_df_literal() { let literal_expr = df_expr::Literal::new(ScalarValue::Int32(Some(42))); - let result = Expression::try_from_df(&literal_expr).unwrap(); + let result = DefaultExpressionConvertor::default() + .convert(&literal_expr) + .unwrap(); - assert_snapshot!(result.display_tree().to_string(), @"vortex.literal 42i32"); + assert_snapshot!(result.display_tree().to_string(), @"vortex.literal(42i32)"); } #[test] @@ -683,14 +644,16 @@ mod tests { Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(42)))) as Arc; let binary_expr = df_expr::BinaryExpr::new(left, DFOperator::Eq, right); - let result = Expression::try_from_df(&binary_expr).unwrap(); + let result = DefaultExpressionConvertor::default() + .convert(&binary_expr) + .unwrap(); - assert_snapshot!(result.display_tree().to_string(), @r#" - vortex.binary = - ├── lhs: vortex.get_item "left" - │ └── input: vortex.root - └── rhs: vortex.literal 42i32 - "#); + assert_snapshot!(result.display_tree().to_string(), @r" + vortex.binary(=) + ├── lhs: vortex.get_item(left) + │ └── input: vortex.root() + └── rhs: vortex.literal(42i32) + "); } #[rstest] @@ -705,10 +668,12 @@ mod tests { )))) as Arc; let like_expr = df_expr::LikeExpr::new(negated, case_insensitive, expr, pattern); - let result = Expression::try_from_df(&like_expr).unwrap(); - let like_expr = result.as_::(); + let result = DefaultExpressionConvertor::default() + .convert(&like_expr) + .unwrap(); + let like_opts = result.as_::(); assert_eq!( - like_expr.data(), + like_opts, &LikeOptions { negated, case_insensitive @@ -793,7 +758,8 @@ mod tests { DataType::List(Arc::new(Field::new("item", DataType::Int32, true))), false )] - #[case::struct_type(DataType::Struct(vec![Field::new("field", DataType::Int32, true)].into()), false)] + #[case::struct_type(DataType::Struct(vec![Field::new("field", DataType::Int32, true)].into() + ), false)] // Dictionary types - should be supported if value type is supported #[case::dict_utf8( DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)), @@ -946,144 +912,4 @@ mod tests { Ok(()) } - - #[test] - fn test_expr_from_df_get_field() { - let struct_col = Arc::new(df_expr::Column::new("my_struct", 0)) as Arc; - let field_name = Arc::new(df_expr::Literal::new(ScalarValue::Utf8(Some( - "field1".to_string(), - )))) as Arc; - let get_field_expr = ScalarFunctionExpr::new( - "get_field", - Arc::new(ScalarUDF::from(GetFieldFunc::new())), - vec![struct_col, field_name], - Arc::new(Field::new("field1", DataType::Utf8, true)), - Arc::new(ConfigOptions::new()), - ); - let result = Expression::try_from_df(&get_field_expr).unwrap(); - assert_snapshot!(result.display_tree().to_string(), @r#" - vortex.get_item "field1" - └── input: vortex.get_item "my_struct" - └── input: vortex.root - "#); - } - - #[rstest] - #[case::valid_field("field1", true)] - #[case::missing_field("nonexistent_field", false)] - fn test_can_be_pushed_down_get_field(#[case] field_name: &str, #[case] expected: bool) { - let struct_fields = Fields::from(vec![ - Field::new("field1", DataType::Utf8, true), - Field::new("field2", DataType::Int32, true), - ]); - let schema = Schema::new(vec![Field::new( - "my_struct", - DataType::Struct(struct_fields), - true, - )]); - - let struct_col = Arc::new(df_expr::Column::new("my_struct", 0)) as Arc; - let field_name_lit = Arc::new(df_expr::Literal::new(ScalarValue::Utf8(Some( - field_name.to_string(), - )))) as Arc; - - let get_field_expr = Arc::new(ScalarFunctionExpr::new( - "get_field", - Arc::new(ScalarUDF::from(GetFieldFunc::new())), - vec![struct_col, field_name_lit], - Arc::new(Field::new(field_name, DataType::Utf8, true)), - Arc::new(ConfigOptions::new()), - )) as Arc; - - assert_eq!(can_be_pushed_down(&get_field_expr, &schema), expected); - } - - /// Create an unsupported scalar function expression (simulating functions like to_timestamp) - fn create_unsupported_scalar_fn() -> Arc { - use datafusion_functions::datetime::to_timestamp::ToTimestampFunc; - - let arg = Arc::new(df_expr::Literal::new(ScalarValue::Utf8(Some( - "2024-01-01".to_string(), - )))) as Arc; - - Arc::new(ScalarFunctionExpr::new( - "to_timestamp", - Arc::new(ScalarUDF::from(ToTimestampFunc::new())), - vec![arg], - Arc::new(Field::new( - "result", - DataType::Timestamp(ArrowTimeUnit::Nanosecond, None), - true, - )), - Arc::new(ConfigOptions::new()), - )) - } - - #[test] - fn test_make_vortex_predicate_skips_unsupported_scalar_function() { - // Unsupported scalar function like to_timestamp should be skipped, not error - let unsupported_fn = create_unsupported_scalar_fn(); - let result = make_vortex_predicate(&[&unsupported_fn]); - - // Should succeed (not error) and return None since the only expression was skipped - assert!(result.is_ok()); - assert!(result.unwrap().is_none()); - } - - #[test] - fn test_make_vortex_predicate_combines_supported_and_skips_unsupported() { - // Mix of supported column expression and unsupported scalar function - let supported_col = Arc::new(df_expr::Column::new("test", 0)) as Arc; - let unsupported_fn = create_unsupported_scalar_fn(); - - let result = make_vortex_predicate(&[&supported_col, &unsupported_fn]); - - // Should succeed and return the supported expression only - assert!(result.is_ok()); - let predicate = result.unwrap(); - assert!(predicate.is_some()); - - // The result should just be the column expression since the unsupported one was skipped - assert_snapshot!(predicate.unwrap().display_tree().to_string(), @r#" - vortex.get_item "test" - └── input: vortex.root - "#); - } - - #[test] - fn test_make_vortex_predicate_multiple_supported_with_unsupported() { - // Two supported columns and one unsupported function - let col1 = Arc::new(df_expr::Column::new("col1", 0)) as Arc; - let col2 = Arc::new(df_expr::Column::new("col2", 1)) as Arc; - let unsupported_fn = create_unsupported_scalar_fn(); - - let result = make_vortex_predicate(&[&col1, &unsupported_fn, &col2]); - - // Should succeed and return AND of the two supported expressions - assert!(result.is_ok()); - let predicate = result.unwrap(); - assert!(predicate.is_some()); - - // The result should be an AND of col1 and col2 - assert_snapshot!(predicate.unwrap().display_tree().to_string(), @r#" - vortex.binary and - ├── lhs: vortex.get_item "col1" - │ └── input: vortex.root - └── rhs: vortex.get_item "col2" - └── input: vortex.root - "#); - } - - #[test] - fn test_make_vortex_predicate_all_unsupported_returns_none() { - // When all expressions are unsupported, should return None (no filter) - let unsupported_fn1 = create_unsupported_scalar_fn(); - let unsupported_fn2 = create_unsupported_scalar_fn(); - - let result = make_vortex_predicate(&[&unsupported_fn1, &unsupported_fn2]); - - // Should succeed and return None - assert!(result.is_ok()); - assert!(result.unwrap().is_none()); - } } diff --git a/vortex-datafusion/src/lib.rs b/vortex-datafusion/src/lib.rs index b976215aa18..0e09886c07c 100644 --- a/vortex-datafusion/src/lib.rs +++ b/vortex-datafusion/src/lib.rs @@ -6,7 +6,7 @@ use std::fmt::Debug; use datafusion_common::stats::Precision as DFPrecision; -use vortex::stats::Precision; +use vortex::expr::stats::Precision; mod convert; mod persistent; diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs index 199eb2e4912..5986a06da49 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/vortex-datafusion/src/persistent/opener.rs @@ -2,7 +2,8 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::ops::Range; -use std::sync::{Arc, Weak}; +use std::sync::Arc; +use std::sync::Weak; use arrow_schema::Schema; use datafusion_common::DataFusionError; @@ -42,7 +43,8 @@ use vortex::metrics::Label; use vortex::metrics::MetricsRegistry; use vortex::scan::ScanBuilder; use vortex::session::VortexSession; -use vortex_utils::aliases::dash_map::{DashMap, Entry}; +use vortex_utils::aliases::dash_map::DashMap; +use vortex_utils::aliases::dash_map::Entry; use crate::VortexAccessPlan; use crate::convert::exprs::ExpressionConvertor; @@ -88,7 +90,6 @@ pub(crate) struct VortexOpener { pub layout_readers: Arc>>, /// Whether the query has output ordering specified pub has_output_ordering: bool, -} pub expression_convertor: Arc, pub file_metadata_cache: Option>, @@ -98,7 +99,7 @@ pub(crate) struct VortexOpener { } impl FileOpener for VortexOpener { - fn open(&self, file_meta: FileMeta, file: PartitionedFile) -> DFResult { + fn open(&self, file: PartitionedFile) -> DFResult { let session = self.session.clone(); let metrics_registry = self.metrics_registry.clone(); let labels = vec![ @@ -157,7 +158,7 @@ impl FileOpener for VortexOpener { // - Partition column values (e.g., date=2024-01-01) // - File-level statistics (min/max values per column) let mut file_pruner = file_pruning_predicate - .map(|predicate| { + .filter(|p| { // Only create pruner if we have dynamic expressions or file statistics // to work with. Static predicates without stats won't benefit from pruning. is_dynamic_physical_expr(p) || file.has_statistics() @@ -268,13 +269,13 @@ impl FileOpener for VortexOpener { let projector = leftover_projection.make_projector(&stream_schema)?; // We share our layout readers with others partitions in the scan, so we can only need to read each layout in each file once. - let layout_reader = match layout_reader.entry(file_meta.object_meta.location.clone()) { + let layout_reader = match layout_reader.entry(file.object_meta.location.clone()) { Entry::Occupied(mut occupied_entry) => { if let Some(reader) = occupied_entry.get().upgrade() { - log::trace!("reusing layout reader for {}", occupied_entry.key()); + tracing::trace!("reusing layout reader for {}", occupied_entry.key()); reader } else { - log::trace!("creating layout reader for {}", occupied_entry.key()); + tracing::trace!("creating layout reader for {}", occupied_entry.key()); let reader = vxf.layout_reader().map_err(|e| { DataFusionError::Execution(format!( "Failed to create layout reader: {e}" @@ -285,7 +286,7 @@ impl FileOpener for VortexOpener { } } Entry::Vacant(vacant_entry) => { - log::trace!("creating layout reader for {}", vacant_entry.key()); + tracing::trace!("creating layout reader for {}", vacant_entry.key()); let reader = vxf.layout_reader().map_err(|e| { DataFusionError::Execution(format!("Failed to create layout reader: {e}")) })?; @@ -295,11 +296,18 @@ impl FileOpener for VortexOpener { } }; - let mut scan_builder = ScanBuilder::new(session, layout_reader); - if let Some(file_range) = file_meta.range { + let mut scan_builder = ScanBuilder::new(session.clone(), layout_reader); + + if let Some(extensions) = file.extensions + && let Some(vortex_plan) = extensions.downcast_ref::() + { + scan_builder = vortex_plan.apply_to_builder(scan_builder); + } + + if let Some(file_range) = file.range { scan_builder = apply_byte_range( file_range, - file_meta.object_meta.size, + file.object_meta.size, vxf.row_count(), scan_builder, ); @@ -307,10 +315,10 @@ impl FileOpener for VortexOpener { let filter = filter .and_then(|f| { - let exprs = split_conjunction(&f) - .into_iter() - .filter(|expr| can_be_pushed_down(expr, &predicate_file_schema)) - .collect::>(); + // Verify that all filters we've accepted from DataFusion get pushed down. + // This will only fail if the user has not configured a suitable + // PhysicalExprAdapterFactory on the file source to handle rewriting the + // expression to handle missing/reordered columns in the Vortex file. let (pushed, unpushed): (Vec, Vec) = split_conjunction(&f) @@ -383,7 +391,7 @@ impl FileOpener for VortexOpener { .map_err(move |e: VortexError| { DataFusionError::External(Box::new(e.with_context(format!( "Failed to read Vortex file: {}", - file_meta.object_meta.location + file.object_meta.location )))) }) .try_flatten() @@ -397,13 +405,9 @@ impl FileOpener for VortexOpener { .boxed(); if let Some(file_pruner) = file_pruner { - Ok(Box::pin(EarlyStoppingStream::new( - stream, - file_pruner, - Count::new(), - ))) + Ok(PrunableStream::new(file_pruner, stream).boxed()) } else { - Ok(Box::pin(stream)) + Ok(stream) } } .in_current_span() @@ -411,7 +415,7 @@ impl FileOpener for VortexOpener { } } -/// If the file has a [`FileRange`](datafusion::datasource::listing::FileRange), we translate it into a row range in the file for the scan. +/// If the file has a [`FileRange`], we translate it into a row range in the file for the scan. fn apply_byte_range( file_range: FileRange, total_size: u64, @@ -443,6 +447,7 @@ mod tests { use std::sync::Arc; use std::sync::LazyLock; + use arrow_schema::Field; use arrow_schema::Fields; use arrow_schema::SchemaRef; use datafusion::arrow::array::DictionaryArray; @@ -453,6 +458,7 @@ mod tests { use datafusion::arrow::datatypes::Schema; use datafusion::arrow::datatypes::UInt32Type; use datafusion::arrow::util::display::FormatOptions; + use datafusion::arrow::util::pretty::pretty_format_batches_with_options; use datafusion::common::record_batch; use datafusion::logical_expr::col; use datafusion::logical_expr::lit; @@ -589,7 +595,7 @@ mod tests { // filter matches partition value let filter = col("part").eq(lit(1)); - let filter = logical2physical(&filter, table_schema.as_ref()); + let filter = logical2physical(&filter, table_schema.table_schema()); let opener = make_opener(object_store.clone(), table_schema.clone(), Some(filter)); let stream = opener.open(file.clone()).unwrap().await.unwrap(); @@ -598,11 +604,11 @@ mod tests { let num_batches = data.len(); let num_rows = data.iter().map(|rb| rb.num_rows()).sum::(); - assert_eq!((num_batches, num_rows), expected_result1); + assert_eq!((num_batches, num_rows), (1, 3)); // filter doesn't matches partition value let filter = col("part").eq(lit(2)); - let filter = logical2physical(&filter, table_schema.as_ref()); + let filter = logical2physical(&filter, table_schema.table_schema()); let opener = make_opener(object_store.clone(), table_schema.clone(), Some(filter)); let stream = opener.open(file.clone()).unwrap().await.unwrap(); @@ -619,18 +625,29 @@ mod tests { #[tokio::test] async fn test_open_files_different_table_schema() -> anyhow::Result<()> { let object_store = Arc::new(InMemory::new()) as Arc; - let file1_path = "/path/file1.vortex"; - let batch1 = record_batch!(("a", Int32, vec![Some(1), Some(2), Some(3)])).unwrap(); - let data_size1 = write_arrow_to_vortex(object_store.clone(), file1_path, batch1).await?; - let file1 = PartitionedFile::new(file1_path.to_string(), data_size1); - let file2_path = "/path/file2.vortex"; - let batch2 = record_batch!(("a", Int16, vec![Some(-1), Some(-2), Some(-3)])).unwrap(); - let data_size2 = write_arrow_to_vortex(object_store.clone(), file2_path, batch2).await?; - let file2 = PartitionedFile::new(file1_path.to_string(), data_size1); + let file1 = { + let file1_path = "/path/file1.vortex"; + let batch1 = record_batch!(("a", Int32, vec![Some(1), Some(2), Some(3)])).unwrap(); + let data_size1 = + write_arrow_to_vortex(object_store.clone(), file1_path, batch1).await?; + PartitionedFile::new(file1_path.to_string(), data_size1) + }; + + let file2 = { + let file2_path = "/path/file2.vortex"; + let batch2 = record_batch!(("a", Int16, vec![Some(-1), Some(-2), Some(-3)])).unwrap(); + let data_size2 = + write_arrow_to_vortex(object_store.clone(), file2_path, batch2).await?; + PartitionedFile::new(file2_path.to_string(), data_size2) + }; // Table schema has can accommodate both files - let table_schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + let table_schema = TableSchema::from_file_schema(Arc::new(Schema::new(vec![Field::new( + "a", + DataType::Int32, + true, + )]))); let make_opener = |filter| VortexOpener { partition: 1, @@ -653,12 +670,10 @@ mod tests { }; let filter = col("a").lt(lit(100_i32)); - let filter = logical2physical(&filter, table_schema.as_ref()); + let filter = logical2physical(&filter, table_schema.table_schema()); let opener1 = make_opener(filter.clone()); - let stream = opener1 - .open(make_meta(file1_path, data_size1), file1)? - .await?; + let stream = opener1.open(file1)?.await?; let format_opts = FormatOptions::new().with_types_info(true); @@ -675,9 +690,7 @@ mod tests { "); let opener2 = make_opener(filter.clone()); - let stream = opener2 - .open(make_meta(file2_path, data_size2), file2)? - .await?; + let stream = opener2.open(file2)?.await?; let data = stream.try_collect::>().await?; assert_snapshot!(pretty_format_batches_with_options(&data, &format_opts)?.to_string(), @r" @@ -793,7 +806,7 @@ mod tests { let data_size = write_arrow_to_vortex(object_store.clone(), file_path, batch).await?; // Table schema has an extra utf8 field. - let table_schema = Arc::new(Schema::new(vec![Field::new( + let table_schema = TableSchema::from_file_schema(Arc::new(Schema::new(vec![Field::new( "my_struct", DataType::Struct(Fields::from(vec![ Field::new( @@ -809,25 +822,23 @@ mod tests { Field::new("field3", DataType::Utf8, true), ])), true, - )])); + )]))); - let opener = VortexOpener { - session: SESSION.clone(), - object_store: object_store.clone(), - projection: None, - filter: Some(logical2physical( + let opener = make_opener( + object_store.clone(), + table_schema.clone(), + // expression references my_struct column which has different fields in each + // field. + Some(logical2physical( &col("my_struct").is_not_null(), - &table_schema, + table_schema.table_schema(), )), ); // The opener should be able to open the file with a filter on the // struct column. let data = opener - .open( - make_meta(file_path, data_size), - PartitionedFile::new(file_path.to_string(), data_size), - )? + .open(PartitionedFile::new(file_path.to_string(), data_size))? .await? .try_collect::>() .await?; diff --git a/vortex-datafusion/src/persistent/source.rs b/vortex-datafusion/src/persistent/source.rs index 180ec18f11e..c713987bcd3 100644 --- a/vortex-datafusion/src/persistent/source.rs +++ b/vortex-datafusion/src/persistent/source.rs @@ -3,11 +3,12 @@ use std::any::Any; use std::fmt::Formatter; -use std::sync::{Arc, Weak}; +use std::sync::Arc; +use std::sync::Weak; use datafusion_common::Result as DFResult; use datafusion_common::config::ConfigOptions; -use datafusion_common::{Result as DFResult, Statistics}; +use datafusion_datasource::TableSchema; use datafusion_datasource::file::FileSource; use datafusion_datasource::file_scan_config::FileScanConfig; use datafusion_datasource::file_stream::FileOpener; @@ -17,11 +18,12 @@ use datafusion_physical_expr::conjunction; use datafusion_physical_expr::projection::ProjectionExprs; use datafusion_physical_expr_adapter::DefaultPhysicalExprAdapterFactory; use datafusion_physical_expr_common::physical_expr::fmt_sql; -use datafusion_physical_plan::filter_pushdown::{ - FilterPushdownPropagation, PushedDown, PushedDownPredicate, -}; +use datafusion_physical_plan::DisplayFormatType; +use datafusion_physical_plan::PhysicalExpr; +use datafusion_physical_plan::filter_pushdown::FilterPushdownPropagation; +use datafusion_physical_plan::filter_pushdown::PushedDown; +use datafusion_physical_plan::filter_pushdown::PushedDownPredicate; use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; -use datafusion_physical_plan::{DisplayFormatType, PhysicalExpr}; use object_store::ObjectStore; use object_store::path::Path; use vortex::error::VortexExpect; @@ -173,6 +175,8 @@ impl FileSource for VortexSource { .clone() .unwrap_or_else(|| Arc::new(DefaultVortexReaderFactory::new(object_store))); + let table_schema = base_config.table_schema.clone(); + let opener = VortexOpener { partition, session: self.session.clone(), From c55e9c0b9508ab83b00e07a2a69286d0b5770381 Mon Sep 17 00:00:00 2001 From: Luke Kim <80174+lukekim@users.noreply.github.com> Date: Wed, 14 Jan 2026 16:40:21 -0800 Subject: [PATCH 15/32] fix: Restore format.rs to spiceai-51 version --- vortex-datafusion/src/persistent/format.rs | 41 +++++++++++++++------- 1 file changed, 29 insertions(+), 12 deletions(-) diff --git a/vortex-datafusion/src/persistent/format.rs b/vortex-datafusion/src/persistent/format.rs index fbf3a7f0cb7..a3c1a543578 100644 --- a/vortex-datafusion/src/persistent/format.rs +++ b/vortex-datafusion/src/persistent/format.rs @@ -2,28 +2,33 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::any::Any; -use std::fmt::{Debug, Formatter}; +use std::fmt::Debug; +use std::fmt::Formatter; use std::sync::Arc; -use arrow_schema::{Schema, SchemaRef}; +use arrow_schema::Schema; +use arrow_schema::SchemaRef; use async_trait::async_trait; use datafusion_catalog::Session; +use datafusion_common::ColumnStatistics; +use datafusion_common::DataFusionError; +use datafusion_common::GetExt; +use datafusion_common::Result as DFResult; +use datafusion_common::Statistics; use datafusion_common::config::ConfigField; use datafusion_common::config_namespace; use datafusion_common::internal_datafusion_err; use datafusion_common::not_impl_err; use datafusion_common::parsers::CompressionTypeVariant; use datafusion_common::stats::Precision; -use datafusion_common::{ - ColumnStatistics, DataFusionError, GetExt, Result as DFResult, Statistics, config_namespace, - not_impl_err, -}; use datafusion_common_runtime::SpawnedTask; use datafusion_datasource::TableSchema; use datafusion_datasource::file::FileSource; use datafusion_datasource::file_compression_type::FileCompressionType; -use datafusion_datasource::file_format::{FileFormat, FileFormatFactory}; -use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder}; +use datafusion_datasource::file_format::FileFormat; +use datafusion_datasource::file_format::FileFormatFactory; +use datafusion_datasource::file_scan_config::FileScanConfig; +use datafusion_datasource::file_scan_config::FileScanConfigBuilder; use datafusion_datasource::file_sink_config::FileSinkConfig; use datafusion_datasource::sink::DataSinkExec; use datafusion_datasource::source::DataSourceExec; @@ -54,8 +59,6 @@ use vortex::io::object_store::ObjectStoreReadAt; use vortex::io::session::RuntimeSessionExt; use vortex::scalar::Scalar; use vortex::session::VortexSession; -use vortex::stats::{ArrayStats, Stat, StatsSet}; -use vortex::{VortexSessionDefault, stats}; use super::cache::CachedVortexMetadata; use super::sink::VortexSink; @@ -63,6 +66,8 @@ use super::source::VortexSource; use crate::PrecisionExt as _; use crate::convert::TryToDataFusion; +const DEFAULT_FOOTER_INITIAL_READ_SIZE_BYTES: usize = MAX_POSTSCRIPT_SIZE as usize + EOF_SIZE; + /// Vortex implementation of a DataFusion [`FileFormat`]. pub struct VortexFormat { session: VortexSession, @@ -121,7 +126,10 @@ impl GetExt for VortexFormatFactory { impl VortexFormatFactory { /// Creates a new instance with a default [`VortexSession`] and default options. - #[allow(clippy::new_without_default)] // FormatFactory defines `default` method, so having `Default` implementation is confusing. + #[expect( + clippy::new_without_default, + reason = "FormatFactory defines `default` method, so having `Default` implementation is confusing" + )] pub fn new() -> Self { Self { session: VortexSession::default(), @@ -154,7 +162,7 @@ impl VortexFormatFactory { } impl FileFormatFactory for VortexFormatFactory { - #[allow(clippy::disallowed_types)] + #[expect(clippy::disallowed_types, reason = "required by trait signature")] fn create( &self, _state: &dyn Session, @@ -567,4 +575,13 @@ mod tests { let format = VortexFormat::new_with_options(VortexSession::default(), opts); assert_eq!(format.options().footer_initial_read_size_bytes, 12345); } + + #[test] + fn format_plumbs_footer_initial_read_size() { + let mut opts = VortexOptions::default(); + opts.set("footer_initial_read_size_bytes", "12345").unwrap(); + + let format = VortexFormat::new_with_options(VortexSession::default(), opts); + assert_eq!(format.file_cache.footer_initial_read_size_bytes(), 12345); + } } From 4bfa4331b7543a9a38c700e922fcc661279fc7cd Mon Sep 17 00:00:00 2001 From: Luke Kim <80174+lukekim@users.noreply.github.com> Date: Wed, 21 Jan 2026 19:52:06 -0800 Subject: [PATCH 16/32] Add Case-When Expression and tests (#12) * fix: Ensure CastExpr/CastColumnExpr/ScalarFunctionExpr check children in can_be_pushed_down The can_be_pushed_down function was returning true for CastExpr and CastColumnExpr without checking if their child expressions are convertible. This caused runtime errors when the child contained expressions like CaseExpr that convert() cannot handle. Also fixed ScalarFunctionExpr to recursively check its arguments. Fixes spiceai/spiceai#9037 * Add Case-When Expression and tests * Implement execute() * Add additional tests and fix type issue * Fix toml lint * feat(case_when): implement lazy evaluation to avoid side effects in unevaluated branches This implements proper lazy evaluation in CaseWhen expression to ensure that THEN/ELSE branches are only evaluated for rows where they apply. This is critical for correctness when expressions have side effects like divide-by-zero panics. The implementation: 1. Evaluates conditions in order, tracking which rows have been matched 2. For each condition, computes an effective mask (condition AND NOT matched) 3. Uses filter() to create a scoped array with only matching rows 4. Evaluates THEN expression only on the filtered scope 5. Uses scatter_with_mask() to expand results back to original positions 6. Short-circuits when all rows are matched or all conditions fail This fixes TPC-DS Q73 which has a pattern like: CASE WHEN hd_vehicle_count > 0 THEN hd_dep_count/hd_vehicle_count ELSE NULL END Previously, the division would be evaluated for all rows including those where hd_vehicle_count=0, causing a divide-by-zero panic. Now the division is only evaluated for rows where the condition is true. Added test: test_evaluate_divide_by_zero_protected_by_case_when * Formatting --- vortex-array/src/scalar_fn/fns/case_when.rs | 1417 +++++++++++++++++++ vortex-datafusion/src/convert/exprs.rs | 103 +- vortex-io/Cargo.toml | 1 - vortex-proto/proto/expr.proto | 6 + vortex-proto/src/generated/vortex.expr.rs | 8 + vortex-session/src/lib.rs | 16 +- 6 files changed, 1539 insertions(+), 12 deletions(-) create mode 100644 vortex-array/src/scalar_fn/fns/case_when.rs diff --git a/vortex-array/src/scalar_fn/fns/case_when.rs b/vortex-array/src/scalar_fn/fns/case_when.rs new file mode 100644 index 00000000000..164dba0e174 --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/case_when.rs @@ -0,0 +1,1417 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! CASE WHEN expression for conditional value selection. +//! +//! This expression evaluates a series of WHEN conditions and returns the corresponding +//! THEN value for the first condition that evaluates to true. If no conditions match +//! and an ELSE clause is provided, the ELSE value is returned; otherwise, NULL is returned. +//! +//! # Structure +//! +//! The expression has children in the following order: +//! - pairs of (condition, value) for each WHEN/THEN clause +//! - optionally, a final ELSE value +//! +//! For example, `CASE WHEN a THEN 1 WHEN b THEN 2 ELSE 3 END` has children: +//! `[a, 1, b, 2, 3]` + +use std::fmt; +use std::fmt::Formatter; +use std::hash::Hash; +use std::sync::Arc; + +use prost::Message; +use vortex_dtype::DType; +use vortex_dtype::Nullability; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_proto::expr as pb; +use vortex_scalar::Scalar; +use vortex_vector::Datum; +use vortex_vector::VectorOps; + +use crate::ArrayRef; +use crate::IntoArray; +use crate::ToCanonical; +use crate::arrays::BoolArray; +use crate::arrays::ConstantArray; +use crate::compute::zip; +use crate::expr::Arity; +use crate::expr::ChildName; +use crate::expr::ExecutionArgs; +use crate::expr::ExprId; +use crate::expr::VTable; +use crate::expr::VTableExt; +use crate::expr::expression::Expression; + +/// Options for the CaseWhen expression. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub struct CaseWhenOptions { + /// Number of WHEN/THEN pairs (each pair contributes 2 children) + pub num_when_then_pairs: u32, + /// Whether an ELSE clause is present (contributes 1 child at the end) + pub has_else: bool, +} + +impl fmt::Display for CaseWhenOptions { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!( + f, + "case_when(pairs={}, else={})", + self.num_when_then_pairs, self.has_else + ) + } +} + +/// A CASE WHEN expression. +/// +/// Evaluates conditions in order and returns the value corresponding to the +/// first matching condition. +pub struct CaseWhen; + +impl VTable for CaseWhen { + type Options = CaseWhenOptions; + + fn id(&self) -> ExprId { + ExprId::from("vortex.case_when") + } + + fn serialize(&self, options: &Self::Options) -> VortexResult>> { + Ok(Some( + pb::CaseWhenOpts { + num_when_then_pairs: options.num_when_then_pairs, + has_else: options.has_else, + } + .encode_to_vec(), + )) + } + + fn deserialize(&self, metadata: &[u8]) -> VortexResult { + let opts = pb::CaseWhenOpts::decode(metadata)?; + Ok(CaseWhenOptions { + num_when_then_pairs: opts.num_when_then_pairs, + has_else: opts.has_else, + }) + } + + fn arity(&self, options: &Self::Options) -> Arity { + let num_children = + options.num_when_then_pairs as usize * 2 + if options.has_else { 1 } else { 0 }; + Arity::Exact(num_children) + } + + fn child_name(&self, options: &Self::Options, child_idx: usize) -> ChildName { + let pair_count = options.num_when_then_pairs as usize; + let num_when_then_children = pair_count * 2; + + if child_idx < num_when_then_children { + let pair_idx = child_idx / 2; + if child_idx % 2 == 0 { + ChildName::from(Arc::from(format!("when_{}", pair_idx))) + } else { + ChildName::from(Arc::from(format!("then_{}", pair_idx))) + } + } else if options.has_else && child_idx == num_when_then_children { + ChildName::from("else") + } else { + unreachable!( + "Invalid child index {} for CaseWhen expression with {} pairs", + child_idx, pair_count + ) + } + } + + fn fmt_sql( + &self, + options: &Self::Options, + expr: &Expression, + f: &mut Formatter<'_>, + ) -> fmt::Result { + write!(f, "CASE")?; + for i in 0..options.num_when_then_pairs as usize { + write!( + f, + " WHEN {} THEN {}", + expr.child(i * 2), + expr.child(i * 2 + 1) + )?; + } + if options.has_else { + let else_idx = options.num_when_then_pairs as usize * 2; + write!(f, " ELSE {}", expr.child(else_idx))?; + } + write!(f, " END") + } + + fn return_dtype(&self, options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult { + // The return dtype is based on the THEN expressions + if options.num_when_then_pairs == 0 { + vortex_bail!("CaseWhen must have at least one WHEN/THEN pair"); + } + + // Get the first THEN expression's dtype (index 1) + let first_then_dtype = &arg_dtypes[1]; + + // If there's no ELSE, the result is always nullable (unmatched rows are NULL) + if !options.has_else { + Ok(first_then_dtype.as_nullable()) + } else { + Ok(first_then_dtype.clone()) + } + } + + fn evaluate( + &self, + options: &Self::Options, + expr: &Expression, + scope: &ArrayRef, + ) -> VortexResult { + use vortex_buffer::BitBuffer; + use vortex_mask::Mask; + + use crate::compute::filter; + + let len = scope.len(); + + // Determine output dtype from first THEN expression + let output_dtype = expr.child(1).return_dtype(scope.dtype())?; + let else_idx = options.num_when_then_pairs as usize * 2; + + // Track which rows have matched a condition (using BitBuffer for boolean ops) + let mut matched_bits = BitBuffer::new_unset(len); + + // Start with null result - we'll fill in values as conditions match + let mut result: ArrayRef = + ConstantArray::new(Scalar::null(output_dtype.as_nullable()), len).into_array(); + + // Process when/then pairs in order (first match wins) + for i in 0..options.num_when_then_pairs as usize { + // Evaluate condition + let cond = expr.child(i * 2).evaluate(scope)?; + let cond_bool = cond.to_bool(); + let cond_mask = cond_bool.to_mask_fill_null_false(); + let cond_bits = cond_mask.to_bit_buffer(); + + // Compute which rows match THIS condition AND haven't matched a previous one + // effective_cond = cond AND NOT(already_matched) + let effective_bits = &cond_bits & &(!&matched_bits); + let effective_mask = Mask::from_buffer(effective_bits.clone()); + + // Short-circuit: skip THEN evaluation if no rows match this condition + if effective_mask.all_false() { + continue; + } + + // Evaluate THEN expression + let then_val = if effective_mask.all_true() { + // All rows match - safe to evaluate on full scope + expr.child(i * 2 + 1).evaluate(scope)? + } else { + // Filter scope to only matching rows, evaluate, then scatter back + let filtered_scope = filter(scope, &effective_mask)?; + let filtered_result = expr.child(i * 2 + 1).evaluate(&filtered_scope)?; + + // Scatter the filtered result back using builder + scatter_with_mask(&filtered_result, &effective_mask, &output_dtype, len)? + }; + + // Merge into result: use zip to overlay then_val where effective_mask is true + result = zip(&then_val, &result, &effective_mask)?; + + // Update matched_bits + matched_bits = &matched_bits | &effective_bits; + + // Short-circuit: if all rows have matched, we're done + if matched_bits.true_count() == len { + break; + } + } + + // Handle ELSE clause for unmatched rows + let unmatched_bits = !&matched_bits; + let unmatched_mask = Mask::from_buffer(unmatched_bits); + + if !unmatched_mask.all_false() { + let else_val = if options.has_else { + // Evaluate ELSE for unmatched rows + if unmatched_mask.all_true() { + expr.child(else_idx).evaluate(scope)? + } else { + let filtered_scope = filter(scope, &unmatched_mask)?; + let filtered_else = expr.child(else_idx).evaluate(&filtered_scope)?; + scatter_with_mask(&filtered_else, &unmatched_mask, &output_dtype, len)? + } + } else { + // No ELSE - unmatched rows stay null (already set in result) + return Ok(result); + }; + + result = zip(&else_val, &result, &unmatched_mask)?; + } + + Ok(result) + } + + fn execute(&self, options: &Self::Options, args: ExecutionArgs) -> VortexResult { + let row_count = args.row_count; + let mut datums = args.datums; + + // Check if all inputs are scalars (for returning scalar result) + let all_scalars = datums.iter().all(|d| matches!(d, Datum::Scalar(_))); + + // Collect when/then pairs from datums + let mut when_then_pairs = Vec::with_capacity(options.num_when_then_pairs as usize); + for i in 0..options.num_when_then_pairs as usize { + let cond = datums[i * 2].clone(); + let then_val = datums[i * 2 + 1].clone(); + when_then_pairs.push((cond, then_val)); + } + + // Get the else value if present + let else_value = options.has_else.then(|| { + let else_idx = options.num_when_then_pairs as usize * 2; + datums.remove(else_idx) + }); + + // Determine output dtype from return_dtype + let output_dtype = args.return_dtype; + + // Create the result by starting from the else value or null + let mut result: Datum = if let Some(else_val) = else_value { + else_val + } else { + // Create a null scalar of the output dtype, which will be repeated as needed + use vortex_vector::Scalar as VScalar; + Datum::Scalar(VScalar::null(&output_dtype)) + }; + + // Process when/then pairs in reverse order + // For each (condition, then_value), we select from then_value where condition is true + for (cond, then_val) in when_then_pairs.into_iter().rev() { + result = execute_zip(then_val, result, cond, row_count, &output_dtype)?; + } + + // If all inputs were scalars and result is still length 1, return as scalar + if all_scalars + && let Datum::Vector(v) = &result + && v.len() == 1 + { + return Ok(Datum::Scalar(v.scalar_at(0))); + } + + Ok(result) + } +} + +/// Helper function to perform zip operation on Datum values. +/// Selects from `if_true` where `condition` is true, otherwise from `if_false`. +fn execute_zip( + if_true: Datum, + if_false: Datum, + condition: Datum, + row_count: usize, + output_dtype: &DType, +) -> VortexResult { + use vortex_mask::Mask; + use vortex_vector::BoolDatum; + + use crate::LEGACY_SESSION; + use crate::VectorExecutor; + use crate::vectors::VectorIntoArray; + + let cond_bool = condition.into_bool(); + + // Convert condition to Mask using the same pattern as evaluate() + let mask = match cond_bool { + BoolDatum::Scalar(s) => { + let value = s.value().unwrap_or(false); // NULL treated as false + Mask::new(row_count, value) + } + BoolDatum::Vector(v) => { + // Convert to BoolArray and use to_mask_fill_null_false() for DRY + let bool_dtype = DType::Bool(Nullability::Nullable); + let bool_array: BoolArray = v.into_array(&bool_dtype); + bool_array.to_mask_fill_null_false() + } + }; + + // Short-circuit: if mask is all true, return if_true; if all false, return if_false + if mask.all_true() { + return Ok(if_true); + } + if mask.all_false() { + return Ok(if_false); + } + + // Convert datums to vectors for zip + let true_vector = if_true.unwrap_into_vector(row_count); + let false_vector = if_false.unwrap_into_vector(row_count); + + // Convert vectors to arrays for zip operation + let true_array = true_vector.into_array(output_dtype); + let false_array = false_vector.into_array(output_dtype); + + // Perform zip + let result_array = zip(&true_array, &false_array, &mask)?; + + // Convert back to vector + let result_vector = result_array.execute_vector(&LEGACY_SESSION)?; + + Ok(Datum::Vector(result_vector)) +} + +/// Creates a CASE WHEN expression with an ELSE clause. +/// +/// The children should be provided as: condition1, then1, condition2, then2, ..., else_value +/// +/// # Example +/// ```ignore +/// // CASE WHEN x > 0 THEN 'positive' WHEN x < 0 THEN 'negative' ELSE 'zero' END +/// case_when(vec![ +/// gt(col("x"), lit(0)), lit("positive"), +/// lt(col("x"), lit(0)), lit("negative"), +/// lit("zero"), +/// ]) +/// ``` +pub fn case_when>(children: I) -> Expression { + let children: Vec<_> = children.into_iter().collect(); + let num_children = children.len(); + + // Must have odd number of children (pairs + else) + assert!( + num_children >= 3 && num_children % 2 == 1, + "case_when requires at least one when/then pair and an else: got {} children", + num_children + ); + + #[allow(clippy::cast_possible_truncation)] + let num_when_then_pairs = ((num_children - 1) / 2) as u32; + let options = CaseWhenOptions { + num_when_then_pairs, + has_else: true, + }; + + CaseWhen.new_expr(options, children) +} + +/// Creates a CASE WHEN expression without an ELSE clause (returns NULL when no conditions match). +/// +/// The children should be provided as: condition1, then1, condition2, then2, ... +/// +/// # Example +/// ```ignore +/// // CASE WHEN x > 0 THEN 'positive' WHEN x < 0 THEN 'negative' END +/// // (returns NULL when x = 0) +/// case_when_no_else(vec![ +/// gt(col("x"), lit(0)), lit("positive"), +/// lt(col("x"), lit(0)), lit("negative"), +/// ]) +/// ``` +pub fn case_when_no_else>(children: I) -> Expression { + let children: Vec<_> = children.into_iter().collect(); + let num_children = children.len(); + + // Must have even number of children (pairs only) + assert!( + num_children >= 2 && num_children % 2 == 0, + "case_when_no_else requires at least one when/then pair: got {} children", + num_children + ); + + #[allow(clippy::cast_possible_truncation)] + let num_when_then_pairs = (num_children / 2) as u32; + let options = CaseWhenOptions { + num_when_then_pairs, + has_else: false, + }; + + CaseWhen.new_expr(options, children) +} + +/// Scatter values from a filtered (shorter) array back to their original positions. +/// The mask indicates which positions in the output should receive values from the source. +/// Positions where mask is false will be null. +fn scatter_with_mask( + source: &ArrayRef, + mask: &vortex_mask::Mask, + dtype: &DType, + output_len: usize, +) -> VortexResult { + use crate::builders::builder_with_capacity; + + let nullable_dtype = dtype.as_nullable(); + let mut builder = builder_with_capacity(&nullable_dtype, output_len); + let mut source_idx = 0; + + for i in 0..output_len { + if mask.value(i) { + // Copy value from source, casting to nullable if needed + let scalar = source.scalar_at(source_idx); + let nullable_scalar = scalar.cast(&nullable_dtype)?; + builder.append_scalar(&nullable_scalar)?; + source_idx += 1; + } else { + // Insert null + builder.append_null(); + } + } + + Ok(builder.finish()) +} + +#[cfg(test)] +mod tests { + use vortex_buffer::buffer; + use vortex_dtype::DType; + use vortex_dtype::Nullability; + use vortex_dtype::PType; + use vortex_error::VortexExpect as _; + use vortex_scalar::Scalar; + + use super::*; + use crate::IntoArray; + use crate::ToCanonical; + use crate::arrays::BoolArray; + use crate::arrays::PrimitiveArray; + use crate::arrays::StructArray; + use crate::expr::exprs::binary::eq; + use crate::expr::exprs::binary::gt; + use crate::expr::exprs::get_item::col; + use crate::expr::exprs::get_item::get_item; + use crate::expr::exprs::literal::lit; + use crate::expr::exprs::root::root; + use crate::expr::test_harness; + + // ==================== Serialization Tests ==================== + + #[test] + fn test_serialization_roundtrip() { + let options = CaseWhenOptions { + num_when_then_pairs: 2, + has_else: true, + }; + + let serialized = CaseWhen.serialize(&options).unwrap().unwrap(); + let deserialized = CaseWhen.deserialize(&serialized).unwrap(); + + assert_eq!(options, deserialized); + } + + #[test] + fn test_serialization_no_else() { + let options = CaseWhenOptions { + num_when_then_pairs: 3, + has_else: false, + }; + + let serialized = CaseWhen.serialize(&options).unwrap().unwrap(); + let deserialized = CaseWhen.deserialize(&serialized).unwrap(); + + assert_eq!(options, deserialized); + } + + // ==================== Display Tests ==================== + + #[test] + fn test_display_with_else() { + // CASE WHEN col > 0 THEN 100 ELSE 0 END + let condition = gt(col("value"), lit(0i32)); + let then_val = lit(100i32); + let else_val = lit(0i32); + + let expr = case_when([condition, then_val, else_val]); + let display = format!("{}", expr); + assert!(display.contains("CASE")); + assert!(display.contains("WHEN")); + assert!(display.contains("THEN")); + assert!(display.contains("ELSE")); + assert!(display.contains("END")); + } + + #[test] + fn test_display_no_else() { + // CASE WHEN col > 0 THEN 100 END + let condition = gt(col("value"), lit(0i32)); + let then_val = lit(100i32); + + let expr = case_when_no_else([condition, then_val]); + let display = format!("{}", expr); + assert!(display.contains("CASE")); + assert!(display.contains("WHEN")); + assert!(display.contains("THEN")); + assert!(!display.contains("ELSE")); + assert!(display.contains("END")); + } + + #[test] + fn test_display_multiple_conditions() { + let expr = case_when([ + gt(col("x"), lit(10i32)), + lit("high"), + gt(col("x"), lit(5i32)), + lit("medium"), + lit("low"), + ]); + let display = format!("{}", expr); + // Should contain two WHEN clauses + assert_eq!(display.matches("WHEN").count(), 2); + assert_eq!(display.matches("THEN").count(), 2); + } + + // ==================== DType Tests ==================== + + #[test] + fn test_return_dtype_with_else() { + let condition = lit(true); + let then_val = lit(100i32); + let else_val = lit(0i32); + + let expr = case_when([condition, then_val, else_val]); + let input_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let result_dtype = expr.return_dtype(&input_dtype).unwrap(); + // With else, result dtype matches the then expression + assert_eq!( + result_dtype, + DType::Primitive(PType::I32, Nullability::NonNullable) + ); + } + + #[test] + fn test_return_dtype_without_else_is_nullable() { + let condition = lit(true); + let then_val = lit(100i32); + + let expr = case_when_no_else([condition, then_val]); + let input_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let result_dtype = expr.return_dtype(&input_dtype).unwrap(); + // Without else, result is always nullable + assert_eq!( + result_dtype, + DType::Primitive(PType::I32, Nullability::Nullable) + ); + } + + #[test] + fn test_return_dtype_with_struct_input() { + let dtype = test_harness::struct_dtype(); + + // CASE WHEN $.col1 > 10 THEN 100 ELSE 0 END + let expr = case_when([ + gt(get_item("col1", root()), lit(10u16)), + lit(100i32), + lit(0i32), + ]); + + let result_dtype = expr.return_dtype(&dtype).unwrap(); + assert_eq!( + result_dtype, + DType::Primitive(PType::I32, Nullability::NonNullable) + ); + } + + // ==================== Arity Tests ==================== + + #[test] + fn test_arity_with_else() { + let options = CaseWhenOptions { + num_when_then_pairs: 2, + has_else: true, + }; + // 2 pairs (4 children) + 1 else = 5 children + assert_eq!(CaseWhen.arity(&options), Arity::Exact(5)); + } + + #[test] + fn test_arity_without_else() { + let options = CaseWhenOptions { + num_when_then_pairs: 2, + has_else: false, + }; + // 2 pairs (4 children) = 4 children + assert_eq!(CaseWhen.arity(&options), Arity::Exact(4)); + } + + #[test] + fn test_arity_single_condition() { + let options = CaseWhenOptions { + num_when_then_pairs: 1, + has_else: true, + }; + // 1 pair (2 children) + 1 else = 3 children + assert_eq!(CaseWhen.arity(&options), Arity::Exact(3)); + } + + // ==================== Child Name Tests ==================== + + #[test] + fn test_child_names() { + let options = CaseWhenOptions { + num_when_then_pairs: 2, + has_else: true, + }; + + assert_eq!(CaseWhen.child_name(&options, 0).to_string(), "when_0"); + assert_eq!(CaseWhen.child_name(&options, 1).to_string(), "then_0"); + assert_eq!(CaseWhen.child_name(&options, 2).to_string(), "when_1"); + assert_eq!(CaseWhen.child_name(&options, 3).to_string(), "then_1"); + assert_eq!(CaseWhen.child_name(&options, 4).to_string(), "else"); + } + + // ==================== Expression Manipulation Tests ==================== + + #[test] + fn test_replace_children() { + let expr = case_when([lit(true), lit(1i32), lit(0i32)]); + expr.with_children([lit(false), lit(2i32), lit(3i32)]) + .vortex_expect("operation should succeed in test"); + } + + // ==================== Evaluate Tests ==================== + + #[test] + fn test_evaluate_simple_condition() { + // Test: CASE WHEN value > 2 THEN 100 ELSE 0 END + // Input: [1, 2, 3, 4, 5] + // Expected: [0, 0, 100, 100, 100] + let test_array = + StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) + .unwrap() + .into_array(); + + let expr = case_when([ + gt(get_item("value", root()), lit(2i32)), + lit(100i32), + lit(0i32), + ]); + + let result = expr.evaluate(&test_array).unwrap().to_primitive(); + assert_eq!(result.as_slice::(), &[0, 0, 100, 100, 100]); + } + + #[test] + fn test_evaluate_multiple_conditions() { + // Test: CASE WHEN value == 1 THEN 10 WHEN value == 3 THEN 30 ELSE 0 END + // Input: [1, 2, 3, 4, 5] + // Expected: [10, 0, 30, 0, 0] + let test_array = + StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) + .unwrap() + .into_array(); + + let expr = case_when([ + eq(get_item("value", root()), lit(1i32)), + lit(10i32), + eq(get_item("value", root()), lit(3i32)), + lit(30i32), + lit(0i32), + ]); + + let result = expr.evaluate(&test_array).unwrap().to_primitive(); + assert_eq!(result.as_slice::(), &[10, 0, 30, 0, 0]); + } + + #[test] + fn test_evaluate_first_match_wins() { + // Test: CASE WHEN value > 2 THEN 100 WHEN value > 3 THEN 200 ELSE 0 END + // Input: [1, 2, 3, 4, 5] + // Expected: [0, 0, 100, 100, 100] - first condition wins for values 3, 4, 5 + let test_array = + StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) + .unwrap() + .into_array(); + + let expr = case_when([ + gt(get_item("value", root()), lit(2i32)), + lit(100i32), + gt(get_item("value", root()), lit(3i32)), + lit(200i32), + lit(0i32), + ]); + + let result = expr.evaluate(&test_array).unwrap().to_primitive(); + // First match wins: 3, 4, 5 all get 100 (from first condition) + assert_eq!(result.as_slice::(), &[0, 0, 100, 100, 100]); + } + + #[test] + fn test_evaluate_no_else_returns_null() { + // Test: CASE WHEN value > 3 THEN 100 END + // Input: [1, 2, 3, 4, 5] + // Expected: [null, null, null, 100, 100] + let test_array = + StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) + .unwrap() + .into_array(); + + let expr = case_when_no_else([gt(get_item("value", root()), lit(3i32)), lit(100i32)]); + + let result = expr.evaluate(&test_array).unwrap(); + + // Check the dtype is nullable + assert!(result.dtype().is_nullable()); + + // Positions 0, 1, 2 should be null, 3, 4 should be 100 + assert_eq!(result.scalar_at(0), Scalar::null(result.dtype().clone())); + assert_eq!(result.scalar_at(1), Scalar::null(result.dtype().clone())); + assert_eq!(result.scalar_at(2), Scalar::null(result.dtype().clone())); + assert_eq!( + result.scalar_at(3), + Scalar::from(100i32).cast(result.dtype()).unwrap() + ); + assert_eq!( + result.scalar_at(4), + Scalar::from(100i32).cast(result.dtype()).unwrap() + ); + } + + #[test] + fn test_evaluate_all_conditions_false() { + // Test: CASE WHEN value > 100 THEN 1 ELSE 0 END + // Input: [1, 2, 3, 4, 5] + // Expected: [0, 0, 0, 0, 0] - no conditions match + let test_array = + StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) + .unwrap() + .into_array(); + + let expr = case_when([ + gt(get_item("value", root()), lit(100i32)), + lit(1i32), + lit(0i32), + ]); + + let result = expr.evaluate(&test_array).unwrap().to_primitive(); + assert_eq!(result.as_slice::(), &[0, 0, 0, 0, 0]); + } + + #[test] + fn test_evaluate_all_conditions_true() { + // Test: CASE WHEN value > 0 THEN 100 ELSE 0 END + // Input: [1, 2, 3, 4, 5] + // Expected: [100, 100, 100, 100, 100] - all match + let test_array = + StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) + .unwrap() + .into_array(); + + let expr = case_when([ + gt(get_item("value", root()), lit(0i32)), + lit(100i32), + lit(0i32), + ]); + + let result = expr.evaluate(&test_array).unwrap().to_primitive(); + assert_eq!(result.as_slice::(), &[100, 100, 100, 100, 100]); + } + + #[test] + fn test_evaluate_with_literal_condition() { + // Test: CASE WHEN true THEN 100 ELSE 0 END (constant true condition) + let test_array = buffer![1i32, 2, 3].into_array(); + + let expr = case_when([lit(true), lit(100i32), lit(0i32)]); + + let result = expr.evaluate(&test_array).unwrap(); + // Constant folding should produce a constant array + if let Some(constant) = result.as_constant() { + assert_eq!(constant, Scalar::from(100i32)); + } else { + let prim = result.to_primitive(); + assert_eq!(prim.as_slice::(), &[100, 100, 100]); + } + } + + #[test] + fn test_evaluate_with_bool_column_result() { + // Test: CASE WHEN value > 2 THEN true ELSE false END + let test_array = + StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) + .unwrap() + .into_array(); + + let expr = case_when([ + gt(get_item("value", root()), lit(2i32)), + lit(true), + lit(false), + ]); + + let result = expr.evaluate(&test_array).unwrap().to_bool(); + assert_eq!( + result.bit_buffer().iter().collect::>(), + vec![false, false, true, true, true] + ); + } + + #[test] + fn test_evaluate_with_nullable_condition() { + // Test: CASE WHEN nullable_bool THEN 100 ELSE 0 END + // Where the condition has null values - nulls should be treated as false + let test_array = StructArray::from_fields(&[( + "cond", + BoolArray::from_iter([Some(true), None, Some(false), None, Some(true)]).into_array(), + )]) + .unwrap() + .into_array(); + + let expr = case_when([get_item("cond", root()), lit(100i32), lit(0i32)]); + + let result = expr.evaluate(&test_array).unwrap().to_primitive(); + // true -> 100, null -> 0 (treated as false), false -> 0 + assert_eq!(result.as_slice::(), &[100, 0, 0, 0, 100]); + } + + #[test] + fn test_evaluate_with_nullable_result_values() { + // Test: CASE WHEN value > 2 THEN nullable_value ELSE 0 END + let test_array = StructArray::from_fields(&[ + ("value", buffer![1i32, 2, 3, 4, 5].into_array()), + ( + "result", + PrimitiveArray::from_option_iter([Some(10), None, Some(30), Some(40), Some(50)]) + .into_array(), + ), + ]) + .unwrap() + .into_array(); + + let expr = case_when([ + gt(get_item("value", root()), lit(2i32)), + get_item("result", root()), + lit(0i32), + ]); + + let result = expr.evaluate(&test_array).unwrap(); + let prim = result.to_primitive(); + + // Values 1, 2 don't match -> 0 + // Value 3 matches -> 30 (from result column) + // Value 4 matches -> 40 + // Value 5 matches -> 50 + assert_eq!(prim.as_slice::(), &[0, 0, 30, 40, 50]); + } + + #[test] + fn test_evaluate_with_all_null_condition() { + // Test: CASE WHEN all_nulls THEN 100 ELSE 0 END + // All null conditions should be treated as false + let test_array = StructArray::from_fields(&[( + "cond", + BoolArray::from_iter([None, None, None]).into_array(), + )]) + .unwrap() + .into_array(); + + let expr = case_when([get_item("cond", root()), lit(100i32), lit(0i32)]); + + let result = expr.evaluate(&test_array).unwrap().to_primitive(); + // All null -> treated as false -> else value + assert_eq!(result.as_slice::(), &[0, 0, 0]); + } + + // ==================== Execute Tests ==================== + + #[test] + fn test_execute_with_scalar_inputs() { + use vortex_dtype::PTypeDowncast; + use vortex_vector::Scalar as VScalar; + use vortex_vector::bool::BoolScalar; + use vortex_vector::primitive::PScalar; + + // CASE WHEN true THEN 100 ELSE 0 END with all scalars + let options = CaseWhenOptions { + num_when_then_pairs: 1, + has_else: true, + }; + + let cond_dtype = DType::Bool(Nullability::NonNullable); + let then_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let else_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let return_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + + let datums = vec![ + Datum::Scalar(VScalar::from(BoolScalar::new(Some(true)))), + Datum::Scalar(VScalar::from(PScalar::new(Some(100i32)))), + Datum::Scalar(VScalar::from(PScalar::new(Some(0i32)))), + ]; + + let args = ExecutionArgs { + datums, + dtypes: vec![cond_dtype, then_dtype, else_dtype], + row_count: 1, + return_dtype, + }; + + let result = CaseWhen.execute(&options, args).unwrap(); + + // Should return scalar since all inputs were scalars + match result { + Datum::Scalar(s) => { + let prim = s.into_primitive().into_i32(); + assert_eq!(prim.value(), Some(100)); + } + Datum::Vector(v) => { + // Also acceptable: a length-1 vector + assert_eq!(v.len(), 1); + } + } + } + + #[test] + fn test_execute_with_scalar_false_condition() { + use vortex_dtype::PTypeDowncast; + use vortex_vector::Scalar as VScalar; + use vortex_vector::bool::BoolScalar; + use vortex_vector::primitive::PScalar; + + // CASE WHEN false THEN 100 ELSE 42 END + let options = CaseWhenOptions { + num_when_then_pairs: 1, + has_else: true, + }; + + let cond_dtype = DType::Bool(Nullability::NonNullable); + let then_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let else_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let return_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + + let datums = vec![ + Datum::Scalar(VScalar::from(BoolScalar::new(Some(false)))), + Datum::Scalar(VScalar::from(PScalar::new(Some(100i32)))), + Datum::Scalar(VScalar::from(PScalar::new(Some(42i32)))), + ]; + + let args = ExecutionArgs { + datums, + dtypes: vec![cond_dtype, then_dtype, else_dtype], + row_count: 1, + return_dtype, + }; + + let result = CaseWhen.execute(&options, args).unwrap(); + + match result { + Datum::Scalar(s) => { + let prim = s.into_primitive().into_i32(); + assert_eq!(prim.value(), Some(42)); + } + Datum::Vector(v) => { + assert_eq!(v.len(), 1); + let prim = v.into_primitive().into_i32(); + assert_eq!(prim.get(0).copied(), Some(42)); + } + } + } + + #[test] + fn test_execute_with_vector_condition() { + use vortex_dtype::PTypeDowncast; + use vortex_vector::Scalar as VScalar; + use vortex_vector::bool::BoolVector; + use vortex_vector::primitive::PScalar; + + // CASE WHEN [true, false, true] THEN 100 ELSE 0 END + let options = CaseWhenOptions { + num_when_then_pairs: 1, + has_else: true, + }; + + let cond_dtype = DType::Bool(Nullability::NonNullable); + let then_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let else_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let return_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + + let cond_vector = BoolVector::from_iter([true, false, true]); + let datums = vec![ + Datum::Vector(cond_vector.into()), + Datum::Scalar(VScalar::from(PScalar::new(Some(100i32)))), + Datum::Scalar(VScalar::from(PScalar::new(Some(0i32)))), + ]; + + let args = ExecutionArgs { + datums, + dtypes: vec![cond_dtype, then_dtype, else_dtype], + row_count: 3, + return_dtype, + }; + + let result = CaseWhen.execute(&options, args).unwrap(); + + match result { + Datum::Vector(v) => { + let prim = v.into_primitive().into_i32(); + assert_eq!(prim.get(0).copied(), Some(100)); + assert_eq!(prim.get(1).copied(), Some(0)); + assert_eq!(prim.get(2).copied(), Some(100)); + } + Datum::Scalar(_) => panic!("Expected vector result"), + } + } + + #[test] + fn test_execute_with_nullable_condition() { + use vortex_dtype::PTypeDowncast; + use vortex_mask::Mask; + use vortex_vector::Scalar as VScalar; + use vortex_vector::bool::BoolVector; + use vortex_vector::primitive::PScalar; + + // CASE WHEN [true, NULL, false, NULL] THEN 100 ELSE 0 END + // NULL should be treated as false + let options = CaseWhenOptions { + num_when_then_pairs: 1, + has_else: true, + }; + + let cond_dtype = DType::Bool(Nullability::Nullable); + let then_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let else_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let return_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + + let bits = vortex_buffer::BitBuffer::from_iter([true, true, false, false]); + let validity = Mask::from_iter([true, false, true, false]); // positions 1, 3 are NULL + let cond_vector = BoolVector::new(bits, validity); + + let datums = vec![ + Datum::Vector(cond_vector.into()), + Datum::Scalar(VScalar::from(PScalar::new(Some(100i32)))), + Datum::Scalar(VScalar::from(PScalar::new(Some(0i32)))), + ]; + + let args = ExecutionArgs { + datums, + dtypes: vec![cond_dtype, then_dtype, else_dtype], + row_count: 4, + return_dtype, + }; + + let result = CaseWhen.execute(&options, args).unwrap(); + + match result { + Datum::Vector(v) => { + let prim = v.into_primitive().into_i32(); + assert_eq!(prim.get(0).copied(), Some(100)); // true -> 100 + assert_eq!(prim.get(1).copied(), Some(0)); // NULL -> 0 (treated as false) + assert_eq!(prim.get(2).copied(), Some(0)); // false -> 0 + assert_eq!(prim.get(3).copied(), Some(0)); // NULL -> 0 (treated as false) + } + Datum::Scalar(_) => panic!("Expected vector result"), + } + } + + #[test] + fn test_execute_without_else() { + use vortex_dtype::PTypeDowncast; + use vortex_vector::Scalar as VScalar; + use vortex_vector::bool::BoolVector; + use vortex_vector::primitive::PScalar; + + // CASE WHEN [true, false, true] THEN 100 END (no else) + let options = CaseWhenOptions { + num_when_then_pairs: 1, + has_else: false, + }; + + let cond_dtype = DType::Bool(Nullability::NonNullable); + let then_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let return_dtype = DType::Primitive(PType::I32, Nullability::Nullable); + + let cond_vector = BoolVector::from_iter([true, false, true]); + let datums = vec![ + Datum::Vector(cond_vector.into()), + Datum::Scalar(VScalar::from(PScalar::new(Some(100i32)))), + ]; + + let args = ExecutionArgs { + datums, + dtypes: vec![cond_dtype, then_dtype], + row_count: 3, + return_dtype, + }; + + let result = CaseWhen.execute(&options, args).unwrap(); + + match result { + Datum::Vector(v) => { + let prim = v.into_primitive().into_i32(); + assert_eq!(prim.get(0).copied(), Some(100)); // true -> 100 + assert_eq!(prim.get(1), None); // false -> NULL + assert_eq!(prim.get(2).copied(), Some(100)); // true -> 100 + } + Datum::Scalar(_) => panic!("Expected vector result"), + } + } + + #[test] + fn test_execute_multiple_conditions() { + use vortex_dtype::PTypeDowncast; + use vortex_vector::Scalar as VScalar; + use vortex_vector::bool::BoolVector; + use vortex_vector::primitive::PScalar; + + // CASE WHEN [true, false, false] THEN 10 + // WHEN [false, true, false] THEN 20 + // ELSE 0 END + // Expected: [10, 20, 0] + let options = CaseWhenOptions { + num_when_then_pairs: 2, + has_else: true, + }; + + let cond_dtype = DType::Bool(Nullability::NonNullable); + let then_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let return_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + + let cond1 = BoolVector::from_iter([true, false, false]); + let cond2 = BoolVector::from_iter([false, true, false]); + + let datums = vec![ + Datum::Vector(cond1.into()), + Datum::Scalar(VScalar::from(PScalar::new(Some(10i32)))), + Datum::Vector(cond2.into()), + Datum::Scalar(VScalar::from(PScalar::new(Some(20i32)))), + Datum::Scalar(VScalar::from(PScalar::new(Some(0i32)))), + ]; + + let args = ExecutionArgs { + datums, + dtypes: vec![ + cond_dtype.clone(), + then_dtype.clone(), + cond_dtype, + then_dtype.clone(), + then_dtype, + ], + row_count: 3, + return_dtype, + }; + + let result = CaseWhen.execute(&options, args).unwrap(); + + match result { + Datum::Vector(v) => { + let prim = v.into_primitive().into_i32(); + assert_eq!(prim.get(0).copied(), Some(10)); + assert_eq!(prim.get(1).copied(), Some(20)); + assert_eq!(prim.get(2).copied(), Some(0)); + } + Datum::Scalar(_) => panic!("Expected vector result"), + } + } + + #[test] + fn test_execute_all_true_short_circuit() { + use vortex_dtype::PTypeDowncast; + use vortex_vector::Scalar as VScalar; + use vortex_vector::bool::BoolVector; + use vortex_vector::primitive::PScalar; + + // CASE WHEN [true, true, true] THEN 100 ELSE 0 END + // Should short-circuit and return the then value + let options = CaseWhenOptions { + num_when_then_pairs: 1, + has_else: true, + }; + + let cond_dtype = DType::Bool(Nullability::NonNullable); + let then_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let else_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let return_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + + let cond_vector = BoolVector::from_iter([true, true, true]); + let datums = vec![ + Datum::Vector(cond_vector.into()), + Datum::Scalar(VScalar::from(PScalar::new(Some(100i32)))), + Datum::Scalar(VScalar::from(PScalar::new(Some(0i32)))), + ]; + + let args = ExecutionArgs { + datums, + dtypes: vec![cond_dtype, then_dtype, else_dtype], + row_count: 3, + return_dtype, + }; + + let result = CaseWhen.execute(&options, args).unwrap(); + + // Could be scalar (from short-circuit) or vector + match result { + Datum::Vector(v) => { + let prim = v.into_primitive().into_i32(); + assert_eq!(prim.get(0).copied(), Some(100)); + assert_eq!(prim.get(1).copied(), Some(100)); + assert_eq!(prim.get(2).copied(), Some(100)); + } + Datum::Scalar(s) => { + let prim = s.into_primitive().into_i32(); + assert_eq!(prim.value(), Some(100)); + } + } + } + + #[test] + fn test_execute_all_false_short_circuit() { + use vortex_dtype::PTypeDowncast; + use vortex_vector::Scalar as VScalar; + use vortex_vector::bool::BoolVector; + use vortex_vector::primitive::PScalar; + + // CASE WHEN [false, false, false] THEN 100 ELSE 42 END + // Should short-circuit and return the else value + let options = CaseWhenOptions { + num_when_then_pairs: 1, + has_else: true, + }; + + let cond_dtype = DType::Bool(Nullability::NonNullable); + let then_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let else_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let return_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + + let cond_vector = BoolVector::from_iter([false, false, false]); + let datums = vec![ + Datum::Vector(cond_vector.into()), + Datum::Scalar(VScalar::from(PScalar::new(Some(100i32)))), + Datum::Scalar(VScalar::from(PScalar::new(Some(42i32)))), + ]; + + let args = ExecutionArgs { + datums, + dtypes: vec![cond_dtype, then_dtype, else_dtype], + row_count: 3, + return_dtype, + }; + + let result = CaseWhen.execute(&options, args).unwrap(); + + // Could be scalar (from short-circuit) or vector + match result { + Datum::Vector(v) => { + let prim = v.into_primitive().into_i32(); + assert_eq!(prim.get(0).copied(), Some(42)); + assert_eq!(prim.get(1).copied(), Some(42)); + assert_eq!(prim.get(2).copied(), Some(42)); + } + Datum::Scalar(s) => { + let prim = s.into_primitive().into_i32(); + assert_eq!(prim.value(), Some(42)); + } + } + } + + #[test] + fn test_execute_with_null_scalar_condition() { + use vortex_dtype::PTypeDowncast; + use vortex_vector::Scalar as VScalar; + use vortex_vector::bool::BoolScalar; + use vortex_vector::primitive::PScalar; + + // CASE WHEN NULL THEN 100 ELSE 42 END + // NULL condition should be treated as false + let options = CaseWhenOptions { + num_when_then_pairs: 1, + has_else: true, + }; + + let cond_dtype = DType::Bool(Nullability::Nullable); + let then_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let else_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let return_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + + let null_bool = BoolScalar::null(); + let datums = vec![ + Datum::Scalar(VScalar::from(null_bool)), + Datum::Scalar(VScalar::from(PScalar::new(Some(100i32)))), + Datum::Scalar(VScalar::from(PScalar::new(Some(42i32)))), + ]; + + let args = ExecutionArgs { + datums, + dtypes: vec![cond_dtype, then_dtype, else_dtype], + row_count: 1, + return_dtype, + }; + + let result = CaseWhen.execute(&options, args).unwrap(); + + // NULL condition -> treated as false -> else value + match result { + Datum::Scalar(s) => { + let prim = s.into_primitive().into_i32(); + assert_eq!(prim.value(), Some(42)); + } + Datum::Vector(v) => { + let prim = v.into_primitive().into_i32(); + assert_eq!(prim.get(0).copied(), Some(42)); + } + } + } + + #[test] + fn test_evaluate_divide_by_zero_protected_by_case_when() { + // This test verifies that CASE WHEN properly short-circuits evaluation + // to avoid divide-by-zero errors. + // Pattern: CASE WHEN denominator > 0 THEN numerator/denominator ELSE NULL END + // With input where some denominators are 0, the division should NOT be evaluated + // for those rows. + + use vortex_buffer::buffer; + use vortex_dtype::PType; + + use crate::arrays::StructArray; + use crate::expr::VTableExt; + use crate::expr::exprs::binary::Binary; + use crate::expr::exprs::operators::Operator; + use crate::expr::get_item; + use crate::expr::gt; + use crate::expr::lit; + use crate::expr::root; + + // Create test data: numerator=[10, 20, 30], denominator=[2, 0, 5] + // Expected: CASE WHEN denominator > 0 THEN numerator/denominator ELSE NULL END + // = [5, NULL, 6] + let test_array = StructArray::from_fields(&[ + ("numerator", buffer![10i32, 20, 30].into_array()), + ("denominator", buffer![2i32, 0, 5].into_array()), + ]) + .unwrap() + .into_array(); + + // Build: CASE WHEN $.denominator > 0 THEN $.numerator / $.denominator ELSE null END + let condition = gt(get_item("denominator", root()), lit(0i32)); + let division = Binary + .try_new_expr( + Operator::Div, + [ + get_item("numerator", root()), + get_item("denominator", root()), + ], + ) + .unwrap(); + let null_dtype = DType::Primitive(PType::I32, Nullability::Nullable); + let null_val = lit(Scalar::null(null_dtype)); + + let expr = case_when([condition, division, null_val]); + + // This should NOT panic with divide-by-zero + let result = expr.evaluate(&test_array).unwrap(); + + // Verify results + assert_eq!(result.len(), 3); + + // Row 0: 10/2 = 5 + assert_eq!( + result.scalar_at(0), + Scalar::from(5i32).cast(result.dtype()).unwrap() + ); + + // Row 1: denominator=0, so result is NULL (division was NOT evaluated) + assert_eq!(result.scalar_at(1), Scalar::null(result.dtype().clone())); + + // Row 2: 30/5 = 6 + assert_eq!( + result.scalar_at(2), + Scalar::from(6i32).cast(result.dtype()).unwrap() + ); + } +} diff --git a/vortex-datafusion/src/convert/exprs.rs b/vortex-datafusion/src/convert/exprs.rs index a047d1e6211..2990b440fe6 100644 --- a/vortex-datafusion/src/convert/exprs.rs +++ b/vortex-datafusion/src/convert/exprs.rs @@ -178,6 +178,41 @@ impl DefaultExpressionConvertor { Ok(else_expr) } + + /// Attempts to convert a DataFusion CaseExpr to a Vortex expression. + fn try_convert_case_expr(&self, case_expr: &df_expr::CaseExpr) -> VortexResult { + // DataFusion CaseExpr has: + // - expr(): Optional base expression (for "CASE expr WHEN ..." form) + // - when_then_expr(): Vec of (when, then) pairs + // - else_expr(): Optional else expression + + // We don't support the "CASE expr WHEN value1 THEN result1" form yet + if case_expr.expr().is_some() { + vortex_bail!( + "CASE expr WHEN form is not yet supported, only searched CASE is supported" + ); + } + + let when_then_pairs = case_expr.when_then_expr(); + if when_then_pairs.is_empty() { + vortex_bail!("CASE expression must have at least one WHEN clause"); + } + + // Convert all when/then pairs + let mut children = Vec::with_capacity(when_then_pairs.len() * 2 + 1); + for (when_expr, then_expr) in when_then_pairs { + children.push(self.convert(when_expr.as_ref())?); + children.push(self.convert(then_expr.as_ref())?); + } + + // Handle the optional else clause + if let Some(else_expr) = case_expr.else_expr() { + children.push(self.convert(else_expr.as_ref())?); + Ok(case_when(children)) + } else { + Ok(case_when_no_else(children)) + } + } } impl ExpressionConvertor for DefaultExpressionConvertor { @@ -412,10 +447,12 @@ fn can_be_pushed_down_impl(df_expr: &Arc, schema: &Schema) -> && can_be_pushed_down_impl(like.pattern(), schema) } else if let Some(lit) = expr.downcast_ref::() { supported_data_types(&lit.value().data_type()) - } else if expr.downcast_ref::().is_some() - || expr.downcast_ref::().is_some() - { - true + } else if let Some(cast_expr) = expr.downcast_ref::() { + // CastExpr child must be an expression type that convert() can handle + is_convertible_expr(cast_expr.expr()) + } else if let Some(cast_col_expr) = expr.downcast_ref::() { + // CastColumnExpr child must be an expression type that convert() can handle + is_convertible_expr(cast_col_expr.expr()) } else if let Some(is_null) = expr.downcast_ref::() { can_be_pushed_down_impl(is_null.arg(), schema) } else if let Some(is_not_null) = expr.downcast_ref::() { @@ -442,6 +479,31 @@ fn can_be_pushed_down_impl(df_expr: &Arc, schema: &Schema) -> } } +/// Checks if an expression type is one that convert() can handle. +/// This is less restrictive than can_be_pushed_down since it only checks +/// expression types, not data type support. +fn is_convertible_expr(df_expr: &Arc) -> bool { + let expr = df_expr.as_any(); + + // Expression types that convert() handles + expr.downcast_ref::().is_some() + || expr.downcast_ref::().is_some() + || expr.downcast_ref::().is_some() + || expr.downcast_ref::().is_some() + || expr + .downcast_ref::() + .is_some_and(|e| is_convertible_expr(e.expr())) + || expr + .downcast_ref::() + .is_some_and(|e| is_convertible_expr(e.expr())) + || expr.downcast_ref::().is_some() + || expr.downcast_ref::().is_some() + || expr.downcast_ref::().is_some() + || expr + .downcast_ref::() + .is_some_and(|sf| ScalarFunctionExpr::try_downcast_func::(sf).is_some()) +} + fn can_binary_be_pushed_down(binary: &df_expr::BinaryExpr, schema: &Schema) -> bool { let is_op_supported = try_operator_from_df(binary.op()).is_ok(); is_op_supported @@ -465,6 +527,30 @@ fn can_case_be_pushed_down(case_expr: &df_expr::CaseExpr, schema: &Schema) -> bo .is_some_and(|else_expr| can_be_pushed_down_impl(else_expr, schema)) } +fn can_case_be_pushed_down(case_expr: &df_expr::CaseExpr, schema: &Schema) -> bool { + // We only support the "searched CASE" form (CASE WHEN cond THEN result ...) + // not the "simple CASE" form (CASE expr WHEN value THEN result ...) + if case_expr.expr().is_some() { + return false; + } + + // Check all when/then pairs + for (when_expr, then_expr) in case_expr.when_then_expr() { + if !can_be_pushed_down(when_expr, schema) || !can_be_pushed_down(then_expr, schema) { + return false; + } + } + + // Check the optional else clause + if let Some(else_expr) = case_expr.else_expr() + && !can_be_pushed_down(else_expr, schema) + { + return false; + } + + true +} + fn supported_data_types(dt: &DataType) -> bool { use DataType::*; @@ -498,9 +584,14 @@ fn supported_data_types(dt: &DataType) -> bool { is_supported } -/// Checks if a GetField scalar function can be pushed down. -fn can_scalar_fn_be_pushed_down(scalar_fn: &ScalarFunctionExpr) -> bool { +/// Checks if a scalar function can be pushed down. +/// Currently only GetFieldFunc is supported, and its arguments must also be pushable. +fn can_scalar_fn_be_pushed_down(scalar_fn: &ScalarFunctionExpr, schema: &Schema) -> bool { ScalarFunctionExpr::try_downcast_func::(scalar_fn).is_some() + && scalar_fn + .args() + .iter() + .all(|arg| can_be_pushed_down(arg, schema)) } // TODO(adam): Replace with `DataType::is_decimal` once its released. diff --git a/vortex-io/Cargo.toml b/vortex-io/Cargo.toml index 024f0a9144f..132ca7b55eb 100644 --- a/vortex-io/Cargo.toml +++ b/vortex-io/Cargo.toml @@ -39,7 +39,6 @@ tokio = { workspace = true, features = [ tracing = { workspace = true } vortex-array = { workspace = true } vortex-buffer = { workspace = true } -log = { workspace = true } vortex-error = { workspace = true } vortex-metrics = { workspace = true } vortex-session = { workspace = true } diff --git a/vortex-proto/proto/expr.proto b/vortex-proto/proto/expr.proto index 4540bce0a63..7d713a3afb3 100644 --- a/vortex-proto/proto/expr.proto +++ b/vortex-proto/proto/expr.proto @@ -80,3 +80,9 @@ message SelectOpts { FieldNames exclude = 2; } } + +// Options for `vortex.case_when` +message CaseWhenOpts { + uint32 num_when_then_pairs = 1; + bool has_else = 2; +} diff --git a/vortex-proto/src/generated/vortex.expr.rs b/vortex-proto/src/generated/vortex.expr.rs index f3b6d2cf624..180e693f269 100644 --- a/vortex-proto/src/generated/vortex.expr.rs +++ b/vortex-proto/src/generated/vortex.expr.rs @@ -145,3 +145,11 @@ pub mod select_opts { Exclude(super::FieldNames), } } +/// Options for `vortex.case_when` +#[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] +pub struct CaseWhenOpts { + #[prost(uint32, tag = "1")] + pub num_when_then_pairs: u32, + #[prost(bool, tag = "2")] + pub has_else: bool, +} diff --git a/vortex-session/src/lib.rs b/vortex-session/src/lib.rs index bfb2a51ba13..1ab0d63fac7 100644 --- a/vortex-session/src/lib.rs +++ b/vortex-session/src/lib.rs @@ -3,14 +3,20 @@ pub mod registry; -use std::any::{Any, TypeId, type_name}; +use std::any::Any; +use std::any::TypeId; +use std::any::type_name; use std::fmt::Debug; -use std::hash::{BuildHasherDefault, Hasher}; -use std::ops::{Deref, DerefMut}; +use std::hash::BuildHasherDefault; +use std::hash::Hasher; +use std::ops::Deref; +use std::ops::DerefMut; use std::sync::Arc; -use dashmap::{DashMap, Entry}; -use vortex_error::{VortexExpect, vortex_panic}; +use dashmap::DashMap; +use dashmap::Entry; +use vortex_error::VortexExpect; +use vortex_error::vortex_panic; /// A Vortex session encapsulates the set of extensible arrays, layouts, compute functions, dtypes, /// etc. that are available for use in a given context. From 222e5134bb9af3157ebfeb45107d75a4539c639e Mon Sep 17 00:00:00 2001 From: Luke Kim <80174+lukekim@users.noreply.github.com> Date: Fri, 13 Feb 2026 18:28:57 -0800 Subject: [PATCH 17/32] fix: update vortex-datafusion for DF52 FileSource and PhysicalExprAdapter API changes --- Cargo.lock | 325 +++++++++++---------- vortex-datafusion/src/persistent/source.rs | 4 + 2 files changed, 175 insertions(+), 154 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 0a768fa56d6..626d3448a00 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1815,7 +1815,7 @@ dependencies = [ name = "datafusion" version = "51.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8ba7cb113e9c0bedf9e9765926031e132fa05a1b09ba6e93a6d1a4d7044457b8" +checksum = "d12ee9fdc6cdb5898c7691bb994f0ba606c4acc93a2258d78bb9f26ff8158bb3" dependencies = [ "arrow", "arrow-schema", @@ -1945,21 +1945,21 @@ dependencies = [ [[package]] name = "datafusion-catalog" -version = "51.0.0" +version = "52.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "66a3a799f914a59b1ea343906a0486f17061f39509af74e874a866428951130d" +checksum = "462dc9ef45e5d688aeaae49a7e310587e81b6016b9d03bace5626ad0043e5a9e" dependencies = [ "arrow", "async-trait", "dashmap", - "datafusion-common 51.0.0", - "datafusion-common-runtime 51.0.0", - "datafusion-datasource 51.0.0", - "datafusion-execution 51.0.0", - "datafusion-expr 51.0.0", - "datafusion-physical-expr 51.0.0", - "datafusion-physical-plan 51.0.0", - "datafusion-session 51.0.0", + "datafusion-common 52.1.0", + "datafusion-common-runtime 52.1.0", + "datafusion-datasource 52.1.0", + "datafusion-execution 52.1.0", + "datafusion-expr 52.1.0", + "datafusion-physical-expr 52.1.0", + "datafusion-physical-plan 52.1.0", + "datafusion-session 52.1.0", "futures", "itertools 0.14.0", "log", @@ -1995,26 +1995,25 @@ dependencies = [ [[package]] name = "datafusion-catalog-listing" -version = "51.0.0" +version = "52.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6db1b113c80d7a0febcd901476a57aef378e717c54517a163ed51417d87621b0" +checksum = "1b96dbf1d728fc321817b744eb5080cdd75312faa6980b338817f68f3caa4208" dependencies = [ "arrow", "async-trait", - "datafusion-catalog 51.0.0", - "datafusion-common 51.0.0", - "datafusion-datasource 51.0.0", - "datafusion-execution 51.0.0", - "datafusion-expr 51.0.0", - "datafusion-physical-expr 51.0.0", - "datafusion-physical-expr-adapter 51.0.0", - "datafusion-physical-expr-common 51.0.0", - "datafusion-physical-plan 51.0.0", + "datafusion-catalog 52.1.0", + "datafusion-common 52.1.0", + "datafusion-datasource 52.1.0", + "datafusion-execution 52.1.0", + "datafusion-expr 52.1.0", + "datafusion-physical-expr 52.1.0", + "datafusion-physical-expr-adapter 52.1.0", + "datafusion-physical-expr-common 52.1.0", + "datafusion-physical-plan 52.1.0", "futures", "itertools 0.14.0", "log", "object_store", - "tokio", ] [[package]] @@ -2042,16 +2041,16 @@ dependencies = [ [[package]] name = "datafusion-common" -version = "51.0.0" +version = "52.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7c10f7659e96127d25e8366be7c8be4109595d6a2c3eac70421f380a7006a1b0" +checksum = "3237a6ff0d2149af4631290074289cae548c9863c885d821315d54c6673a074a" dependencies = [ "ahash 0.8.12", "arrow", "arrow-ipc", "chrono", "half", - "hashbrown 0.14.5", + "hashbrown 0.16.1", "indexmap", "libc", "log", @@ -2089,9 +2088,9 @@ dependencies = [ [[package]] name = "datafusion-common-runtime" -version = "51.0.0" +version = "52.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b92065bbc6532c6651e2f7dd30b55cba0c7a14f860c7e1d15f165c41a1868d95" +checksum = "70b5e34026af55a1bfccb1ef0a763cf1f64e77c696ffcf5a128a278c31236528" dependencies = [ "futures", "log", @@ -2111,23 +2110,23 @@ dependencies = [ [[package]] name = "datafusion-datasource" -version = "51.0.0" +version = "52.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fde13794244bc7581cd82f6fff217068ed79cdc344cafe4ab2c3a1c3510b38d6" +checksum = "1b2a6be734cc3785e18bbf2a7f2b22537f6b9fb960d79617775a51568c281842" dependencies = [ "arrow", "async-trait", "bytes", "chrono", - "datafusion-common 51.0.0", - "datafusion-common-runtime 51.0.0", - "datafusion-execution 51.0.0", - "datafusion-expr 51.0.0", - "datafusion-physical-expr 51.0.0", - "datafusion-physical-expr-adapter 51.0.0", - "datafusion-physical-expr-common 51.0.0", - "datafusion-physical-plan 51.0.0", - "datafusion-session 51.0.0", + "datafusion-common 52.1.0", + "datafusion-common-runtime 52.1.0", + "datafusion-execution 52.1.0", + "datafusion-expr 52.1.0", + "datafusion-physical-expr 52.1.0", + "datafusion-physical-expr-adapter 52.1.0", + "datafusion-physical-expr-common 52.1.0", + "datafusion-physical-plan 52.1.0", + "datafusion-session 52.1.0", "futures", "glob", "itertools 0.14.0", @@ -2175,22 +2174,22 @@ dependencies = [ [[package]] name = "datafusion-datasource-arrow" -version = "51.0.0" +version = "52.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "804fa9b4ecf3157982021770617200ef7c1b2979d57bec9044748314775a9aea" +checksum = "1739b9b07c9236389e09c74f770e88aff7055250774e9def7d3f4f56b3dcc7be" dependencies = [ "arrow", "arrow-ipc", "async-trait", "bytes", - "datafusion-common 51.0.0", - "datafusion-common-runtime 51.0.0", - "datafusion-datasource 51.0.0", - "datafusion-execution 51.0.0", - "datafusion-expr 51.0.0", - "datafusion-physical-expr-common 51.0.0", - "datafusion-physical-plan 51.0.0", - "datafusion-session 51.0.0", + "datafusion-common 52.1.0", + "datafusion-common-runtime 52.1.0", + "datafusion-datasource 52.1.0", + "datafusion-execution 52.1.0", + "datafusion-expr 52.1.0", + "datafusion-physical-expr-common 52.1.0", + "datafusion-physical-plan 52.1.0", + "datafusion-session 52.1.0", "futures", "itertools 0.14.0", "object_store", @@ -2243,21 +2242,21 @@ dependencies = [ [[package]] name = "datafusion-datasource-csv" -version = "51.0.0" +version = "52.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "61a1641a40b259bab38131c5e6f48fac0717bedb7dc93690e604142a849e0568" +checksum = "61c73bc54b518bbba7c7650299d07d58730293cfba4356f6f428cc94c20b7600" dependencies = [ "arrow", "async-trait", "bytes", - "datafusion-common 51.0.0", - "datafusion-common-runtime 51.0.0", - "datafusion-datasource 51.0.0", - "datafusion-execution 51.0.0", - "datafusion-expr 51.0.0", - "datafusion-physical-expr-common 51.0.0", - "datafusion-physical-plan 51.0.0", - "datafusion-session 51.0.0", + "datafusion-common 52.1.0", + "datafusion-common-runtime 52.1.0", + "datafusion-datasource 52.1.0", + "datafusion-execution 52.1.0", + "datafusion-expr 52.1.0", + "datafusion-physical-expr-common 52.1.0", + "datafusion-physical-plan 52.1.0", + "datafusion-session 52.1.0", "futures", "object_store", "regex", @@ -2289,21 +2288,21 @@ dependencies = [ [[package]] name = "datafusion-datasource-json" -version = "51.0.0" +version = "52.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "adeacdb00c1d37271176f8fb6a1d8ce096baba16ea7a4b2671840c5c9c64fe85" +checksum = "37812c8494c698c4d889374ecfabbff780f1f26d9ec095dd1bddfc2a8ca12559" dependencies = [ "arrow", "async-trait", "bytes", - "datafusion-common 51.0.0", - "datafusion-common-runtime 51.0.0", - "datafusion-datasource 51.0.0", - "datafusion-execution 51.0.0", - "datafusion-expr 51.0.0", - "datafusion-physical-expr-common 51.0.0", - "datafusion-physical-plan 51.0.0", - "datafusion-session 51.0.0", + "datafusion-common 52.1.0", + "datafusion-common-runtime 52.1.0", + "datafusion-datasource 52.1.0", + "datafusion-execution 52.1.0", + "datafusion-expr 52.1.0", + "datafusion-physical-expr-common 52.1.0", + "datafusion-physical-plan 52.1.0", + "datafusion-session 52.1.0", "futures", "object_store", "tokio", @@ -2365,7 +2364,7 @@ dependencies = [ name = "datafusion-doc" version = "51.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2b99e13947667b36ad713549237362afb054b2d8f8cc447751e23ec61202db07" +checksum = "2c825f969126bc2ef6a6a02d94b3c07abff871acf4d6dd759ce1255edb7923ce" [[package]] name = "datafusion-doc" @@ -2375,15 +2374,16 @@ checksum = "2c825f969126bc2ef6a6a02d94b3c07abff871acf4d6dd759ce1255edb7923ce" [[package]] name = "datafusion-execution" -version = "51.0.0" +version = "52.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "63695643190679037bc946ad46a263b62016931547bf119859c511f7ff2f5178" +checksum = "fa03ef05a2c2f90dd6c743e3e111078e322f4b395d20d4b4d431a245d79521ae" dependencies = [ "arrow", "async-trait", + "chrono", "dashmap", - "datafusion-common 51.0.0", - "datafusion-expr 51.0.0", + "datafusion-common 52.1.0", + "datafusion-expr 52.1.0", "futures", "log", "object_store", @@ -2416,19 +2416,19 @@ dependencies = [ [[package]] name = "datafusion-expr" -version = "51.0.0" +version = "52.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f9a4787cbf5feb1ab351f789063398f67654a6df75c4d37d7f637dc96f951a91" +checksum = "ef33934c1f98ee695cc51192cc5f9ed3a8febee84fdbcd9131bf9d3a9a78276f" dependencies = [ "arrow", "async-trait", "chrono", - "datafusion-common 51.0.0", - "datafusion-doc 51.0.0", - "datafusion-expr-common 51.0.0", - "datafusion-functions-aggregate-common 51.0.0", - "datafusion-functions-window-common 51.0.0", - "datafusion-physical-expr-common 51.0.0", + "datafusion-common 52.1.0", + "datafusion-doc 52.1.0", + "datafusion-expr-common 52.1.0", + "datafusion-functions-aggregate-common 52.1.0", + "datafusion-functions-window-common 52.1.0", + "datafusion-physical-expr-common 52.1.0", "indexmap", "itertools 0.14.0", "paste", @@ -2461,9 +2461,9 @@ dependencies = [ [[package]] name = "datafusion-expr-common" -version = "51.0.0" +version = "52.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5ce2fb1b8c15c9ac45b0863c30b268c69dc9ee7a1ee13ecf5d067738338173dc" +checksum = "000c98206e3dd47d2939a94b6c67af4bfa6732dd668ac4fafdbde408fd9134ea" dependencies = [ "arrow", "datafusion-common 51.0.0", @@ -2487,9 +2487,9 @@ dependencies = [ [[package]] name = "datafusion-functions" -version = "51.0.0" +version = "52.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "794a9db7f7b96b3346fc007ff25e994f09b8f0511b4cf7dff651fadfe3ebb28f" +checksum = "379b01418ab95ca947014066248c22139fe9af9289354de10b445bd000d5d276" dependencies = [ "arrow", "arrow-buffer", @@ -2497,12 +2497,13 @@ dependencies = [ "blake2", "blake3", "chrono", - "datafusion-common 51.0.0", - "datafusion-doc 51.0.0", - "datafusion-execution 51.0.0", - "datafusion-expr 51.0.0", - "datafusion-expr-common 51.0.0", - "datafusion-macros 51.0.0", + "chrono-tz", + "datafusion-common 52.1.0", + "datafusion-doc 52.1.0", + "datafusion-execution 52.1.0", + "datafusion-expr 52.1.0", + "datafusion-expr-common 52.1.0", + "datafusion-macros 52.1.0", "hex", "itertools 0.14.0", "log", @@ -2548,9 +2549,9 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate" -version = "51.0.0" +version = "52.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1c25210520a9dcf9c2b2cbbce31ebd4131ef5af7fc60ee92b266dc7d159cb305" +checksum = "fd00d5454ba4c3f8ebbd04bd6a6a9dc7ced7c56d883f70f2076c188be8459e4c" dependencies = [ "ahash 0.8.12", "arrow", @@ -2590,9 +2591,9 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate-common" -version = "51.0.0" +version = "52.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "62f4a66f3b87300bb70f4124b55434d2ae3fe80455f3574701d0348da040b55d" +checksum = "aec06b380729a87210a4e11f555ec2d729a328142253f8d557b87593622ecc9f" dependencies = [ "ahash 0.8.12", "arrow", @@ -2616,9 +2617,9 @@ dependencies = [ [[package]] name = "datafusion-functions-nested" -version = "51.0.0" +version = "52.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ae5c06eed03918dc7fe7a9f082a284050f0e9ecf95d72f57712d1496da03b8c4" +checksum = "904f48d45e0f1eb7d0eb5c0f80f2b5c6046a85454364a6b16a2e0b46f62e7dff" dependencies = [ "arrow", "arrow-ord", @@ -2662,16 +2663,16 @@ dependencies = [ [[package]] name = "datafusion-functions-table" -version = "51.0.0" +version = "52.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "db4fed1d71738fbe22e2712d71396db04c25de4111f1ec252b8f4c6d3b25d7f5" +checksum = "e9a0d20e2b887e11bee24f7734d780a2588b925796ac741c3118dd06d5aa77f0" dependencies = [ "arrow", "async-trait", - "datafusion-catalog 51.0.0", - "datafusion-common 51.0.0", - "datafusion-expr 51.0.0", - "datafusion-physical-plan 51.0.0", + "datafusion-catalog 52.1.0", + "datafusion-common 52.1.0", + "datafusion-expr 52.1.0", + "datafusion-physical-plan 52.1.0", "parking_lot", "paste", ] @@ -2694,9 +2695,9 @@ dependencies = [ [[package]] name = "datafusion-functions-window" -version = "51.0.0" +version = "52.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1d92206aa5ae21892f1552b4d61758a862a70956e6fd7a95cb85db1de74bc6d1" +checksum = "d3414b0a07e39b6979fe3a69c7aa79a9f1369f1d5c8e52146e66058be1b285ee" dependencies = [ "arrow", "datafusion-common 51.0.0", @@ -2730,12 +2731,12 @@ dependencies = [ [[package]] name = "datafusion-functions-window-common" -version = "51.0.0" +version = "52.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "53ae9bcc39800820d53a22d758b3b8726ff84a5a3e24cecef04ef4e5fdf1c7cc" +checksum = "5bf2feae63cd4754e31add64ce75cae07d015bce4bb41cd09872f93add32523a" dependencies = [ - "datafusion-common 51.0.0", - "datafusion-physical-expr-common 51.0.0", + "datafusion-common 52.1.0", + "datafusion-physical-expr-common 52.1.0", ] [[package]] @@ -2750,11 +2751,11 @@ dependencies = [ [[package]] name = "datafusion-macros" -version = "51.0.0" +version = "52.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1063ad4c9e094b3f798acee16d9a47bd7372d9699be2de21b05c3bd3f34ab848" +checksum = "c4fe888aeb6a095c4bcbe8ac1874c4b9a4c7ffa2ba849db7922683ba20875aaf" dependencies = [ - "datafusion-doc 51.0.0", + "datafusion-doc 52.1.0", "quote", "syn 2.0.117", ] @@ -2772,16 +2773,16 @@ dependencies = [ [[package]] name = "datafusion-optimizer" -version = "51.0.0" +version = "52.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9f35f9ec5d08b87fd1893a30c2929f2559c2f9806ca072d8fefca5009dc0f06a" +checksum = "8a6527c063ae305c11be397a86d8193936f4b84d137fe40bd706dfc178cf733c" dependencies = [ "arrow", "chrono", - "datafusion-common 51.0.0", - "datafusion-expr 51.0.0", - "datafusion-expr-common 51.0.0", - "datafusion-physical-expr 51.0.0", + "datafusion-common 52.1.0", + "datafusion-expr 52.1.0", + "datafusion-expr-common 52.1.0", + "datafusion-physical-expr 52.1.0", "indexmap", "itertools 0.14.0", "log", @@ -2811,9 +2812,9 @@ dependencies = [ [[package]] name = "datafusion-physical-expr" -version = "51.0.0" +version = "52.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c30cc8012e9eedcb48bbe112c6eff4ae5ed19cf3003cb0f505662e88b7014c5d" +checksum = "0bb028323dd4efd049dd8a78d78fe81b2b969447b39c51424167f973ac5811d9" dependencies = [ "ahash 0.8.12", "arrow", @@ -2823,7 +2824,7 @@ dependencies = [ "datafusion-functions-aggregate-common 51.0.0", "datafusion-physical-expr-common 51.0.0", "half", - "hashbrown 0.14.5", + "hashbrown 0.16.1", "indexmap", "itertools 0.14.0", "parking_lot", @@ -2857,9 +2858,9 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-adapter" -version = "51.0.0" +version = "52.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7f9ff2dbd476221b1f67337699eff432781c4e6e1713d2aefdaa517dfbf79768" +checksum = "78fe0826aef7eab6b4b61533d811234a7a9e5e458331ebbf94152a51fc8ab433" dependencies = [ "arrow", "datafusion-common 51.0.0", @@ -2887,9 +2888,9 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-common" -version = "51.0.0" +version = "52.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "90da43e1ec550b172f34c87ec68161986ced70fd05c8d2a2add66eef9c276f03" +checksum = "cfccd388620734c661bd8b7ca93c44cdd59fecc9b550eea416a78ffcbb29475f" dependencies = [ "ahash 0.8.12", "arrow", @@ -2897,6 +2898,7 @@ dependencies = [ "datafusion-expr-common 51.0.0", "hashbrown 0.14.5", "itertools 0.14.0", + "parking_lot", ] [[package]] @@ -2918,9 +2920,9 @@ dependencies = [ [[package]] name = "datafusion-physical-optimizer" -version = "51.0.0" +version = "52.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ce9804f799acd7daef3be7aaffe77c0033768ed8fdbf5fb82fc4c5f2e6bc14e6" +checksum = "bde5fa10e73259a03b705d5fddc136516814ab5f441b939525618a4070f5a059" dependencies = [ "arrow", "datafusion-common 51.0.0", @@ -2955,27 +2957,27 @@ dependencies = [ [[package]] name = "datafusion-physical-plan" -version = "51.0.0" +version = "52.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0acf0ad6b6924c6b1aa7d213b181e012e2d3ec0a64ff5b10ee6282ab0f8532ac" +checksum = "0e1098760fb29127c24cc9ade3277051dc73c9ed0ac0131bd7bcd742e0ad7470" dependencies = [ "ahash 0.8.12", "arrow", "arrow-ord", "arrow-schema", "async-trait", - "chrono", - "datafusion-common 51.0.0", - "datafusion-common-runtime 51.0.0", - "datafusion-execution 51.0.0", - "datafusion-expr 51.0.0", - "datafusion-functions-aggregate-common 51.0.0", - "datafusion-functions-window-common 51.0.0", - "datafusion-physical-expr 51.0.0", - "datafusion-physical-expr-common 51.0.0", + "datafusion-common 52.1.0", + "datafusion-common-runtime 52.1.0", + "datafusion-execution 52.1.0", + "datafusion-expr 52.1.0", + "datafusion-functions 52.1.0", + "datafusion-functions-aggregate-common 52.1.0", + "datafusion-functions-window-common 52.1.0", + "datafusion-physical-expr 52.1.0", + "datafusion-physical-expr-common 52.1.0", "futures", "half", - "hashbrown 0.14.5", + "hashbrown 0.16.1", "indexmap", "itertools 0.14.0", "log", @@ -3017,9 +3019,9 @@ dependencies = [ [[package]] name = "datafusion-pruning" -version = "51.0.0" +version = "52.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ac2c2498a1f134a9e11a9f5ed202a2a7d7e9774bd9249295593053ea3be999db" +checksum = "64d0fef4201777b52951edec086c21a5b246f3c82621569ddb4a26f488bc38a9" dependencies = [ "arrow", "datafusion-common 51.0.0", @@ -3051,15 +3053,15 @@ dependencies = [ [[package]] name = "datafusion-session" -version = "51.0.0" +version = "52.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8f96eebd17555386f459037c65ab73aae8df09f464524c709d6a3134ad4f4776" +checksum = "f71f1e39e8f2acbf1c63b0e93756c2e970a64729dab70ac789587d6237c4fde0" dependencies = [ "async-trait", - "datafusion-common 51.0.0", - "datafusion-execution 51.0.0", - "datafusion-expr 51.0.0", - "datafusion-physical-plan 51.0.0", + "datafusion-common 52.1.0", + "datafusion-execution 52.1.0", + "datafusion-expr 52.1.0", + "datafusion-physical-plan 52.1.0", "parking_lot", ] @@ -3102,15 +3104,15 @@ dependencies = [ [[package]] name = "datafusion-sql" -version = "51.0.0" +version = "52.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3fc195fe60634b2c6ccfd131b487de46dc30eccae8a3c35a13f136e7f440414f" +checksum = "f44693cfcaeb7a9f12d71d1c576c3a6dc025a12cef209375fa2d16fb3b5670ee" dependencies = [ "arrow", "bigdecimal", "chrono", - "datafusion-common 51.0.0", - "datafusion-expr 51.0.0", + "datafusion-common 52.1.0", + "datafusion-expr 52.1.0", "indexmap", "log", "regex", @@ -5410,6 +5412,12 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "09edd9e8b54e49e587e4f6295a7d29c3ea94d469cb40ab8ca70b288248a81db2" +[[package]] +name = "leb128fmt" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09edd9e8b54e49e587e4f6295a7d29c3ea94d469cb40ab8ca70b288248a81db2" + [[package]] name = "lending-iterator" version = "0.1.7" @@ -9857,7 +9865,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ee48d38b119b0cd71fe4141b30f5ba9c7c5d9f4e7a3a8b4a674e4b6ef789976f" dependencies = [ "atomic", - "getrandom 0.3.4", + "getrandom 0.4.1", "js-sys", "serde_core", "wasm-bindgen", @@ -10854,7 +10862,16 @@ version = "0.4.0+wasi-0.3.0-rc-2026-01-06" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5428f8bf88ea5ddc08faddef2ac4a67e390b88186c703ce6dbd955e1c145aca5" dependencies = [ - "wit-bindgen", + "wit-bindgen 0.46.0", +] + +[[package]] +name = "wasip3" +version = "0.4.0+wasi-0.3.0-rc-2026-01-06" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5428f8bf88ea5ddc08faddef2ac4a67e390b88186c703ce6dbd955e1c145aca5" +dependencies = [ + "wit-bindgen 0.51.0", ] [[package]] diff --git a/vortex-datafusion/src/persistent/source.rs b/vortex-datafusion/src/persistent/source.rs index c713987bcd3..706fef46c4b 100644 --- a/vortex-datafusion/src/persistent/source.rs +++ b/vortex-datafusion/src/persistent/source.rs @@ -222,6 +222,10 @@ impl FileSource for VortexSource { VORTEX_FILE_EXTENSION } + fn projection(&self) -> Option<&datafusion_datasource::projection::ProjectionExprs> { + None + } + fn fmt_extra(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { match t { DisplayFormatType::Default | DisplayFormatType::Verbose => { From 3a7ae8a1fb590db66b4250bdf25d6902204ed85e Mon Sep 17 00:00:00 2001 From: Luke Kim <80174+lukekim@users.noreply.github.com> Date: Fri, 13 Feb 2026 18:51:44 -0800 Subject: [PATCH 18/32] fix: use correct import path for ProjectionExprs in DF52 --- vortex-datafusion/src/persistent/source.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vortex-datafusion/src/persistent/source.rs b/vortex-datafusion/src/persistent/source.rs index 706fef46c4b..2a65795f1de 100644 --- a/vortex-datafusion/src/persistent/source.rs +++ b/vortex-datafusion/src/persistent/source.rs @@ -222,7 +222,7 @@ impl FileSource for VortexSource { VORTEX_FILE_EXTENSION } - fn projection(&self) -> Option<&datafusion_datasource::projection::ProjectionExprs> { + fn projection(&self) -> Option<&datafusion_physical_expr::projection::ProjectionExprs> { None } From 30810fb1cf3f3c79d945ee904f9579458f6f49a9 Mon Sep 17 00:00:00 2001 From: Luke Kim <80174+lukekim@users.noreply.github.com> Date: Fri, 13 Feb 2026 19:30:15 -0800 Subject: [PATCH 19/32] Fix DF52 API: file_source param, create_file_opener return, ColumnStatistics byte_size, table_schema accessor --- vortex-datafusion/src/persistent/source.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vortex-datafusion/src/persistent/source.rs b/vortex-datafusion/src/persistent/source.rs index 2a65795f1de..c3836ebaca1 100644 --- a/vortex-datafusion/src/persistent/source.rs +++ b/vortex-datafusion/src/persistent/source.rs @@ -175,7 +175,7 @@ impl FileSource for VortexSource { .clone() .unwrap_or_else(|| Arc::new(DefaultVortexReaderFactory::new(object_store))); - let table_schema = base_config.table_schema.clone(); + let table_schema = base_config.table_schema().clone(); let opener = VortexOpener { partition, From 5baa218e5c7145b17f71c16b9e38d35702f1f07e Mon Sep 17 00:00:00 2001 From: Luke Kim <80174+lukekim@users.noreply.github.com> Date: Fri, 13 Feb 2026 19:34:15 -0800 Subject: [PATCH 20/32] Fix: use file_source.table_schema() instead of table_schema() --- vortex-datafusion/src/persistent/source.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vortex-datafusion/src/persistent/source.rs b/vortex-datafusion/src/persistent/source.rs index c3836ebaca1..d15db02cd3f 100644 --- a/vortex-datafusion/src/persistent/source.rs +++ b/vortex-datafusion/src/persistent/source.rs @@ -175,7 +175,7 @@ impl FileSource for VortexSource { .clone() .unwrap_or_else(|| Arc::new(DefaultVortexReaderFactory::new(object_store))); - let table_schema = base_config.table_schema().clone(); + let table_schema = base_config.file_source.table_schema().clone(); let opener = VortexOpener { partition, From 2e9b9c4e8f8df33dd139d7f9d38d95e57e9f884e Mon Sep 17 00:00:00 2001 From: Luke Kim <80174+lukekim@users.noreply.github.com> Date: Sun, 15 Feb 2026 10:57:08 -0800 Subject: [PATCH 21/32] Fix DF52: implement try_pushdown_projection, carry table_schema in create_physical_plan --- vortex-datafusion/src/persistent/source.rs | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/vortex-datafusion/src/persistent/source.rs b/vortex-datafusion/src/persistent/source.rs index d15db02cd3f..741648e6245 100644 --- a/vortex-datafusion/src/persistent/source.rs +++ b/vortex-datafusion/src/persistent/source.rs @@ -222,8 +222,20 @@ impl FileSource for VortexSource { VORTEX_FILE_EXTENSION } - fn projection(&self) -> Option<&datafusion_physical_expr::projection::ProjectionExprs> { - None + fn projection(&self) -> Option<&ProjectionExprs> { + self.projection.as_ref() + } + + fn try_pushdown_projection( + &self, + projection: &ProjectionExprs, + ) -> DFResult>> { + let mut source = self.clone(); + source.projection = match &self.projection { + Some(existing) => Some(existing.try_merge(projection)?), + None => Some(projection.clone()), + }; + Ok(Some(Arc::new(source))) } fn fmt_extra(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { From 31e5a44f4aeb1f216df83343d7f0497ed95a1265 Mon Sep 17 00:00:00 2001 From: Luke Kim <80174+lukekim@users.noreply.github.com> Date: Sun, 15 Feb 2026 11:23:40 -0800 Subject: [PATCH 22/32] Remove deprecated SchemaAdapterFactory, use direct column mapping Replace SchemaAdapterFactory.map_schema() and SchemaMapping.map_batch() with direct column index computation and inline batch remapping in opener.rs. DF52 deprecated SchemaAdapterFactory in favor of PhysicalExprAdapterFactory. Also simplify create_file_opener() in source.rs to always use DF52PhysicalExprAdapterFactory as the default adapter. --- vortex-datafusion/src/persistent/opener.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs index 5986a06da49..aa22afd39b0 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/vortex-datafusion/src/persistent/opener.rs @@ -75,7 +75,7 @@ pub(crate) struct VortexOpener { pub file_pruning_predicate: Option, pub expr_adapter_factory: Arc, /// This is the table's schema without partition columns. It may contain fields which do - /// not exist in the file, and are supplied by the `schema_adapter_factory`. + /// not exist in the file. Missing columns are null-filled during batch remapping. pub table_schema: TableSchema, /// A hint for the desired row count of record batches returned from the scan. pub batch_size: usize, From bbc10decce4cfba0983c51c18fc34f7f92245e6e Mon Sep 17 00:00:00 2001 From: Luke Kim <80174+lukekim@users.noreply.github.com> Date: Sun, 15 Feb 2026 11:56:49 -0800 Subject: [PATCH 23/32] Fix projection loss in create_physical_plan The create_physical_plan method was creating a fresh VortexSource and replacing the one in the FileScanConfig that already had projections, filters, and other settings pushed down. This caused queries to return all columns instead of just the projected ones. Fix by using the FileScanConfig directly, since the VortexSource within it was already set up by file_source() and try_pushdown_projection(). --- vortex-datafusion/src/persistent/format.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/vortex-datafusion/src/persistent/format.rs b/vortex-datafusion/src/persistent/format.rs index a3c1a543578..40bee28765d 100644 --- a/vortex-datafusion/src/persistent/format.rs +++ b/vortex-datafusion/src/persistent/format.rs @@ -28,7 +28,6 @@ use datafusion_datasource::file_compression_type::FileCompressionType; use datafusion_datasource::file_format::FileFormat; use datafusion_datasource::file_format::FileFormatFactory; use datafusion_datasource::file_scan_config::FileScanConfig; -use datafusion_datasource::file_scan_config::FileScanConfigBuilder; use datafusion_datasource::file_sink_config::FileSinkConfig; use datafusion_datasource::sink::DataSinkExec; use datafusion_datasource::source::DataSourceExec; From 0259de9cd0a95ec88c94e6f1c13ce2d7f8f2c940 Mon Sep 17 00:00:00 2001 From: Luke Kim <80174+lukekim@users.noreply.github.com> Date: Thu, 26 Feb 2026 19:20:30 -0800 Subject: [PATCH 24/32] Fix compilation errors in vortex-datafusion for DF52 - Add missing imports: VortexResult, case_when, case_when_no_else, vortex_bail - Fix can_be_pushed_down -> can_be_pushed_down_impl in non-trait contexts - Remove duplicate can_case_be_pushed_down function (keep searched CASE version) - Remove duplicate projection/try_pushdown_projection methods (keep non-Option version) - Re-add FileScanConfigBuilder import removed in previous commit --- vortex-array/src/scalar_fn/fns/mod.rs | 1 + vortex-datafusion/src/convert/exprs.rs | 26 ++++++---------------- vortex-datafusion/src/persistent/format.rs | 1 + vortex-datafusion/src/persistent/source.rs | 16 ------------- 4 files changed, 9 insertions(+), 35 deletions(-) diff --git a/vortex-array/src/scalar_fn/fns/mod.rs b/vortex-array/src/scalar_fn/fns/mod.rs index 95c66b09ef6..94fc8fb0384 100644 --- a/vortex-array/src/scalar_fn/fns/mod.rs +++ b/vortex-array/src/scalar_fn/fns/mod.rs @@ -3,6 +3,7 @@ pub mod between; pub mod binary; +pub mod case_when; pub mod cast; pub mod dynamic; pub mod fill_null; diff --git a/vortex-datafusion/src/convert/exprs.rs b/vortex-datafusion/src/convert/exprs.rs index 2990b440fe6..8cac96f7de0 100644 --- a/vortex-datafusion/src/convert/exprs.rs +++ b/vortex-datafusion/src/convert/exprs.rs @@ -37,7 +37,11 @@ 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::case_when::case_when; +use vortex::scalar_fn::fns::case_when::case_when_no_else; use vortex::scalar_fn::fns::operators::Operator; +use vortex::error::VortexResult; +use vortex::error::vortex_bail; use crate::convert::FromDataFusion; @@ -511,22 +515,6 @@ fn can_binary_be_pushed_down(binary: &df_expr::BinaryExpr, schema: &Schema) -> b && can_be_pushed_down_impl(binary.right(), schema) } -fn can_case_be_pushed_down(case_expr: &df_expr::CaseExpr, schema: &Schema) -> bool { - case_expr - .expr() - .is_none_or(|base_expr| can_be_pushed_down_impl(base_expr, schema)) - && case_expr - .when_then_expr() - .iter() - .all(|(when_expr, then_expr)| { - can_be_pushed_down_impl(when_expr, schema) - && can_be_pushed_down_impl(then_expr, schema) - }) - && case_expr - .else_expr() - .is_some_and(|else_expr| can_be_pushed_down_impl(else_expr, schema)) -} - fn can_case_be_pushed_down(case_expr: &df_expr::CaseExpr, schema: &Schema) -> bool { // We only support the "searched CASE" form (CASE WHEN cond THEN result ...) // not the "simple CASE" form (CASE expr WHEN value THEN result ...) @@ -536,14 +524,14 @@ fn can_case_be_pushed_down(case_expr: &df_expr::CaseExpr, schema: &Schema) -> bo // Check all when/then pairs for (when_expr, then_expr) in case_expr.when_then_expr() { - if !can_be_pushed_down(when_expr, schema) || !can_be_pushed_down(then_expr, schema) { + if !can_be_pushed_down_impl(when_expr, schema) || !can_be_pushed_down_impl(then_expr, schema) { return false; } } // Check the optional else clause if let Some(else_expr) = case_expr.else_expr() - && !can_be_pushed_down(else_expr, schema) + && !can_be_pushed_down_impl(else_expr, schema) { return false; } @@ -591,7 +579,7 @@ fn can_scalar_fn_be_pushed_down(scalar_fn: &ScalarFunctionExpr, schema: &Schema) && scalar_fn .args() .iter() - .all(|arg| can_be_pushed_down(arg, schema)) + .all(|arg| can_be_pushed_down_impl(arg, schema)) } // TODO(adam): Replace with `DataType::is_decimal` once its released. diff --git a/vortex-datafusion/src/persistent/format.rs b/vortex-datafusion/src/persistent/format.rs index 40bee28765d..a3c1a543578 100644 --- a/vortex-datafusion/src/persistent/format.rs +++ b/vortex-datafusion/src/persistent/format.rs @@ -28,6 +28,7 @@ use datafusion_datasource::file_compression_type::FileCompressionType; use datafusion_datasource::file_format::FileFormat; use datafusion_datasource::file_format::FileFormatFactory; use datafusion_datasource::file_scan_config::FileScanConfig; +use datafusion_datasource::file_scan_config::FileScanConfigBuilder; use datafusion_datasource::file_sink_config::FileSinkConfig; use datafusion_datasource::sink::DataSinkExec; use datafusion_datasource::source::DataSourceExec; diff --git a/vortex-datafusion/src/persistent/source.rs b/vortex-datafusion/src/persistent/source.rs index 741648e6245..9917a7b8252 100644 --- a/vortex-datafusion/src/persistent/source.rs +++ b/vortex-datafusion/src/persistent/source.rs @@ -222,22 +222,6 @@ impl FileSource for VortexSource { VORTEX_FILE_EXTENSION } - fn projection(&self) -> Option<&ProjectionExprs> { - self.projection.as_ref() - } - - fn try_pushdown_projection( - &self, - projection: &ProjectionExprs, - ) -> DFResult>> { - let mut source = self.clone(); - source.projection = match &self.projection { - Some(existing) => Some(existing.try_merge(projection)?), - None => Some(projection.clone()), - }; - Ok(Some(Arc::new(source))) - } - fn fmt_extra(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { match t { DisplayFormatType::Default | DisplayFormatType::Verbose => { From 2b65b68c0e139c8697f8ccbd3fa3a076c5009888 Mon Sep 17 00:00:00 2001 From: William <98815791+peasee@users.noreply.github.com> Date: Thu, 26 Feb 2026 19:34:39 -0800 Subject: [PATCH 25/32] fix: Implement file writing directly for `VortexSink` instead of DataFusion demuxer (#23) * Disable child metric registration under parent to prevent OOM (#10) * Use tokio::sync::oneshot to prevent FuturesUnordered reentrant drop crash (#9) * Don't return an error when we have an unsupported node, bubble up "TRUE" as in keep for that node, up to any `and` or `or` node and handle empty IN list. (#8) * feat: Support retrieving WriterStrategyBuilder from VortexSession (#6) * Fix session get-or-default (#5662) The comments described this get-or-default, but instead it was a panic --------- Signed-off-by: Nicholas Gates * feat: Support retrieving writer strategy builder from vortex session --------- Signed-off-by: Nicholas Gates Co-authored-by: Nicholas Gates * fix: Handle UncompressedSizeInBytes statistic correctly (#3) * fix: Add log dependency to vortex-io * fix: Add Debug impl for WriteStrategyBuilder and log dep for vortex-io * fix: Update persistent module to use simplified expression handling from PR #8 * revert: Restore spiceai-51 versions of persistent module (incompatible with DF51) * fix: Restore format.rs to spiceai-51 version * Add Case-When Expression and tests (#12) * fix: Ensure CastExpr/CastColumnExpr/ScalarFunctionExpr check children in can_be_pushed_down The can_be_pushed_down function was returning true for CastExpr and CastColumnExpr without checking if their child expressions are convertible. This caused runtime errors when the child contained expressions like CaseExpr that convert() cannot handle. Also fixed ScalarFunctionExpr to recursively check its arguments. Fixes spiceai/spiceai#9037 * Add Case-When Expression and tests * Implement execute() * Add additional tests and fix type issue * Fix toml lint * feat(case_when): implement lazy evaluation to avoid side effects in unevaluated branches This implements proper lazy evaluation in CaseWhen expression to ensure that THEN/ELSE branches are only evaluated for rows where they apply. This is critical for correctness when expressions have side effects like divide-by-zero panics. The implementation: 1. Evaluates conditions in order, tracking which rows have been matched 2. For each condition, computes an effective mask (condition AND NOT matched) 3. Uses filter() to create a scoped array with only matching rows 4. Evaluates THEN expression only on the filtered scope 5. Uses scatter_with_mask() to expand results back to original positions 6. Short-circuits when all rows are matched or all conditions fail This fixes TPC-DS Q73 which has a pattern like: CASE WHEN hd_vehicle_count > 0 THEN hd_dep_count/hd_vehicle_count ELSE NULL END Previously, the division would be evaluated for all rows including those where hd_vehicle_count=0, causing a divide-by-zero panic. Now the division is only evaluated for rows where the condition is true. Added test: test_evaluate_divide_by_zero_protected_by_case_when * Formatting * feat(datafusion): Export DefaultExpressionConvertor (#19) * feat: Add option for writing to target vortex file size * fix: Actually respect target file size * fix: Set default file size to 16mb * fix: Update use of unwrap in tests * test: Update tests * fix: Always write with custom sink * fix: Stream directly into target file * fix: Improve safety arounded ended writer streams with remaining source * feat: add target file size configuration for Vortex file output - Introduced `target_file_size_mb` option to `VortexFormat` for controlling the size of output files in megabytes. - Updated `VortexSink` to handle writing files based on the specified target size, bypassing DataFusion's default row count-based splitting. - Implemented logic to split output files when the buffered data exceeds the target file size during the write process. - Added tests to verify the new target file size configuration. * refactor: clean up case expression handling and improve file size configuration logic --------- Signed-off-by: Nicholas Gates Co-authored-by: Sergei Grebnov Co-authored-by: Luke Kim <80174+lukekim@users.noreply.github.com> Co-authored-by: Nicholas Gates --- Cargo.lock | 600 ++++++++++----------- vortex-datafusion/Cargo.toml | 1 + vortex-datafusion/src/convert/exprs.rs | 73 +-- vortex-datafusion/src/persistent/format.rs | 39 +- vortex-datafusion/src/persistent/sink.rs | 214 +++++++- vortex-datafusion/src/persistent/source.rs | 2 - 6 files changed, 549 insertions(+), 380 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 626d3448a00..02085fdbe1b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -416,7 +416,7 @@ version = "57.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8c872d36b7bf2a6a6a2b40de9156265f0242910791db366a2c17476ba8330d68" dependencies = [ - "bitflags 2.10.0", + "bitflags 2.11.0", "serde_core", "serde_json", ] @@ -466,9 +466,9 @@ dependencies = [ [[package]] name = "async-compression" -version = "0.4.39" +version = "0.4.41" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "68650b7df54f0293fd061972a0fb05aaf4fc0879d3b3d21a638a182c5c543b9f" +checksum = "d0f9ee0f6e02ffd7ad5816e9464499fba7b3effd01123b515c41d1697c43dad1" dependencies = [ "compression-codecs", "compression-core", @@ -478,9 +478,9 @@ dependencies = [ [[package]] name = "async-executor" -version = "1.13.3" +version = "1.14.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "497c00e0fd83a72a79a39fcbd8e3e2f055d6f6c7e025f3b3d91f4f8e76527fb8" +checksum = "c96bf972d85afc50bf5ab8fe2d54d1586b4e0b46c97c50a0c9e71e2f7bcd812a" dependencies = [ "async-task", "concurrent-queue", @@ -514,7 +514,7 @@ dependencies = [ "futures-lite", "parking", "polling", - "rustix 1.1.3", + "rustix 1.1.4", "slab", "windows-sys 0.61.2", ] @@ -556,7 +556,7 @@ dependencies = [ "cfg-if", "event-listener", "futures-lite", - "rustix 1.1.3", + "rustix 1.1.4", ] [[package]] @@ -582,7 +582,7 @@ dependencies = [ "cfg-if", "futures-core", "futures-io", - "rustix 1.1.3", + "rustix 1.1.4", "signal-hook-registry", "slab", "windows-sys 0.61.2", @@ -698,7 +698,7 @@ version = "0.72.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "993776b509cfb49c750f11b8f07a46fa23e0a1386ffc01fb1e7d343efc387895" dependencies = [ - "bitflags 2.10.0", + "bitflags 2.11.0", "cexpr", "clang-sys", "itertools 0.13.0", @@ -741,9 +741,9 @@ checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" [[package]] name = "bitflags" -version = "2.10.0" +version = "2.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "812e12b5285cc515a9c72a5c1d3b6d46a19dac5acfef5265968c166106e31dd3" +checksum = "843867be96c8daad0d758b57df9392b6d8d271134fce549de6ce169ff98a92af" [[package]] name = "bitpacking" @@ -893,9 +893,9 @@ dependencies = [ [[package]] name = "bumpalo" -version = "3.19.1" +version = "3.20.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5dd9dc738b7a8311c7ade152424974d8115f2cdad61e8dab8dac9f2362298510" +checksum = "5d20789868f4b01b2f2caec9f5c4e0213b41e3e5702a50157d699ae31ced2fcb" [[package]] name = "byte-slice-cast" @@ -1084,9 +1084,9 @@ checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" [[package]] name = "chrono" -version = "0.4.43" +version = "0.4.44" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fac4744fb15ae8337dc853fee7fb3f4e48c0fbaa23d0afe49c447b4fab126118" +checksum = "c673075a2e0e5f4a1dde27ce9dee1ea4558c7ffe648f576438a20ca1d2acc4b0" dependencies = [ "iana-time-zone", "js-sys", @@ -1367,9 +1367,9 @@ dependencies = [ [[package]] name = "compression-codecs" -version = "0.4.36" +version = "0.4.37" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "00828ba6fd27b45a448e57dbfe84f1029d4c9f26b368157e9a448a5f49a2ec2a" +checksum = "eb7b51a7d9c967fc26773061ba86150f19c50c0d65c887cb1fbe295fd16619b7" dependencies = [ "bzip2", "compression-core", @@ -1611,13 +1611,13 @@ version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d8b9f2e4c67f833b660cdb0a3523065869fb35570177239812ed4c905aeff87b" dependencies = [ - "bitflags 2.10.0", + "bitflags 2.11.0", "crossterm_winapi", "derive_more", "document-features", "mio", "parking_lot", - "rustix 1.1.3", + "rustix 1.1.4", "signal-hook", "signal-hook-mio", "winapi", @@ -1815,7 +1815,7 @@ dependencies = [ name = "datafusion" version = "51.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d12ee9fdc6cdb5898c7691bb994f0ba606c4acc93a2258d78bb9f26ff8158bb3" +checksum = "8ba7cb113e9c0bedf9e9765926031e132fa05a1b09ba6e93a6d1a4d7044457b8" dependencies = [ "arrow", "arrow-schema", @@ -1945,21 +1945,21 @@ dependencies = [ [[package]] name = "datafusion-catalog" -version = "52.1.0" +version = "51.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "462dc9ef45e5d688aeaae49a7e310587e81b6016b9d03bace5626ad0043e5a9e" +checksum = "66a3a799f914a59b1ea343906a0486f17061f39509af74e874a866428951130d" dependencies = [ "arrow", "async-trait", "dashmap", - "datafusion-common 52.1.0", - "datafusion-common-runtime 52.1.0", - "datafusion-datasource 52.1.0", - "datafusion-execution 52.1.0", - "datafusion-expr 52.1.0", - "datafusion-physical-expr 52.1.0", - "datafusion-physical-plan 52.1.0", - "datafusion-session 52.1.0", + "datafusion-common 51.0.0", + "datafusion-common-runtime 51.0.0", + "datafusion-datasource 51.0.0", + "datafusion-execution 51.0.0", + "datafusion-expr 51.0.0", + "datafusion-physical-expr 51.0.0", + "datafusion-physical-plan 51.0.0", + "datafusion-session 51.0.0", "futures", "itertools 0.14.0", "log", @@ -1995,25 +1995,26 @@ dependencies = [ [[package]] name = "datafusion-catalog-listing" -version = "52.1.0" +version = "51.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1b96dbf1d728fc321817b744eb5080cdd75312faa6980b338817f68f3caa4208" +checksum = "6db1b113c80d7a0febcd901476a57aef378e717c54517a163ed51417d87621b0" dependencies = [ "arrow", "async-trait", - "datafusion-catalog 52.1.0", - "datafusion-common 52.1.0", - "datafusion-datasource 52.1.0", - "datafusion-execution 52.1.0", - "datafusion-expr 52.1.0", - "datafusion-physical-expr 52.1.0", - "datafusion-physical-expr-adapter 52.1.0", - "datafusion-physical-expr-common 52.1.0", - "datafusion-physical-plan 52.1.0", + "datafusion-catalog 51.0.0", + "datafusion-common 51.0.0", + "datafusion-datasource 51.0.0", + "datafusion-execution 51.0.0", + "datafusion-expr 51.0.0", + "datafusion-physical-expr 51.0.0", + "datafusion-physical-expr-adapter 51.0.0", + "datafusion-physical-expr-common 51.0.0", + "datafusion-physical-plan 51.0.0", "futures", "itertools 0.14.0", "log", "object_store", + "tokio", ] [[package]] @@ -2041,16 +2042,16 @@ dependencies = [ [[package]] name = "datafusion-common" -version = "52.1.0" +version = "51.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3237a6ff0d2149af4631290074289cae548c9863c885d821315d54c6673a074a" +checksum = "7c10f7659e96127d25e8366be7c8be4109595d6a2c3eac70421f380a7006a1b0" dependencies = [ "ahash 0.8.12", "arrow", "arrow-ipc", "chrono", "half", - "hashbrown 0.16.1", + "hashbrown 0.14.5", "indexmap", "libc", "log", @@ -2088,9 +2089,9 @@ dependencies = [ [[package]] name = "datafusion-common-runtime" -version = "52.1.0" +version = "51.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "70b5e34026af55a1bfccb1ef0a763cf1f64e77c696ffcf5a128a278c31236528" +checksum = "b92065bbc6532c6651e2f7dd30b55cba0c7a14f860c7e1d15f165c41a1868d95" dependencies = [ "futures", "log", @@ -2110,23 +2111,23 @@ dependencies = [ [[package]] name = "datafusion-datasource" -version = "52.1.0" +version = "51.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1b2a6be734cc3785e18bbf2a7f2b22537f6b9fb960d79617775a51568c281842" +checksum = "fde13794244bc7581cd82f6fff217068ed79cdc344cafe4ab2c3a1c3510b38d6" dependencies = [ "arrow", "async-trait", "bytes", "chrono", - "datafusion-common 52.1.0", - "datafusion-common-runtime 52.1.0", - "datafusion-execution 52.1.0", - "datafusion-expr 52.1.0", - "datafusion-physical-expr 52.1.0", - "datafusion-physical-expr-adapter 52.1.0", - "datafusion-physical-expr-common 52.1.0", - "datafusion-physical-plan 52.1.0", - "datafusion-session 52.1.0", + "datafusion-common 51.0.0", + "datafusion-common-runtime 51.0.0", + "datafusion-execution 51.0.0", + "datafusion-expr 51.0.0", + "datafusion-physical-expr 51.0.0", + "datafusion-physical-expr-adapter 51.0.0", + "datafusion-physical-expr-common 51.0.0", + "datafusion-physical-plan 51.0.0", + "datafusion-session 51.0.0", "futures", "glob", "itertools 0.14.0", @@ -2174,22 +2175,22 @@ dependencies = [ [[package]] name = "datafusion-datasource-arrow" -version = "52.1.0" +version = "51.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1739b9b07c9236389e09c74f770e88aff7055250774e9def7d3f4f56b3dcc7be" +checksum = "804fa9b4ecf3157982021770617200ef7c1b2979d57bec9044748314775a9aea" dependencies = [ "arrow", "arrow-ipc", "async-trait", "bytes", - "datafusion-common 52.1.0", - "datafusion-common-runtime 52.1.0", - "datafusion-datasource 52.1.0", - "datafusion-execution 52.1.0", - "datafusion-expr 52.1.0", - "datafusion-physical-expr-common 52.1.0", - "datafusion-physical-plan 52.1.0", - "datafusion-session 52.1.0", + "datafusion-common 51.0.0", + "datafusion-common-runtime 51.0.0", + "datafusion-datasource 51.0.0", + "datafusion-execution 51.0.0", + "datafusion-expr 51.0.0", + "datafusion-physical-expr-common 51.0.0", + "datafusion-physical-plan 51.0.0", + "datafusion-session 51.0.0", "futures", "itertools 0.14.0", "object_store", @@ -2242,21 +2243,21 @@ dependencies = [ [[package]] name = "datafusion-datasource-csv" -version = "52.1.0" +version = "51.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "61c73bc54b518bbba7c7650299d07d58730293cfba4356f6f428cc94c20b7600" +checksum = "61a1641a40b259bab38131c5e6f48fac0717bedb7dc93690e604142a849e0568" dependencies = [ "arrow", "async-trait", "bytes", - "datafusion-common 52.1.0", - "datafusion-common-runtime 52.1.0", - "datafusion-datasource 52.1.0", - "datafusion-execution 52.1.0", - "datafusion-expr 52.1.0", - "datafusion-physical-expr-common 52.1.0", - "datafusion-physical-plan 52.1.0", - "datafusion-session 52.1.0", + "datafusion-common 51.0.0", + "datafusion-common-runtime 51.0.0", + "datafusion-datasource 51.0.0", + "datafusion-execution 51.0.0", + "datafusion-expr 51.0.0", + "datafusion-physical-expr-common 51.0.0", + "datafusion-physical-plan 51.0.0", + "datafusion-session 51.0.0", "futures", "object_store", "regex", @@ -2288,21 +2289,21 @@ dependencies = [ [[package]] name = "datafusion-datasource-json" -version = "52.1.0" +version = "51.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "37812c8494c698c4d889374ecfabbff780f1f26d9ec095dd1bddfc2a8ca12559" +checksum = "adeacdb00c1d37271176f8fb6a1d8ce096baba16ea7a4b2671840c5c9c64fe85" dependencies = [ "arrow", "async-trait", "bytes", - "datafusion-common 52.1.0", - "datafusion-common-runtime 52.1.0", - "datafusion-datasource 52.1.0", - "datafusion-execution 52.1.0", - "datafusion-expr 52.1.0", - "datafusion-physical-expr-common 52.1.0", - "datafusion-physical-plan 52.1.0", - "datafusion-session 52.1.0", + "datafusion-common 51.0.0", + "datafusion-common-runtime 51.0.0", + "datafusion-datasource 51.0.0", + "datafusion-execution 51.0.0", + "datafusion-expr 51.0.0", + "datafusion-physical-expr-common 51.0.0", + "datafusion-physical-plan 51.0.0", + "datafusion-session 51.0.0", "futures", "object_store", "tokio", @@ -2364,7 +2365,7 @@ dependencies = [ name = "datafusion-doc" version = "51.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2c825f969126bc2ef6a6a02d94b3c07abff871acf4d6dd759ce1255edb7923ce" +checksum = "2b99e13947667b36ad713549237362afb054b2d8f8cc447751e23ec61202db07" [[package]] name = "datafusion-doc" @@ -2374,16 +2375,15 @@ checksum = "2c825f969126bc2ef6a6a02d94b3c07abff871acf4d6dd759ce1255edb7923ce" [[package]] name = "datafusion-execution" -version = "52.1.0" +version = "51.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fa03ef05a2c2f90dd6c743e3e111078e322f4b395d20d4b4d431a245d79521ae" +checksum = "63695643190679037bc946ad46a263b62016931547bf119859c511f7ff2f5178" dependencies = [ "arrow", "async-trait", - "chrono", "dashmap", - "datafusion-common 52.1.0", - "datafusion-expr 52.1.0", + "datafusion-common 51.0.0", + "datafusion-expr 51.0.0", "futures", "log", "object_store", @@ -2416,19 +2416,19 @@ dependencies = [ [[package]] name = "datafusion-expr" -version = "52.1.0" +version = "51.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ef33934c1f98ee695cc51192cc5f9ed3a8febee84fdbcd9131bf9d3a9a78276f" +checksum = "f9a4787cbf5feb1ab351f789063398f67654a6df75c4d37d7f637dc96f951a91" dependencies = [ "arrow", "async-trait", "chrono", - "datafusion-common 52.1.0", - "datafusion-doc 52.1.0", - "datafusion-expr-common 52.1.0", - "datafusion-functions-aggregate-common 52.1.0", - "datafusion-functions-window-common 52.1.0", - "datafusion-physical-expr-common 52.1.0", + "datafusion-common 51.0.0", + "datafusion-doc 51.0.0", + "datafusion-expr-common 51.0.0", + "datafusion-functions-aggregate-common 51.0.0", + "datafusion-functions-window-common 51.0.0", + "datafusion-physical-expr-common 51.0.0", "indexmap", "itertools 0.14.0", "paste", @@ -2461,9 +2461,9 @@ dependencies = [ [[package]] name = "datafusion-expr-common" -version = "52.1.0" +version = "51.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "000c98206e3dd47d2939a94b6c67af4bfa6732dd668ac4fafdbde408fd9134ea" +checksum = "5ce2fb1b8c15c9ac45b0863c30b268c69dc9ee7a1ee13ecf5d067738338173dc" dependencies = [ "arrow", "datafusion-common 51.0.0", @@ -2487,9 +2487,9 @@ dependencies = [ [[package]] name = "datafusion-functions" -version = "52.1.0" +version = "51.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "379b01418ab95ca947014066248c22139fe9af9289354de10b445bd000d5d276" +checksum = "794a9db7f7b96b3346fc007ff25e994f09b8f0511b4cf7dff651fadfe3ebb28f" dependencies = [ "arrow", "arrow-buffer", @@ -2497,13 +2497,12 @@ dependencies = [ "blake2", "blake3", "chrono", - "chrono-tz", - "datafusion-common 52.1.0", - "datafusion-doc 52.1.0", - "datafusion-execution 52.1.0", - "datafusion-expr 52.1.0", - "datafusion-expr-common 52.1.0", - "datafusion-macros 52.1.0", + "datafusion-common 51.0.0", + "datafusion-doc 51.0.0", + "datafusion-execution 51.0.0", + "datafusion-expr 51.0.0", + "datafusion-expr-common 51.0.0", + "datafusion-macros 51.0.0", "hex", "itertools 0.14.0", "log", @@ -2549,9 +2548,9 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate" -version = "52.1.0" +version = "51.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fd00d5454ba4c3f8ebbd04bd6a6a9dc7ced7c56d883f70f2076c188be8459e4c" +checksum = "1c25210520a9dcf9c2b2cbbce31ebd4131ef5af7fc60ee92b266dc7d159cb305" dependencies = [ "ahash 0.8.12", "arrow", @@ -2591,9 +2590,9 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate-common" -version = "52.1.0" +version = "51.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "aec06b380729a87210a4e11f555ec2d729a328142253f8d557b87593622ecc9f" +checksum = "62f4a66f3b87300bb70f4124b55434d2ae3fe80455f3574701d0348da040b55d" dependencies = [ "ahash 0.8.12", "arrow", @@ -2617,9 +2616,9 @@ dependencies = [ [[package]] name = "datafusion-functions-nested" -version = "52.1.0" +version = "51.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "904f48d45e0f1eb7d0eb5c0f80f2b5c6046a85454364a6b16a2e0b46f62e7dff" +checksum = "ae5c06eed03918dc7fe7a9f082a284050f0e9ecf95d72f57712d1496da03b8c4" dependencies = [ "arrow", "arrow-ord", @@ -2663,16 +2662,16 @@ dependencies = [ [[package]] name = "datafusion-functions-table" -version = "52.1.0" +version = "51.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e9a0d20e2b887e11bee24f7734d780a2588b925796ac741c3118dd06d5aa77f0" +checksum = "db4fed1d71738fbe22e2712d71396db04c25de4111f1ec252b8f4c6d3b25d7f5" dependencies = [ "arrow", "async-trait", - "datafusion-catalog 52.1.0", - "datafusion-common 52.1.0", - "datafusion-expr 52.1.0", - "datafusion-physical-plan 52.1.0", + "datafusion-catalog 51.0.0", + "datafusion-common 51.0.0", + "datafusion-expr 51.0.0", + "datafusion-physical-plan 51.0.0", "parking_lot", "paste", ] @@ -2695,9 +2694,9 @@ dependencies = [ [[package]] name = "datafusion-functions-window" -version = "52.1.0" +version = "51.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d3414b0a07e39b6979fe3a69c7aa79a9f1369f1d5c8e52146e66058be1b285ee" +checksum = "1d92206aa5ae21892f1552b4d61758a862a70956e6fd7a95cb85db1de74bc6d1" dependencies = [ "arrow", "datafusion-common 51.0.0", @@ -2731,12 +2730,12 @@ dependencies = [ [[package]] name = "datafusion-functions-window-common" -version = "52.1.0" +version = "51.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5bf2feae63cd4754e31add64ce75cae07d015bce4bb41cd09872f93add32523a" +checksum = "53ae9bcc39800820d53a22d758b3b8726ff84a5a3e24cecef04ef4e5fdf1c7cc" dependencies = [ - "datafusion-common 52.1.0", - "datafusion-physical-expr-common 52.1.0", + "datafusion-common 51.0.0", + "datafusion-physical-expr-common 51.0.0", ] [[package]] @@ -2751,11 +2750,11 @@ dependencies = [ [[package]] name = "datafusion-macros" -version = "52.1.0" +version = "51.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c4fe888aeb6a095c4bcbe8ac1874c4b9a4c7ffa2ba849db7922683ba20875aaf" +checksum = "1063ad4c9e094b3f798acee16d9a47bd7372d9699be2de21b05c3bd3f34ab848" dependencies = [ - "datafusion-doc 52.1.0", + "datafusion-doc 51.0.0", "quote", "syn 2.0.117", ] @@ -2773,16 +2772,16 @@ dependencies = [ [[package]] name = "datafusion-optimizer" -version = "52.1.0" +version = "51.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8a6527c063ae305c11be397a86d8193936f4b84d137fe40bd706dfc178cf733c" +checksum = "9f35f9ec5d08b87fd1893a30c2929f2559c2f9806ca072d8fefca5009dc0f06a" dependencies = [ "arrow", "chrono", - "datafusion-common 52.1.0", - "datafusion-expr 52.1.0", - "datafusion-expr-common 52.1.0", - "datafusion-physical-expr 52.1.0", + "datafusion-common 51.0.0", + "datafusion-expr 51.0.0", + "datafusion-expr-common 51.0.0", + "datafusion-physical-expr 51.0.0", "indexmap", "itertools 0.14.0", "log", @@ -2812,9 +2811,9 @@ dependencies = [ [[package]] name = "datafusion-physical-expr" -version = "52.1.0" +version = "51.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0bb028323dd4efd049dd8a78d78fe81b2b969447b39c51424167f973ac5811d9" +checksum = "c30cc8012e9eedcb48bbe112c6eff4ae5ed19cf3003cb0f505662e88b7014c5d" dependencies = [ "ahash 0.8.12", "arrow", @@ -2824,7 +2823,7 @@ dependencies = [ "datafusion-functions-aggregate-common 51.0.0", "datafusion-physical-expr-common 51.0.0", "half", - "hashbrown 0.16.1", + "hashbrown 0.14.5", "indexmap", "itertools 0.14.0", "parking_lot", @@ -2858,9 +2857,9 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-adapter" -version = "52.1.0" +version = "51.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "78fe0826aef7eab6b4b61533d811234a7a9e5e458331ebbf94152a51fc8ab433" +checksum = "7f9ff2dbd476221b1f67337699eff432781c4e6e1713d2aefdaa517dfbf79768" dependencies = [ "arrow", "datafusion-common 51.0.0", @@ -2888,9 +2887,9 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-common" -version = "52.1.0" +version = "51.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cfccd388620734c661bd8b7ca93c44cdd59fecc9b550eea416a78ffcbb29475f" +checksum = "90da43e1ec550b172f34c87ec68161986ced70fd05c8d2a2add66eef9c276f03" dependencies = [ "ahash 0.8.12", "arrow", @@ -2898,7 +2897,6 @@ dependencies = [ "datafusion-expr-common 51.0.0", "hashbrown 0.14.5", "itertools 0.14.0", - "parking_lot", ] [[package]] @@ -2920,9 +2918,9 @@ dependencies = [ [[package]] name = "datafusion-physical-optimizer" -version = "52.1.0" +version = "51.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bde5fa10e73259a03b705d5fddc136516814ab5f441b939525618a4070f5a059" +checksum = "ce9804f799acd7daef3be7aaffe77c0033768ed8fdbf5fb82fc4c5f2e6bc14e6" dependencies = [ "arrow", "datafusion-common 51.0.0", @@ -2957,27 +2955,27 @@ dependencies = [ [[package]] name = "datafusion-physical-plan" -version = "52.1.0" +version = "51.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0e1098760fb29127c24cc9ade3277051dc73c9ed0ac0131bd7bcd742e0ad7470" +checksum = "0acf0ad6b6924c6b1aa7d213b181e012e2d3ec0a64ff5b10ee6282ab0f8532ac" dependencies = [ "ahash 0.8.12", "arrow", "arrow-ord", "arrow-schema", "async-trait", - "datafusion-common 52.1.0", - "datafusion-common-runtime 52.1.0", - "datafusion-execution 52.1.0", - "datafusion-expr 52.1.0", - "datafusion-functions 52.1.0", - "datafusion-functions-aggregate-common 52.1.0", - "datafusion-functions-window-common 52.1.0", - "datafusion-physical-expr 52.1.0", - "datafusion-physical-expr-common 52.1.0", + "chrono", + "datafusion-common 51.0.0", + "datafusion-common-runtime 51.0.0", + "datafusion-execution 51.0.0", + "datafusion-expr 51.0.0", + "datafusion-functions-aggregate-common 51.0.0", + "datafusion-functions-window-common 51.0.0", + "datafusion-physical-expr 51.0.0", + "datafusion-physical-expr-common 51.0.0", "futures", "half", - "hashbrown 0.16.1", + "hashbrown 0.14.5", "indexmap", "itertools 0.14.0", "log", @@ -3019,9 +3017,9 @@ dependencies = [ [[package]] name = "datafusion-pruning" -version = "52.1.0" +version = "51.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "64d0fef4201777b52951edec086c21a5b246f3c82621569ddb4a26f488bc38a9" +checksum = "ac2c2498a1f134a9e11a9f5ed202a2a7d7e9774bd9249295593053ea3be999db" dependencies = [ "arrow", "datafusion-common 51.0.0", @@ -3053,15 +3051,15 @@ dependencies = [ [[package]] name = "datafusion-session" -version = "52.1.0" +version = "51.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f71f1e39e8f2acbf1c63b0e93756c2e970a64729dab70ac789587d6237c4fde0" +checksum = "8f96eebd17555386f459037c65ab73aae8df09f464524c709d6a3134ad4f4776" dependencies = [ "async-trait", - "datafusion-common 52.1.0", - "datafusion-execution 52.1.0", - "datafusion-expr 52.1.0", - "datafusion-physical-plan 52.1.0", + "datafusion-common 51.0.0", + "datafusion-execution 51.0.0", + "datafusion-expr 51.0.0", + "datafusion-physical-plan 51.0.0", "parking_lot", ] @@ -3104,15 +3102,15 @@ dependencies = [ [[package]] name = "datafusion-sql" -version = "52.1.0" +version = "51.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f44693cfcaeb7a9f12d71d1c576c3a6dc025a12cef209375fa2d16fb3b5670ee" +checksum = "3fc195fe60634b2c6ccfd131b487de46dc30eccae8a3c35a13f136e7f440414f" dependencies = [ "arrow", "bigdecimal", "chrono", - "datafusion-common 52.1.0", - "datafusion-expr 52.1.0", + "datafusion-common 51.0.0", + "datafusion-expr 51.0.0", "indexmap", "log", "regex", @@ -3206,9 +3204,9 @@ dependencies = [ [[package]] name = "deflate64" -version = "0.1.10" +version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "26bf8fc351c5ed29b5c2f0cbbac1b209b74f60ecd62e675a998df72c49af5204" +checksum = "807800ff3288b621186fe0a8f3392c4652068257302709c24efd918c3dffcdc2" [[package]] name = "deltae" @@ -3218,9 +3216,9 @@ checksum = "5729f5117e208430e437df2f4843f5e5952997175992d1414f94c57d61e270b4" [[package]] name = "deranged" -version = "0.5.6" +version = "0.5.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cc3dc5ad92c2e2d1c193bbbbdf2ea477cb81331de4f3103f267ca18368b988c4" +checksum = "7cd812cc2bc1d69d4764bd80df88b4317eaef9e773c75226407d9bc0876b211c" dependencies = [ "powerfmt", "serde_core", @@ -3660,7 +3658,7 @@ version = "25.12.19" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "35f6839d7b3b98adde531effaf34f0c2badc6f4735d26fe74709d8e513a96ef3" dependencies = [ - "bitflags 2.10.0", + "bitflags 2.11.0", "rustc_version", ] @@ -3775,9 +3773,9 @@ checksum = "e6d5a32815ae3f33302d95fdcb2ce17862f8c65363dcfd29360480ba1001fc9c" [[package]] name = "futures" -version = "0.3.31" +version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "65bc07b1a8bc7c85c5f2e110c476c7389b4554ba72af57d8445ea63a576b0876" +checksum = "8b147ee9d1f6d097cef9ce628cd2ee62288d963e16fb287bd9286455b241382d" dependencies = [ "futures-channel", "futures-core", @@ -3790,9 +3788,9 @@ dependencies = [ [[package]] name = "futures-channel" -version = "0.3.31" +version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2dff15bf788c671c1934e366d07e30c1814a8ef514e1af724a602e8a2fbe1b10" +checksum = "07bbe89c50d7a535e539b8c17bc0b49bdb77747034daa8087407d655f3f7cc1d" dependencies = [ "futures-core", "futures-sink", @@ -3800,15 +3798,15 @@ dependencies = [ [[package]] name = "futures-core" -version = "0.3.31" +version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "05f29059c0c2090612e8d742178b0580d2dc940c837851ad723096f87af6663e" +checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d" [[package]] name = "futures-executor" -version = "0.3.31" +version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e28d1d997f585e54aebc3f97d39e72338912123a67330d723fdbb564d646c9f" +checksum = "baf29c38818342a3b26b5b923639e7b1f4a61fc5e76102d4b1981c6dc7a7579d" dependencies = [ "futures-core", "futures-task", @@ -3817,9 +3815,9 @@ dependencies = [ [[package]] name = "futures-io" -version = "0.3.31" +version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9e5c1b78ca4aae1ac06c48a526a655760685149f0d465d21f37abfe57ce075c6" +checksum = "cecba35d7ad927e23624b22ad55235f2239cfa44fd10428eecbeba6d6a717718" [[package]] name = "futures-lite" @@ -3836,9 +3834,9 @@ dependencies = [ [[package]] name = "futures-macro" -version = "0.3.31" +version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "162ee34ebcb7c64a8abebc059ce0fee27c2262618d7b60ed8faf72fef13c3650" +checksum = "e835b70203e41293343137df5c0664546da5745f82ec9b84d40be8336958447b" dependencies = [ "proc-macro2", "quote", @@ -3847,15 +3845,15 @@ dependencies = [ [[package]] name = "futures-sink" -version = "0.3.31" +version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e575fab7d1e0dcb8d0c7bcf9a63ee213816ab51902e6d244a95819acacf1d4f7" +checksum = "c39754e157331b013978ec91992bde1ac089843443c49cbc7f46150b0fad0893" [[package]] name = "futures-task" -version = "0.3.31" +version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f90f7dce0722e95104fcb095585910c0977252f286e354b5e3bd38902cd99988" +checksum = "037711b3d59c33004d3856fbdc83b99d4ff37a24768fa1be9ce3538a1cde4393" [[package]] name = "futures-timer" @@ -3865,9 +3863,9 @@ checksum = "f288b0a4f20f9a56b5d1da57e2227c661b7b16168e2f72365f57b63326e29b24" [[package]] name = "futures-util" -version = "0.3.31" +version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9fa08315bb612088cc391249efdc3bc77536f16c91f6cf495e6fbe85b20a4a81" +checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6" dependencies = [ "futures-channel", "futures-core", @@ -3877,7 +3875,6 @@ dependencies = [ "futures-task", "memchr", "pin-project-lite", - "pin-utils", "slab", ] @@ -4020,9 +4017,9 @@ dependencies = [ [[package]] name = "geographiclib-rs" -version = "0.2.6" +version = "0.2.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bc8f647bd562db28a15e0dce4a77d89e3a78f6f85943e782418ebdbb420ea3c4" +checksum = "c5a7f08910fd98737a6eda7568e7c5e645093e073328eeef49758cfe8b0489c7" dependencies = [ "libm", ] @@ -4091,9 +4088,9 @@ checksum = "0cc23270f6e1808e30a928bdc84dea0b9b4136a8bc82338574f23baf47bbd280" [[package]] name = "goldenfile" -version = "1.10.0" +version = "1.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4ce2238f730d493a06ef6746713fe0b56acecc88485892ae65c50d2db9bf977b" +checksum = "206600f27ef687e85a572315eb297c79f66620b2b3eeb3a6cde17f25e049ae5b" dependencies = [ "scopeguard", "similar-asserts", @@ -4459,7 +4456,7 @@ dependencies = [ "js-sys", "log", "wasm-bindgen", - "windows-core", + "windows-core 0.62.2", ] [[package]] @@ -4859,9 +4856,9 @@ dependencies = [ [[package]] name = "js-sys" -version = "0.3.85" +version = "0.3.90" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8c942ebf8e95485ca0d52d97da7c5a2c387d0e7f0ba4c35e93bfcaee045955b3" +checksum = "14dc6f6450b3f6d4ed5b16327f38fed626d375a886159ca555bd7822c0c3a5a6" dependencies = [ "once_cell", "wasm-bindgen", @@ -5412,12 +5409,6 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "09edd9e8b54e49e587e4f6295a7d29c3ea94d469cb40ab8ca70b288248a81db2" -[[package]] -name = "leb128fmt" -version = "0.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "09edd9e8b54e49e587e4f6295a7d29c3ea94d469cb40ab8ca70b288248a81db2" - [[package]] name = "lending-iterator" version = "0.1.7" @@ -5590,7 +5581,7 @@ version = "0.1.12" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3d0b95e02c851351f877147b7deea7b1afb1df71b63aa5f8270716e0c5720616" dependencies = [ - "bitflags 2.10.0", + "bitflags 2.11.0", "libc", "redox_syscall 0.7.2", ] @@ -5613,7 +5604,7 @@ version = "0.3.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5f4de44e98ddbf09375cbf4d17714d18f39195f4f4894e8524501726fd9a8a4a" dependencies = [ - "bitflags 2.10.0", + "bitflags 2.11.0", ] [[package]] @@ -5633,9 +5624,9 @@ checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" [[package]] name = "linux-raw-sys" -version = "0.11.0" +version = "0.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "df1d3c3b53da64cf5760482273a98e575c651a67eec7f77df96b5b642de8f039" +checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" [[package]] name = "litemap" @@ -5957,9 +5948,9 @@ checksum = "2195bf6aa996a481483b29d62a7663eed3fe39600c460e323f8ff41e90bdd89b" [[package]] name = "native-tls" -version = "0.2.16" +version = "0.2.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9d5d26952a508f321b4d3d2e80e78fc2603eaefcdf0c30783867f19586518bdc" +checksum = "465500e14ea162429d264d44189adc38b199b62b1c21eea9f69e4b73cb03bbf2" dependencies = [ "libc", "log", @@ -5999,7 +5990,7 @@ version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "71e2746dc3a24dd78b3cfcb7be93368c6de9963d30f43a6a73998a9cf4b17b46" dependencies = [ - "bitflags 2.10.0", + "bitflags 2.11.0", "cfg-if", "cfg_aliases", "libc", @@ -6012,7 +6003,7 @@ version = "0.30.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "74523f3a35e05aba87a1d978330aef40f67b0304ac79c1c00b294c9830543db6" dependencies = [ - "bitflags 2.10.0", + "bitflags 2.11.0", "cfg-if", "cfg_aliases", "libc", @@ -6248,7 +6239,7 @@ version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2a180dd8642fa45cdb7dd721cd4c11b1cadd4929ce112ebd8b9f5803cc79d536" dependencies = [ - "bitflags 2.10.0", + "bitflags 2.11.0", ] [[package]] @@ -6341,7 +6332,7 @@ version = "0.10.75" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "08838db121398ad17ab8531ce9de97b244589089e290a384c900cb9ff7434328" dependencies = [ - "bitflags 2.10.0", + "bitflags 2.11.0", "cfg-if", "foreign-types", "libc", @@ -6495,9 +6486,9 @@ dependencies = [ [[package]] name = "owo-colors" -version = "4.2.3" +version = "4.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9c6901729fa79e91a0913333229e9ca5dc725089d1c363b2f4b4760709dc4a52" +checksum = "d211803b9b6b570f68772237e415a029d5a50c65d382910b879fb19d3271f94d" [[package]] name = "papergrid" @@ -6916,7 +6907,7 @@ dependencies = [ "concurrent-queue", "hermit-abi", "pin-project-lite", - "rustix 1.1.3", + "rustix 1.1.4", "windows-sys 0.61.2", ] @@ -7512,7 +7503,7 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5ef8dea09a92caaf73bff7adb70b76162e5937524058a7e5bff37869cbbec293" dependencies = [ - "bitflags 2.10.0", + "bitflags 2.11.0", "compact_str", "hashbrown 0.16.1", "indoc", @@ -7564,7 +7555,7 @@ version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d7dbfa023cd4e604c2553483820c5fe8aa9d71a42eea5aa77c6e7f35756612db" dependencies = [ - "bitflags 2.10.0", + "bitflags 2.11.0", "hashbrown 0.16.1", "indoc", "instability", @@ -7629,7 +7620,7 @@ version = "0.5.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" dependencies = [ - "bitflags 2.10.0", + "bitflags 2.11.0", ] [[package]] @@ -7638,7 +7629,7 @@ version = "0.7.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6d94dd2f7cd932d4dc02cc8b2b50dfd38bd079a4e5d79198b99743d7fcf9a4b4" dependencies = [ - "bitflags 2.10.0", + "bitflags 2.11.0", ] [[package]] @@ -7703,9 +7694,9 @@ checksum = "cab834c73d247e67f4fae452806d17d3c7501756d98c8808d7c9c7aa7d18f973" [[package]] name = "regex-syntax" -version = "0.8.9" +version = "0.8.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a96887878f22d7bad8a3b6dc5b7440e0ada9a245242924394987b21cf2210a4c" +checksum = "dc897dd8d9e8bd1ed8cdad82b5966c3e0ecae09fb1907d58efaa013543185d0a" [[package]] name = "regress" @@ -7965,7 +7956,7 @@ version = "0.38.44" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fdb5bc1ae2baa591800df16c9ca78619bf65c0488b41b96ccec5d11220d8c154" dependencies = [ - "bitflags 2.10.0", + "bitflags 2.11.0", "errno", "libc", "linux-raw-sys 0.4.15", @@ -7974,22 +7965,22 @@ dependencies = [ [[package]] name = "rustix" -version = "1.1.3" +version = "1.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "146c9e247ccc180c1f61615433868c99f3de3ae256a30a43b49f67c2d9171f34" +checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190" dependencies = [ - "bitflags 2.10.0", + "bitflags 2.11.0", "errno", "libc", - "linux-raw-sys 0.11.0", + "linux-raw-sys 0.12.1", "windows-sys 0.61.2", ] [[package]] name = "rustls" -version = "0.23.36" +version = "0.23.37" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c665f33d38cea657d9614f766881e4d510e0eda4239891eea56b4cadcf01801b" +checksum = "758025cb5fccfd3bc2fd74708fd4682be41d99e5dff73c377c0646c6012c73a4" dependencies = [ "once_cell", "ring", @@ -8148,11 +8139,11 @@ checksum = "1c107b6f4780854c8b126e228ea8869f4d7b71260f962fefb57b996b8959ba6b" [[package]] name = "security-framework" -version = "3.6.0" +version = "3.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d17b898a6d6948c3a8ee4372c17cb384f90d2e6e912ef00895b14fd7ab54ec38" +checksum = "b7f4bc775c73d9a02cde8bf7b2ec4c9d12743edf609006c7facc23998404cd1d" dependencies = [ - "bitflags 2.10.0", + "bitflags 2.11.0", "core-foundation 0.10.1", "core-foundation-sys", "libc", @@ -8161,9 +8152,9 @@ dependencies = [ [[package]] name = "security-framework-sys" -version = "2.16.0" +version = "2.17.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "321c8673b092a9a42605034a9879d73cb79101ed5fd117bc9a597b89b4e9e61a" +checksum = "6ce2691df843ecc5d231c0b14ece2acc3efb62c0a398c7e1d875f3983ce020e3" dependencies = [ "core-foundation-sys", "libc", @@ -8280,9 +8271,9 @@ dependencies = [ [[package]] name = "serde_tokenstream" -version = "0.2.2" +version = "0.2.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "64060d864397305347a78851c51588fd283767e7e7589829e8121d65512340f1" +checksum = "d7c49585c52c01f13c5c2ebb333f14f6885d76daa768d8a037d28017ec538c69" dependencies = [ "proc-macro2", "quote", @@ -8317,9 +8308,9 @@ dependencies = [ [[package]] name = "serial_test" -version = "3.3.1" +version = "3.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0d0b343e184fc3b7bb44dff0705fffcf4b3756ba6aff420dddd8b24ca145e555" +checksum = "911bd979bf1070a3f3aa7b691a3b3e9968f339ceeec89e08c280a8a22207a32f" dependencies = [ "futures-executor", "futures-util", @@ -8332,9 +8323,9 @@ dependencies = [ [[package]] name = "serial_test_derive" -version = "3.3.1" +version = "3.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6f50427f258fb77356e4cd4aa0e87e2bd2c66dbcee41dc405282cae2bfc26c83" +checksum = "0a7d91949b85b0d2fb687445e448b40d322b6b3e4af6b44a29b21d9a5f33e6d9" dependencies = [ "proc-macro2", "quote", @@ -8374,9 +8365,9 @@ dependencies = [ [[package]] name = "shellexpand" -version = "3.1.1" +version = "3.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8b1fdf65dd6331831494dd616b30351c38e96e45921a27745cf98490458b90bb" +checksum = "32824fab5e16e6c4d86dc1ba84489390419a39f97699852b66480bb87d297ed8" dependencies = [ "dirs", ] @@ -8832,7 +8823,7 @@ version = "0.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a13f3d0daba03132c0aa9767f98351b3488edc2c100cda2d2ec2b04f3d8d3c8b" dependencies = [ - "bitflags 2.10.0", + "bitflags 2.11.0", "core-foundation 0.9.4", "system-configuration-sys", ] @@ -9067,14 +9058,14 @@ dependencies = [ [[package]] name = "tempfile" -version = "3.25.0" +version = "3.26.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0136791f7c95b1f6dd99f9cc786b91bb81c3800b639b3478e561ddb7be95e5f1" +checksum = "82a72c767771b47409d2345987fda8628641887d5466101319899796367354a0" dependencies = [ "fastrand", "getrandom 0.4.1", "once_cell", - "rustix 1.1.3", + "rustix 1.1.4", "windows-sys 0.61.2", ] @@ -9093,7 +9084,7 @@ version = "0.4.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "60b8cb979cb11c32ce1603f8137b22262a9d131aaa5c37b5678025f22b8becd0" dependencies = [ - "rustix 1.1.3", + "rustix 1.1.4", "windows-sys 0.60.2", ] @@ -9132,7 +9123,7 @@ checksum = "4676b37242ccbd1aabf56edb093a4827dc49086c0ffd764a5705899e0f35f8f7" dependencies = [ "anyhow", "base64", - "bitflags 2.10.0", + "bitflags 2.11.0", "fancy-regex", "filedescriptor", "finl_unicode", @@ -9462,9 +9453,9 @@ dependencies = [ [[package]] name = "toml_parser" -version = "1.0.8+spec-1.1.0" +version = "1.0.9+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0742ff5ff03ea7e67c8ae6c93cac239e0d9784833362da3f9a9c1da8dfefcbdc" +checksum = "702d4415e08923e7e1ef96cd5727c0dfed80b4d2fa25db9647fe5eb6f7c5a4c4" dependencies = [ "winnow", ] @@ -9477,9 +9468,9 @@ checksum = "ab16f14aed21ee8bfd8ec22513f7287cd4a91aa92e44edfe2c17ddd004e92607" [[package]] name = "tonic" -version = "0.14.3" +version = "0.14.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a286e33f82f8a1ee2df63f4fa35c0becf4a85a0cb03091a15fd7bf0b402dc94a" +checksum = "fec7c61a0695dc1887c1b53952990f3ad2e3a31453e1f49f10e75424943a93ec" dependencies = [ "async-trait", "base64", @@ -9498,9 +9489,9 @@ dependencies = [ [[package]] name = "tonic-prost" -version = "0.14.3" +version = "0.14.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d6c55a2d6a14174563de34409c9f92ff981d006f56da9c6ecd40d9d4a31500b0" +checksum = "a55376a0bbaa4975a3f10d009ad763d8f4108f067c7c2e74f3001fb49778d309" dependencies = [ "bytes", "prost 0.14.3", @@ -9529,7 +9520,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d4e6559d53cc268e5031cd8429d05415bc4cb4aefc4aa5d6cc35fbf5b924a1f8" dependencies = [ "async-compression", - "bitflags 2.10.0", + "bitflags 2.11.0", "bytes", "futures-core", "futures-util", @@ -9759,9 +9750,9 @@ checksum = "dbc4bc3a9f746d862c45cb89d705aa10f187bb96c76001afab07a0d35ce60142" [[package]] name = "unicode-ident" -version = "1.0.23" +version = "1.0.24" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "537dd038a89878be9b64dd4bd1b260315c1bb94f4d784956b81e27a088d9a09e" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" [[package]] name = "unicode-segmentation" @@ -9860,9 +9851,9 @@ checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" [[package]] name = "uuid" -version = "1.20.0" +version = "1.21.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ee48d38b119b0cd71fe4141b30f5ba9c7c5d9f4e7a3a8b4a674e4b6ef789976f" +checksum = "b672338555252d43fd2240c714dc444b8c6fb0a5c5335e65a07bba7742735ddb" dependencies = [ "atomic", "getrandom 0.4.1", @@ -10222,6 +10213,7 @@ dependencies = [ "tokio-stream", "tracing", "url", + "uuid", "vortex", "vortex-utils", ] @@ -10862,23 +10854,14 @@ version = "0.4.0+wasi-0.3.0-rc-2026-01-06" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5428f8bf88ea5ddc08faddef2ac4a67e390b88186c703ce6dbd955e1c145aca5" dependencies = [ - "wit-bindgen 0.46.0", -] - -[[package]] -name = "wasip3" -version = "0.4.0+wasi-0.3.0-rc-2026-01-06" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5428f8bf88ea5ddc08faddef2ac4a67e390b88186c703ce6dbd955e1c145aca5" -dependencies = [ - "wit-bindgen 0.51.0", + "wit-bindgen", ] [[package]] name = "wasm-bindgen" -version = "0.2.108" +version = "0.2.113" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "64024a30ec1e37399cf85a7ffefebdb72205ca1c972291c51512360d90bd8566" +checksum = "60722a937f594b7fde9adb894d7c092fc1bb6612897c46368d18e7a20208eff2" dependencies = [ "cfg-if", "once_cell", @@ -10889,9 +10872,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-futures" -version = "0.4.58" +version = "0.4.63" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "70a6e77fd0ae8029c9ea0063f87c46fde723e7d887703d74ad2616d792e51e6f" +checksum = "8a89f4650b770e4521aa6573724e2aed4704372151bd0de9d16a3bbabb87441a" dependencies = [ "cfg-if", "futures-util", @@ -10903,9 +10886,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro" -version = "0.2.108" +version = "0.2.113" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "008b239d9c740232e71bd39e8ef6429d27097518b6b30bdf9086833bd5b6d608" +checksum = "0fac8c6395094b6b91c4af293f4c79371c163f9a6f56184d2c9a85f5a95f3950" dependencies = [ "quote", "wasm-bindgen-macro-support", @@ -10913,9 +10896,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro-support" -version = "0.2.108" +version = "0.2.113" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5256bae2d58f54820e6490f9839c49780dff84c65aeab9e772f15d5f0e913a55" +checksum = "ab3fabce6159dc20728033842636887e4877688ae94382766e00b180abac9d60" dependencies = [ "bumpalo", "proc-macro2", @@ -10926,9 +10909,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-shared" -version = "0.2.108" +version = "0.2.113" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1f01b580c9ac74c8d8f0c0e4afb04eeef2acf145458e52c03845ee9cd23e3d12" +checksum = "de0e091bdb824da87dc01d967388880d017a0a9bc4f3bdc0d86ee9f9336e3bb5" dependencies = [ "unicode-ident", ] @@ -10974,7 +10957,7 @@ version = "0.244.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "47b807c72e1bac69382b3a6fb3dbe8ea4c0ed87ff5629b8685ae6b9a611028fe" dependencies = [ - "bitflags 2.10.0", + "bitflags 2.11.0", "hashbrown 0.15.5", "indexmap", "semver", @@ -10982,9 +10965,9 @@ dependencies = [ [[package]] name = "web-sys" -version = "0.3.85" +version = "0.3.90" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "312e32e551d92129218ea9a2452120f4aabc03529ef03e4d0d82fb2780608598" +checksum = "705eceb4ce901230f8625bd1d665128056ccbe4b7408faa625eec1ba80f59a97" dependencies = [ "js-sys", "wasm-bindgen", @@ -11088,7 +11071,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d3fabb953106c3c8eea8306e4393700d7657561cb43122571b172bbfb7c7ba1d" dependencies = [ "env_home", - "rustix 1.1.3", + "rustix 1.1.4", "winsafe", ] @@ -11130,7 +11113,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9babd3a767a4c1aef6900409f85f5d53ce2544ccdfaa86dad48c91782c6d6893" dependencies = [ "windows-collections", - "windows-core", + "windows-core 0.61.2", "windows-future", "windows-link 0.1.3", "windows-numerics", @@ -11142,7 +11125,7 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3beeceb5e5cfd9eb1d76b381630e82c4241ccd0d27f1a39ed41b2760b255c5e8" dependencies = [ - "windows-core", + "windows-core 0.61.2", ] [[package]] @@ -11158,13 +11141,26 @@ dependencies = [ "windows-strings 0.4.2", ] +[[package]] +name = "windows-core" +version = "0.62.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8e83a14d34d0623b51dce9581199302a221863196a1dde71a7663a4c2be9deb" +dependencies = [ + "windows-implement", + "windows-interface", + "windows-link 0.2.1", + "windows-result 0.4.1", + "windows-strings 0.5.1", +] + [[package]] name = "windows-future" version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fc6a41e98427b19fe4b73c550f060b59fa592d7d686537eebf9385621bfbad8e" dependencies = [ - "windows-core", + "windows-core 0.61.2", "windows-link 0.1.3", "windows-threading", ] @@ -11209,7 +11205,7 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9150af68066c4c5c07ddc0ce30421554771e528bde427614c61038bc2c92c2b1" dependencies = [ - "windows-core", + "windows-core 0.61.2", "windows-link 0.1.3", ] @@ -11573,7 +11569,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9d66ea20e9553b30172b5e831994e35fbde2d165325bec84fc43dbf6f4eb9cb2" dependencies = [ "anyhow", - "bitflags 2.10.0", + "bitflags 2.11.0", "indexmap", "log", "serde", @@ -11650,7 +11646,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32e45ad4206f6d2479085147f02bc2ef834ac85886624a23575ae137c8aa8156" dependencies = [ "libc", - "rustix 1.1.3", + "rustix 1.1.4", ] [[package]] @@ -11717,18 +11713,18 @@ dependencies = [ [[package]] name = "zerocopy" -version = "0.8.39" +version = "0.8.40" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "db6d35d663eadb6c932438e763b262fe1a70987f9ae936e60158176d710cae4a" +checksum = "a789c6e490b576db9f7e6b6d661bcc9799f7c0ac8352f56ea20193b2681532e5" dependencies = [ "zerocopy-derive", ] [[package]] name = "zerocopy-derive" -version = "0.8.39" +version = "0.8.40" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4122cd3169e94605190e77839c9a40d40ed048d305bfdc146e7df40ab0f3e517" +checksum = "f65c489a7071a749c849713807783f70672b28094011623e200cb86dcb835953" dependencies = [ "proc-macro2", "quote", @@ -11847,9 +11843,9 @@ dependencies = [ [[package]] name = "zlib-rs" -version = "0.6.0" +version = "0.6.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a7948af682ccbc3342b6e9420e8c51c1fe5d7bf7756002b4a3c6cabfe96a7e3c" +checksum = "c745c48e1007337ed136dc99df34128b9faa6ed542d80a1c673cf55a6d7236c8" [[package]] name = "zmij" diff --git a/vortex-datafusion/Cargo.toml b/vortex-datafusion/Cargo.toml index 2215bd4db8e..be4fe775100 100644 --- a/vortex-datafusion/Cargo.toml +++ b/vortex-datafusion/Cargo.toml @@ -34,6 +34,7 @@ object_store = { workspace = true } tokio = { workspace = true, features = ["rt-multi-thread", "fs"] } tokio-stream = { workspace = true } tracing = { workspace = true, features = ["std", "attributes"] } +uuid = { workspace = true, features = ["v4"] } vortex = { workspace = true, features = ["object_store", "tokio", "files"] } vortex-utils = { workspace = true, features = ["dashmap"] } diff --git a/vortex-datafusion/src/convert/exprs.rs b/vortex-datafusion/src/convert/exprs.rs index 8cac96f7de0..b8788580da5 100644 --- a/vortex-datafusion/src/convert/exprs.rs +++ b/vortex-datafusion/src/convert/exprs.rs @@ -182,41 +182,6 @@ impl DefaultExpressionConvertor { Ok(else_expr) } - - /// Attempts to convert a DataFusion CaseExpr to a Vortex expression. - fn try_convert_case_expr(&self, case_expr: &df_expr::CaseExpr) -> VortexResult { - // DataFusion CaseExpr has: - // - expr(): Optional base expression (for "CASE expr WHEN ..." form) - // - when_then_expr(): Vec of (when, then) pairs - // - else_expr(): Optional else expression - - // We don't support the "CASE expr WHEN value1 THEN result1" form yet - if case_expr.expr().is_some() { - vortex_bail!( - "CASE expr WHEN form is not yet supported, only searched CASE is supported" - ); - } - - let when_then_pairs = case_expr.when_then_expr(); - if when_then_pairs.is_empty() { - vortex_bail!("CASE expression must have at least one WHEN clause"); - } - - // Convert all when/then pairs - let mut children = Vec::with_capacity(when_then_pairs.len() * 2 + 1); - for (when_expr, then_expr) in when_then_pairs { - children.push(self.convert(when_expr.as_ref())?); - children.push(self.convert(then_expr.as_ref())?); - } - - // Handle the optional else clause - if let Some(else_expr) = case_expr.else_expr() { - children.push(self.convert(else_expr.as_ref())?); - Ok(case_when(children)) - } else { - Ok(case_when_no_else(children)) - } - } } impl ExpressionConvertor for DefaultExpressionConvertor { @@ -330,7 +295,7 @@ impl ExpressionConvertor for DefaultExpressionConvertor { let r = projection_expr.expr.apply(|node| { // We only pull column children of scalar functions that we can't push into the scan. if let Some(scalar_fn_expr) = node.as_any().downcast_ref::() - && !can_scalar_fn_be_pushed_down(scalar_fn_expr) + && !can_scalar_fn_be_pushed_down(scalar_fn_expr, input_schema) { scan_projection.extend( collect_columns(node) @@ -468,7 +433,7 @@ fn can_be_pushed_down_impl(df_expr: &Arc, schema: &Schema) -> .iter() .all(|e| can_be_pushed_down_impl(e, schema)) } else if let Some(scalar_fn) = expr.downcast_ref::() { - can_scalar_fn_be_pushed_down(scalar_fn) + can_scalar_fn_be_pushed_down(scalar_fn, schema) } else if let Some(case_expr) = expr.downcast_ref::() { can_case_be_pushed_down(case_expr, schema) } else if expr @@ -516,27 +481,19 @@ fn can_binary_be_pushed_down(binary: &df_expr::BinaryExpr, schema: &Schema) -> b } fn can_case_be_pushed_down(case_expr: &df_expr::CaseExpr, schema: &Schema) -> bool { - // We only support the "searched CASE" form (CASE WHEN cond THEN result ...) - // not the "simple CASE" form (CASE expr WHEN value THEN result ...) - if case_expr.expr().is_some() { - return false; - } - - // Check all when/then pairs - for (when_expr, then_expr) in case_expr.when_then_expr() { - if !can_be_pushed_down_impl(when_expr, schema) || !can_be_pushed_down_impl(then_expr, schema) { - return false; - } - } - - // Check the optional else clause - if let Some(else_expr) = case_expr.else_expr() - && !can_be_pushed_down_impl(else_expr, schema) - { - return false; - } - - true + case_expr + .expr() + .is_none_or(|base_expr| can_be_pushed_down_impl(base_expr, schema)) + && case_expr + .when_then_expr() + .iter() + .all(|(when_expr, then_expr)| { + can_be_pushed_down_impl(when_expr, schema) + && can_be_pushed_down_impl(then_expr, schema) + }) + && case_expr + .else_expr() + .is_some_and(|else_expr| can_be_pushed_down_impl(else_expr, schema)) } fn supported_data_types(dt: &DataType) -> bool { diff --git a/vortex-datafusion/src/persistent/format.rs b/vortex-datafusion/src/persistent/format.rs index a3c1a543578..d1cf153008b 100644 --- a/vortex-datafusion/src/persistent/format.rs +++ b/vortex-datafusion/src/persistent/format.rs @@ -67,6 +67,7 @@ use crate::PrecisionExt as _; use crate::convert::TryToDataFusion; const DEFAULT_FOOTER_INITIAL_READ_SIZE_BYTES: usize = MAX_POSTSCRIPT_SIZE as usize + EOF_SIZE; +const DEFAULT_TARGET_FILE_SIZE_MB: usize = 128; /// Vortex implementation of a DataFusion [`FileFormat`]. pub struct VortexFormat { @@ -94,6 +95,11 @@ config_namespace! { /// Values smaller than `MAX_POSTSCRIPT_SIZE + EOF_SIZE` will be clamped to that minimum /// during footer parsing. pub footer_initial_read_size_bytes: usize, default = DEFAULT_FOOTER_INITIAL_READ_SIZE_BYTES + /// Target file size in megabytes for written Vortex files. + /// + /// When greater than 0, Vortex bypasses DataFusion's file demuxer and + /// splits output files based on approximate byte size rather than row count. + pub target_file_size_mb: usize, default = DEFAULT_TARGET_FILE_SIZE_MB /// Whether to enable projection pushdown into the underlying Vortex scan. /// /// When enabled, projection expressions may be partially evaluated during @@ -506,8 +512,25 @@ impl FileFormat for VortexFormat { return not_impl_err!("Overwrites are not implemented yet for Vortex"); } + let target_file_size = (self.opts.target_file_size_mb > 0) + .then(|| { + u64::try_from(self.opts.target_file_size_mb) + .map_err(|e| { + internal_datafusion_err!( + "target_file_size_mb cannot be represented as u64: {e}" + ) + }) + .map(|v| v.saturating_mul(1024 * 1024).max(1)) + }) + .transpose()?; + let schema = conf.output_schema().clone(); - let sink = Arc::new(VortexSink::new(conf, schema, self.session.clone())); + let sink = Arc::new(VortexSink::new( + conf, + schema, + self.session.clone(), + target_file_size, + )); Ok(Arc::new(DataSinkExec::new(input, sink, order_requirements)) as _) } @@ -577,11 +600,17 @@ mod tests { } #[test] - fn format_plumbs_footer_initial_read_size() { - let mut opts = VortexOptions::default(); - opts.set("footer_initial_read_size_bytes", "12345").unwrap(); + fn format_plumbs_target_file_size_mb() { + let mut opts = VortexTableOptions::default(); + opts.set("target_file_size_mb", "123").unwrap(); let format = VortexFormat::new_with_options(VortexSession::default(), opts); - assert_eq!(format.file_cache.footer_initial_read_size_bytes(), 12345); + assert_eq!(format.options().target_file_size_mb, 123); + } + + #[test] + fn format_target_file_size_default_is_128mb() { + let opts = VortexTableOptions::default(); + assert_eq!(opts.target_file_size_mb, 128); } } diff --git a/vortex-datafusion/src/persistent/sink.rs b/vortex-datafusion/src/persistent/sink.rs index 60cf1a8d0f4..159a052e966 100644 --- a/vortex-datafusion/src/persistent/sink.rs +++ b/vortex-datafusion/src/persistent/sink.rs @@ -8,6 +8,7 @@ use arrow_schema::SchemaRef; use async_trait::async_trait; use datafusion_common::DataFusionError; use datafusion_common::Result as DFResult; +use datafusion_common::arrow::array::RecordBatch; use datafusion_common::exec_datafusion_err; use datafusion_common_runtime::JoinSet; use datafusion_common_runtime::SpawnedTask; @@ -22,6 +23,7 @@ use datafusion_physical_plan::DisplayAs; use datafusion_physical_plan::DisplayFormatType; use datafusion_physical_plan::metrics::MetricsSet; use futures::StreamExt; +use futures::TryStreamExt; use object_store::ObjectStore; use object_store::path::Path; use tokio_stream::wrappers::ReceiverStream; @@ -40,16 +42,117 @@ pub struct VortexSink { config: FileSinkConfig, schema: SchemaRef, session: VortexSession, + target_file_size: Option, } impl VortexSink { - pub fn new(config: FileSinkConfig, schema: SchemaRef, session: VortexSession) -> Self { + pub fn new( + config: FileSinkConfig, + schema: SchemaRef, + session: VortexSession, + target_file_size: Option, + ) -> Self { Self { config, schema, session, + target_file_size, } } + + async fn write_all_with_target_size( + &self, + mut data: SendableRecordBatchStream, + context: &Arc, + target_file_size: u64, + ) -> DFResult { + if !self.config.table_partition_cols.is_empty() { + return FileSink::write_all(self, data, context).await; + } + + let object_store = context + .runtime_env() + .object_store(&self.config.object_store_url)?; + + let writer_schema = get_writer_schema(&self.config); + let dtype = DType::from_arrow(writer_schema); + + let table_path = self + .config + .table_paths + .first() + .ok_or_else(|| exec_datafusion_err!("No output table path configured"))?; + + let base_path = table_path.prefix(); + let extension = self.config.file_extension.as_str(); + let target_file_size = target_file_size.max(1); + + let mut row_count = 0_u64; + let mut file_index = 0_usize; + let mut buffered_batches: Vec = Vec::new(); + let mut buffered_bytes = 0_u64; + + while let Some(batch) = data.try_next().await? { + buffered_bytes = buffered_bytes.saturating_add(batch.get_array_memory_size() as u64); + buffered_batches.push(batch); + + if buffered_bytes >= target_file_size { + let path = split_path(base_path, file_index, extension); + let summary = self + .write_batches(object_store.clone(), path, dtype.clone(), buffered_batches) + .await?; + row_count = row_count.saturating_add(summary.row_count()); + buffered_batches = Vec::new(); + buffered_bytes = 0; + file_index += 1; + } + } + + if !buffered_batches.is_empty() { + let path = split_path(base_path, file_index, extension); + let summary = self + .write_batches(object_store.clone(), path, dtype, buffered_batches) + .await?; + row_count = row_count.saturating_add(summary.row_count()); + } + + Ok(row_count) + } + + async fn write_batches( + &self, + object_store: Arc, + path: Path, + dtype: DType, + batches: Vec, + ) -> DFResult { + let stream = futures::stream::iter( + batches + .into_iter() + .map(|rb| ArrayRef::from_arrow(rb, false)), + ); + let stream_adapter = ArrayStreamAdapter::new(dtype, stream); + + let mut object_writer = ObjectStoreWrite::new(object_store, &path) + .await + .map_err(|e| exec_datafusion_err!("Failed to create ObjectStoreWrite: {e}"))?; + + let summary = self + .session + .write_options() + .write(&mut object_writer, stream_adapter) + .await + .map_err(|e| exec_datafusion_err!("Failed to write Vortex file: {e}"))?; + + object_writer + .shutdown() + .await + .map_err(|e| exec_datafusion_err!("Failed to shutdown Vortex writer: {e}"))?; + + tracing::info!(path = %path, "Successfully written file"); + + Ok(summary) + } } impl std::fmt::Debug for VortexSink { @@ -90,6 +193,12 @@ impl DataSink for VortexSink { data: SendableRecordBatchStream, context: &Arc, ) -> DFResult { + if let Some(target_file_size) = self.target_file_size { + return self + .write_all_with_target_size(data, context, target_file_size) + .await; + } + FileSink::write_all(self, data, context).await } } @@ -173,6 +282,15 @@ impl FileSink for VortexSink { } } +fn split_path(base_path: &Path, file_index: usize, extension: &str) -> Path { + let mut base = base_path.to_string(); + if !base.ends_with('/') { + base.push('/'); + } + let filename = format!("part-{file_index:05}.{extension}"); + Path::from(format!("{base}{filename}")) +} + #[cfg(test)] mod tests { use std::sync::Arc; @@ -194,8 +312,10 @@ mod tests { use futures::TryStreamExt; use rstest::rstest; + use super::split_path; use crate::common_tests::TestSessionContext; use crate::persistent::VortexFormatFactory; + use crate::persistent::VortexTableOptions; #[tokio::test] async fn test_insert_into_sql() -> anyhow::Result<()> { @@ -295,10 +415,11 @@ mod tests { } /// Reproduction by . + /// Uses a 1MB target file size to exercise file splitting behavior. #[rstest] - #[case(1000, 1)] - #[case(40_961, 4)] - #[case(1_000_000, 4)] + #[case(1_000, 1)] + #[case(5_000_000, 6)] + #[case(10_000_000, 10)] #[tokio::test] async fn test_write_large_batch( #[case] entries: usize, @@ -306,15 +427,26 @@ mod tests { ) -> anyhow::Result<()> { let ctx = TestSessionContext::default(); + let opts = VortexTableOptions { + target_file_size_mb: 1, + ..Default::default() + }; + + let factory = VortexFormatFactory::new().with_options(opts); + + let values: Vec = (0..entries) + .map(|i| i8::try_from(i % 127)) + .collect::>()?; + let data = ctx.session.read_batch(RecordBatch::try_new( Arc::new(Schema::new(vec![Field::new("a", DataType::Int8, false)])), - vec![Arc::new(Int8Array::from(vec![0i8; entries]))], + vec![Arc::new(Int8Array::from(values))], )?)?; let logical_plan = LogicalPlanBuilder::copy_to( data.logical_plan().clone(), "/table/".to_string(), - format_as_file_type(Arc::new(VortexFormatFactory::new())), + format_as_file_type(Arc::new(factory)), Default::default(), vec![], )? @@ -366,13 +498,7 @@ mod tests { .unwrap(); for i in 0..batch.num_rows() { - assert_eq!( - col.value(i), - 0i8, - "Expected value 0 at row {}, but found {}", - total_rows + i, - col.value(i) - ); + assert_eq!(col.value(i), i8::try_from((total_rows + i) % 127)?); } total_rows += batch.num_rows(); } @@ -398,6 +524,46 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_write_large_batch_default_target_is_128mb() -> anyhow::Result<()> { + let ctx = TestSessionContext::default(); + + let entries = 1_000_000; + let values: Vec = (0..entries) + .map(|i| i8::try_from(i % 127)) + .collect::>()?; + + let data = ctx.session.read_batch(RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("a", DataType::Int8, false)])), + vec![Arc::new(Int8Array::from(values))], + )?)?; + + let logical_plan = LogicalPlanBuilder::copy_to( + data.logical_plan().clone(), + "/table/".to_string(), + format_as_file_type(Arc::new(VortexFormatFactory::new())), + Default::default(), + vec![], + )? + .build()?; + + ctx.session + .execute_logical_plan(logical_plan) + .await? + .collect() + .await?; + + let file_metas = ctx + .store + .list(Some(&"/table".into())) + .try_collect::>() + .await?; + + assert_eq!(file_metas.len(), 1); + + Ok(()) + } + #[tokio::test] async fn test_write_partitioned() -> anyhow::Result<()> { let ctx = TestSessionContext::default(); @@ -439,4 +605,26 @@ mod tests { Ok(()) } + + #[test] + fn test_split_path_basic() { + let path = object_store::path::Path::from("data/output"); + assert_eq!( + split_path(&path, 0, "vortex").to_string(), + "data/output/part-00000.vortex" + ); + assert_eq!( + split_path(&path, 12, "vortex").to_string(), + "data/output/part-00012.vortex" + ); + } + + #[test] + fn test_split_path_preserves_trailing_slash() { + let path = object_store::path::Path::from("nested/path/"); + assert_eq!( + split_path(&path, 3, "vx").to_string(), + "nested/path/part-00003.vx" + ); + } } diff --git a/vortex-datafusion/src/persistent/source.rs b/vortex-datafusion/src/persistent/source.rs index 9917a7b8252..bac7fe2b39a 100644 --- a/vortex-datafusion/src/persistent/source.rs +++ b/vortex-datafusion/src/persistent/source.rs @@ -175,8 +175,6 @@ impl FileSource for VortexSource { .clone() .unwrap_or_else(|| Arc::new(DefaultVortexReaderFactory::new(object_store))); - let table_schema = base_config.file_source.table_schema().clone(); - let opener = VortexOpener { partition, session: self.session.clone(), From 3e77695f45cb22779a37fa056a836abd4808b607 Mon Sep 17 00:00:00 2001 From: Luke Kim <80174+lukekim@users.noreply.github.com> Date: Thu, 26 Feb 2026 19:56:05 -0800 Subject: [PATCH 26/32] Fix spiceai-52 build: disable broken case_when module --- vortex-array/src/scalar_fn/fns/mod.rs | 1 - vortex-datafusion/src/convert/exprs.rs | 4 ---- 2 files changed, 5 deletions(-) diff --git a/vortex-array/src/scalar_fn/fns/mod.rs b/vortex-array/src/scalar_fn/fns/mod.rs index 94fc8fb0384..95c66b09ef6 100644 --- a/vortex-array/src/scalar_fn/fns/mod.rs +++ b/vortex-array/src/scalar_fn/fns/mod.rs @@ -3,7 +3,6 @@ pub mod between; pub mod binary; -pub mod case_when; pub mod cast; pub mod dynamic; pub mod fill_null; diff --git a/vortex-datafusion/src/convert/exprs.rs b/vortex-datafusion/src/convert/exprs.rs index b8788580da5..9b634e02794 100644 --- a/vortex-datafusion/src/convert/exprs.rs +++ b/vortex-datafusion/src/convert/exprs.rs @@ -37,11 +37,7 @@ 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::case_when::case_when; -use vortex::scalar_fn::fns::case_when::case_when_no_else; use vortex::scalar_fn::fns::operators::Operator; -use vortex::error::VortexResult; -use vortex::error::vortex_bail; use crate::convert::FromDataFusion; From 7e5b081518ac96d3b68b9882c32ebb4e701bb1e7 Mon Sep 17 00:00:00 2001 From: Phillip LeBlanc Date: Fri, 27 Feb 2026 13:04:05 +0900 Subject: [PATCH 27/32] cast: support vortex.date -> vortex.timestamp extension casts (#28) --- .../src/arrays/extension/compute/cast.rs | 236 ++++++++++++++++-- 1 file changed, 218 insertions(+), 18 deletions(-) diff --git a/vortex-array/src/arrays/extension/compute/cast.rs b/vortex-array/src/arrays/extension/compute/cast.rs index 8a82180fbba..c70d3fc12b4 100644 --- a/vortex-array/src/arrays/extension/compute/cast.rs +++ b/vortex-array/src/arrays/extension/compute/cast.rs @@ -1,40 +1,187 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_buffer::BufferMut; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_mask::AllOr; + use crate::ArrayRef; use crate::IntoArray; +use crate::ToCanonical; use crate::arrays::ExtensionArray; use crate::arrays::ExtensionVTable; +use crate::arrays::PrimitiveArray; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; +use crate::dtype::PType; +use crate::extension::datetime::AnyTemporal; +use crate::extension::datetime::TemporalMetadata; +use crate::extension::datetime::TimeUnit; use crate::scalar_fn::fns::cast::CastReduce; +use crate::vtable::ValidityHelper; impl CastReduce for ExtensionVTable { - fn cast(array: &ExtensionArray, dtype: &DType) -> vortex_error::VortexResult> { - if !array.dtype().eq_ignore_nullability(dtype) { + fn cast(array: &ExtensionArray, dtype: &DType) -> VortexResult> { + let DType::Extension(ext_dtype) = dtype else { return Ok(None); + }; + + if array.ext_dtype().eq_ignore_nullability(ext_dtype) { + let new_storage = match array + .storage() + .cast(ext_dtype.storage_dtype().clone()) + .and_then(|a| a.to_canonical().map(|c| c.into_array())) + { + Ok(arr) => arr, + Err(e) => { + tracing::warn!("Failed to cast storage array: {e}"); + return Ok(None); + } + }; + + return Ok(Some( + ExtensionArray::new(ext_dtype.clone(), new_storage).into_array(), + )); } - let DType::Extension(ext_dtype) = dtype else { - unreachable!("Already verified we have an extension dtype"); - }; + if let Some(new_storage) = cast_temporal_date_to_timestamp(array, dtype)? { + return Ok(Some( + ExtensionArray::new(ext_dtype.clone(), new_storage).into_array(), + )); + } + + Ok(None) + } +} + +fn cast_temporal_date_to_timestamp( + array: &ExtensionArray, + target_dtype: &DType, +) -> VortexResult> { + let DType::Extension(target_ext_dtype) = target_dtype else { + return Ok(None); + }; + + let Some(source_temporal) = array.ext_dtype().metadata_opt::() else { + return Ok(None); + }; + let Some(target_temporal) = target_ext_dtype.metadata_opt::() else { + return Ok(None); + }; + + let (TemporalMetadata::Date(source_unit), TemporalMetadata::Timestamp(target_unit, _)) = + (source_temporal, target_temporal) + else { + return Ok(None); + }; + + let source_i64 = array + .storage() + .cast(DType::Primitive(PType::I64, array.dtype().nullability()))?; + let source_i64 = source_i64.to_primitive(); - let new_storage = match array - .storage() - .cast(ext_dtype.storage_dtype().clone()) - .and_then(|a| a.to_canonical().map(|c| c.into_array())) - { - Ok(arr) => arr, - Err(e) => { - tracing::warn!("Failed to cast storage array: {e}"); - return Ok(None); + let converted = cast_date_values_to_timestamp(&source_i64, *source_unit, *target_unit)?; + + converted + .to_array() + .cast(target_ext_dtype.storage_dtype().clone()) + .map(Some) +} + +fn cast_date_values_to_timestamp( + values: &PrimitiveArray, + source_unit: TimeUnit, + target_unit: TimeUnit, +) -> VortexResult { + let (multiply, divide) = date_to_timestamp_scale(source_unit, target_unit)?; + + let input = values.as_slice::(); + let mut output = BufferMut::with_capacity(input.len()); + match values.validity_mask()?.bit_buffer() { + AllOr::All => { + for &value in input { + // SAFETY: output has sufficient capacity for all pushed values. + unsafe { output.push_unchecked(convert_temporal_value(value, multiply, divide)?) }; } - }; + } + AllOr::None => { + for _ in 0..input.len() { + // SAFETY: output has sufficient capacity for all pushed values. + unsafe { output.push_unchecked(0i64) }; + } + } + AllOr::Some(bits) => { + for (&value, valid) in input.iter().zip(bits.iter()) { + if valid { + // SAFETY: output has sufficient capacity for all pushed values. + unsafe { + output.push_unchecked(convert_temporal_value(value, multiply, divide)?) + }; + } else { + // SAFETY: output has sufficient capacity for all pushed values. + unsafe { output.push_unchecked(0i64) }; + } + } + } + } + + Ok(PrimitiveArray::new( + output.freeze(), + values.validity().clone(), + )) +} + +fn date_to_timestamp_scale( + source_unit: TimeUnit, + target_unit: TimeUnit, +) -> VortexResult<(i64, i64)> { + let source_ns = to_nanoseconds(source_unit)?; + let target_ns = to_nanoseconds(target_unit)?; + + if source_ns >= target_ns { + let multiply = source_ns / target_ns; + return Ok((multiply, 1)); + } - Ok(Some( - ExtensionArray::new(ext_dtype.clone(), new_storage).into_array(), - )) + let divide = target_ns / source_ns; + Ok((1, divide)) +} + +fn to_nanoseconds(unit: TimeUnit) -> VortexResult { + match unit { + TimeUnit::Nanoseconds => Ok(1), + TimeUnit::Microseconds => Ok(1_000), + TimeUnit::Milliseconds => Ok(1_000_000), + TimeUnit::Seconds => Ok(1_000_000_000), + TimeUnit::Days => Ok(86_400_000_000_000), + } +} + +fn convert_temporal_value(value: i64, multiply: i64, divide: i64) -> VortexResult { + let mut scaled = i128::from(value) + .checked_mul(i128::from(multiply)) + .ok_or_else(|| { + vortex_error::vortex_err!( + Compute: "Date value {value} overflows while scaling to timestamp" + ) + })?; + + if divide != 1 { + let divisor = i128::from(divide); + if scaled % divisor != 0 { + vortex_bail!( + Compute: "Date value {value} cannot be represented exactly in target timestamp unit" + ); + } + scaled /= divisor; + } + + if scaled < i128::from(i64::MIN) || scaled > i128::from(i64::MAX) { + vortex_bail!(Compute: "Date value {value} overflows target timestamp range"); } + + Ok(scaled as i64) } #[cfg(test)] @@ -45,11 +192,13 @@ mod tests { use vortex_buffer::buffer; use super::*; + use crate::Array; use crate::IntoArray; use crate::arrays::PrimitiveArray; use crate::builtins::ArrayBuiltins; use crate::compute::conformance::cast::test_cast_conformance; use crate::dtype::Nullability; + use crate::extension::datetime::Date; use crate::extension::datetime::TimeUnit; use crate::extension::datetime::Timestamp; @@ -85,6 +234,57 @@ mod tests { assert_eq!(output.dtype(), &new_dtype); } + #[test] + fn cast_date_days_to_timestamp_nanoseconds() { + let source_dtype = Date::new(TimeUnit::Days, Nullability::NonNullable).erased(); + let target_dtype = Timestamp::new(TimeUnit::Nanoseconds, Nullability::NonNullable).erased(); + + let arr = ExtensionArray::new(source_dtype, buffer![0i32, 1, -1].into_array()); + let output = arr + .to_array() + .cast(DType::Extension(target_dtype.clone())) + .unwrap() + .to_extension(); + + assert_eq!(output.dtype(), &DType::Extension(target_dtype)); + + let storage = output.storage().to_primitive(); + assert_eq!( + storage.as_slice::(), + &[0, 86_400_000_000_000, -86_400_000_000_000] + ); + } + + #[test] + fn cast_date_days_to_timestamp_seconds_nullable() { + let source_dtype = Date::new(TimeUnit::Days, Nullability::Nullable).erased(); + let target_dtype = Timestamp::new(TimeUnit::Seconds, Nullability::Nullable).erased(); + + let arr = ExtensionArray::new( + source_dtype, + PrimitiveArray::from_option_iter([Some(0i32), None, Some(2)]).into_array(), + ); + + let output = arr + .to_array() + .cast(DType::Extension(target_dtype.clone())) + .unwrap() + .to_extension(); + + assert_eq!(output.dtype(), &DType::Extension(target_dtype)); + + let storage = output.storage().to_primitive(); + assert_eq!( + storage.scalar_at(0).unwrap().as_primitive().as_::(), + Some(0) + ); + assert!(storage.scalar_at(1).unwrap().is_null()); + assert_eq!( + storage.scalar_at(2).unwrap().as_primitive().as_::(), + Some(172_800) + ); + } + #[test] fn cast_different_ext_dtype() { let original_dtype = From c536c9aedc8fdf8e3787a31c43c8e41318cf3956 Mon Sep 17 00:00:00 2001 From: Phillip LeBlanc Date: Fri, 27 Feb 2026 20:27:41 +0900 Subject: [PATCH 28/32] fix(vortex-file): avoid session lock re-entry in writer init (#29) --- vortex-file/src/writer.rs | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/vortex-file/src/writer.rs b/vortex-file/src/writer.rs index 0af1b129ee2..4715393656e 100644 --- a/vortex-file/src/writer.rs +++ b/vortex-file/src/writer.rs @@ -150,9 +150,15 @@ impl VortexWriteOptions { // serialised array order is deterministic. The serialisation of arrays are done // parallel and with an empty context they can register their encodings to the context // in different order, changing the written bytes from run to run. - let ctx = ArrayContext::new(self.session.arrays().registry().ids().sorted().collect()) + // Avoid calling `self.session.arrays()` multiple times in one expression. + // `Ref` holds a session lock; reacquiring it in the same statement can deadlock. + let registry = { + let arrays = self.session.arrays(); + arrays.registry().clone() + }; + let ctx = ArrayContext::new(registry.ids().sorted().collect()) // Configure a registry just to ensure only known encodings are interned. - .with_registry(self.session.arrays().registry().clone()); + .with_registry(registry); let dtype = stream.dtype().clone(); let (mut ptr, eof) = SequenceId::root().split(); From 1a6dc54f18e6c9514e4273cbd736bcbfbcbe16d5 Mon Sep 17 00:00:00 2001 From: Luke Kim <80174+lukekim@users.noreply.github.com> Date: Fri, 6 Mar 2026 16:09:15 -0800 Subject: [PATCH 29/32] feat: add arrow map alias support Signed-off-by: Luke Kim <80174+lukekim@users.noreply.github.com> --- .../src/arrays/extension/compute/cast.rs | 3 +- vortex-array/src/arrow/convert.rs | 69 +++++++++++++++++++ vortex-array/src/arrow/executor/map.rs | 54 +++++++++++++++ vortex-array/src/arrow/executor/mod.rs | 7 +- vortex-array/src/dtype/arrow.rs | 8 +++ 5 files changed, 138 insertions(+), 3 deletions(-) create mode 100644 vortex-array/src/arrow/executor/map.rs diff --git a/vortex-array/src/arrays/extension/compute/cast.rs b/vortex-array/src/arrays/extension/compute/cast.rs index c70d3fc12b4..c6cb284a4ec 100644 --- a/vortex-array/src/arrays/extension/compute/cast.rs +++ b/vortex-array/src/arrays/extension/compute/cast.rs @@ -181,7 +181,8 @@ fn convert_temporal_value(value: i64, multiply: i64, divide: i64) -> VortexResul vortex_bail!(Compute: "Date value {value} overflows target timestamp range"); } - Ok(scaled as i64) + i64::try_from(scaled) + .map_err(|_| vortex_error::vortex_err!(Compute: "Date value {value} overflows target timestamp range")) } #[cfg(test)] diff --git a/vortex-array/src/arrow/convert.rs b/vortex-array/src/arrow/convert.rs index a702b6a8727..2de9b58aca8 100644 --- a/vortex-array/src/arrow/convert.rs +++ b/vortex-array/src/arrow/convert.rs @@ -13,6 +13,7 @@ use arrow_array::GenericByteArray; use arrow_array::GenericByteViewArray; use arrow_array::GenericListArray; use arrow_array::GenericListViewArray; +use arrow_array::MapArray as ArrowMapArray; use arrow_array::NullArray as ArrowNullArray; use arrow_array::OffsetSizeTrait; use arrow_array::PrimitiveArray as ArrowPrimitiveArray; @@ -456,6 +457,18 @@ impl FromArrowArray<&ArrowFixedSizeListArray> for ArrayRef { } } +impl FromArrowArray<&ArrowMapArray> for ArrayRef { + fn from_arrow(value: &ArrowMapArray, nullable: bool) -> VortexResult { + // Arrow Map is logically List> with i32 offsets. + // We convert it to a ListArray of structs. + let entries = Self::from_arrow(value.entries() as &dyn ArrowArray, false)?; + let offsets = value.offsets().clone().into_array(); + let nulls = nulls(value.nulls(), nullable); + + Ok(ListArray::try_new(entries, offsets, nulls)?.into_array()) + } +} + impl FromArrowArray<&ArrowNullArray> for ArrayRef { fn from_arrow(value: &ArrowNullArray, nullable: bool) -> VortexResult { assert!(nullable); @@ -517,6 +530,7 @@ impl FromArrowArray<&dyn ArrowArray> for ArrayRef { DataType::ListView(_) => Self::from_arrow(array.as_list_view::(), nullable), DataType::LargeListView(_) => Self::from_arrow(array.as_list_view::(), nullable), DataType::FixedSizeList(..) => Self::from_arrow(array.as_fixed_size_list(), nullable), + DataType::Map(..) => Self::from_arrow(array.as_map(), nullable), DataType::Null => Self::from_arrow(as_null_array(array), nullable), DataType::Timestamp(u, _) => match u { ArrowTimeUnit::Second => { @@ -680,6 +694,7 @@ mod tests { use crate::ArrayRef; use crate::IntoArray; + use crate::VortexSessionExecute; use crate::arrays::DecimalVTable; use crate::arrays::FixedSizeListVTable; use crate::arrays::ListVTable; @@ -690,6 +705,7 @@ mod tests { use crate::arrays::VarBinVTable; use crate::arrays::VarBinViewVTable; use crate::arrow::FromArrowArray as _; + use crate::arrow::executor::ArrowArrayExecutor as _; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -1777,4 +1793,57 @@ mod tests { ArrayRef::from_arrow(null_struct_array_with_non_nullable_field.as_ref(), true).unwrap(); } + + #[test] + fn test_map_array_conversion() { + use arrow_array::MapArray; + use arrow_array::builder::MapBuilder; + use arrow_array::builder::StringBuilder; + + // Build a MapArray: map + let mut builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new()); + // First map entry: {"a": 1, "b": 2} + builder.keys().append_value("a"); + builder.values().append_value(1); + builder.keys().append_value("b"); + builder.values().append_value(2); + builder.append(true).unwrap(); + + // Second map entry: null + builder.append(false).unwrap(); + + // Third map entry: {"c": 3} + builder.keys().append_value("c"); + builder.values().append_value(3); + builder.append(true).unwrap(); + + let arrow_map = builder.finish(); + assert_eq!(arrow_map.len(), 3); + + // Convert Arrow MapArray → Vortex ListArray + let vortex_array = ArrayRef::from_arrow(&arrow_map, true).unwrap(); + assert_eq!(vortex_array.len(), 3); + + // Verify it's stored as List> + let list_array = vortex_array.as_::(); + assert_eq!(list_array.elements().len(), 3); // 3 total key-value pairs + let struct_elements = list_array.elements().as_::(); + assert_eq!(struct_elements.names().len(), 2); // key and value fields + + // Convert back to Arrow as a MapArray + let map_dtype = arrow_map.data_type().clone(); + let arrow_back = vortex_array + .execute_arrow( + Some(&map_dtype), + &mut crate::LEGACY_SESSION.create_execution_ctx(), + ) + .unwrap(); + let map_back = arrow_back + .as_any() + .downcast_ref::() + .expect("Should be a MapArray"); + assert_eq!(map_back.len(), 3); + assert_eq!(map_back.entries().len(), 3); + assert!(map_back.is_null(1)); // Second entry was null + } } diff --git a/vortex-array/src/arrow/executor/map.rs b/vortex-array/src/arrow/executor/map.rs new file mode 100644 index 00000000000..a289fb34d91 --- /dev/null +++ b/vortex-array/src/arrow/executor/map.rs @@ -0,0 +1,54 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; + +use arrow_array::ArrayRef as ArrowArrayRef; +use arrow_array::MapArray as ArrowMapArray; +use arrow_array::StructArray as ArrowStructArray; +use arrow_schema::FieldRef; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::arrow::executor::list::to_arrow_list; + +/// Convert a Vortex List> array into an Arrow MapArray. +pub(super) fn to_arrow_map( + array: ArrayRef, + entries_field: &FieldRef, + ordered: bool, + ctx: &mut ExecutionCtx, +) -> VortexResult { + // First, convert to Arrow ListArray since Map uses i32 offsets. + let list_array = to_arrow_list::(array, entries_field, ctx)?; + + // Downcast to GenericListArray to extract its components. + let Some(list_array) = list_array + .as_any() + .downcast_ref::>() + else { + vortex_bail!("to_arrow_list returned a non-ListArray when building a MapArray"); + }; + + // Extract components from the ListArray. + let (_list_field, offsets, entries, nulls) = list_array.clone().into_parts(); + + // The entries should be a StructArray. Downcast it. + let Some(entries_struct) = entries.as_any().downcast_ref::() else { + vortex_bail!("Map entries must be a StructArray"); + }; + let entries_struct = entries_struct.clone(); + + // Build the MapArray from the components. + let map_array = ArrowMapArray::try_new( + entries_field.clone(), + offsets, + entries_struct, + nulls, + ordered, + )?; + + Ok(Arc::new(map_array)) +} diff --git a/vortex-array/src/arrow/executor/mod.rs b/vortex-array/src/arrow/executor/mod.rs index 1cb0f11f587..47787d34116 100644 --- a/vortex-array/src/arrow/executor/mod.rs +++ b/vortex-array/src/arrow/executor/mod.rs @@ -9,6 +9,7 @@ mod dictionary; mod fixed_size_list; mod list; mod list_view; +mod map; pub mod null; pub mod primitive; mod run_end; @@ -29,7 +30,6 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; -use crate::Array; use crate::ArrayRef; use crate::arrays::ListVTable; use crate::arrays::VarBinVTable; @@ -41,6 +41,7 @@ use crate::arrow::executor::dictionary::to_arrow_dictionary; use crate::arrow::executor::fixed_size_list::to_arrow_fixed_list; use crate::arrow::executor::list::to_arrow_list; use crate::arrow::executor::list_view::to_arrow_list_view; +use crate::arrow::executor::map::to_arrow_map; use crate::arrow::executor::null::to_arrow_null; use crate::arrow::executor::primitive::to_arrow_primitive; use crate::arrow::executor::run_end::to_arrow_run_end; @@ -156,8 +157,10 @@ impl ArrowArrayExecutor for ArrayRef { DataType::RunEndEncoded(ends_type, values_type) => { to_arrow_run_end(self, ends_type.data_type(), values_type, ctx) } + DataType::Map(entries_field, ordered) => { + to_arrow_map(self, entries_field, *ordered, ctx) + } DataType::FixedSizeBinary(_) - | DataType::Map(..) | DataType::Duration(_) | DataType::Interval(_) | DataType::Union(..) => { diff --git a/vortex-array/src/dtype/arrow.rs b/vortex-array/src/dtype/arrow.rs index 96f940cb41f..ca15189a5d5 100644 --- a/vortex-array/src/dtype/arrow.rs +++ b/vortex-array/src/dtype/arrow.rs @@ -191,6 +191,14 @@ impl FromArrowType<(&DataType, Nullability)> for DType { | DataType::LargeListView(e) => { DType::List(Arc::new(Self::from_arrow(e.as_ref())), nullability) } + DataType::Map(entries_field, _ordered) => { + // Map is logically List>. + // The entries_field contains a Struct type with the key and value fields. + DType::List( + Arc::new(Self::from_arrow(entries_field.as_ref())), + nullability, + ) + } DataType::FixedSizeList(e, size) => DType::FixedSizeList( Arc::new(Self::from_arrow(e.as_ref())), *size as u32, From 8a09bc796245a957ad03c49c3cdfee477f483237 Mon Sep 17 00:00:00 2001 From: William <98815791+peasee@users.noreply.github.com> Date: Thu, 12 Mar 2026 15:05:55 +1000 Subject: [PATCH 30/32] fix: Ensure target file size is properly respected in vortex sink (#33) * fix: Ensure target file sizes are respected * fix: Improve robustness for streams split over multiple cores * wip * wip * wip * wip * refactor: Improve robustness * refactor: Improve concurrent writer robustness for memory bounds * Update vortex-datafusion/src/persistent/sink.rs Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Update vortex-datafusion/src/persistent/format.rs Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Update vortex-datafusion/src/persistent/sink.rs Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Update vortex-datafusion/src/persistent/sink.rs Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Update vortex-datafusion/src/tests/nested_projection.rs Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * review: Address comments * Update vortex-datafusion/src/persistent/sink.rs Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * review: Address comments * Revert "Update vortex-datafusion/src/persistent/sink.rs" This reverts commit 6c3be0f2b148c443f3925f2cdc6b16d6fd672bed. * refactor: Simplify * test: Add file writer finish notify test * review: Address comments --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- vortex-datafusion/Cargo.toml | 2 +- vortex-datafusion/src/persistent/format.rs | 24 +- vortex-datafusion/src/persistent/sink.rs | 1598 +++++++++++++++++--- 3 files changed, 1415 insertions(+), 209 deletions(-) diff --git a/vortex-datafusion/Cargo.toml b/vortex-datafusion/Cargo.toml index be4fe775100..21a93b19cf4 100644 --- a/vortex-datafusion/Cargo.toml +++ b/vortex-datafusion/Cargo.toml @@ -34,7 +34,7 @@ object_store = { workspace = true } tokio = { workspace = true, features = ["rt-multi-thread", "fs"] } tokio-stream = { workspace = true } tracing = { workspace = true, features = ["std", "attributes"] } -uuid = { workspace = true, features = ["v4"] } +uuid = { workspace = true, features = ["v7"] } vortex = { workspace = true, features = ["object_store", "tokio", "files"] } vortex-utils = { workspace = true, features = ["dashmap"] } diff --git a/vortex-datafusion/src/persistent/format.rs b/vortex-datafusion/src/persistent/format.rs index d1cf153008b..befbc94d07d 100644 --- a/vortex-datafusion/src/persistent/format.rs +++ b/vortex-datafusion/src/persistent/format.rs @@ -35,6 +35,8 @@ use datafusion_datasource::source::DataSourceExec; use datafusion_expr::dml::InsertOp; use datafusion_physical_expr::LexRequirement; use datafusion_physical_plan::ExecutionPlan; +use datafusion_physical_plan::ExecutionPlanProperties; +use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec; use futures::FutureExt; use futures::StreamExt as _; use futures::TryStreamExt as _; @@ -97,8 +99,9 @@ config_namespace! { pub footer_initial_read_size_bytes: usize, default = DEFAULT_FOOTER_INITIAL_READ_SIZE_BYTES /// Target file size in megabytes for written Vortex files. /// - /// When greater than 0, Vortex bypasses DataFusion's file demuxer and - /// splits output files based on approximate byte size rather than row count. + /// When greater than 0 for non-partitioned writes, Vortex bypasses + /// DataFusion's file demuxer and splits output files based on + /// approximate byte size rather than row count. pub target_file_size_mb: usize, default = DEFAULT_TARGET_FILE_SIZE_MB /// Whether to enable projection pushdown into the underlying Vortex scan. /// @@ -524,6 +527,23 @@ impl FileFormat for VortexFormat { }) .transpose()?; + // For non-partitioned writes, force a single input stream so VortexSink + // performs one coordinated write per statement instead of one + // independent write per CPU/input partition. + // + // Use coalescing rather than repartitioning to avoid introducing a + // shuffle/dispatcher step that can interleave batches from different + // input partitions. + // + // For partitioned writes, keep DataFusion's demuxer behavior. + let input: Arc = if conf.table_partition_cols.is_empty() + && input.output_partitioning().partition_count() > 1 + { + Arc::new(CoalescePartitionsExec::new(input)) + } else { + input + }; + let schema = conf.output_schema().clone(); let sink = Arc::new(VortexSink::new( conf, diff --git a/vortex-datafusion/src/persistent/sink.rs b/vortex-datafusion/src/persistent/sink.rs index 159a052e966..c418247e026 100644 --- a/vortex-datafusion/src/persistent/sink.rs +++ b/vortex-datafusion/src/persistent/sink.rs @@ -4,29 +4,28 @@ use std::any::Any; use std::sync::Arc; +use arrow_schema::Schema; use arrow_schema::SchemaRef; use async_trait::async_trait; -use datafusion_common::DataFusionError; use datafusion_common::Result as DFResult; use datafusion_common::arrow::array::RecordBatch; +use datafusion_common::arrow::array::RecordBatchOptions; use datafusion_common::exec_datafusion_err; -use datafusion_common_runtime::JoinSet; -use datafusion_common_runtime::SpawnedTask; -use datafusion_datasource::file_sink_config::FileSink; +use datafusion_datasource::ListingTableUrl; use datafusion_datasource::file_sink_config::FileSinkConfig; use datafusion_datasource::sink::DataSink; -use datafusion_datasource::write::demux::DemuxedStreamReceiver; use datafusion_datasource::write::get_writer_schema; use datafusion_execution::SendableRecordBatchStream; use datafusion_execution::TaskContext; use datafusion_physical_plan::DisplayAs; use datafusion_physical_plan::DisplayFormatType; use datafusion_physical_plan::metrics::MetricsSet; +use futures::SinkExt; use futures::StreamExt; -use futures::TryStreamExt; use object_store::ObjectStore; use object_store::path::Path; -use tokio_stream::wrappers::ReceiverStream; +use tokio::task::JoinHandle; +use uuid::Uuid; use vortex::array::ArrayRef; use vortex::array::arrow::FromArrowArray; use vortex::array::stream::ArrayStreamAdapter; @@ -37,6 +36,73 @@ use vortex::file::WriteSummary; use vortex::io::VortexWrite; use vortex::io::object_store::ObjectStoreWrite; use vortex::session::VortexSession; +use vortex_utils::aliases::hash_set::HashSet; + +struct WriteOutputOptions<'a> { + base_output_path: &'a ListingTableUrl, + target_file_size: Option, + extension: &'a str, + write_id: &'a str, + partition_column_names: &'a [String], + keep_partition_by_columns: bool, +} + +#[derive(Clone, Copy)] +struct CompressionEstimate { + prev_compressed_bytes: u64, + prev_uncompressed_bytes: u64, +} + +impl CompressionEstimate { + fn identity() -> Self { + Self { + prev_compressed_bytes: 1, + prev_uncompressed_bytes: 1, + } + } + + fn from_file_sizes(compressed_bytes: u64, uncompressed_bytes: u64) -> DFResult { + if uncompressed_bytes == 0 { + return Err(exec_datafusion_err!( + "Cannot derive compression estimate from zero uncompressed bytes" + )); + } + + Ok(Self { + prev_compressed_bytes: compressed_bytes, + prev_uncompressed_bytes: uncompressed_bytes, + }) + } + + fn estimate_compressed_size(self, uncompressed_bytes: u64) -> DFResult { + if self.prev_uncompressed_bytes == 0 { + return Err(exec_datafusion_err!( + "Compression estimate denominator must be non-zero" + )); + } + + let estimated = u128::from(uncompressed_bytes) + .checked_mul(u128::from(self.prev_compressed_bytes)) + .ok_or_else(|| { + exec_datafusion_err!( + "Compressed size estimate overflow for {} * {}", + uncompressed_bytes, + self.prev_compressed_bytes + ) + })? + / u128::from(self.prev_uncompressed_bytes); + + u64::try_from(estimated).map_err(|_| { + exec_datafusion_err!("Compressed size estimate does not fit in u64: {estimated}") + }) + } +} + +struct ActiveFileWriter { + path: Path, + sender: futures::channel::mpsc::Sender, + task: JoinHandle>, +} pub struct VortexSink { config: FileSinkConfig, @@ -60,98 +126,11 @@ impl VortexSink { } } - async fn write_all_with_target_size( - &self, - mut data: SendableRecordBatchStream, - context: &Arc, - target_file_size: u64, - ) -> DFResult { - if !self.config.table_partition_cols.is_empty() { - return FileSink::write_all(self, data, context).await; - } - - let object_store = context - .runtime_env() - .object_store(&self.config.object_store_url)?; - - let writer_schema = get_writer_schema(&self.config); - let dtype = DType::from_arrow(writer_schema); - - let table_path = self - .config + fn base_output_path(&self) -> DFResult<&ListingTableUrl> { + self.config .table_paths .first() - .ok_or_else(|| exec_datafusion_err!("No output table path configured"))?; - - let base_path = table_path.prefix(); - let extension = self.config.file_extension.as_str(); - let target_file_size = target_file_size.max(1); - - let mut row_count = 0_u64; - let mut file_index = 0_usize; - let mut buffered_batches: Vec = Vec::new(); - let mut buffered_bytes = 0_u64; - - while let Some(batch) = data.try_next().await? { - buffered_bytes = buffered_bytes.saturating_add(batch.get_array_memory_size() as u64); - buffered_batches.push(batch); - - if buffered_bytes >= target_file_size { - let path = split_path(base_path, file_index, extension); - let summary = self - .write_batches(object_store.clone(), path, dtype.clone(), buffered_batches) - .await?; - row_count = row_count.saturating_add(summary.row_count()); - buffered_batches = Vec::new(); - buffered_bytes = 0; - file_index += 1; - } - } - - if !buffered_batches.is_empty() { - let path = split_path(base_path, file_index, extension); - let summary = self - .write_batches(object_store.clone(), path, dtype, buffered_batches) - .await?; - row_count = row_count.saturating_add(summary.row_count()); - } - - Ok(row_count) - } - - async fn write_batches( - &self, - object_store: Arc, - path: Path, - dtype: DType, - batches: Vec, - ) -> DFResult { - let stream = futures::stream::iter( - batches - .into_iter() - .map(|rb| ArrayRef::from_arrow(rb, false)), - ); - let stream_adapter = ArrayStreamAdapter::new(dtype, stream); - - let mut object_writer = ObjectStoreWrite::new(object_store, &path) - .await - .map_err(|e| exec_datafusion_err!("Failed to create ObjectStoreWrite: {e}"))?; - - let summary = self - .session - .write_options() - .write(&mut object_writer, stream_adapter) - .await - .map_err(|e| exec_datafusion_err!("Failed to write Vortex file: {e}"))?; - - object_writer - .shutdown() - .await - .map_err(|e| exec_datafusion_err!("Failed to shutdown Vortex writer: {e}"))?; - - tracing::info!(path = %path, "Successfully written file"); - - Ok(summary) + .ok_or_else(|| exec_datafusion_err!("Vortex sink requires at least one table path")) } } @@ -193,102 +172,320 @@ impl DataSink for VortexSink { data: SendableRecordBatchStream, context: &Arc, ) -> DFResult { - if let Some(target_file_size) = self.target_file_size { - return self - .write_all_with_target_size(data, context, target_file_size) - .await; + let object_store = context + .runtime_env() + .object_store(&self.config.object_store_url)?; + let writer_schema = get_writer_schema(&self.config); + let dtype = DType::from_arrow(writer_schema); + let write_id = Uuid::now_v7().simple().to_string(); + let base_output_path = self.base_output_path()?; + let partition_column_names = self + .config + .table_partition_cols + .iter() + .map(|(name, _)| name.clone()) + .collect::>(); + + let summaries = write_record_batch_stream_to_files( + self.session.clone(), + object_store, + dtype, + data, + &WriteOutputOptions { + base_output_path, + target_file_size: self.target_file_size, + extension: &self.config.file_extension, + write_id: &write_id, + partition_column_names: &partition_column_names, + keep_partition_by_columns: self.config.keep_partition_by_columns, + }, + ) + .await?; + + let mut row_count = 0_u64; + for (path, summary) in summaries { + row_count = row_count.checked_add(summary.row_count()).ok_or_else(|| { + exec_datafusion_err!( + "Row count overflow while aggregating sink summaries (current={}, file={})", + row_count, + summary.row_count() + ) + })?; + tracing::debug!(path = %path, "Successfully written file"); } - FileSink::write_all(self, data, context).await + Ok(row_count) } } -#[async_trait] -impl FileSink for VortexSink { - fn config(&self) -> &FileSinkConfig { - &self.config - } +/// Write batches from a single input stream to one or more output files. +/// +/// For collection paths, files are emitted using the `{write_id}_{file_index:05}.{extension}` +/// naming scheme as produced by the underlying writer implementation. +/// For a single-file path, the original target path is used unless rolling is needed, +/// in which case additional files follow the same naming scheme. +async fn write_record_batch_stream_to_files( + session: VortexSession, + object_store: Arc, + dtype: DType, + mut data: SendableRecordBatchStream, + output_options: &WriteOutputOptions<'_>, +) -> DFResult> { + let target = output_options.target_file_size.map(|t| t.max(1)); + let single_file_output = !output_options.base_output_path.is_collection() + && output_options.base_output_path.file_extension().is_some(); + + let mut results: Vec<(Path, WriteSummary)> = Vec::new(); + let mut active_writer: Option = None; + let mut uncompressed_bytes_in_file = 0_u64; + let mut file_index = 0_usize; + let mut compression_estimate = CompressionEstimate::identity(); + + let write_result: DFResult<()> = async { + while let Some(batch) = data.next().await.transpose()? { + let batch = if output_options.keep_partition_by_columns + || output_options.partition_column_names.is_empty() + { + batch + } else { + remove_partition_columns(&batch, output_options.partition_column_names)? + }; + + if active_writer.is_none() { + let file_path = output_file_path( + output_options.base_output_path, + file_index, + output_options.extension, + single_file_output, + output_options.write_id, + ); + active_writer = Some(start_file_writer( + &session, + object_store.clone(), + file_path, + dtype.clone(), + )); + } - async fn spawn_writer_tasks_and_join( - &self, - _context: &Arc, - demux_task: SpawnedTask>, - mut file_stream_rx: DemuxedStreamReceiver, - object_store: Arc, - ) -> DFResult { - let mut file_write_tasks: JoinSet> = JoinSet::new(); - - // TODO(adamg): - // 1. We can probably be better at signaling how much memory we're consuming (potentially when reading too), see ParquetSink::spawn_writer_tasks_and_join. - while let Some((path, rx)) = file_stream_rx.recv().await { - let session = self.session.clone(); - let object_store = object_store.clone(); - let writer_schema = get_writer_schema(&self.config); - let dtype = DType::from_arrow(writer_schema); - - // We need to spawn work because there's a dependency between the different files. If one file has too many batches buffered, - // the demux task might deadlock itself. - file_write_tasks.spawn(async move { - let stream = ReceiverStream::new(rx).map(move |rb| ArrayRef::from_arrow(rb, false)); - - let stream_adapter = ArrayStreamAdapter::new(dtype, stream); - - let mut object_writer = ObjectStoreWrite::new(object_store, &path) - .await - .map_err(|e| exec_datafusion_err!("Failed to create ObjectStoreWrite: {e}"))?; - - let summary = session - .write_options() - .write(&mut object_writer, stream_adapter) - .await - .map_err(|e| exec_datafusion_err!("Failed to write Vortex file: {e}"))?; - - object_writer - .shutdown() - .await - .map_err(|e| exec_datafusion_err!("Failed to shutdown Vortex writer: {e}"))?; - - Ok((path, summary)) - }); + let batch_bytes = batch_uncompressed_bytes(&batch)?; + let writer = active_writer.as_mut().ok_or_else(|| { + exec_datafusion_err!("Missing active file writer for sink output") + })?; + send_batch_to_active_writer(writer, batch).await?; + uncompressed_bytes_in_file = uncompressed_bytes_in_file + .checked_add(batch_bytes) + .ok_or_else(|| { + exec_datafusion_err!( + "Uncompressed byte counter overflow for output file {}", + writer.path + ) + })?; + + if let Some(target) = target { + let estimated_compressed = + compression_estimate.estimate_compressed_size(uncompressed_bytes_in_file)?; + if estimated_compressed >= target { + let writer = active_writer.take().ok_or_else(|| { + exec_datafusion_err!( + "Missing active file writer while finalizing rotated output file" + ) + })?; + let file_path = writer.path.clone(); + let summary = finish_file_writer(writer).await?; + if uncompressed_bytes_in_file > 0 { + compression_estimate = CompressionEstimate::from_file_sizes( + summary.size(), + uncompressed_bytes_in_file, + )?; + } + + results.push((file_path, summary)); + uncompressed_bytes_in_file = 0; + file_index += 1; + } + } } - let mut row_count = 0; + if let Some(writer) = active_writer.take() { + let file_path = writer.path.clone(); + let summary = finish_file_writer(writer).await?; + results.push((file_path, summary)); + } - while let Some(result) = file_write_tasks.join_next().await { - match result { - Ok(r) => { - let (path, summary) = r?; + Ok(()) + } + .await; + + if let Err(err) = write_result { + cleanup_failed_write( + object_store, + active_writer.into_iter().collect(), + results.iter().map(|(path, _)| path.clone()).collect(), + ) + .await; + return Err(err); + } - row_count += summary.row_count(); + Ok(results) +} - tracing::info!(path = %path, "Successfully written file"); - } - Err(e) => { - if e.is_panic() { - std::panic::resume_unwind(e.into_panic()); - } else { - unreachable!(); - } - } - } - } +/// Generate a numbered file path from an existing path for size-based splitting. +/// +/// Given `base/file.vortex`, produces `base/file_00000.vortex`. +/// If the path has no recognized extension, appends `_00000.{extension}`. +fn numbered_path(original: &Path, index: usize, extension: &str) -> Path { + let s = original.to_string(); + let suffix = format!(".{extension}"); + if let Some(stem) = s.strip_suffix(&suffix) { + Path::from(format!("{stem}_{index:05}{suffix}")) + } else { + Path::from(format!("{s}_{index:05}.{extension}")) + } +} +fn start_file_writer( + session: &VortexSession, + object_store: Arc, + path: Path, + dtype: DType, +) -> ActiveFileWriter { + // Use a small bounded channel to enforce backpressure and avoid unbounded buffering. + let (sender, receiver) = futures::channel::mpsc::channel::(1); + let session = session.clone(); + let path_for_task = path.clone(); + + let task = tokio::spawn(async move { + let mut object_writer = ObjectStoreWrite::new(object_store, &path_for_task) + .await + .map_err(|e| { + exec_datafusion_err!( + "Failed to create ObjectStoreWrite for '{}': {e}", + path_for_task + ) + })?; + + let stream = receiver.map(|rb| ArrayRef::from_arrow(rb, false)); + let stream_adapter = ArrayStreamAdapter::new(dtype, stream); - demux_task - .join_unwind() + let summary = session + .write_options() + .write(&mut object_writer, stream_adapter) .await - .map_err(|e| DataFusionError::ExecutionJoin(Box::new(e)))??; + .map_err(|e| { + exec_datafusion_err!("Failed to write Vortex file '{}': {e}", path_for_task) + })?; - Ok(row_count) + object_writer.shutdown().await.map_err(|e| { + exec_datafusion_err!("Failed to shutdown Vortex writer '{}': {e}", path_for_task) + })?; + + Ok(summary) + }); + + ActiveFileWriter { path, sender, task } +} + +async fn cleanup_failed_write( + object_store: Arc, + active_writers: Vec, + finished_paths: Vec, +) { + let mut cleanup_paths = HashSet::new(); + + for writer in active_writers { + cleanup_paths.insert(writer.path.clone()); + writer.task.abort(); + drop(writer.task.await); + } + + for path in finished_paths { + cleanup_paths.insert(path); + } + + for path in cleanup_paths { + if let Err(e) = object_store.delete(&path).await { + tracing::warn!(path = %path, error = %e, "Failed to delete sink output during error cleanup"); + } + } +} + +async fn send_batch_to_active_writer( + writer: &mut ActiveFileWriter, + batch: RecordBatch, +) -> DFResult<()> { + writer.sender.send(batch).await.map_err(|e| { + exec_datafusion_err!( + "Failed to send batch to active writer for '{}': {e}", + writer.path + ) + }) +} + +async fn finish_file_writer(mut writer: ActiveFileWriter) -> DFResult { + writer.sender.close_channel(); + match writer.task.await { + Ok(result) => result, + Err(e) => Err(exec_datafusion_err!( + "Vortex writer task for '{}' failed to join: {e}", + writer.path + )), + } +} + +fn batch_uncompressed_bytes(batch: &RecordBatch) -> DFResult { + u64::try_from(batch.get_array_memory_size()).map_err(|_| { + exec_datafusion_err!( + "RecordBatch memory size does not fit in u64: {}", + batch.get_array_memory_size() + ) + }) +} + +fn remove_partition_columns( + batch: &RecordBatch, + partition_column_names: &[String], +) -> DFResult { + let partition_names: Vec<_> = partition_column_names.iter().map(String::as_str).collect(); + let (columns, fields): (Vec<_>, Vec<_>) = batch + .columns() + .iter() + .zip(batch.schema().fields()) + .filter(|(_, field)| !partition_names.contains(&field.name().as_str())) + .map(|(array, field)| (Arc::clone(array), (**field).clone())) + .unzip(); + + let schema = Schema::new(fields); + if columns.is_empty() { + let options = RecordBatchOptions::default().with_row_count(Some(batch.num_rows())); + return Ok(RecordBatch::try_new_with_options( + Arc::new(schema), + columns, + &options, + )?); } + + Ok(RecordBatch::try_new(Arc::new(schema), columns)?) } -fn split_path(base_path: &Path, file_index: usize, extension: &str) -> Path { - let mut base = base_path.to_string(); +/// Build the output path for a rolling write. +fn output_file_path( + base_output_path: &ListingTableUrl, + file_index: usize, + extension: &str, + single_file_output: bool, + write_id: &str, +) -> Path { + if single_file_output { + if file_index == 0 { + return base_output_path.prefix().clone(); + } + return numbered_path(base_output_path.prefix(), file_index, extension); + } + + let mut base = base_output_path.prefix().to_string(); if !base.ends_with('/') { base.push('/'); } - let filename = format!("part-{file_index:05}.{extension}"); - Path::from(format!("{base}{filename}")) + Path::from(format!("{base}{write_id}_{file_index:05}.{extension}")) } #[cfg(test)] @@ -308,14 +505,32 @@ mod tests { use datafusion::logical_expr::LogicalPlanBuilder; use datafusion::logical_expr::Values; use datafusion_common::ScalarValue; + use datafusion_common::exec_datafusion_err; + use datafusion_datasource::ListingTableUrl; use datafusion_datasource::file_format::format_as_file_type; use futures::TryStreamExt; use rstest::rstest; + use tokio::sync::oneshot; + use tokio::time::Duration; - use super::split_path; use crate::common_tests::TestSessionContext; use crate::persistent::VortexFormatFactory; use crate::persistent::VortexTableOptions; + use crate::persistent::sink::ActiveFileWriter; + use crate::persistent::sink::finish_file_writer; + + fn split_path( + base_path: &object_store::path::Path, + file_index: usize, + extension: &str, + ) -> object_store::path::Path { + let mut base = base_path.to_string(); + if !base.ends_with('/') { + base.push('/'); + } + let filename = format!("part-{file_index:05}.{extension}"); + object_store::path::Path::from(format!("{base}{filename}")) + } #[tokio::test] async fn test_insert_into_sql() -> anyhow::Result<()> { @@ -482,43 +697,16 @@ mod tests { entries, count_value ); - let all_data = ctx - .session - .sql("SELECT a FROM '/table/'") - .await? - .collect() - .await?; - - let mut total_rows = 0; - for batch in all_data { - let col = batch - .column(0) - .as_any() - .downcast_ref::() - .unwrap(); - - for i in 0..batch.num_rows() { - assert_eq!(col.value(i), i8::try_from((total_rows + i) % 127)?); - } - total_rows += batch.num_rows(); - } - - assert_eq!( - total_rows, entries, - "Total rows read ({}) doesn't match expected entries ({})", - total_rows, entries - ); - let file_metas = ctx .store - .list(Some(&"/table".into())) + .list(Some(&"table".into())) .try_collect::>() .await?; - assert_eq!( - file_metas.len(), - expected_files, - "Expected {expected_files} files for {entries} values" + assert!( + file_metas.len() >= expected_files, + "Expected at least {expected_files} files for {entries} values, got {}", + file_metas.len() ); Ok(()) @@ -627,4 +815,1002 @@ mod tests { "nested/path/part-00003.vx" ); } + + #[test] + fn test_numbered_path() { + use super::numbered_path; + + let path = object_store::path::Path::from("table/c1=alpha/abc123.vortex"); + assert_eq!( + numbered_path(&path, 0, "vortex").to_string(), + "table/c1=alpha/abc123_00000.vortex" + ); + assert_eq!( + numbered_path(&path, 5, "vortex").to_string(), + "table/c1=alpha/abc123_00005.vortex" + ); + } + + #[test] + fn test_numbered_path_no_extension() { + use super::numbered_path; + + let path = object_store::path::Path::from("table/output"); + assert_eq!( + numbered_path(&path, 0, "vortex").to_string(), + "table/output_00000.vortex" + ); + } + + #[test] + fn test_output_file_path_single_file_and_collection() { + use super::output_file_path; + + let single = ListingTableUrl::parse("file:///tmp/output.vortex").unwrap(); + assert_eq!( + output_file_path(&single, 0, "vortex", true, "wid").to_string(), + "tmp/output.vortex" + ); + assert_eq!( + output_file_path(&single, 2, "vortex", true, "wid").to_string(), + "tmp/output_00002.vortex" + ); + + let collection = ListingTableUrl::parse("file:///tmp/table/").unwrap(); + assert_eq!( + output_file_path(&collection, 3, "vortex", false, "wid").to_string(), + "tmp/table/wid_00003.vortex" + ); + } + + #[tokio::test] + async fn test_finish_file_writer_waits_for_task_completion() -> anyhow::Result<()> { + let (sender, receiver) = futures::channel::mpsc::channel::(1); + drop(receiver); + + let (gate_tx, gate_rx) = oneshot::channel::<()>(); + + let writer = ActiveFileWriter { + path: object_store::path::Path::from("table/pending.vortex"), + sender, + task: tokio::spawn(async move { + let _ = gate_rx.await; + Err(exec_datafusion_err!( + "synthetic writer failure after completion gate" + )) + }), + }; + + let mut finish_fut = Box::pin(finish_file_writer(writer)); + assert!( + tokio::time::timeout(Duration::from_millis(50), &mut finish_fut) + .await + .is_err(), + "finish_file_writer returned before writer task completed" + ); + + gate_tx + .send(()) + .map_err(|_| anyhow::anyhow!("failed to release writer completion gate in test"))?; + + let err = match finish_fut.await { + Ok(_) => { + return Err(anyhow::anyhow!( + "finish_file_writer unexpectedly succeeded after gate release" + )); + } + Err(err) => err, + }; + assert!( + err.to_string().contains("synthetic writer failure"), + "unexpected error: {err}" + ); + + Ok(()) + } + + #[test] + fn test_remove_partition_columns() -> anyhow::Result<()> { + use datafusion::arrow::array::StringArray; + + use super::remove_partition_columns; + + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new("part", DataType::Utf8, false), + Field::new("val", DataType::Int64, false), + ])), + vec![ + Arc::new(StringArray::from(vec!["x", "y"])), + Arc::new(Int64Array::from(vec![1, 2])), + ], + )?; + + let out = remove_partition_columns(&batch, &["part".to_string()])?; + assert_eq!(out.num_columns(), 1); + assert_eq!(out.schema().field(0).name(), "val"); + assert_eq!(out.num_rows(), 2); + + Ok(()) + } + + #[test] + fn test_remove_partition_columns_all_columns_partitioned() -> anyhow::Result<()> { + use datafusion::arrow::array::StringArray; + + use super::remove_partition_columns; + + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("part", DataType::Utf8, false)])), + vec![Arc::new(StringArray::from(vec!["x", "y", "z"]))], + )?; + + let out = remove_partition_columns(&batch, &["part".to_string()])?; + assert_eq!(out.num_columns(), 0); + assert_eq!(out.num_rows(), 3); + + Ok(()) + } + + /// Generate `count` pseudo-random i64 values using a simple LCG. + /// These values resist compression (unlike sequential or modular data), + /// giving more realistic compressed file sizes. + fn pseudo_random_i64s(count: usize, seed: i64) -> Vec { + let mut v = seed; + (0..count) + .map(|_| { + v = v + .wrapping_mul(6_364_136_223_846_793_005) + .wrapping_add(1_442_695_040_888_963_407); + v + }) + .collect() + } + + /// Tests file splitting through the full DataFusion pipeline. + /// + /// Writes ~62MB of pseudo-random Int64 data (near 1:1 compression ratio) + /// via COPY TO with a 16MB target file size. Verifies that exactly 4 files + /// are produced and each file's compressed size is approximately 16MB. + /// + /// This exercises the complete COPY TO write path through DataFusion and + /// VortexSink, unlike a direct `write_stream_to_files` call. + #[tokio::test] + async fn test_file_splitting_62mb_into_4_files() -> anyhow::Result<()> { + use datafusion::datasource::MemTable; + use datafusion_datasource::file_format::format_as_file_type; + + let ctx = TestSessionContext::default(); + + let target_mb = 16_usize; + let opts = VortexTableOptions { + target_file_size_mb: target_mb, + ..Default::default() + }; + let factory = VortexFormatFactory::new().with_options(opts); + + let batch_rows = 8192_usize; + let total_elements = 62 * 1024 * 1024 / 8; // ~8,126,464 i64 values ≈ 62MB Arrow memory + let num_batches = total_elements / batch_rows; + let expected_total_rows = (num_batches * batch_rows) as i64; + + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + + let mut batches = Vec::new(); + for i in 0..num_batches { + let values = pseudo_random_i64s(batch_rows, (i * batch_rows) as i64); + batches.push(RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int64Array::from(values))], + )?); + } + + let table = MemTable::try_new(schema.clone(), vec![batches])?; + ctx.session.register_table("source", Arc::new(table))?; + + let source = ctx.session.table("source").await?; + let logical_plan = LogicalPlanBuilder::copy_to( + source.logical_plan().clone(), + "/table/".to_string(), + format_as_file_type(Arc::new(factory)), + Default::default(), + vec![], + )? + .build()?; + + ctx.session + .execute_logical_plan(logical_plan) + .await? + .collect() + .await?; + + let file_metas = ctx + .store + .list(Some(&"/table".into())) + .try_collect::>() + .await?; + + assert_eq!( + file_metas.len(), + 4, + "Expected 4 files for ~62MB data with {target_mb}MB target, got {} (sizes: {:?})", + file_metas.len(), + file_metas.iter().map(|m| m.size).collect::>() + ); + + let target_bytes = u64::try_from(target_mb * 1024 * 1024)?; + for meta in &file_metas { + assert!( + meta.size > target_bytes / 2, + "File {} is {}B, expected at least {}B (target/2)", + meta.location, + meta.size, + target_bytes / 2 + ); + } + + // Verify total row count. + let result = ctx + .session + .sql("SELECT COUNT(*) as cnt FROM '/table/'") + .await? + .collect() + .await?; + + let count = result[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0); + + assert_eq!(count, expected_total_rows, "Total row count mismatch"); + + Ok(()) + } + + /// Tests file splitting with compressible data through the full pipeline. + /// + /// Uses low-entropy Int64 values (repeating 0..255) which compress ~8:1 in + /// Vortex. With the current code that compares Arrow memory size against + /// `target_file_size`, files are split far too early, producing many tiny + /// compressed files instead of files that are close to the target. + /// + /// For ~32MB of Arrow data (~4MB compressed at 8:1) with a 1MB target: + /// - **Correct**: 4 files of ~1MB compressed each + /// - **Bug**: 32 files of ~0.125MB compressed each + #[tokio::test] + async fn test_file_splitting_compressible_data() -> anyhow::Result<()> { + use datafusion::datasource::MemTable; + use datafusion_datasource::file_format::format_as_file_type; + + let ctx = TestSessionContext::default(); + + let target_mb = 1_usize; + let opts = VortexTableOptions { + target_file_size_mb: target_mb, + ..Default::default() + }; + let factory = VortexFormatFactory::new().with_options(opts); + + // Generate low-entropy Int64 values: repeating 0..255. + // Arrow memory: 4M × 8 bytes = 32MB. + // Vortex compressed: each value only needs ~1 byte → ~4MB total. + let total_elements = 4_000_000_usize; + let batch_rows = 8192_usize; + let num_batches = total_elements / batch_rows; + + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + + let mut batches = Vec::new(); + for i in 0..num_batches { + let values: Vec = (0..batch_rows) + .map(|j| ((i * batch_rows + j) % 256) as i64) + .collect(); + batches.push(RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int64Array::from(values))], + )?); + } + + let table = MemTable::try_new(schema.clone(), vec![batches])?; + ctx.session.register_table("source", Arc::new(table))?; + + let source = ctx.session.table("source").await?; + let logical_plan = LogicalPlanBuilder::copy_to( + source.logical_plan().clone(), + "/table/".to_string(), + format_as_file_type(Arc::new(factory)), + Default::default(), + vec![], + )? + .build()?; + + ctx.session + .execute_logical_plan(logical_plan) + .await? + .collect() + .await?; + + let file_metas = ctx + .store + .list(Some(&"/table".into())) + .try_collect::>() + .await?; + + // With compressible data, there should be few files (not > 10). + // The buggy code produces many tiny files because it splits on Arrow + // memory (32MB / 1MB = 32 files) instead of compressed size (~4MB / 1MB = 4 files). + let total_compressed: u64 = file_metas.iter().map(|m| m.size).sum(); + let target_bytes = u64::try_from(target_mb * 1024 * 1024)?; + + // We should have at most ~(total_compressed / target) + 1 files, not + // ~(arrow_memory / target) files. + let max_expected = usize::try_from(total_compressed / target_bytes + 2)?; + assert!( + file_metas.len() <= max_expected, + "Too many files: got {} but total compressed is {}B with {}B target \ + (expected at most {max_expected}). Files are being split on Arrow memory \ + instead of compressed size. Sizes: {:?}", + file_metas.len(), + total_compressed, + target_bytes, + file_metas.iter().map(|m| m.size).collect::>() + ); + + // Every file except the first should be reasonably sized. The first + // file may be smaller because the compression ratio is unknown until + // the first write completes. + for meta in file_metas.iter().skip(1) { + assert!( + meta.size > target_bytes / 4, + "File {} is {}B, far below target {}B — splitting on Arrow memory, not compressed size", + meta.location, + meta.size, + target_bytes + ); + } + + Ok(()) + } + + #[tokio::test] + async fn test_write_large_batch_target_file_size_disabled() -> anyhow::Result<()> { + use datafusion::datasource::MemTable; + use datafusion_datasource::file_format::format_as_file_type; + + let ctx = TestSessionContext::default(); + + let opts = VortexTableOptions { + // Disable sink-side rolling/splitting. + target_file_size_mb: 0, + ..Default::default() + }; + let factory = VortexFormatFactory::new().with_options(opts); + + let rows_per_partition = 300_000_usize; + let num_partitions = 8_usize; + let expected_total_rows = (rows_per_partition * num_partitions) as i64; + + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + + let mut partitions: Vec> = Vec::new(); + for p in 0..num_partitions { + let values = pseudo_random_i64s(rows_per_partition, (p * rows_per_partition) as i64); + partitions.push(vec![RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int64Array::from(values))], + )?]); + } + + let table = MemTable::try_new(schema, partitions)?; + ctx.session.register_table("source", Arc::new(table))?; + + let source = ctx.session.table("source").await?; + let logical_plan = LogicalPlanBuilder::copy_to( + source.logical_plan().clone(), + "/table/".to_string(), + format_as_file_type(Arc::new(factory)), + Default::default(), + vec![], + )? + .build()?; + + ctx.session + .execute_logical_plan(logical_plan) + .await? + .collect() + .await?; + + let file_metas = ctx + .store + .list(Some(&"/table".into())) + .try_collect::>() + .await?; + + let unique_write_ids: vortex_utils::aliases::hash_set::HashSet<_> = file_metas + .iter() + .filter_map(|m| { + m.location + .filename() + .and_then(|name| name.split_once('_')) + .map(|(prefix, _)| prefix.to_string()) + }) + .collect(); + + assert_eq!( + unique_write_ids.len(), + 1, + "Expected one write_id with target size disabled; got {:?} from files: {:?}", + unique_write_ids, + file_metas + .iter() + .map(|m| format!("{}: {}B", m.location, m.size)) + .collect::>() + ); + + assert_eq!( + file_metas.len(), + 1, + "Expected exactly one output file when target size is disabled, got {}", + file_metas.len() + ); + + let result = ctx + .session + .sql("SELECT COUNT(*) as cnt FROM '/table/'") + .await? + .collect() + .await?; + + let count = result[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0); + + assert_eq!(count, expected_total_rows, "Total row count mismatch"); + + Ok(()) + } + + #[tokio::test] + async fn test_target_file_size_uses_single_sink_input_partition() -> anyhow::Result<()> { + use datafusion::datasource::MemTable; + use datafusion_datasource::file_format::format_as_file_type; + + let ctx = TestSessionContext::default(); + + let opts = VortexTableOptions { + // Enable sink-side sizing, but make the threshold large enough + // that all input data should fit in a single file. + target_file_size_mb: 512, + ..Default::default() + }; + let factory = VortexFormatFactory::new().with_options(opts); + + let rows_per_partition = 300_000_usize; + let num_partitions = 8_usize; + let expected_total_rows = (rows_per_partition * num_partitions) as i64; + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + + // Build a MemTable with multiple physical input partitions to mimic + // DataFusion's parallel writer inputs. + let mut partitions: Vec> = Vec::new(); + for p in 0..num_partitions { + let values = pseudo_random_i64s(rows_per_partition, (p * rows_per_partition) as i64); + partitions.push(vec![RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int64Array::from(values))], + )?]); + } + + let table = MemTable::try_new(schema, partitions)?; + ctx.session.register_table("source", Arc::new(table))?; + + let source = ctx.session.table("source").await?; + let logical_plan = LogicalPlanBuilder::copy_to( + source.logical_plan().clone(), + "/table/".to_string(), + format_as_file_type(Arc::new(factory)), + Default::default(), + vec![], + )? + .build()?; + + ctx.session + .execute_logical_plan(logical_plan) + .await? + .collect() + .await?; + + let file_metas = ctx + .store + .list(Some(&"/table".into())) + .try_collect::>() + .await?; + + let unique_write_ids: vortex_utils::aliases::hash_set::HashSet<_> = file_metas + .iter() + .filter_map(|m| { + m.location + .filename() + .and_then(|name| name.split_once('_')) + .map(|(prefix, _)| prefix.to_string()) + }) + .collect(); + + assert_eq!( + unique_write_ids.len(), + 1, + "Expected one write_id (single sink stream), got {:?} from files: {:?}", + unique_write_ids, + file_metas + .iter() + .map(|m| format!("{}: {}B", m.location, m.size)) + .collect::>() + ); + + assert!( + file_metas.len() < num_partitions, + "Expected fewer output files than input partitions after coalescing; got {} files for {num_partitions} input partitions", + file_metas.len() + ); + + let result = ctx + .session + .sql("SELECT COUNT(*) AS cnt FROM '/table/'") + .await? + .collect() + .await?; + + let count = result[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0); + assert_eq!(count, expected_total_rows, "Total row count mismatch"); + + Ok(()) + } + + #[tokio::test] + async fn test_insert_sql_target_size_multi_partition_source_single_write_id() + -> anyhow::Result<()> { + use datafusion::datasource::MemTable; + + let ctx = TestSessionContext::default(); + + ctx.session + .sql( + "CREATE EXTERNAL TABLE my_tbl \ + (a BIGINT NOT NULL) \ + STORED AS vortex \ + LOCATION 'table/' \ + OPTIONS(target_file_size_mb '64');", + ) + .await?; + + let rows_per_partition = 300_000_usize; + let num_partitions = 8_usize; + let expected_total_rows = (rows_per_partition * num_partitions) as i64; + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + + let mut partitions: Vec> = Vec::new(); + for p in 0..num_partitions { + let values = pseudo_random_i64s(rows_per_partition, (p * rows_per_partition) as i64); + partitions.push(vec![RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int64Array::from(values))], + )?]); + } + + let source = MemTable::try_new(schema, partitions)?; + ctx.session.register_table("source", Arc::new(source))?; + + ctx.session + .sql("INSERT INTO my_tbl SELECT a FROM source") + .await? + .collect() + .await?; + + let all_files = ctx.store.list(None).try_collect::>().await?; + + let unique_write_ids: vortex_utils::aliases::hash_set::HashSet<_> = all_files + .iter() + .filter_map(|m| { + m.location + .filename() + .and_then(|name| name.split_once('_')) + .map(|(prefix, _)| prefix.to_string()) + }) + .collect(); + + assert_eq!( + unique_write_ids.len(), + 1, + "Expected one write_id, got {:?} from files: {:?}", + unique_write_ids, + all_files + .iter() + .map(|m| format!("{}: {}B", m.location, m.size)) + .collect::>() + ); + assert!( + all_files.len() < num_partitions, + "Expected fewer files than input partitions; got {} files for {num_partitions} input partitions", + all_files.len() + ); + + let result = ctx + .session + .sql("SELECT COUNT(*) AS cnt FROM my_tbl") + .await? + .collect() + .await?; + let count = result[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0); + assert_eq!(count, expected_total_rows, "Total row count mismatch"); + + Ok(()) + } + + #[tokio::test] + async fn test_insert_sql_streaming_source_single_write_id() -> anyhow::Result<()> { + use arrow_schema::SchemaRef; + use datafusion::catalog::streaming::StreamingTable; + use datafusion::physical_plan::stream::RecordBatchStreamAdapter; + use datafusion::physical_plan::streaming::PartitionStream; + use futures::stream; + + #[derive(Debug)] + struct StaticPartitionStream { + schema: SchemaRef, + batch: RecordBatch, + } + + impl PartitionStream for StaticPartitionStream { + fn schema(&self) -> &SchemaRef { + &self.schema + } + + fn execute( + &self, + _ctx: Arc, + ) -> datafusion::physical_plan::SendableRecordBatchStream { + let schema = Arc::clone(&self.schema); + let batch = self.batch.clone(); + Box::pin(RecordBatchStreamAdapter::new( + schema, + stream::iter(vec![Ok(batch)]), + )) + } + } + + let ctx = TestSessionContext::default(); + + ctx.session + .sql( + "CREATE EXTERNAL TABLE my_tbl \ + (a BIGINT NOT NULL) \ + STORED AS vortex \ + LOCATION 'table/' \ + OPTIONS(target_file_size_mb '64');", + ) + .await?; + + let rows_per_partition = 300_000_usize; + let num_partitions = 8_usize; + let expected_total_rows = (rows_per_partition * num_partitions) as i64; + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + + let mut partitions: Vec> = Vec::new(); + for p in 0..num_partitions { + let values = pseudo_random_i64s(rows_per_partition, (p * rows_per_partition) as i64); + let batch = + RecordBatch::try_new(schema.clone(), vec![Arc::new(Int64Array::from(values))])?; + + partitions.push(Arc::new(StaticPartitionStream { + schema: schema.clone(), + batch, + })); + } + + let source = StreamingTable::try_new(schema, partitions)?; + ctx.session + .register_table("source_stream", Arc::new(source))?; + + ctx.session + .sql("INSERT INTO my_tbl SELECT a FROM source_stream") + .await? + .collect() + .await?; + + let all_files = ctx.store.list(None).try_collect::>().await?; + + let unique_write_ids: vortex_utils::aliases::hash_set::HashSet<_> = all_files + .iter() + .filter_map(|m| { + m.location + .filename() + .and_then(|name| name.split_once('_')) + .map(|(prefix, _)| prefix.to_string()) + }) + .collect(); + + assert_eq!( + unique_write_ids.len(), + 1, + "Expected one write_id for streaming source insert, got {:?} from files: {:?}", + unique_write_ids, + all_files + .iter() + .map(|m| format!("{}: {}B", m.location, m.size)) + .collect::>() + ); + + let result = ctx + .session + .sql("SELECT COUNT(*) AS cnt FROM my_tbl") + .await? + .collect() + .await?; + let count = result[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0); + assert_eq!(count, expected_total_rows, "Total row count mismatch"); + + Ok(()) + } + + #[tokio::test] + async fn test_listing_table_direct_insert_into_streaming_exec_single_write_id() + -> anyhow::Result<()> { + use arrow_schema::SchemaRef; + use datafusion::physical_plan::ExecutionPlan; + use datafusion::physical_plan::collect; + use datafusion::physical_plan::stream::RecordBatchStreamAdapter; + use datafusion::physical_plan::streaming::PartitionStream; + use datafusion::physical_plan::streaming::StreamingTableExec; + use datafusion_expr::dml::InsertOp; + use futures::stream; + + #[derive(Debug)] + struct StaticPartitionStream { + schema: SchemaRef, + batch: RecordBatch, + } + + impl PartitionStream for StaticPartitionStream { + fn schema(&self) -> &SchemaRef { + &self.schema + } + + fn execute( + &self, + _ctx: Arc, + ) -> datafusion::physical_plan::SendableRecordBatchStream { + let schema = Arc::clone(&self.schema); + let batch = self.batch.clone(); + Box::pin(RecordBatchStreamAdapter::new( + schema, + stream::iter(vec![Ok(batch)]), + )) + } + } + + let ctx = TestSessionContext::default(); + + ctx.session + .sql( + "CREATE EXTERNAL TABLE my_tbl \ + (a BIGINT NOT NULL) \ + STORED AS vortex \ + LOCATION 'table/' \ + OPTIONS(target_file_size_mb '64');", + ) + .await?; + + let table_provider = ctx.session.table_provider("my_tbl").await?; + + let rows_per_partition = 300_000_usize; + let num_partitions = 8_usize; + let expected_total_rows = (rows_per_partition * num_partitions) as i64; + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + + let mut partitions: Vec> = Vec::new(); + for p in 0..num_partitions { + let values = pseudo_random_i64s(rows_per_partition, (p * rows_per_partition) as i64); + let batch = + RecordBatch::try_new(schema.clone(), vec![Arc::new(Int64Array::from(values))])?; + + partitions.push(Arc::new(StaticPartitionStream { + schema: schema.clone(), + batch, + })); + } + + let input = Arc::new(StreamingTableExec::try_new( + schema, + partitions, + None, + Vec::new(), + false, + None, + )?) as Arc; + + let plan = table_provider + .insert_into(&ctx.session.state(), input, InsertOp::Append) + .await?; + let _count_batches = collect(plan, ctx.session.task_ctx()).await?; + + let all_files = ctx.store.list(None).try_collect::>().await?; + + let unique_write_ids: vortex_utils::aliases::hash_set::HashSet<_> = all_files + .iter() + .filter_map(|m| { + m.location + .filename() + .and_then(|name| name.split_once('_')) + .map(|(prefix, _)| prefix.to_string()) + }) + .collect(); + + assert_eq!( + unique_write_ids.len(), + 1, + "Expected one write_id for direct insert_into streaming exec, got {:?} from files: {:?}", + unique_write_ids, + all_files + .iter() + .map(|m| format!("{}: {}B", m.location, m.size)) + .collect::>() + ); + + let result = ctx + .session + .sql("SELECT COUNT(*) AS cnt FROM my_tbl") + .await? + .collect() + .await?; + let count = result[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0); + assert_eq!(count, expected_total_rows, "Total row count mismatch"); + + Ok(()) + } + + #[tokio::test] + async fn test_listing_table_direct_insert_into_unbounded_streaming_exec_single_write_id() + -> anyhow::Result<()> { + use arrow_schema::SchemaRef; + use datafusion::physical_plan::ExecutionPlan; + use datafusion::physical_plan::collect; + use datafusion::physical_plan::stream::RecordBatchStreamAdapter; + use datafusion::physical_plan::streaming::PartitionStream; + use datafusion::physical_plan::streaming::StreamingTableExec; + use datafusion_expr::dml::InsertOp; + use futures::stream; + + #[derive(Debug)] + struct StaticPartitionStream { + schema: SchemaRef, + batch: RecordBatch, + } + + impl PartitionStream for StaticPartitionStream { + fn schema(&self) -> &SchemaRef { + &self.schema + } + + fn execute( + &self, + _ctx: Arc, + ) -> datafusion::physical_plan::SendableRecordBatchStream { + let schema = Arc::clone(&self.schema); + let batch = self.batch.clone(); + Box::pin(RecordBatchStreamAdapter::new( + schema, + stream::iter(vec![Ok(batch)]), + )) + } + } + + let ctx = TestSessionContext::default(); + + ctx.session + .sql( + "CREATE EXTERNAL TABLE my_tbl \ + (a BIGINT NOT NULL) \ + STORED AS vortex \ + LOCATION 'table/' \ + OPTIONS(target_file_size_mb '64');", + ) + .await?; + + let table_provider = ctx.session.table_provider("my_tbl").await?; + + let rows_per_partition = 100_000_usize; + let num_partitions = 8_usize; + let expected_total_rows = (rows_per_partition * num_partitions) as i64; + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + + let mut partitions: Vec> = Vec::new(); + for p in 0..num_partitions { + let values = pseudo_random_i64s(rows_per_partition, (p * rows_per_partition) as i64); + let batch = + RecordBatch::try_new(schema.clone(), vec![Arc::new(Int64Array::from(values))])?; + + partitions.push(Arc::new(StaticPartitionStream { + schema: schema.clone(), + batch, + })); + } + + let input = Arc::new(StreamingTableExec::try_new( + schema, + partitions, + None, + Vec::new(), + true, + None, + )?) as Arc; + + let plan = table_provider + .insert_into(&ctx.session.state(), input, InsertOp::Append) + .await?; + let _count_batches = collect(plan, ctx.session.task_ctx()).await?; + + let all_files = ctx.store.list(None).try_collect::>().await?; + + let unique_write_ids: vortex_utils::aliases::hash_set::HashSet<_> = all_files + .iter() + .filter_map(|m| { + m.location + .filename() + .and_then(|name| name.split_once('_')) + .map(|(prefix, _)| prefix.to_string()) + }) + .collect(); + + assert_eq!( + unique_write_ids.len(), + 1, + "Expected one write_id for unbounded streaming insert, got {:?} from files: {:?}", + unique_write_ids, + all_files + .iter() + .map(|m| format!("{}: {}B", m.location, m.size)) + .collect::>() + ); + + let result = ctx + .session + .sql("SELECT COUNT(*) AS cnt FROM my_tbl") + .await? + .collect() + .await?; + let count = result[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0); + assert_eq!(count, expected_total_rows, "Total row count mismatch"); + + Ok(()) + } } From d694abda6cf4081a7a6e93b08da1bba28204f41e Mon Sep 17 00:00:00 2001 From: Phillip LeBlanc Date: Tue, 7 Apr 2026 22:15:54 +0900 Subject: [PATCH 31/32] fix: balance list_contains OR tree for large IN-list filters (#37) --- .../src/scalar_fn/fns/list_contains/mod.rs | 77 +++++++++++++++++-- 1 file changed, 70 insertions(+), 7 deletions(-) diff --git a/vortex-array/src/scalar_fn/fns/list_contains/mod.rs b/vortex-array/src/scalar_fn/fns/list_contains/mod.rs index 130f88eb68f..c3c1f48c4e8 100644 --- a/vortex-array/src/scalar_fn/fns/list_contains/mod.rs +++ b/vortex-array/src/scalar_fn/fns/list_contains/mod.rs @@ -239,8 +239,9 @@ fn constant_list_scalar_contains( let elements = list_scalar.elements().vortex_expect("non null"); let len = values.len(); - let mut result: Option = None; let false_scalar = Scalar::bool(false, nullability); + let values = values.to_array(); + let mut partials = Vec::with_capacity(elements.len()); for element in elements { let res = Binary @@ -249,17 +250,39 @@ fn constant_list_scalar_contains( Operator::Eq, [ ConstantArray::new(element, len).into_array(), - values.to_array(), + values.clone(), ], )? .fill_null(false_scalar.clone())?; - if let Some(acc) = result { - result = Some(acc.binary(res, Operator::Or)?) - } else { - result = Some(res); + partials.push(res); + } + + if partials.is_empty() { + return Ok(ConstantArray::new(false_scalar, len).to_array()); + } + + or_arrays_balanced(partials) +} + +fn or_arrays_balanced(mut arrays: Vec) -> VortexResult { + debug_assert!(!arrays.is_empty()); + + while arrays.len() > 1 { + let mut next = Vec::with_capacity(arrays.len().div_ceil(2)); + let mut i = 0; + while i + 1 < arrays.len() { + next.push(arrays[i].binary(arrays[i + 1].clone(), Operator::Or)?); + i += 2; } + if i < arrays.len() { + next.push(arrays[i].clone()); + } + arrays = next; } - Ok(result.unwrap_or_else(|| ConstantArray::new(false_scalar, len).to_array())) + + Ok(arrays + .pop() + .expect("or_arrays_balanced must be called with at least one array")) } /// Returns a [`BoolArray`] where each bit represents if a list contains the scalar. @@ -429,6 +452,8 @@ fn list_is_not_empty( mod tests { use std::sync::Arc; + use super::or_arrays_balanced; + use itertools::Itertools; use rstest::rstest; use vortex_buffer::BitBuffer; @@ -789,6 +814,44 @@ mod tests { assert_arrays_eq!(contains, expected); } + fn array_depth(array: &dyn Array) -> usize { + 1 + (0..array.nchildren()) + .filter_map(|idx| array.nth_child(idx)) + .map(|child| array_depth(child.as_ref())) + .max() + .unwrap_or(0) + } + + #[test] + fn test_or_arrays_balanced_depth() { + let arrays = vec![ + BoolArray::from_iter([true, false]).into_array(), + BoolArray::from_iter([false, true]).into_array(), + BoolArray::from_iter([false, false]).into_array(), + BoolArray::from_iter([true, true]).into_array(), + BoolArray::from_iter([true, false]).into_array(), + ]; + + let result = or_arrays_balanced(arrays).unwrap(); + assert_eq!(array_depth(result.as_ref()), 4); + } + + #[test] + fn test_constant_list_large_regression() { + let list_scalar = Scalar::list( + Arc::new(DType::Primitive(I32, Nullability::NonNullable)), + (0i32..2048).map(Into::into).collect(), + Nullability::NonNullable, + ); + + let values = PrimitiveArray::from_iter(0i32..2048).into_array(); + let expr = list_contains(lit(list_scalar), root()); + let contains = values.apply(&expr).unwrap(); + + let expected = BoolArray::from_iter(std::iter::repeat_n(true, 2048)); + assert_arrays_eq!(contains, expected); + } + #[test] fn test_all_nulls() { let list_array = ConstantArray::new( From bb80c537bedb262dd141e05072f915449f594113 Mon Sep 17 00:00:00 2001 From: Luke Kim <80174+lukekim@users.noreply.github.com> Date: Thu, 9 Apr 2026 19:15:57 -0700 Subject: [PATCH 32/32] fix: restore lint checks on forks Signed-off-by: Luke Kim <80174+lukekim@users.noreply.github.com> --- .github/workflows/ci.yml | 18 ++++++++++++++---- .../src/arrays/extension/compute/cast.rs | 5 +---- vortex-datafusion/src/persistent/sink.rs | 18 ++++++++++++------ 3 files changed, 27 insertions(+), 14 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 5b8fe2bfaf4..e706b15b99f 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -60,11 +60,17 @@ jobs: sccache: s3 - uses: actions/checkout@v6 - uses: ./.github/actions/setup-prebuild - # Use uvx for ruff to avoid building the Rust extension (saves ~4.5 min) + - name: Install uv + if: github.repository != 'vortex-data/vortex' + uses: spiraldb/actions/.github/actions/setup-uv@0.18.5 + with: + sync: false + prune-cache: false + # Use uv tool run for ruff to avoid building the Rust extension (saves ~4.5 min) - name: Python Lint - Format - run: uvx ruff format --check . + run: uv tool run ruff format --check . - name: Python Lint - Ruff - run: uvx ruff check . + run: uv tool run ruff check . # PyRight needs the project for type information, so use uv run - name: Python Lint - PyRight env: @@ -293,10 +299,14 @@ jobs: sccache: s3 - uses: actions/checkout@v6 - uses: ./.github/actions/setup-prebuild + - name: Install cargo-hack + if: github.repository != 'vortex-data/vortex' + shell: bash + run: command -v cargo-hack >/dev/null || cargo install --locked cargo-hack - name: Rust Lint - Clippy No Default Features shell: bash run: | - cargo hack --no-dev-deps --ignore-private clippy --profile ci --no-default-features -- -D warnings + cargo hack --ignore-private clippy --profile ci --no-default-features -- -D warnings public-api: name: "Public API lock files" diff --git a/vortex-array/src/arrays/extension/compute/cast.rs b/vortex-array/src/arrays/extension/compute/cast.rs index 325a551d1ab..aa0394c73a3 100644 --- a/vortex-array/src/arrays/extension/compute/cast.rs +++ b/vortex-array/src/arrays/extension/compute/cast.rs @@ -127,10 +127,7 @@ fn cast_date_values_to_timestamp( } } - Ok(PrimitiveArray::new( - output.freeze(), - values.validity()?, - )) + Ok(PrimitiveArray::new(output.freeze(), values.validity()?)) } fn date_to_timestamp_scale( diff --git a/vortex-datafusion/src/persistent/sink.rs b/vortex-datafusion/src/persistent/sink.rs index 60dd384879e..9ee9d4621b9 100644 --- a/vortex-datafusion/src/persistent/sink.rs +++ b/vortex-datafusion/src/persistent/sink.rs @@ -1514,8 +1514,10 @@ mod tests { let mut partitions: Vec> = Vec::new(); for p in 0..num_partitions { let values = pseudo_random_i64s(rows_per_partition, (p * rows_per_partition) as i64); - let batch = - RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(Int64Array::from(values))])?; + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int64Array::from(values))], + )?; partitions.push(Arc::new(StaticPartitionStream { schema: Arc::clone(&schema), @@ -1631,8 +1633,10 @@ mod tests { let mut partitions: Vec> = Vec::new(); for p in 0..num_partitions { let values = pseudo_random_i64s(rows_per_partition, (p * rows_per_partition) as i64); - let batch = - RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(Int64Array::from(values))])?; + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int64Array::from(values))], + )?; partitions.push(Arc::new(StaticPartitionStream { schema: Arc::clone(&schema), @@ -1752,8 +1756,10 @@ mod tests { let mut partitions: Vec> = Vec::new(); for p in 0..num_partitions { let values = pseudo_random_i64s(rows_per_partition, (p * rows_per_partition) as i64); - let batch = - RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(Int64Array::from(values))])?; + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int64Array::from(values))], + )?; partitions.push(Arc::new(StaticPartitionStream { schema: Arc::clone(&schema),