diff --git a/datafusion/common/src/hash_utils.rs b/datafusion/common/src/hash_utils.rs index cfe57999689b1..d10634e1d3745 100644 --- a/datafusion/common/src/hash_utils.rs +++ b/datafusion/common/src/hash_utils.rs @@ -19,15 +19,11 @@ use arrow::array::types::{IntervalDayTime, IntervalMonthDayNano}; use arrow::array::*; -#[cfg(not(feature = "force_hash_collisions"))] use arrow::compute::take; use arrow::datatypes::*; -#[cfg(not(feature = "force_hash_collisions"))] use arrow::{downcast_dictionary_array, downcast_primitive_array}; use foldhash::fast::FixedState; -#[cfg(not(feature = "force_hash_collisions"))] use itertools::Itertools; -#[cfg(not(feature = "force_hash_collisions"))] use std::collections::HashMap; use std::hash::{BuildHasher, Hash, Hasher}; @@ -80,7 +76,6 @@ impl HashState for foldhash::quality::FixedState { } } -#[cfg(not(feature = "force_hash_collisions"))] use crate::cast::{ as_binary_view_array, as_boolean_array, as_fixed_size_list_array, as_generic_binary_array, as_large_list_array, as_large_list_view_array, @@ -207,7 +202,6 @@ where build_hasher::with_hashes_with_hasher(arrays, hash_builder, callback) } -#[cfg(not(feature = "force_hash_collisions"))] fn hash_null( random_state: &S, hashes_buffer: &'_ mut [u64], @@ -275,7 +269,6 @@ macro_rules! hash_float_value { } hash_float_value!((half::f16, u16), (f32, u32), (f64, u64)); -#[cfg(not(feature = "force_hash_collisions"))] trait ChildHashing { fn create_hashes(&self, arrays: I, hashes_buffer: &mut [u64]) -> Result<()> where @@ -283,26 +276,23 @@ trait ChildHashing { T: AsDynArray; } -#[cfg(not(feature = "force_hash_collisions"))] struct HashStateChildHashing<'a, S> { hash_state: &'a S, } -#[cfg(not(feature = "force_hash_collisions"))] impl ChildHashing for HashStateChildHashing<'_, S> { fn create_hashes(&self, arrays: I, hashes_buffer: &mut [u64]) -> Result<()> where I: IntoIterator, T: AsDynArray, { - create_hashes(arrays, self.hash_state, hashes_buffer).map(|_| ()) + create_hashes_for_partitioning(arrays, self.hash_state, hashes_buffer).map(|_| ()) } } /// Builds hash values of PrimitiveArray and writes them into `hashes_buffer` /// If `rehash==true` this folds the existing hash into the hasher state /// and hashes only the new value (avoiding a separate combine step). -#[cfg(not(feature = "force_hash_collisions"))] fn hash_array_primitive( array: &PrimitiveArray, random_state: &impl HashState, @@ -347,7 +337,6 @@ fn hash_array_primitive( /// Hashes one array into the `hashes_buffer` /// If `rehash==true` this combines the previous hash value in the buffer /// with the new hash using `combine_hashes` -#[cfg(not(feature = "force_hash_collisions"))] fn hash_array( array: &T, random_state: &impl HashState, @@ -396,7 +385,6 @@ fn hash_array( /// HAS_NULLS: do we have to check null in the inner loop /// HAS_BUFFERS: if true, array has external buffers; if false, all strings are inlined/ less then 12 bytes /// REHASH: if true, combining with existing hash, otherwise initializing -#[cfg(not(feature = "force_hash_collisions"))] #[inline(never)] fn hash_string_view_array_inner< T: ByteViewType, @@ -457,7 +445,6 @@ fn hash_string_view_array_inner< /// Builds hash values for array views and writes them into `hashes_buffer` /// If `rehash==true` this combines the previous hash value in the buffer /// with the new hash using `combine_hashes` -#[cfg(not(feature = "force_hash_collisions"))] fn hash_generic_byte_view_array( array: &GenericByteViewArray, random_state: &impl HashState, @@ -523,7 +510,6 @@ fn hash_generic_byte_view_array( /// - `HAS_NULL_KEYS`: Whether to check for null dictionary keys /// - `HAS_NULL_VALUES`: Whether to check for null dictionary values /// - `MULTI_COL`: Whether to combine with existing hash (true) or initialize (false) -#[cfg(not(feature = "force_hash_collisions"))] #[inline(never)] fn hash_dictionary_scatter< K: ArrowDictionaryKeyType, @@ -563,7 +549,6 @@ fn hash_dictionary_scatter< } } -#[cfg(not(feature = "force_hash_collisions"))] fn dispatch_dictionary_scatter( array: &DictionaryArray, dict_hashes: &[u64], @@ -618,7 +603,6 @@ fn dispatch_dictionary_scatter( } /// Hash the values in a dictionary array. -#[cfg(not(feature = "force_hash_collisions"))] fn hash_dictionary( array: &DictionaryArray, random_state: &impl HashState, @@ -630,7 +614,7 @@ fn hash_dictionary( // redundant hashing for large dictionary elements (e.g. strings) let dict_values = array.values(); let mut dict_hashes = vec![0; dict_values.len()]; - create_hashes([dict_values], random_state, &mut dict_hashes)?; + create_hashes_for_partitioning([dict_values], random_state, &mut dict_hashes)?; dispatch_dictionary_scatter(array, &dict_hashes, hashes_buffer, multi_col); Ok(()) } @@ -649,7 +633,6 @@ fn hash_dictionary_with_child_hashing( Ok(()) } -#[cfg(not(feature = "force_hash_collisions"))] fn hash_struct_array( array: &StructArray, child_hashing: &impl ChildHashing, @@ -678,8 +661,6 @@ fn hash_struct_array( Ok(()) } -// only adding this `cfg` b/c this function is only used with this `cfg` -#[cfg(not(feature = "force_hash_collisions"))] fn hash_map_array( array: &MapArray, child_hashing: &impl ChildHashing, @@ -730,7 +711,6 @@ fn hash_map_array( Ok(()) } -#[cfg(not(feature = "force_hash_collisions"))] fn hash_list_array( array: &GenericListArray, child_hashing: &impl ChildHashing, @@ -780,7 +760,6 @@ where Ok(()) } -#[cfg(not(feature = "force_hash_collisions"))] fn hash_list_view_array( array: &GenericListViewArray, child_hashing: &impl ChildHashing, @@ -819,7 +798,6 @@ where Ok(()) } -#[cfg(not(feature = "force_hash_collisions"))] fn hash_union_array( array: &UnionArray, child_hashing: &impl ChildHashing, @@ -850,7 +828,6 @@ fn hash_union_array( /// For sparse unions with 3+ types, the optimized take/scatter approach in /// `hash_sparse_union_array` is more efficient, but for 1-2 types or dense unions, /// this simpler approach is preferred. -#[cfg(not(feature = "force_hash_collisions"))] fn hash_union_array_default( array: &UnionArray, union_fields: &UnionFields, @@ -891,7 +868,6 @@ fn hash_union_array_default( /// /// For 1-2 types, the overhead of take/scatter outweighs the benefit, so we use /// the default approach of hashing all children (same as dense unions). -#[cfg(not(feature = "force_hash_collisions"))] fn hash_sparse_union_array( array: &UnionArray, union_fields: &UnionFields, @@ -947,7 +923,6 @@ fn hash_sparse_union_array( Ok(()) } -#[cfg(not(feature = "force_hash_collisions"))] fn hash_fixed_list_array( array: &FixedSizeListArray, child_hashing: &impl ChildHashing, @@ -982,7 +957,6 @@ fn hash_fixed_list_array( /// Inner hash function for RunArray #[inline(never)] -#[cfg(not(feature = "force_hash_collisions"))] fn hash_run_array_inner< R: RunEndIndexType, C: ChildHashing + ?Sized, @@ -1051,7 +1025,6 @@ fn hash_run_array_inner< Ok(()) } -#[cfg(not(feature = "force_hash_collisions"))] fn hash_run_array( array: &RunArray, child_hashing: &impl ChildHashing, @@ -1080,8 +1053,7 @@ fn hash_run_array( /// Internal helper function that hashes a single array and either initializes or combines /// the hash values in the buffer. -#[cfg(not(feature = "force_hash_collisions"))] -fn hash_single_array( +fn hash_single_array_for_partitioning( array: &dyn Array, random_state: &impl HashState, hashes_buffer: &mut [u64], @@ -1181,6 +1153,16 @@ fn hash_single_array( Ok(()) } +#[cfg(not(feature = "force_hash_collisions"))] +fn hash_single_array( + array: &dyn Array, + random_state: &impl HashState, + hashes_buffer: &mut [u64], + rehash: bool, +) -> Result<()> { + hash_single_array_for_partitioning(array, random_state, hashes_buffer, rehash) +} + /// Test version of `hash_single_array` that forces all hashes to collide to zero. #[cfg(feature = "force_hash_collisions")] fn hash_single_array( @@ -1253,6 +1235,31 @@ where Ok(hashes_buffer) } +/// Creates hashes for partition routing even when collision-forcing tests are enabled. +/// +/// The `force_hash_collisions` feature intentionally collapses hashes used by hash +/// tables. Partition routing must remain independent so a grace hash join can split +/// the input into bounded partitions before exercising those colliding hash tables. +pub fn create_hashes_for_partitioning<'a, I, T>( + arrays: I, + random_state: &impl HashState, + hashes_buffer: &'a mut [u64], +) -> Result<&'a mut [u64]> +where + I: IntoIterator, + T: AsDynArray, +{ + for (i, array) in arrays.into_iter().enumerate() { + hash_single_array_for_partitioning( + array.as_dyn_array(), + random_state, + hashes_buffer, + i >= 1, + )?; + } + Ok(hashes_buffer) +} + /// Creates hash values for every row using a caller-provided hash builder. /// /// The number of rows to hash is determined by `hashes_buffer.len()`. @@ -1308,6 +1315,23 @@ mod tests { } } + #[cfg(feature = "force_hash_collisions")] + #[test] + fn partition_hashes_are_not_forced_to_collide() -> Result<()> { + let array: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 1])); + let random_state = RandomState::with_seed(0); + + let mut collision_hashes = vec![0; array.len()]; + create_hashes([&array], &random_state, &mut collision_hashes)?; + assert_eq!(collision_hashes, vec![0; array.len()]); + + let mut partition_hashes = vec![0; array.len()]; + create_hashes_for_partitioning([&array], &random_state, &mut partition_hashes)?; + assert_eq!(partition_hashes[0], partition_hashes[2]); + assert_ne!(partition_hashes[0], partition_hashes[1]); + Ok(()) + } + #[test] fn create_hashes_for_decimal_array() -> Result<()> { let array = vec![1, 2, 3, 4] diff --git a/datafusion/physical-expr/src/expressions/binary.rs b/datafusion/physical-expr/src/expressions/binary.rs index 909a601bc6c91..732b4d3cfdd55 100644 --- a/datafusion/physical-expr/src/expressions/binary.rs +++ b/datafusion/physical-expr/src/expressions/binary.rs @@ -4798,7 +4798,7 @@ mod tests { Field::new("b", DataType::Decimal128(10, 2), true), ])); let expect = Arc::new(create_decimal_array( - &[Some(1000000), None, Some(1008196), Some(1000000)], + &[Some(1000000), None, Some(1008197), Some(1000000)], 16, 4, )) as ArrayRef; diff --git a/datafusion/physical-plan/src/joins/cross_join.rs b/datafusion/physical-plan/src/joins/cross_join.rs index 8a477c1021d1b..b75e9f327b34c 100644 --- a/datafusion/physical-plan/src/joins/cross_join.rs +++ b/datafusion/physical-plan/src/joins/cross_join.rs @@ -452,6 +452,7 @@ impl ExecutionPlan for CrossJoinExec { let (new_left, new_right) = new_join_children( &projection_as_columns, + projection.schema().as_ref(), far_right_left_col_ind, far_left_right_col_ind, self.left(), diff --git a/datafusion/physical-plan/src/joins/hash_join/spill.rs b/datafusion/physical-plan/src/joins/hash_join/spill.rs index decbf6ddb43ff..1b0c159ac86bd 100644 --- a/datafusion/physical-plan/src/joins/hash_join/spill.rs +++ b/datafusion/physical-plan/src/joins/hash_join/spill.rs @@ -33,7 +33,7 @@ //! and per-partition write buffers; held until the spill join completes. //! - `HashJoinSpillPartition[p.k]`: per partition-pair; covers the loaded //! build batches plus the hash table built from them (moved into the -//! pair's [`JoinLeftData`]); dropped when the pair finishes. +//! pair's in-memory join state); dropped when the pair finishes. //! //! The scatter hash uses seeds distinct from both `RepartitionExec`'s //! `(0,0,0,0)` routing seeds and the join hash map's `HASH_JOIN_SEED` @@ -45,7 +45,7 @@ use std::sync::Arc; use std::sync::atomic::AtomicUsize; use std::task::{Context, Poll}; -use crate::hash_utils::create_hashes; +use crate::hash_utils::create_hashes_for_partitioning; use crate::joins::PartitionMode; use crate::joins::SharedBitmapBuilder; use crate::joins::hash_join::exec::{ @@ -429,7 +429,11 @@ impl SideScatter { let keys = evaluate_expressions_to_arrays(&self.on_exprs, batch)?; self.hashes_buffer.clear(); self.hashes_buffer.resize(num_rows, 0); - create_hashes(&keys, &self.random_state, &mut self.hashes_buffer)?; + create_hashes_for_partitioning( + &keys, + &self.random_state, + &mut self.hashes_buffer, + )?; let partition_count = self.writers.len() as u64; let mut indices: Vec> = vec![Vec::new(); self.writers.len()]; @@ -1711,6 +1715,7 @@ fn shared_build_loader( #[cfg(test)] mod tests { use super::*; + use crate::hash_utils::create_hashes; use crate::joins::HashJoinExec; use crate::joins::PartitionMode; use crate::metrics::SpillMetrics; diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs index b48905500d546..977e0d0a0814c 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -662,6 +662,7 @@ impl ExecutionPlan for SortMergeJoinExec { let (new_left, new_right) = new_join_children( &projection_as_columns, + projection.schema().as_ref(), far_right_left_col_ind, far_left_right_col_ind, self.children()[0], diff --git a/datafusion/physical-plan/src/projection.rs b/datafusion/physical-plan/src/projection.rs index cf362cdee55d3..11ccf4d1e419f 100644 --- a/datafusion/physical-plan/src/projection.rs +++ b/datafusion/physical-plan/src/projection.rs @@ -744,9 +744,10 @@ pub fn try_embed_projection( }); } // Old projection may contain some alias or expression such as `a + 1` and `CAST('true' AS BOOLEAN)`, but our projection_exprs in hash join just contain column, so we need to create the new projection to keep the original projection. - let new_projection = Arc::new(ProjectionExec::try_new( + let new_projection = Arc::new(ProjectionExec::try_new_with_schema_metadata( new_projection_exprs, Arc::clone(&new_execution_plan) as _, + projection.schema().as_ref(), )?); if is_projection_removable(&new_projection) { // Residual is identity — embedding fully absorbed the projection. @@ -872,8 +873,11 @@ pub fn try_pushdown_through_join_with_column_indices( } let mut left_proj: Vec<(Column, String)> = Vec::new(); let mut right_proj: Vec<(Column, String)> = Vec::new(); + let mut left_fields = Vec::new(); + let mut right_fields = Vec::new(); + let projection_schema = projection.schema(); let mut seen_right = false; - for (col, alias) in &projection_as_columns { + for (projection_index, (col, alias)) in projection_as_columns.iter().enumerate() { let Some(origin) = column_indices.get(col.index()) else { return plan_err!( "Projection column {} is outside the {}-entry column index mapping", @@ -889,10 +893,14 @@ pub fn try_pushdown_through_join_with_column_indices( return Ok(None); } left_proj.push((Column::new(col.name(), origin.index), alias.clone())); + left_fields + .push(Arc::clone(&projection_schema.fields()[projection_index])); } JoinSide::Right => { seen_right = true; right_proj.push((Column::new(col.name(), origin.index), alias.clone())); + right_fields + .push(Arc::clone(&projection_schema.fields()[projection_index])); } // Synthetic column (e.g. mark): belongs to neither child. // Phase 2 declines; Phase 3 keeps it at the join output instead. @@ -922,8 +930,16 @@ pub fn try_pushdown_through_join_with_column_indices( return Ok(None); }; - let (new_left, new_right) = - new_join_children_from_groups(&left_proj, &right_proj, join_left, join_right)?; + let left_schema = Schema::new(left_fields); + let right_schema = Schema::new(right_fields); + let (new_left, new_right) = new_join_children_from_groups( + &left_proj, + &right_proj, + &left_schema, + &right_schema, + join_left, + join_right, + )?; Ok(Some(JoinData { projected_left_child: new_left, @@ -968,6 +984,7 @@ fn is_projection_removable(projection: &ProjectionExec) -> bool { }; col.name() == proj_expr.alias && col.index() == idx }) && exprs.len() == projection.input().schema().fields().len() + && projection.schema() == projection.input().schema() } /// Given the expression set of a projection, checks if the projection causes @@ -1006,8 +1023,12 @@ pub fn make_with_child( projection: &ProjectionExec, child: &Arc, ) -> Result> { - ProjectionExec::try_new(projection.expr().to_vec(), Arc::clone(child)) - .map(|e| Arc::new(e) as _) + ProjectionExec::try_new_with_schema_metadata( + projection.expr().to_vec(), + Arc::clone(child), + projection.schema().as_ref(), + ) + .map(|e| Arc::new(e) as _) } /// Returns `true` if all the expressions in the argument are `Column`s. @@ -1072,12 +1093,16 @@ pub fn physical_to_column_exprs( /// of the original children of the join. pub fn new_join_children( projection_as_columns: &[(Column, String)], + projection_schema: &Schema, far_right_left_col_ind: i32, far_left_right_col_ind: i32, left_child: &Arc, right_child: &Arc, ) -> Result<(ProjectionExec, ProjectionExec)> { - let new_left = ProjectionExec::try_new( + let left_schema = Schema::new( + projection_schema.fields()[0..=far_right_left_col_ind as usize].to_vec(), + ); + let new_left = ProjectionExec::try_new_with_schema_metadata( projection_as_columns[0..=far_right_left_col_ind as _] .iter() .map(|(col, alias)| ProjectionExpr { @@ -1085,9 +1110,13 @@ pub fn new_join_children( alias: alias.clone(), }), Arc::clone(left_child), + &left_schema, )?; let left_size = left_child.schema().fields().len() as i32; - let new_right = ProjectionExec::try_new( + let right_schema = Schema::new( + projection_schema.fields()[far_left_right_col_ind as usize..].to_vec(), + ); + let new_right = ProjectionExec::try_new_with_schema_metadata( projection_as_columns[far_left_right_col_ind as _..] .iter() .map(|(col, alias)| { @@ -1102,6 +1131,7 @@ pub fn new_join_children( } }), Arc::clone(right_child), + &right_schema, )?; Ok((new_left, new_right)) @@ -1116,22 +1146,26 @@ pub fn new_join_children( fn new_join_children_from_groups( left_proj: &[(Column, String)], right_proj: &[(Column, String)], + left_schema: &Schema, + right_schema: &Schema, left_child: &Arc, right_child: &Arc, ) -> Result<(ProjectionExec, ProjectionExec)> { - let build = |cols: &[(Column, String)], child: &Arc| { - ProjectionExec::try_new( - cols.iter().map(|(col, alias)| ProjectionExpr { - expr: Arc::new(Column::new(col.name(), col.index())) as _, - alias: alias.clone(), - }), - Arc::clone(child), - ) - }; + let build = + |cols: &[(Column, String)], schema: &Schema, child: &Arc| { + ProjectionExec::try_new_with_schema_metadata( + cols.iter().map(|(col, alias)| ProjectionExpr { + expr: Arc::new(Column::new(col.name(), col.index())) as _, + alias: alias.clone(), + }), + Arc::clone(child), + schema, + ) + }; Ok(( - build(left_proj, left_child)?, - build(right_proj, right_child)?, + build(left_proj, left_schema, left_child)?, + build(right_proj, right_schema, right_child)?, )) } @@ -1313,8 +1347,32 @@ fn try_collapse_projection_chain( } // To unify 3 or more sequential projections: - let unified: Arc = - Arc::new(ProjectionExec::try_new(current_exprs, current_input)?); + let unified_projection = ProjectionExec::try_new(current_exprs, current_input)?; + let metadata_fields = unified_projection + .schema() + .fields() + .iter() + .zip(outer.expr()) + .zip(outer.schema().fields()) + .map(|((unified_field, outer_expr), outer_field)| { + if outer_expr.expr.is::() { + unified_field + .as_ref() + .clone() + .with_metadata(outer_field.metadata().clone()) + } else { + unified_field.as_ref().clone() + } + }) + .collect::>(); + let metadata_schema = + Schema::new_with_metadata(metadata_fields, outer.schema().metadata().clone()); + let unified_projection = ProjectionExec::try_new_with_schema_metadata( + unified_projection.expr().to_vec(), + Arc::clone(unified_projection.input()), + &metadata_schema, + )?; + let unified: Arc = Arc::new(unified_projection); remove_unnecessary_projections(unified).data().map(Some) } @@ -1441,6 +1499,7 @@ mod tests { use crate::statistics::{StatisticsArgs, StatisticsContext}; use crate::test; use crate::test::exec::StatisticsExec; + use crate::union::UnionExec; use arrow::datatypes::{DataType, Field, Schema}; use datafusion_common::ScalarValue; @@ -1491,6 +1550,165 @@ mod tests { Ok(()) } + #[test] + fn test_projection_pushdown_preserves_output_metadata() -> Result<()> { + let input_schema = Arc::new(Schema::new(vec![ + Field::new("input", DataType::Int32, false).with_metadata(HashMap::from([( + "source".to_string(), + "input".to_string(), + )])), + Field::new("unused", DataType::Int32, false), + ])); + let input: Arc = UnionExec::try_new(vec![ + Arc::new(EmptyExec::new(Arc::clone(&input_schema))), + Arc::new(EmptyExec::new(input_schema)), + ])?; + let projected_schema = + Schema::new(vec![Field::new("input", DataType::Int32, false)]); + let projection = ProjectionExec::try_new_with_schema_metadata( + [ProjectionExpr { + expr: Arc::new(Column::new("input", 0)), + alias: "input".to_string(), + }], + input, + &projected_schema, + )?; + + assert!(!is_projection_removable(&projection)); + let plan: Arc = Arc::new(projection); + let optimized = remove_unnecessary_projections(Arc::clone(&plan))?; + assert!(optimized.transformed); + assert_eq!(optimized.data.schema(), plan.schema()); + Ok(()) + } + + #[test] + fn test_projection_chain_does_not_restore_stripped_column_metadata() -> Result<()> { + let source_metadata = + HashMap::from([("PARQUET:field_id".to_string(), "6".to_string())]); + let input_schema = Arc::new(Schema::new(vec![ + Field::new("sales", DataType::Int64, true).with_metadata(source_metadata), + ])); + let input: Arc = Arc::new(EmptyExec::new(input_schema)); + let inner: Arc = Arc::new(ProjectionExec::try_new( + [ProjectionExpr { + expr: Arc::new(Column::new("sales", 0)), + alias: "inner_sales".to_string(), + }], + input, + )?); + let projected_schema = + Schema::new(vec![Field::new("sales", DataType::Int64, true)]); + let outer = ProjectionExec::try_new_with_schema_metadata( + [ProjectionExpr { + expr: Arc::new(Column::new("inner_sales", 0)), + alias: "sales".to_string(), + }], + inner, + &projected_schema, + )?; + + let Some(collapsed) = try_collapse_projection_chain(&outer)? else { + return internal_err!("projection chain should collapse"); + }; + assert_eq!(collapsed.schema(), outer.schema()); + assert!(collapsed.schema().field(0).metadata().is_empty()); + Ok(()) + } + + #[test] + fn test_join_projection_pushdown_preserves_output_metadata() -> Result<()> { + let source_metadata = + HashMap::from([("source".to_string(), "iceberg".to_string())]); + let left_schema = Arc::new(Schema::new(vec![ + Field::new("left_keep", DataType::Int32, false) + .with_metadata(source_metadata.clone()), + Field::new("left_unused", DataType::Int32, false), + ])); + let right_schema = Arc::new(Schema::new(vec![ + Field::new("right_keep", DataType::Int32, false) + .with_metadata(source_metadata.clone()), + Field::new("right_unused", DataType::Int32, false), + ])); + let left: Arc = + Arc::new(EmptyExec::new(Arc::clone(&left_schema))); + let right: Arc = + Arc::new(EmptyExec::new(Arc::clone(&right_schema))); + let join_schema = Arc::new(Schema::new(vec![ + left_schema.field(0).clone(), + left_schema.field(1).clone(), + right_schema.field(0).clone(), + right_schema.field(1).clone(), + ])); + let join: Arc = + Arc::new(EmptyExec::new(Arc::clone(&join_schema))); + let projected_schema = Schema::new(vec![ + Field::new("left_keep", DataType::Int32, false), + Field::new("right_keep", DataType::Int32, false), + ]); + let projection = ProjectionExec::try_new_with_schema_metadata( + [ + ProjectionExpr { + expr: Arc::new(Column::new("left_keep", 0)), + alias: "left_keep".to_string(), + }, + ProjectionExpr { + expr: Arc::new(Column::new("right_keep", 2)), + alias: "right_keep".to_string(), + }, + ], + join, + &projected_schema, + )?; + let column_indices = [ + ColumnIndex { + index: 0, + side: JoinSide::Left, + }, + ColumnIndex { + index: 1, + side: JoinSide::Left, + }, + ColumnIndex { + index: 0, + side: JoinSide::Right, + }, + ColumnIndex { + index: 1, + side: JoinSide::Right, + }, + ]; + + let pushed = try_pushdown_through_join_with_column_indices( + &projection, + &left, + &right, + &[], + &join_schema, + None, + &column_indices, + )? + .expect("projection should be pushed through the join"); + + assert!( + pushed + .projected_left_child + .schema() + .field(0) + .metadata() + .is_empty() + ); + assert!( + pushed + .projected_right_child + .schema() + .field(0) + .metadata() + .is_empty() + ); + Ok(()) + } + #[test] fn test_collect_column_indices() -> Result<()> { let expr = Arc::new(BinaryExpr::new( diff --git a/datafusion/sql/src/parser.rs b/datafusion/sql/src/parser.rs index 86a00ca767a4c..efd78ffabe35a 100644 --- a/datafusion/sql/src/parser.rs +++ b/datafusion/sql/src/parser.rs @@ -925,28 +925,57 @@ impl<'a> DFParser<'a> { /// Parse a SQL `CREATE` statement handling `CREATE EXTERNAL TABLE` pub fn parse_create(&mut self) -> Result { // TODO: Change sql parser to take in `or_replace: bool` inside parse_create() - if self - .parser - .parse_keywords(&[Keyword::OR, Keyword::REPLACE, Keyword::EXTERNAL]) + if self.external_is_followed_by_table(2) + && self.parser.parse_keywords(&[ + Keyword::OR, + Keyword::REPLACE, + Keyword::EXTERNAL, + ]) { self.parse_create_external_table(false, true) - } else if self.parser.parse_keywords(&[ - Keyword::OR, - Keyword::REPLACE, - Keyword::UNBOUNDED, - Keyword::EXTERNAL, - ]) { + } else if self.external_is_followed_by_table(3) + && self.parser.parse_keywords(&[ + Keyword::OR, + Keyword::REPLACE, + Keyword::UNBOUNDED, + Keyword::EXTERNAL, + ]) + { self.parse_create_external_table(true, true) - } else if self.parser.parse_keyword(Keyword::EXTERNAL) { + } else if self.external_is_followed_by_table(0) + && self.parser.parse_keyword(Keyword::EXTERNAL) + { self.parse_create_external_table(false, false) - } else if self - .parser - .parse_keywords(&[Keyword::UNBOUNDED, Keyword::EXTERNAL]) + } else if self.external_is_followed_by_table(1) + && self + .parser + .parse_keywords(&[Keyword::UNBOUNDED, Keyword::EXTERNAL]) { self.parse_create_external_table(true, false) } else { - Ok(Statement::Statement(Box::from(self.parser.parse_create()?))) + // Let the configured dialect see the complete CREATE statement. Dialects such as + // Snowflake extend CREATE with object types (for example ICEBERG TABLE) that are not + // handled by sqlparser's generic `parse_create` entry point. + self.parser.prev_token(); + self.parse_and_handle_statement() + } + } + + fn external_is_followed_by_table(&self, external_offset: usize) -> bool { + let next_keyword = match self.parser.peek_nth_token(external_offset + 1).token { + Token::Word(word) => word.keyword, + _ => return false, + }; + if next_keyword == Keyword::TABLE { + return true; } + if !matches!(next_keyword, Keyword::TEMP | Keyword::TEMPORARY) { + return false; + } + matches!( + self.parser.peek_nth_token(external_offset + 2).token, + Token::Word(word) if word.keyword == Keyword::TABLE + ) } fn parse_partitions(&mut self) -> Result, DataFusionError> { @@ -1392,6 +1421,13 @@ mod tests { }); expect_parse_ok(sql, expected)?; + let sql = "CREATE EXTERNAL TEMPORARY TABLE t STORED AS CSV LOCATION 'foo.csv'"; + let expected = Statement::CreateExternalTable(CreateExternalTable { + temporary: true, + ..make_create_external_table("foo.csv") + }); + expect_parse_ok(sql, expected)?; + // positive case: literal comma remains part of a single path let sql = "CREATE EXTERNAL TABLE t(c1 int) STORED AS CSV LOCATION 'foo,bar.csv'"; let expected = Statement::CreateExternalTable(CreateExternalTable { @@ -1857,6 +1893,50 @@ mod tests { Ok(()) } + #[test] + fn delegate_create_iceberg_table_to_snowflake_dialect() -> Result<(), DataFusionError> + { + let sql = "CREATE ICEBERG TABLE t (id INT) BASE_LOCATION = 't'"; + let dialect = SnowflakeDialect; + let statements = DFParser::parse_sql_with_dialect(sql, &dialect)?; + + assert_eq!(statements.len(), 1); + let Statement::Statement(statement) = &statements[0] else { + panic!("expected a sqlparser statement, got {:?}", statements[0]); + }; + let sqlparser::ast::Statement::CreateTable(table) = statement.as_ref() else { + panic!("expected CREATE TABLE, got {statement:?}"); + }; + assert!(table.iceberg); + assert_eq!(table.base_location.as_deref(), Some("t")); + Ok(()) + } + + #[test] + fn delegate_create_stage_to_snowflake_dialect() -> Result<(), DataFusionError> { + let sql = + "CREATE OR REPLACE STAGE stage URL='s3://data.csv' FILE_FORMAT=(TYPE=csv)"; + let dialect = SnowflakeDialect; + let statements = DFParser::parse_sql_with_dialect(sql, &dialect)?; + + assert_eq!(statements.len(), 1); + let Statement::Statement(statement) = &statements[0] else { + panic!("expected a sqlparser statement, got {:?}", statements[0]); + }; + assert_eq!(statement.to_string(), sql); + Ok(()) + } + + #[test] + fn parse_standalone_begin_with_snowflake_dialect() -> Result<(), DataFusionError> { + let dialect = SnowflakeDialect; + let statements = DFParser::parse_sql_with_dialect("BEGIN", &dialect)?; + + assert_eq!(statements.len(), 1); + assert!(matches!(statements[0], Statement::Statement(_))); + Ok(()) + } + #[test] fn explain_copy_to_table_to_table() -> Result<(), DataFusionError> { let cases = vec![ diff --git a/datafusion/sqllogictest/test_files/decimal.slt b/datafusion/sqllogictest/test_files/decimal.slt index 4335ec06685f2..77e2aa13d7f6c 100644 --- a/datafusion/sqllogictest/test_files/decimal.slt +++ b/datafusion/sqllogictest/test_files/decimal.slt @@ -121,7 +121,7 @@ Decimal128(20, 6) 0.00055 query TR select arrow_typeof(avg(c1)), avg(c1) from decimal_simple; ---- -Decimal128(14, 10) 0.0000366666 +Decimal128(14, 10) 0.0000366667 query TR @@ -399,19 +399,19 @@ select c1/c5 from decimal_simple; ---- 0.5 0.641025641 -0.7142857142 +0.7142857143 0.7352941176 0.8 0.8571428571 -0.909090909 -0.909090909 +0.9090909091 +0.9090909091 0.9375 0.9615384615 1 1 1.0526315789 -1.5151515151 -2.7272727272 +1.5151515152 +2.7272727273 query T @@ -678,7 +678,7 @@ select * from decimal256_simple where c1 > c5; query TR select arrow_typeof(avg(c1)), avg(c1) from decimal256_simple; ---- -Decimal256(54, 10) 0.0000366666 +Decimal256(54, 10) 0.0000366667 query TR select arrow_typeof(min(c1)), min(c1) from decimal256_simple where c4=false; diff --git a/datafusion/sqllogictest/test_files/table_functions.slt b/datafusion/sqllogictest/test_files/table_functions.slt index e67d898d71475..72a9297b1af79 100644 --- a/datafusion/sqllogictest/test_files/table_functions.slt +++ b/datafusion/sqllogictest/test_files/table_functions.slt @@ -620,12 +620,16 @@ LIMIT 10; 1 1 -# Test that unsupported function argument types are properly reported -# rather than being silently dropped (which previously caused a misleading -# "requires 1 to 3 arguments" error instead) +# Named table-function arguments are preserved by the SQL planner. -statement error DataFusion error: Error during planning: Unsupported function argument type: start => 1 +query I SELECT * FROM generate_series(start => 1, stop => 5) +---- +1 +2 +3 +4 +5 statement error DataFusion error: Error during planning: Unsupported function argument type: \* SELECT * FROM generate_series(*) diff --git a/datafusion/sqllogictest/test_files/tpch/answers/q1.slt.part b/datafusion/sqllogictest/test_files/tpch/answers/q1.slt.part index bd8761bbb7fb4..9cdf63ac91ecc 100644 --- a/datafusion/sqllogictest/test_files/tpch/answers/q1.slt.part +++ b/datafusion/sqllogictest/test_files/tpch/answers/q1.slt.part @@ -39,7 +39,7 @@ order by l_returnflag, l_linestatus; ---- -A F 3774200 5320753880.69 5054096266.6828 5256751331.449234 25.537587 36002.123829 0.050144 147790 +A F 3774200 5320753880.69 5054096266.6828 5256751331.449234 25.537587 36002.123829 0.050145 147790 N F 95257 133737795.84 127132372.6512 132286291.229445 25.300664 35521.326916 0.049394 3765 -N O 7459297 10512270008.9 9986238338.3847 10385578376.585467 25.545537 36000.924688 0.050095 292000 -R F 3785523 5337950526.47 5071818532.942 5274405503.049367 25.525943 35994.029214 0.049989 148301 +N O 7459297 10512270008.9 9986238338.3847 10385578376.585467 25.545538 36000.924688 0.050096 292000 +R F 3785523 5337950526.47 5071818532.942 5274405503.049367 25.525944 35994.029214 0.049989 148301