diff --git a/datafusion/common/src/lib.rs b/datafusion/common/src/lib.rs index 2f6d9848b6e55..2eebfe4963057 100644 --- a/datafusion/common/src/lib.rs +++ b/datafusion/common/src/lib.rs @@ -99,7 +99,7 @@ pub use schema_reference::SchemaReference; pub use spans::{Location, Span, Spans}; pub use stats::{ColumnStatistics, Statistics}; pub use table_reference::{ResolvedTableReference, TableReference}; -pub use unnest::{RecursionUnnestOption, UnnestOptions}; +pub use unnest::{NullHandling, RecursionUnnestOption, UnnestOptions}; pub use utils::project_schema; // These are hidden from docs purely to avoid polluting the public view of what this crate exports. diff --git a/datafusion/common/src/unnest.rs b/datafusion/common/src/unnest.rs index db48edd061605..58aed390ace78 100644 --- a/datafusion/common/src/unnest.rs +++ b/datafusion/common/src/unnest.rs @@ -19,23 +19,38 @@ use crate::Column; +/// How [`UnnestOptions`] handles `NULL` and empty list values in the input column. +/// +/// The variants enumerate the three observable behaviors so that callers do +/// not have to compose multiple boolean flags to express what they want. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, PartialOrd, Hash)] +pub enum NullHandling { + /// Drop rows where the input list is `NULL` or empty. Matches the + /// default behavior of systems such as DuckDB and ClickHouse. + Drop, + /// Preserve `NULL` input rows as a single output row containing `NULL`. + /// Empty lists still produce zero output rows. This is the default and + /// matches DataFusion's historical `preserve_nulls = true` behavior. + #[default] + Preserve, + /// Like [`Self::Preserve`], and additionally treat an empty list + /// identically to a `NULL` list, producing a single output row + /// containing `NULL`. + PreserveAndExpandEmpty, +} + /// Options for unnesting a column that contains a list type, /// replicating values in the other, non nested rows. /// /// Conceptually this operation is like joining each row with all the /// values in the list column. /// -/// If `preserve_nulls` is false, nulls and empty lists -/// from the input column are not carried through to the output. This -/// is the default behavior for other systems such as ClickHouse and -/// DuckDB -/// -/// If `preserve_nulls` is true (the default), nulls from the input -/// column are carried through to the output. +/// The behavior with `NULL` and empty input lists is controlled by +/// [`NullHandling`]. See its variants for full details. /// /// # Examples /// -/// ## `Unnest(c1)`, preserve_nulls: false +/// ## `Unnest(c1)`, null_handling: NullHandling::Drop /// ```text /// ┌─────────┐ ┌─────┐ ┌─────────┐ ┌─────┐ /// │ {1, 2} │ │ A │ Unnest │ 1 │ │ A │ @@ -49,7 +64,7 @@ use crate::Column; /// c1 c2 /// ``` /// -/// ## `Unnest(c1)`, preserve_nulls: true +/// ## `Unnest(c1)`, null_handling: NullHandling::Preserve /// ```text /// ┌─────────┐ ┌─────┐ ┌─────────┐ ┌─────┐ /// │ {1, 2} │ │ A │ Unnest │ 1 │ │ A │ @@ -63,13 +78,30 @@ use crate::Column; /// c1 c2 c1 c2 /// ``` /// +/// ## `Unnest(c1)`, null_handling: NullHandling::PreserveAndExpandEmpty +/// ```text +/// ┌─────────┐ ┌─────┐ ┌─────────┐ ┌─────┐ +/// │ {1, 2} │ │ A │ Unnest │ 1 │ │ A │ +/// ├─────────┤ ├─────┤ ├─────────┤ ├─────┤ +/// │ null │ │ B │ │ 2 │ │ A │ +/// ├─────────┤ ├─────┤ ────────────▶ ├─────────┤ ├─────┤ +/// │ {} │ │ D │ │ null │ │ B │ +/// ├─────────┤ ├─────┤ ├─────────┤ ├─────┤ +/// │ {3} │ │ E │ │ null │ │ D │ +/// └─────────┘ └─────┘ ├─────────┤ ├─────┤ +/// c1 c2 │ 3 │ │ E │ +/// └─────────┘ └─────┘ +/// c1 c2 +/// ``` +/// /// `recursions` instruct how a column should be unnested (e.g unnesting a column multiple /// time, with depth = 1 and depth = 2). Any unnested column not being mentioned inside this /// options is inferred to be unnested with depth = 1 #[derive(Debug, Clone, PartialEq, PartialOrd, Hash, Eq)] pub struct UnnestOptions { - /// Should nulls in the input be preserved? Defaults to true - pub preserve_nulls: bool, + /// How to handle `NULL` and empty list values in the input column. + /// Defaults to [`NullHandling::Preserve`]. + pub null_handling: NullHandling, /// If specific columns need to be unnested multiple times (e.g at different depth), /// declare them here. Any unnested columns not being mentioned inside this option /// will be unnested with depth = 1 @@ -88,8 +120,7 @@ pub struct RecursionUnnestOption { impl Default for UnnestOptions { fn default() -> Self { Self { - // default to true to maintain backwards compatible behavior - preserve_nulls: true, + null_handling: NullHandling::Preserve, recursions: vec![], } } @@ -101,13 +132,41 @@ impl UnnestOptions { Default::default() } - /// Set the behavior with nulls in the input as described on - /// [`Self`] - pub fn with_preserve_nulls(mut self, preserve_nulls: bool) -> Self { - self.preserve_nulls = preserve_nulls; + /// Set the [`NullHandling`] mode used when unnesting `NULL` or empty + /// input lists. + pub fn with_null_handling(mut self, null_handling: NullHandling) -> Self { + self.null_handling = null_handling; self } + /// Backward-compatible setter that maps the previous boolean + /// `preserve_nulls` flag onto [`NullHandling`]. + /// + /// `true` maps to [`NullHandling::Preserve`]; `false` maps to + /// [`NullHandling::Drop`]. To opt into the new empty-list-preserving + /// mode, call [`Self::with_null_handling`] directly with + /// [`NullHandling::PreserveAndExpandEmpty`]. + pub fn with_preserve_nulls(self, preserve_nulls: bool) -> Self { + let null_handling = if preserve_nulls { + NullHandling::Preserve + } else { + NullHandling::Drop + }; + self.with_null_handling(null_handling) + } + + /// Returns true if `NULL` input rows produce a single output row + /// containing `NULL`. + pub fn preserve_nulls(&self) -> bool { + !matches!(self.null_handling, NullHandling::Drop) + } + + /// Returns true if empty input lists should produce a single + /// output row containing `NULL`. + pub fn expand_empty_as_null(&self) -> bool { + matches!(self.null_handling, NullHandling::PreserveAndExpandEmpty) + } + /// Set the recursions for the unnest operation pub fn with_recursions(mut self, recursion: RecursionUnnestOption) -> Self { self.recursions.push(recursion); diff --git a/datafusion/core/tests/dataframe/mod.rs b/datafusion/core/tests/dataframe/mod.rs index b9fecb5fdd732..73a9177ab738a 100644 --- a/datafusion/core/tests/dataframe/mod.rs +++ b/datafusion/core/tests/dataframe/mod.rs @@ -4370,6 +4370,28 @@ async fn unnest_column_nulls() -> Result<()> { ); let options = UnnestOptions::new().with_preserve_nulls(false); + let results = df + .clone() + .unnest_columns_with_options(&["list"], options)? + .collect() + .await?; + assert_snapshot!( + batches_to_string(&results), + @r" + +------+----+ + | list | id | + +------+----+ + | 1 | A | + | 2 | A | + | 3 | D | + +------+----+ + " + ); + + // Outer-unnest semantics: NULL and empty lists both produce a single + // output row containing NULL. + let options = UnnestOptions::new() + .with_null_handling(datafusion_common::NullHandling::PreserveAndExpandEmpty); let results = df .unnest_columns_with_options(&["list"], options)? .collect() @@ -4382,6 +4404,8 @@ async fn unnest_column_nulls() -> Result<()> { +------+----+ | 1 | A | | 2 | A | + | | B | + | | C | | 3 | D | +------+----+ " @@ -4390,6 +4414,133 @@ async fn unnest_column_nulls() -> Result<()> { Ok(()) } +/// Outer-unnest on a list-of-struct column. Verifies that +/// (a) struct elements unnest into flattened sub-columns and +/// (b) NULL and empty lists both still produce a single output row whose +/// struct sub-columns are all NULL. +#[tokio::test] +async fn unnest_outer_list_of_struct() -> Result<()> { + use arrow::array::{Int32Array, StructArray}; + + // Per-row sub-list lengths: 2, 1, 0 (empty), 0 (null) + let names = StringArray::from(vec!["alice", "bob", "carol"]); + let ages = Int32Array::from(vec![30, 40, 50]); + let struct_values = StructArray::from(vec![ + ( + Arc::new(Field::new("name", DataType::Utf8, true)), + Arc::new(names) as ArrayRef, + ), + ( + Arc::new(Field::new("age", DataType::Int32, true)), + Arc::new(ages) as ArrayRef, + ), + ]); + let struct_field = + Arc::new(Field::new("item", struct_values.data_type().clone(), true)); + let offsets = arrow::buffer::OffsetBuffer::::from_lengths([2, 1, 0, 0]); + let validity = arrow::buffer::NullBuffer::from(vec![true, true, true, false]); + let people = ListArray::new( + struct_field, + offsets, + Arc::new(struct_values), + Some(validity), + ); + let group = Int32Array::from(vec![1, 2, 3, 4]); + + let batch = RecordBatch::try_from_iter(vec![ + ("people", Arc::new(people) as ArrayRef), + ("group", Arc::new(group) as ArrayRef), + ])?; + + let ctx = SessionContext::new(); + ctx.register_batch("teams", batch)?; + let df = ctx.table("teams").await?; + + let options = UnnestOptions::new() + .with_null_handling(datafusion_common::NullHandling::PreserveAndExpandEmpty); + let results = df + // Unnest the list, then expand the resulting struct rows into columns. + .unnest_columns_with_options(&["people"], options.clone())? + .unnest_columns_with_options(&["people"], options)? + .collect() + .await?; + assert_snapshot!( + batches_to_string(&results), + @r" + +-------------+------------+-------+ + | people.name | people.age | group | + +-------------+------------+-------+ + | alice | 30 | 1 | + | bob | 40 | 1 | + | carol | 50 | 2 | + | | | 3 | + | | | 4 | + +-------------+------------+-------+ + " + ); + + Ok(()) +} + +/// Outer-unnest applied to a `FixedSizeList` column. For fixed-size lists, +/// every non-null row has the fixed length, so "empty" never occurs — +/// `PreserveAndExpandEmpty` should behave identically to `Preserve` here. +/// The test pins that equivalence so we notice if it ever diverges. +#[tokio::test] +async fn unnest_outer_fixed_size_list() -> Result<()> { + let batch = get_fixed_list_batch()?; + let ctx = SessionContext::new(); + ctx.register_batch("shapes", batch)?; + let df = ctx.table("shapes").await?; + + let preserve_results = df + .clone() + .unnest_columns_with_options( + &["tags"], + UnnestOptions::new().with_preserve_nulls(true), + )? + .collect() + .await?; + let outer_results = df + .unnest_columns_with_options( + &["tags"], + UnnestOptions::new().with_null_handling( + datafusion_common::NullHandling::PreserveAndExpandEmpty, + ), + )? + .collect() + .await?; + assert_eq!( + batches_to_sort_string(&preserve_results), + batches_to_sort_string(&outer_results), + "FixedSizeList has no empty case, so PreserveAndExpandEmpty must \ + match Preserve exactly" + ); + + // And the snapshot itself, to make the expected shape explicit. + assert_snapshot!( + batches_to_sort_string(&outer_results), + @r" + +----------+-------+ + | shape_id | tags | + +----------+-------+ + | 1 | | + | 2 | tag21 | + | 2 | tag22 | + | 3 | tag31 | + | 3 | tag32 | + | 4 | | + | 5 | tag51 | + | 5 | tag52 | + | 6 | tag61 | + | 6 | tag62 | + +----------+-------+ + " + ); + + Ok(()) +} + #[tokio::test] async fn unnest_fixed_list() -> Result<()> { let batch = get_fixed_list_batch()?; diff --git a/datafusion/expr/src/expr.rs b/datafusion/expr/src/expr.rs index b6ffd74ea2ecf..f9c0662e682e8 100644 --- a/datafusion/expr/src/expr.rs +++ b/datafusion/expr/src/expr.rs @@ -671,22 +671,43 @@ pub fn intersect_metadata_for_union<'a>( } /// UNNEST expression. +/// +/// When `outer` is `true`, the unnest should preserve `NULL` and empty input +/// lists by emitting a single `NULL` output row for each. When `false` (the +/// historical default), the behavior is identical to the plain `UNNEST(col)` +/// SQL form: `NULL` and empty input lists are dropped from the output. #[derive(Clone, PartialEq, Eq, PartialOrd, Hash, Debug)] pub struct Unnest { pub expr: Box, + /// Outer-unnest behavior: also expand empty input lists into a single + /// `NULL` output row (in addition to preserving `NULL` input rows). + pub outer: bool, } impl Unnest { - /// Create a new Unnest expression. + /// Create a new Unnest expression with default (non-outer) semantics. pub fn new(expr: Expr) -> Self { Self { expr: Box::new(expr), + outer: false, } } - /// Create a new Unnest expression. + /// Create a new Unnest expression with default (non-outer) semantics. pub fn new_boxed(boxed: Box) -> Self { - Self { expr: boxed } + Self { + expr: boxed, + outer: false, + } + } + + /// Create a new Unnest expression with outer-unnest semantics: `NULL` + /// and empty input lists each produce a single `NULL` output row. + pub fn new_outer(expr: Expr) -> Self { + Self { + expr: Box::new(expr), + outer: true, + } } } @@ -2431,11 +2452,19 @@ impl NormalizeEq for Expr { | (Expr::IsNotTrue(self_expr), Expr::IsNotTrue(other_expr)) | (Expr::IsNotFalse(self_expr), Expr::IsNotFalse(other_expr)) | (Expr::IsNotUnknown(self_expr), Expr::IsNotUnknown(other_expr)) - | (Expr::Negative(self_expr), Expr::Negative(other_expr)) - | ( - Expr::Unnest(Unnest { expr: self_expr }), - Expr::Unnest(Unnest { expr: other_expr }), - ) => self_expr.normalize_eq(other_expr), + | (Expr::Negative(self_expr), Expr::Negative(other_expr)) => { + self_expr.normalize_eq(other_expr) + } + ( + Expr::Unnest(Unnest { + expr: self_expr, + outer: self_outer, + }), + Expr::Unnest(Unnest { + expr: other_expr, + outer: other_outer, + }), + ) => self_outer == other_outer && self_expr.normalize_eq(other_expr), ( Expr::Between(Between { expr: self_expr, @@ -2883,7 +2912,9 @@ impl HashNode for Expr { field.hash(state); column.hash(state); } - Expr::Unnest(Unnest { expr: _expr }) => {} + Expr::Unnest(Unnest { expr: _expr, outer }) => { + outer.hash(state); + } Expr::HigherOrderFunction(HigherOrderFunction { func, args: _args }) => { func.hash(state); } @@ -3123,8 +3154,9 @@ impl Display for SchemaDisplay<'_> { } Expr::Negative(expr) => write!(f, "(- {})", SchemaDisplay(expr)), Expr::Not(expr) => write!(f, "NOT {}", SchemaDisplay(expr)), - Expr::Unnest(Unnest { expr }) => { - write!(f, "UNNEST({})", SchemaDisplay(expr)) + Expr::Unnest(Unnest { expr, outer }) => { + let name = if *outer { "UNNEST_OUTER" } else { "UNNEST" }; + write!(f, "{name}({})", SchemaDisplay(expr)) } Expr::ScalarFunction(ScalarFunction { func, args }) => { match func.schema_name(args) { @@ -3398,8 +3430,9 @@ impl Display for SqlDisplay<'_> { } Expr::Negative(expr) => write!(f, "(- {})", SqlDisplay(expr)), Expr::Not(expr) => write!(f, "NOT {}", SqlDisplay(expr)), - Expr::Unnest(Unnest { expr }) => { - write!(f, "UNNEST({})", SqlDisplay(expr)) + Expr::Unnest(Unnest { expr, outer }) => { + let name = if *outer { "UNNEST_OUTER" } else { "UNNEST" }; + write!(f, "{name}({})", SqlDisplay(expr)) } Expr::SimilarTo(Like { negated, @@ -3752,7 +3785,7 @@ impl Display for Expr { } }, Expr::Placeholder(Placeholder { id, .. }) => write!(f, "{id}"), - Expr::Unnest(Unnest { expr }) => { + Expr::Unnest(Unnest { expr, .. }) => { write!(f, "{UNNEST_COLUMN_PREFIX}({expr})") } Expr::HigherOrderFunction(fun) => { diff --git a/datafusion/expr/src/expr_fn.rs b/datafusion/expr/src/expr_fn.rs index 9d711113e4f74..b1a5a12d155ce 100644 --- a/datafusion/expr/src/expr_fn.rs +++ b/datafusion/expr/src/expr_fn.rs @@ -386,10 +386,11 @@ pub fn when(when: Expr, then: Expr) -> CaseBuilder { CaseBuilder::new(None, vec![when], vec![then], None) } -/// Create a Unnest expression +/// Create a Unnest expression with default (non-outer) semantics. pub fn unnest(expr: Expr) -> Expr { Expr::Unnest(Unnest { expr: Box::new(expr), + outer: false, }) } diff --git a/datafusion/expr/src/expr_rewriter/mod.rs b/datafusion/expr/src/expr_rewriter/mod.rs index a9a0c156538f9..7a6ac3fc8b062 100644 --- a/datafusion/expr/src/expr_rewriter/mod.rs +++ b/datafusion/expr/src/expr_rewriter/mod.rs @@ -87,13 +87,16 @@ pub fn normalize_col_with_schemas_and_ambiguity_check( using_columns: &[HashSet], ) -> Result { // Normalize column inside Unnest - if let Expr::Unnest(Unnest { expr }) = expr { + if let Expr::Unnest(Unnest { expr, outer }) = expr { let e = normalize_col_with_schemas_and_ambiguity_check( expr.as_ref().clone(), schemas, using_columns, )?; - return Ok(Expr::Unnest(Unnest { expr: Box::new(e) })); + return Ok(Expr::Unnest(Unnest { + expr: Box::new(e), + outer, + })); } expr.transform(|expr| { diff --git a/datafusion/expr/src/expr_schema.rs b/datafusion/expr/src/expr_schema.rs index 039bbad65a660..ec367de846d63 100644 --- a/datafusion/expr/src/expr_schema.rs +++ b/datafusion/expr/src/expr_schema.rs @@ -157,7 +157,7 @@ impl ExprSchemable for Expr { Expr::Cast(Cast { field, .. }) | Expr::TryCast(TryCast { field, .. }) => { Ok(field.data_type().clone()) } - Expr::Unnest(Unnest { expr }) => { + Expr::Unnest(Unnest { expr, .. }) => { let arg_data_type = expr.get_type(schema)?; // Unnest's output type is the inner type of the list match arg_data_type { diff --git a/datafusion/expr/src/tree_node.rs b/datafusion/expr/src/tree_node.rs index 010441b5a25d1..941fd22ea179f 100644 --- a/datafusion/expr/src/tree_node.rs +++ b/datafusion/expr/src/tree_node.rs @@ -49,7 +49,7 @@ impl TreeNode for Expr { ) -> Result { match self { Expr::Alias(Alias { expr, .. }) - | Expr::Unnest(Unnest { expr }) + | Expr::Unnest(Unnest { expr, .. }) | Expr::Not(expr) | Expr::IsNotNull(expr) | Expr::IsTrue(expr) @@ -150,9 +150,9 @@ impl TreeNode for Expr { quantifier, }) }), - Expr::Unnest(Unnest { expr, .. }) => expr + Expr::Unnest(Unnest { expr, outer }) => expr .map_elements(f)? - .update_data(|expr| Expr::Unnest(Unnest { expr })), + .update_data(|expr| Expr::Unnest(Unnest { expr, outer })), Expr::Alias(Alias { expr, relation, diff --git a/datafusion/physical-plan/src/unnest.rs b/datafusion/physical-plan/src/unnest.rs index fbe849229941a..c4b157b94e26e 100644 --- a/datafusion/physical-plan/src/unnest.rs +++ b/datafusion/physical-plan/src/unnest.rs @@ -888,14 +888,21 @@ fn build_batch( /// l2: [4,5], [], null, [6, 7] /// ``` /// -/// If `preserve_nulls` is false, the longest length array will be: +/// With [`datafusion_common::NullHandling::Drop`], the longest length array will be: /// /// ```ignore /// longest_length: [3, 0, 0, 2] /// ``` /// -/// whereas if `preserve_nulls` is true, the longest length array will be: +/// With [`datafusion_common::NullHandling::Preserve`] (the default), the longest length array +/// will be: /// +/// ```ignore +/// longest_length: [3, 1, 1, 2] +/// ``` +/// +/// With [`datafusion_common::NullHandling::PreserveAndExpandEmpty`], empty input lists are +/// also bumped to length 1 so they produce a single `NULL` output row: /// /// ```ignore /// longest_length: [3, 1, 1, 2] @@ -904,12 +911,16 @@ fn find_longest_length( list_arrays: &[ArrayRef], options: &UnnestOptions, ) -> Result { - // The length of a NULL list - let null_length = if options.preserve_nulls { + // The length to substitute for a NULL input list. + let null_length = if options.preserve_nulls() { Scalar::new(Int64Array::from_value(1, 1)) } else { Scalar::new(Int64Array::from_value(0, 1)) }; + let expand_empty = options.expand_empty_as_null(); + // Reused scalars for the empty-list rewrite when expand_empty is set. + let zero = Scalar::new(Int64Array::from_value(0, 1)); + let one = Scalar::new(Int64Array::from_value(1, 1)); let list_lengths: Vec = list_arrays .iter() .map(|list_array| { @@ -918,6 +929,12 @@ fn find_longest_length( length_array = cast(&length_array, &DataType::Int64)?; length_array = zip(&is_not_null(&length_array)?, &length_array, &null_length)?; + if expand_empty { + // Bump empty lists (length 0) to length 1 so they + // produce a single output row padded with NULL. + let is_zero = arrow_ord::cmp::eq(&length_array, &zero)?; + length_array = zip(&is_zero, &one, &length_array)?; + } Ok(length_array) }) .collect::>()?; @@ -1187,6 +1204,7 @@ mod tests { }; use arrow::buffer::{NullBuffer, OffsetBuffer}; use arrow::datatypes::{Field, Int32Type}; + use datafusion_common::NullHandling; use datafusion_common::test_util::batches_to_string; use insta::assert_snapshot; @@ -1369,7 +1387,7 @@ mod tests { list_type_columns.as_ref(), &HashSet::default(), &UnnestOptions { - preserve_nulls: true, + null_handling: NullHandling::Preserve, recursions: vec![], }, )? @@ -1405,6 +1423,369 @@ mod tests { Ok(()) } + #[test] + fn test_build_batch_preserve_and_expand_empty() -> Result<()> { + // c1: [A, B, C], [], NULL, [D], NULL, [NULL, F] c2: 1, 2, 3, 4, 5, 6 + // Expected for `NullHandling::PreserveAndExpandEmpty`: + // [A, B, C] -> three rows with c2 = 1, 1, 1 + // [] -> one row with c2 = 2 and unnested value NULL + // NULL -> one row with c2 = 3 and unnested value NULL + // [D] -> one row with c2 = 4 + // NULL -> one row with c2 = 5 and unnested value NULL + // [NULL, F] -> two rows with c2 = 6, 6 + let list_array = Arc::new(make_generic_array::()) as ArrayRef; + let other = + Arc::new(arrow::array::Int32Array::from(vec![1, 2, 3, 4, 5, 6])) as ArrayRef; + let in_schema = Arc::new(Schema::new(vec![ + Field::new( + "c1", + DataType::List(Arc::new(Field::new_list_field(DataType::Utf8, true))), + true, + ), + Field::new("c2", DataType::Int32, true), + ])); + let out_schema = Arc::new(Schema::new(vec![ + Field::new("c1_unnested", DataType::Utf8, true), + Field::new("c2", DataType::Int32, true), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&in_schema), + vec![Arc::clone(&list_array), Arc::clone(&other)], + )?; + let list_type_columns = vec![ListUnnest { + index_in_input_schema: 0, + depth: 1, + }]; + + let ret = build_batch( + &batch, + &out_schema, + &list_type_columns, + &HashSet::default(), + &UnnestOptions { + null_handling: NullHandling::PreserveAndExpandEmpty, + recursions: vec![], + }, + )? + .unwrap(); + + assert_snapshot!(batches_to_string(&[ret]), + @r" + +-------------+----+ + | c1_unnested | c2 | + +-------------+----+ + | A | 1 | + | B | 1 | + | C | 1 | + | | 2 | + | | 3 | + | D | 4 | + | | 5 | + | | 6 | + | F | 6 | + +-------------+----+ + "); + Ok(()) + } + + // PreserveAndExpandEmpty must work for LargeListArray (i64 offsets) too, + // not just the i32-offset ListArray exercised above. + #[test] + fn test_build_batch_preserve_and_expand_empty_largelist() -> Result<()> { + let list_array = Arc::new(make_generic_array::()) as ArrayRef; + let other = + Arc::new(arrow::array::Int32Array::from(vec![1, 2, 3, 4, 5, 6])) as ArrayRef; + let in_schema = Arc::new(Schema::new(vec![ + Field::new( + "c1", + DataType::LargeList(Arc::new(Field::new_list_field( + DataType::Utf8, + true, + ))), + true, + ), + Field::new("c2", DataType::Int32, true), + ])); + let out_schema = Arc::new(Schema::new(vec![ + Field::new("c1_unnested", DataType::Utf8, true), + Field::new("c2", DataType::Int32, true), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&in_schema), + vec![Arc::clone(&list_array), Arc::clone(&other)], + )?; + let list_type_columns = vec![ListUnnest { + index_in_input_schema: 0, + depth: 1, + }]; + + let ret = build_batch( + &batch, + &out_schema, + &list_type_columns, + &HashSet::default(), + &UnnestOptions { + null_handling: NullHandling::PreserveAndExpandEmpty, + recursions: vec![], + }, + )? + .unwrap(); + + // Same expected shape as the ListArray case — exercises the LargeList + // code path in unnest_list_array. + assert_snapshot!(batches_to_string(&[ret]), + @r" + +-------------+----+ + | c1_unnested | c2 | + +-------------+----+ + | A | 1 | + | B | 1 | + | C | 1 | + | | 2 | + | | 3 | + | D | 4 | + | | 5 | + | | 6 | + | F | 6 | + +-------------+----+ + "); + Ok(()) + } + + // When two list columns are unnested together, `find_longest_length` + // takes the per-row max. PreserveAndExpandEmpty must bump zeros to ones + // in each input column independently, then the row-wise max picks up + // the right value. + #[test] + fn test_build_batch_preserve_and_expand_empty_multi_column() -> Result<()> { + // col_a: [1, 2], [], NULL, [3] + // col_b: ['x'], ['y'],['z'], NULL + let col_a = ListArray::from_iter_primitive::(vec![ + Some(vec![Some(1), Some(2)]), + Some(vec![]), + None, + Some(vec![Some(3)]), + ]); + let col_b = { + let mut b = + arrow::array::ListBuilder::new(arrow::array::StringBuilder::new()); + b.values().append_value("x"); + b.append(true); + b.values().append_value("y"); + b.append(true); + b.values().append_value("z"); + b.append(true); + b.append(false); + b.finish() + }; + let id = + Arc::new(arrow::array::Int32Array::from(vec![10, 20, 30, 40])) as ArrayRef; + + let in_schema = Arc::new(Schema::new(vec![ + Field::new( + "a", + DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true))), + true, + ), + Field::new( + "b", + DataType::List(Arc::new(Field::new_list_field(DataType::Utf8, true))), + true, + ), + Field::new("id", DataType::Int32, true), + ])); + let out_schema = Arc::new(Schema::new(vec![ + Field::new("a_unnested", DataType::Int32, true), + Field::new("b_unnested", DataType::Utf8, true), + Field::new("id", DataType::Int32, true), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&in_schema), + vec![ + Arc::new(col_a) as ArrayRef, + Arc::new(col_b) as ArrayRef, + Arc::clone(&id), + ], + )?; + let list_type_columns = vec![ + ListUnnest { + index_in_input_schema: 0, + depth: 1, + }, + ListUnnest { + index_in_input_schema: 1, + depth: 1, + }, + ]; + + let ret = build_batch( + &batch, + &out_schema, + &list_type_columns, + &HashSet::default(), + &UnnestOptions { + null_handling: NullHandling::PreserveAndExpandEmpty, + recursions: vec![], + }, + )? + .unwrap(); + + // Row 0: longest = max(len([1,2])=2, len(['x'])=1) = 2 → a=[1,2], b=['x',NULL] + // Row 1: a=[] bumped to len 1, b=['y'] len 1 → a=[NULL], b=['y'] + // Row 2: a=NULL bumped to len 1, b=['z'] len 1 → a=[NULL], b=['z'] + // Row 3: a=[3] len 1, b=NULL bumped to len 1 → a=[3], b=[NULL] + assert_snapshot!(batches_to_string(&[ret]), + @r" + +------------+------------+----+ + | a_unnested | b_unnested | id | + +------------+------------+----+ + | 1 | x | 10 | + | 2 | | 10 | + | | y | 20 | + | | z | 30 | + | 3 | | 40 | + +------------+------------+----+ + "); + Ok(()) + } + + // PreserveAndExpandEmpty must propagate through recursive depth-2 + // unnesting: an outer NULL or empty produces one NULL output row at + // each level. Adapted from `test_build_batch_list_arr_recursive`. + #[test] + fn test_build_batch_preserve_and_expand_empty_recursive() -> Result<()> { + // col1 | col2 + // [[1,2,3],null,[4,5]] | ['a','b'] + // [[7,8,9,10], null, [11,12,13]] | ['c','d'] + // null | ['e'] + let list_arr1 = ListArray::from_iter_primitive::(vec![ + Some(vec![Some(1), Some(2), Some(3)]), + None, + Some(vec![Some(4), Some(5)]), + Some(vec![Some(7), Some(8), Some(9), Some(10)]), + None, + Some(vec![Some(11), Some(12), Some(13)]), + ]); + let list_arr1_ref = Arc::new(list_arr1) as ArrayRef; + let offsets = OffsetBuffer::from_lengths([3, 3, 0]); + let mut nulls = NullBufferBuilder::new(3); + nulls.append_non_null(); + nulls.append_non_null(); + nulls.append_null(); + let col1_field = Field::new_list_field( + DataType::List(Arc::new(Field::new_list_field( + list_arr1_ref.data_type().to_owned(), + true, + ))), + true, + ); + let col1 = ListArray::new( + Arc::new(Field::new_list_field( + list_arr1_ref.data_type().to_owned(), + true, + )), + offsets, + list_arr1_ref, + nulls.finish(), + ); + + let list_arr2 = StringArray::from(vec![ + Some("a"), + Some("b"), + Some("c"), + Some("d"), + Some("e"), + ]); + let offsets = OffsetBuffer::from_lengths([2, 2, 1]); + let mut nulls = NullBufferBuilder::new(3); + nulls.append_n_non_nulls(3); + let col2_field = Field::new( + "col2", + DataType::List(Arc::new(Field::new_list_field(DataType::Utf8, true))), + true, + ); + let col2 = GenericListArray::::new( + Arc::new(Field::new_list_field(DataType::Utf8, true)), + OffsetBuffer::new(offsets.into()), + Arc::new(list_arr2), + nulls.finish(), + ); + let schema = Arc::new(Schema::new(vec![col1_field, col2_field])); + let out_schema = Arc::new(Schema::new(vec![ + Field::new( + "col1_unnest_placeholder_depth_1", + DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true))), + true, + ), + Field::new("col1_unnest_placeholder_depth_2", DataType::Int32, true), + Field::new("col2_unnest_placeholder_depth_1", DataType::Utf8, true), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(col1) as ArrayRef, Arc::new(col2) as ArrayRef], + )?; + let list_type_columns = vec![ + ListUnnest { + index_in_input_schema: 0, + depth: 1, + }, + ListUnnest { + index_in_input_schema: 0, + depth: 2, + }, + ListUnnest { + index_in_input_schema: 1, + depth: 1, + }, + ]; + + let ret = build_batch( + &batch, + &out_schema, + &list_type_columns, + &HashSet::default(), + &UnnestOptions { + null_handling: NullHandling::PreserveAndExpandEmpty, + recursions: vec![], + }, + )? + .unwrap(); + + // The third input row (col1 = null, col2 = ['e']) now produces a + // NULL row for the depth-1 col1 placeholder *and* the depth-2 one, + // instead of being dropped at depth 1 and again at depth 2 the way + // it would be under `Drop`. Inner NULLs inside [...null...] sub- + // lists are still padded with NULL as before. + assert_snapshot!(batches_to_string(&[ret]), + @r" + +---------------------------------+---------------------------------+---------------------------------+ + | col1_unnest_placeholder_depth_1 | col1_unnest_placeholder_depth_2 | col2_unnest_placeholder_depth_1 | + +---------------------------------+---------------------------------+---------------------------------+ + | [1, 2, 3] | 1 | a | + | | 2 | b | + | [4, 5] | 3 | | + | [1, 2, 3] | | a | + | | | b | + | [4, 5] | | | + | [1, 2, 3] | 4 | a | + | | 5 | b | + | [4, 5] | | | + | [7, 8, 9, 10] | 7 | c | + | | 8 | d | + | [11, 12, 13] | 9 | | + | | 10 | | + | [7, 8, 9, 10] | | c | + | | | d | + | [11, 12, 13] | | | + | [7, 8, 9, 10] | 11 | c | + | | 12 | d | + | [11, 12, 13] | 13 | | + | | | e | + +---------------------------------+---------------------------------+---------------------------------+ + "); + Ok(()) + } + #[test] fn test_unnest_list_array() -> Result<()> { // [A, B, C], [], NULL, [D], NULL, [NULL, F] @@ -1452,11 +1833,11 @@ mod tests { fn verify_longest_length( list_arrays: &[ArrayRef], - preserve_nulls: bool, + null_handling: NullHandling, expected: Vec, ) -> Result<()> { let options = UnnestOptions { - preserve_nulls, + null_handling, recursions: vec![], }; let longest_length = find_longest_length(list_arrays, &options)?; @@ -1476,20 +1857,55 @@ mod tests { // Test with single ListArray // [A, B, C], [], NULL, [D], NULL, [NULL, F] let list_array = Arc::new(make_generic_array::()) as ArrayRef; - verify_longest_length(&[Arc::clone(&list_array)], false, vec![3, 0, 0, 1, 0, 2])?; - verify_longest_length(&[Arc::clone(&list_array)], true, vec![3, 0, 1, 1, 1, 2])?; + verify_longest_length( + &[Arc::clone(&list_array)], + NullHandling::Drop, + vec![3, 0, 0, 1, 0, 2], + )?; + verify_longest_length( + &[Arc::clone(&list_array)], + NullHandling::Preserve, + vec![3, 0, 1, 1, 1, 2], + )?; + // PreserveAndExpandEmpty also treats empty lists as a NULL row. + verify_longest_length( + &[Arc::clone(&list_array)], + NullHandling::PreserveAndExpandEmpty, + vec![3, 1, 1, 1, 1, 2], + )?; // Test with single LargeListArray // [A, B, C], [], NULL, [D], NULL, [NULL, F] let list_array = Arc::new(make_generic_array::()) as ArrayRef; - verify_longest_length(&[Arc::clone(&list_array)], false, vec![3, 0, 0, 1, 0, 2])?; - verify_longest_length(&[Arc::clone(&list_array)], true, vec![3, 0, 1, 1, 1, 2])?; + verify_longest_length( + &[Arc::clone(&list_array)], + NullHandling::Drop, + vec![3, 0, 0, 1, 0, 2], + )?; + verify_longest_length( + &[Arc::clone(&list_array)], + NullHandling::Preserve, + vec![3, 0, 1, 1, 1, 2], + )?; + verify_longest_length( + &[Arc::clone(&list_array)], + NullHandling::PreserveAndExpandEmpty, + vec![3, 1, 1, 1, 1, 2], + )?; // Test with single FixedSizeListArray // [A, B], NULL, [C, D], NULL, [NULL, F], [NULL, NULL] let list_array = Arc::new(make_fixed_list()) as ArrayRef; - verify_longest_length(&[Arc::clone(&list_array)], false, vec![2, 0, 2, 0, 2, 2])?; - verify_longest_length(&[Arc::clone(&list_array)], true, vec![2, 1, 2, 1, 2, 2])?; + verify_longest_length( + &[Arc::clone(&list_array)], + NullHandling::Drop, + vec![2, 0, 2, 0, 2, 2], + )?; + verify_longest_length( + &[Arc::clone(&list_array)], + NullHandling::Preserve, + vec![2, 1, 2, 1, 2, 2], + )?; // Test with multiple list arrays // [A, B, C], [], NULL, [D], NULL, [NULL, F] @@ -1497,8 +1913,17 @@ mod tests { let list1 = Arc::new(make_generic_array::()) as ArrayRef; let list2 = Arc::new(make_fixed_list()) as ArrayRef; let list_arrays = vec![Arc::clone(&list1), Arc::clone(&list2)]; - verify_longest_length(&list_arrays, false, vec![3, 0, 2, 1, 2, 2])?; - verify_longest_length(&list_arrays, true, vec![3, 1, 2, 1, 2, 2])?; + verify_longest_length(&list_arrays, NullHandling::Drop, vec![3, 0, 2, 1, 2, 2])?; + verify_longest_length( + &list_arrays, + NullHandling::Preserve, + vec![3, 1, 2, 1, 2, 2], + )?; + verify_longest_length( + &list_arrays, + NullHandling::PreserveAndExpandEmpty, + vec![3, 1, 2, 1, 2, 2], + )?; Ok(()) } diff --git a/datafusion/proto-models/proto/datafusion.proto b/datafusion/proto-models/proto/datafusion.proto index 205cf89abed1b..84bf69e26276f 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/datafusion/proto-models/proto/datafusion.proto @@ -403,7 +403,22 @@ message ColumnUnnestListRecursion { } message UnnestOptions { - bool preserve_nulls = 1; + // Reserved for the historical `bool preserve_nulls = 1;` field. + // Use `null_handling` instead. + reserved 1; + reserved "preserve_nulls"; + + enum NullHandling { + // Preserve nulls; empty lists produce no rows. The historical default. + PRESERVE = 0; + // Drop both null and empty lists from the output. + DROP = 1; + // Preserve nulls, and additionally expand empty lists into a single + // NULL output row (outer-unnest semantics). + PRESERVE_AND_EXPAND_EMPTY = 2; + } + + NullHandling null_handling = 3; repeated RecursionUnnestOption recursions = 2; } @@ -618,6 +633,9 @@ message NegativeNode { message Unnest { repeated LogicalExprNode exprs = 1; + // When true, this Unnest expression has outer-unnest semantics: NULL and + // empty input lists both produce a single NULL output row. + bool outer = 2; } message InListNode { diff --git a/datafusion/proto-models/src/generated/pbjson.rs b/datafusion/proto-models/src/generated/pbjson.rs index d23f8eee5fd2c..5124ff018bd0c 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -26825,10 +26825,16 @@ impl serde::Serialize for Unnest { if !self.exprs.is_empty() { len += 1; } + if self.outer { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion.Unnest", len)?; if !self.exprs.is_empty() { struct_ser.serialize_field("exprs", &self.exprs)?; } + if self.outer { + struct_ser.serialize_field("outer", &self.outer)?; + } struct_ser.end() } } @@ -26840,11 +26846,13 @@ impl<'de> serde::Deserialize<'de> for Unnest { { const FIELDS: &[&str] = &[ "exprs", + "outer", ]; #[allow(clippy::enum_variant_names)] enum GeneratedField { Exprs, + Outer, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -26867,6 +26875,7 @@ impl<'de> serde::Deserialize<'de> for Unnest { { match value { "exprs" => Ok(GeneratedField::Exprs), + "outer" => Ok(GeneratedField::Outer), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -26887,6 +26896,7 @@ impl<'de> serde::Deserialize<'de> for Unnest { V: serde::de::MapAccess<'de>, { let mut exprs__ = None; + let mut outer__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::Exprs => { @@ -26895,10 +26905,17 @@ impl<'de> serde::Deserialize<'de> for Unnest { } exprs__ = Some(map_.next_value()?); } + GeneratedField::Outer => { + if outer__.is_some() { + return Err(serde::de::Error::duplicate_field("outer")); + } + outer__ = Some(map_.next_value()?); + } } } Ok(Unnest { exprs: exprs__.unwrap_or_default(), + outer: outer__.unwrap_or_default(), }) } } @@ -27280,15 +27297,17 @@ impl serde::Serialize for UnnestOptions { { use serde::ser::SerializeStruct; let mut len = 0; - if self.preserve_nulls { + if self.null_handling != 0 { len += 1; } if !self.recursions.is_empty() { len += 1; } let mut struct_ser = serializer.serialize_struct("datafusion.UnnestOptions", len)?; - if self.preserve_nulls { - struct_ser.serialize_field("preserveNulls", &self.preserve_nulls)?; + if self.null_handling != 0 { + let v = unnest_options::NullHandling::try_from(self.null_handling) + .map_err(|_| serde::ser::Error::custom(format!("Invalid variant {}", self.null_handling)))?; + struct_ser.serialize_field("nullHandling", &v)?; } if !self.recursions.is_empty() { struct_ser.serialize_field("recursions", &self.recursions)?; @@ -27303,14 +27322,14 @@ impl<'de> serde::Deserialize<'de> for UnnestOptions { D: serde::Deserializer<'de>, { const FIELDS: &[&str] = &[ - "preserve_nulls", - "preserveNulls", + "null_handling", + "nullHandling", "recursions", ]; #[allow(clippy::enum_variant_names)] enum GeneratedField { - PreserveNulls, + NullHandling, Recursions, } impl<'de> serde::Deserialize<'de> for GeneratedField { @@ -27333,7 +27352,7 @@ impl<'de> serde::Deserialize<'de> for UnnestOptions { E: serde::de::Error, { match value { - "preserveNulls" | "preserve_nulls" => Ok(GeneratedField::PreserveNulls), + "nullHandling" | "null_handling" => Ok(GeneratedField::NullHandling), "recursions" => Ok(GeneratedField::Recursions), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } @@ -27354,15 +27373,15 @@ impl<'de> serde::Deserialize<'de> for UnnestOptions { where V: serde::de::MapAccess<'de>, { - let mut preserve_nulls__ = None; + let mut null_handling__ = None; let mut recursions__ = None; while let Some(k) = map_.next_key()? { match k { - GeneratedField::PreserveNulls => { - if preserve_nulls__.is_some() { - return Err(serde::de::Error::duplicate_field("preserveNulls")); + GeneratedField::NullHandling => { + if null_handling__.is_some() { + return Err(serde::de::Error::duplicate_field("nullHandling")); } - preserve_nulls__ = Some(map_.next_value()?); + null_handling__ = Some(map_.next_value::()? as i32); } GeneratedField::Recursions => { if recursions__.is_some() { @@ -27373,7 +27392,7 @@ impl<'de> serde::Deserialize<'de> for UnnestOptions { } } Ok(UnnestOptions { - preserve_nulls: preserve_nulls__.unwrap_or_default(), + null_handling: null_handling__.unwrap_or_default(), recursions: recursions__.unwrap_or_default(), }) } @@ -27381,6 +27400,80 @@ impl<'de> serde::Deserialize<'de> for UnnestOptions { deserializer.deserialize_struct("datafusion.UnnestOptions", FIELDS, GeneratedVisitor) } } +impl serde::Serialize for unnest_options::NullHandling { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + let variant = match self { + Self::Preserve => "PRESERVE", + Self::Drop => "DROP", + Self::PreserveAndExpandEmpty => "PRESERVE_AND_EXPAND_EMPTY", + }; + serializer.serialize_str(variant) + } +} +impl<'de> serde::Deserialize<'de> for unnest_options::NullHandling { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "PRESERVE", + "DROP", + "PRESERVE_AND_EXPAND_EMPTY", + ]; + + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = unnest_options::NullHandling; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + fn visit_i64(self, v: i64) -> std::result::Result + where + E: serde::de::Error, + { + i32::try_from(v) + .ok() + .and_then(|x| x.try_into().ok()) + .ok_or_else(|| { + serde::de::Error::invalid_value(serde::de::Unexpected::Signed(v), &self) + }) + } + + fn visit_u64(self, v: u64) -> std::result::Result + where + E: serde::de::Error, + { + i32::try_from(v) + .ok() + .and_then(|x| x.try_into().ok()) + .ok_or_else(|| { + serde::de::Error::invalid_value(serde::de::Unexpected::Unsigned(v), &self) + }) + } + + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "PRESERVE" => Ok(unnest_options::NullHandling::Preserve), + "DROP" => Ok(unnest_options::NullHandling::Drop), + "PRESERVE_AND_EXPAND_EMPTY" => Ok(unnest_options::NullHandling::PreserveAndExpandEmpty), + _ => Err(serde::de::Error::unknown_variant(value, FIELDS)), + } + } + } + deserializer.deserialize_any(GeneratedVisitor) + } +} impl serde::Serialize for ValuesNode { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result diff --git a/datafusion/proto-models/src/generated/prost.rs b/datafusion/proto-models/src/generated/prost.rs index 6baabbf37a41c..6c99a84e631f6 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -670,11 +670,57 @@ pub struct ColumnUnnestListRecursion { } #[derive(Clone, PartialEq, ::prost::Message)] pub struct UnnestOptions { - #[prost(bool, tag = "1")] - pub preserve_nulls: bool, + #[prost(enumeration = "unnest_options::NullHandling", tag = "3")] + pub null_handling: i32, #[prost(message, repeated, tag = "2")] pub recursions: ::prost::alloc::vec::Vec, } +/// Nested message and enum types in `UnnestOptions`. +pub mod unnest_options { + #[derive( + Clone, + Copy, + Debug, + PartialEq, + Eq, + Hash, + PartialOrd, + Ord, + ::prost::Enumeration + )] + #[repr(i32)] + pub enum NullHandling { + /// Preserve nulls; empty lists produce no rows. The historical default. + Preserve = 0, + /// Drop both null and empty lists from the output. + Drop = 1, + /// Preserve nulls, and additionally expand empty lists into a single + /// NULL output row (Spark `explode_outer` semantics). + PreserveAndExpandEmpty = 2, + } + impl NullHandling { + /// String value of the enum field names used in the ProtoBuf definition. + /// + /// The values are not transformed in any way and thus are considered stable + /// (if the ProtoBuf definition does not change) and safe for programmatic use. + pub fn as_str_name(&self) -> &'static str { + match self { + Self::Preserve => "PRESERVE", + Self::Drop => "DROP", + Self::PreserveAndExpandEmpty => "PRESERVE_AND_EXPAND_EMPTY", + } + } + /// Creates an enum from field names used in the ProtoBuf definition. + pub fn from_str_name(value: &str) -> ::core::option::Option { + match value { + "PRESERVE" => Some(Self::Preserve), + "DROP" => Some(Self::Drop), + "PRESERVE_AND_EXPAND_EMPTY" => Some(Self::PreserveAndExpandEmpty), + _ => None, + } + } + } +} #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] pub struct RecursionUnnestOption { #[prost(message, optional, tag = "1")] @@ -963,6 +1009,10 @@ pub struct NegativeNode { pub struct Unnest { #[prost(message, repeated, tag = "1")] pub exprs: ::prost::alloc::vec::Vec, + /// When true, this Unnest expression has Spark `explode_outer` semantics: + /// NULL and empty input lists both produce a single NULL output row. + #[prost(bool, tag = "2")] + pub outer: bool, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct InListNode { diff --git a/datafusion/proto/src/logical_plan/from_proto.rs b/datafusion/proto/src/logical_plan/from_proto.rs index 6d9a73e06ff45..00cc7f6a9d835 100644 --- a/datafusion/proto/src/logical_plan/from_proto.rs +++ b/datafusion/proto/src/logical_plan/from_proto.rs @@ -64,8 +64,20 @@ use super::{AsLogicalPlan, LogicalExtensionCodec}; impl FromProto<&protobuf::UnnestOptions> for UnnestOptions { fn from_proto(opts: &protobuf::UnnestOptions) -> Self { + use datafusion_common::NullHandling; + use protobuf::unnest_options::NullHandling as ProtoNullHandling; + let null_handling = match ProtoNullHandling::try_from(opts.null_handling) { + Ok(ProtoNullHandling::Preserve) => NullHandling::Preserve, + Ok(ProtoNullHandling::Drop) => NullHandling::Drop, + Ok(ProtoNullHandling::PreserveAndExpandEmpty) => { + NullHandling::PreserveAndExpandEmpty + } + // Unknown enum values fall back to the default (Preserve), which + // matches DataFusion's historical behavior. + Err(_) => NullHandling::Preserve, + }; Self { - preserve_nulls: opts.preserve_nulls, + null_handling, recursions: opts .recursions .iter() @@ -681,7 +693,10 @@ pub fn parse_expr( if exprs.len() != 1 { return Err(proto_error("Unnest must have exactly one expression")); } - Ok(Expr::Unnest(Unnest::new(exprs.swap_remove(0)))) + Ok(Expr::Unnest(Unnest { + expr: Box::new(exprs.swap_remove(0)), + outer: unnest.outer, + })) } ExprType::InList(in_list) => Ok(Expr::InList(InList::new( Box::new(parse_required_expr( diff --git a/datafusion/proto/src/logical_plan/to_proto.rs b/datafusion/proto/src/logical_plan/to_proto.rs index 23ce254e99a40..89de342ff00b7 100644 --- a/datafusion/proto/src/logical_plan/to_proto.rs +++ b/datafusion/proto/src/logical_plan/to_proto.rs @@ -57,8 +57,17 @@ use crate::protobuf::LogicalPlanNode; impl FromProto<&UnnestOptions> for protobuf::UnnestOptions { fn from_proto(opts: &UnnestOptions) -> Self { + use datafusion_common::NullHandling; + use protobuf::unnest_options::NullHandling as ProtoNullHandling; + let null_handling = match opts.null_handling { + NullHandling::Preserve => ProtoNullHandling::Preserve, + NullHandling::Drop => ProtoNullHandling::Drop, + NullHandling::PreserveAndExpandEmpty => { + ProtoNullHandling::PreserveAndExpandEmpty + } + } as i32; Self { - preserve_nulls: opts.preserve_nulls, + null_handling, recursions: opts .recursions .iter() @@ -571,9 +580,10 @@ pub fn serialize_expr( expr_type: Some(ExprType::Negative(expr)), } } - Expr::Unnest(Unnest { expr }) => { + Expr::Unnest(Unnest { expr, outer }) => { let expr = protobuf::Unnest { exprs: vec![serialize_expr(expr.as_ref(), codec)?], + outer: *outer, }; protobuf::LogicalExprNode { expr_type: Some(ExprType::Unnest(expr)), diff --git a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs index 74f7253386764..0e77aa76f4a4d 100644 --- a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs @@ -2736,6 +2736,18 @@ fn roundtrip_inlist() { fn roundtrip_unnest() { let test_expr = Expr::Unnest(Unnest { expr: Box::new(col("col")), + outer: false, + }); + + let ctx = SessionContext::new(); + roundtrip_expr_test(test_expr, ctx); +} + +#[test] +fn roundtrip_unnest_outer() { + let test_expr = Expr::Unnest(Unnest { + expr: Box::new(col("col")), + outer: true, }); let ctx = SessionContext::new(); diff --git a/datafusion/sql/src/expr/function.rs b/datafusion/sql/src/expr/function.rs index 701485eee733c..e6bee31fbf106 100644 --- a/datafusion/sql/src/expr/function.rs +++ b/datafusion/sql/src/expr/function.rs @@ -546,15 +546,25 @@ impl SqlToRel<'_, S> { } } - // Build Unnest expression - if name.eq("unnest") { + // Build Unnest expression. + // + // `unnest(col)` drops `NULL` and empty input lists (default SQL + // semantics, matching DuckDB/PostgreSQL). `unnest_outer(col)` sets + // `outer = true` so the downstream planner picks + // `NullHandling::PreserveAndExpandEmpty`, which preserves `NULL` + // and empty input lists as a single `NULL` output row. + if name.eq("unnest") || name.eq("unnest_outer") { + let outer = name.eq("unnest_outer"); let mut exprs = self.function_args_to_expr(args, schema, planner_context)?; if exprs.len() != 1 { - return plan_err!("unnest() requires exactly one argument"); + return plan_err!("{name}() requires exactly one argument"); } let expr = exprs.swap_remove(0); Self::check_unnest_arg(&expr, schema)?; - return Ok(Expr::Unnest(Unnest::new(expr))); + return Ok(Expr::Unnest(Unnest { + expr: Box::new(expr), + outer, + })); } if !order_by.is_empty() && is_function_window { diff --git a/datafusion/sql/src/select.rs b/datafusion/sql/src/select.rs index ba7353c424f4e..bdab013144462 100644 --- a/datafusion/sql/src/select.rs +++ b/datafusion/sql/src/select.rs @@ -33,9 +33,10 @@ use arrow::datatypes::DataType; use datafusion_common::error::DataFusionErrorBuilder; use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; use datafusion_common::{Column, DFSchema, DFSchemaRef, Result, not_impl_err, plan_err}; -use datafusion_common::{RecursionUnnestOption, UnnestOptions}; +use datafusion_common::{NullHandling, RecursionUnnestOption, UnnestOptions}; use datafusion_expr::ExprSchemable; use datafusion_expr::builder::get_struct_unnested_columns; +use datafusion_expr::expr::Unnest as UnnestExpr; use datafusion_expr::expr::{PlannedReplaceSelectItem, WildcardOptions}; use datafusion_expr::expr_rewriter::{ normalize_col, normalize_col_with_schemas_and_ambiguity_check, normalize_sorts, @@ -665,8 +666,15 @@ impl SqlToRel<'_, S> { }); } - // Set preserve_nulls to false to ensure compatibility with DuckDB and PostgreSQL - let mut unnest_options = UnnestOptions::new().with_preserve_nulls(false); + // The default SQL `UNNEST` matches DuckDB/PostgreSQL: drop both + // NULL and empty input lists. Outer-unnest (modelled as + // `Unnest { outer: true }`) overrides that and selects + // `NullHandling::PreserveAndExpandEmpty`. Mixing the two in a + // single SELECT is a planning error because `UnnestOptions` is + // per-`UnnestExec`, not per-column. + let null_handling = collect_unnest_null_handling(&intermediate_expr_groups)?; + let mut unnest_options = + UnnestOptions::new().with_null_handling(null_handling); let mut unnest_col_vec = vec![]; for (col, maybe_list_unnest) in unnest_columns.into_iter() { @@ -1451,3 +1459,45 @@ fn has_unnest_expr_recursively(expr: &Expr) -> bool { }); has_unnest } + +/// Walk `select_exprs`, observe every [`Expr::Unnest`] inside them, and +/// derive the [`NullHandling`] mode for the resulting [`UnnestOptions`]. +/// +/// * No unnest with `outer = true` → [`NullHandling::Drop`] (default SQL +/// `UNNEST(...)` semantics, matching DuckDB/PostgreSQL). +/// * Every unnest with `outer = true` → [`NullHandling::PreserveAndExpandEmpty`] +/// (outer-unnest semantics: `NULL` and empty input lists each produce a +/// single `NULL` output row). +/// * A mix of `outer = true` and `outer = false` in one SELECT → planning +/// error, because `UnnestOptions` applies per `Unnest` plan node, not +/// per output column. +fn collect_unnest_null_handling(expr_groups: &[Vec]) -> Result { + let mut saw_outer = false; + let mut saw_inner = false; + for group in expr_groups { + for expr in group { + expr.apply(|e| { + if let Expr::Unnest(UnnestExpr { outer, .. }) = e { + if *outer { + saw_outer = true; + } else { + saw_inner = true; + } + } + Ok(TreeNodeRecursion::Continue) + })?; + } + } + if saw_outer && saw_inner { + return plan_err!( + "Cannot mix `unnest(...)` with `unnest_outer(...)` in the same \ + SELECT — the unnest operator carries a single null-handling \ + mode. Split the query so each unnest projection uses one mode." + ); + } + Ok(if saw_outer { + NullHandling::PreserveAndExpandEmpty + } else { + NullHandling::Drop + }) +} diff --git a/datafusion/sql/src/unparser/expr.rs b/datafusion/sql/src/unparser/expr.rs index 89560b23791a3..9403e15406344 100644 --- a/datafusion/sql/src/unparser/expr.rs +++ b/datafusion/sql/src/unparser/expr.rs @@ -2452,6 +2452,7 @@ mod tests { name: "array_col".to_string(), spans: Spans::new(), })), + outer: false, }), r#"UNNEST("table".array_col)"#, ), diff --git a/datafusion/sql/src/unparser/plan.rs b/datafusion/sql/src/unparser/plan.rs index 9fe97a8291b6a..f4b60176cfba9 100644 --- a/datafusion/sql/src/unparser/plan.rs +++ b/datafusion/sql/src/unparser/plan.rs @@ -2047,7 +2047,7 @@ impl Unparser<'_> { let mut flatten = FlattenRelationBuilder::default(); flatten.input_expr(input_expr); - flatten.outer(unnest.options.preserve_nulls); + flatten.outer(unnest.options.preserve_nulls()); Ok(Some(flatten)) } diff --git a/datafusion/sqllogictest/test_files/unnest.slt b/datafusion/sqllogictest/test_files/unnest.slt index 5cca3cbfe461f..a3385b81d70d1 100644 --- a/datafusion/sqllogictest/test_files/unnest.slt +++ b/datafusion/sqllogictest/test_files/unnest.slt @@ -1463,3 +1463,186 @@ FROM list_struct_table; statement ok DROP TABLE list_struct_table; + +#################################### +# `unnest_outer` Tests +# +# `unnest_outer(col)` is the outer-unnest peer to `unnest(col)`. Rows whose +# input list is `NULL` or empty produce a single output row containing +# `NULL`; rows with values are exploded element-by-element the same way as +# plain `unnest`. +# +# Column types on the tables below are inferred from the `VALUES` rows. +# DataFusion's SQL parser does not accept PostgreSQL `TYPE[]` array-column +# syntax inside `CREATE TABLE ... AS VALUES`, so tables are declared as +# CTAS over a `VALUES` subquery with aliased column names. +#################################### + +## unnest vs unnest_outer on an integer list + +statement ok +CREATE TABLE int_lists AS +SELECT column1 AS id, column2 AS xs FROM (VALUES + (1, [10, 20, 30]), + (4, [40]), + (5, [NULL, 50]), + (2, arrow_cast(make_array(), 'List(Int64)')), + (3, NULL) +); + +## Plain `unnest`: drops both NULL and empty input rows. +## Inner NULL elements survive. +query II +SELECT id, unnest(xs) AS x FROM int_lists ORDER BY id, x; +---- +1 10 +1 20 +1 30 +4 40 +5 50 +5 NULL + +## `unnest_outer`: NULL and empty input lists each produce one NULL row. +query II +SELECT id, unnest_outer(xs) AS x FROM int_lists ORDER BY id, x; +---- +1 10 +1 20 +1 30 +2 NULL +3 NULL +4 40 +5 50 +5 NULL + +## String list with inner NULLs: inner NULL elements must survive (they are +## not the same as "empty"), while NULL and empty input lists become a +## single NULL output row. + +statement ok +CREATE TABLE str_lists AS +SELECT column1 AS id, column2 AS tags FROM (VALUES + ('A', ['x', 'y']), + ('B', ['p', NULL, 'q']), + ('C', arrow_cast(make_array(), 'List(Utf8)')), + ('D', NULL) +); + +query TT +SELECT id, unnest_outer(tags) AS tag FROM str_lists ORDER BY id, tag; +---- +A x +A y +B p +B q +B NULL +C NULL +D NULL + +## Mixed list lengths — verify row-wise expansion. + +statement ok +CREATE TABLE varied_lists AS +SELECT column1 AS id, column2 AS xs FROM (VALUES + (1, [1, 2, 3, 4]), + (2, [5]), + (3, arrow_cast(make_array(), 'List(Int64)')), + (4, NULL) +); + +query II +SELECT id, unnest_outer(xs) AS x FROM varied_lists ORDER BY id, x; +---- +1 1 +1 2 +1 3 +1 4 +2 5 +3 NULL +4 NULL + +## Aliased output column. +query II +SELECT id, unnest_outer(xs) AS unwrapped FROM int_lists WHERE id = 2; +---- +2 NULL + +## Mixing `unnest` and `unnest_outer` in one SELECT is a planning error. +## `UnnestOptions` is per-`UnnestExec`, so we refuse to silently pick one mode. + +statement error DataFusion error: Error during planning: Cannot mix `unnest\(\.\.\.\)` with `unnest_outer\(\.\.\.\)` in the same SELECT +SELECT unnest(xs), unnest_outer(xs) FROM int_lists; + +## Chained `unnest` → `unnest_outer` via subquery. +## `unnest(xs)` (inner) drops NULL and empty outer rows, then +## `unnest_outer(ys)` (outer) preserves NULL and empty sub-lists from the +## inner unnest. + +statement ok +CREATE TABLE nested_lists AS +SELECT column1 AS id, column2 AS xs FROM (VALUES + (100, [[1, 2, 3], NULL, [4, 5]]), + (200, [[7], arrow_cast(make_array(), 'List(Int64)')]), + (300, NULL) +); + +query II +SELECT id, unnest_outer(ys) AS y +FROM (SELECT id, unnest(xs) AS ys FROM nested_lists) +ORDER BY id, y; +---- +100 1 +100 2 +100 3 +100 4 +100 5 +100 NULL +200 7 +200 NULL + +statement ok +DROP TABLE nested_lists; + +## `unnest_outer` agrees with `unnest` when no NULL or empty rows exist. + +statement ok +CREATE TABLE dense_lists AS +SELECT column1 AS id, column2 AS xs FROM (VALUES + (1, [10, 20]), + (2, [30]), + (3, [40, 50, 60]) +); + +query II +SELECT id, unnest(xs) AS x FROM dense_lists ORDER BY id, x; +---- +1 10 +1 20 +2 30 +3 40 +3 50 +3 60 + +query II +SELECT id, unnest_outer(xs) AS x FROM dense_lists ORDER BY id, x; +---- +1 10 +1 20 +2 30 +3 40 +3 50 +3 60 + +## Cleanup + +statement ok +DROP TABLE int_lists; + +statement ok +DROP TABLE str_lists; + +statement ok +DROP TABLE varied_lists; + +statement ok +DROP TABLE dense_lists; diff --git a/docs/source/user-guide/sql/select.md b/docs/source/user-guide/sql/select.md index ea96f6ae4528d..af442de6597c1 100644 --- a/docs/source/user-guide/sql/select.md +++ b/docs/source/user-guide/sql/select.md @@ -279,6 +279,29 @@ SELECT id, UNNEST(items) FROM orders; items (implicit lateral references such as `FROM orders AS t, UNNEST(t.items)` are not currently supported). +### `unnest_outer` + +`unnest_outer(col)` is the outer-unnest peer to `UNNEST(col)`. The two differ +only in how `NULL` and empty input lists are handled: + +| Form | `NULL` input list | Empty input list | +| ------------------- | ----------------- | ---------------- | +| `UNNEST(col)` | dropped | dropped | +| `unnest_outer(col)` | one `NULL` row | one `NULL` row | + +```sql +SELECT id, unnest_outer(tags) AS tag FROM rows; +``` + +An input row with an empty `tags` array or `NULL` `tags` produces one output +row whose `tag` is `NULL`, instead of being dropped. This is analogous to the +outer variant offered by other engines (Spark `explode_outer`, Hive `EXPLODE OUTER`, Snowflake `FLATTEN(OUTER => true)`). + +`unnest_outer` cannot be mixed with `unnest` in the same `SELECT` — the +unnest plan node carries a single null-handling mode for all its output +columns, so a mix would be ambiguous. The planner returns an error in that +case. + ## WHERE clause ```text