diff --git a/datafusion/functions/src/core/getfield.rs b/datafusion/functions/src/core/getfield.rs index a0f024bbc7ea4..c7003a372eeaf 100644 --- a/datafusion/functions/src/core/getfield.rs +++ b/datafusion/functions/src/core/getfield.rs @@ -18,9 +18,10 @@ use std::sync::{Arc, OnceLock}; use arrow::array::{ - Array, Capacities, MutableArrayData, Scalar, cast::AsArray, make_array, + Array, ArrayRef, Capacities, MutableArrayData, Scalar, cast::AsArray, make_array, make_comparator, }; +use arrow::buffer::NullBuffer; use arrow::compute::SortOptions; use arrow::datatypes::{DataType, Field, FieldRef}; @@ -125,6 +126,10 @@ fn process_map_array( let matches = keys.values(); for entry in 0..map_array.len() { + if map_array.is_null(entry) { + mutable.try_extend_nulls(1)?; + continue; + } let start = offsets[entry] as usize; let end = offsets[entry + 1] as usize; @@ -160,6 +165,10 @@ fn process_map_with_nested_key( MutableArrayData::with_capacities(vec![&original_data], true, capacity); for entry in 0..map_array.len() { + if map_array.is_null(entry) { + mutable.try_extend_nulls(1)?; + continue; + } let start = map_array.value_offsets()[entry] as usize; let end = map_array.value_offsets()[entry + 1] as usize; @@ -182,6 +191,37 @@ fn process_map_with_nested_key( Ok(ColumnarValue::Array(data)) } +/// Apply a struct's nulls to one of its fields. +fn apply_parent_nulls(col: &ArrayRef, parent_nulls: &NullBuffer) -> Result { + // NullArray is already entirely null and cannot have a validity bitmap. + // If we have 0 parent nulls, we can also avoid extra work. + if col.data_type().is_null() || parent_nulls.null_count() == 0 { + return Ok(Arc::clone(col)); + } + + let data = col.to_data(); + match col.data_type() { + DataType::Union(_, _) => { + // Unions represent nulls in their children. Rebuild the array so null + // parents become null union values in both sparse and dense layouts. + let mut mutable = MutableArrayData::new(vec![&data], true, data.len()); + let mut end = 0; + for (start, valid_end) in parent_nulls.valid_slices() { + mutable.try_extend_nulls(start - end)?; + mutable.try_extend(0, start, valid_end)?; + end = valid_end; + } + mutable.try_extend_nulls(data.len() - end)?; + + Ok(make_array(mutable.freeze())) + } + _ => { + let nulls = NullBuffer::union(col.nulls(), Some(parent_nulls)); + Ok(make_array(data.into_builder().nulls(nulls).build()?)) + } + } +} + /// Extract a single field from a struct or map array fn extract_single_field(base: ColumnarValue, name: ScalarValue) -> Result { let arrays = ColumnarValue::values_to_arrays(&[base])?; @@ -204,9 +244,12 @@ fn extract_single_field(base: ColumnarValue, name: ScalarValue) -> Result { // The lookup key is a single scalar. `eq` does not support nested @@ -220,9 +263,13 @@ fn extract_single_field(base: ColumnarValue, name: ScalarValue) -> Result { let as_struct_array = as_struct_array(&array)?; - match as_struct_array.column_by_name(&k) { - None => exec_err!("Field {k} not found in struct"), - Some(col) => Ok(ColumnarValue::Array(Arc::clone(col))), + let nulls = as_struct_array.nulls(); + match (as_struct_array.column_by_name(&k), nulls) { + (None, _) => exec_err!("Field {k} not found in struct"), + (Some(col), None) => Ok(ColumnarValue::Array(Arc::clone(col))), + (Some(col), Some(parent_nulls)) => { + Ok(ColumnarValue::Array(apply_parent_nulls(col, parent_nulls)?)) + } } } (DataType::Struct(_), name, _) => exec_err!( @@ -658,9 +705,9 @@ mod tests { use super::*; use arrow::array::{ ArrayRef, Int32Array, Int32Builder, ListArray, ListBuilder, MapBuilder, - StructArray, + StructArray, UnionArray, }; - use arrow::datatypes::{Fields, Int32Type}; + use arrow::datatypes::{Fields, Int32Type, UnionFields}; #[test] fn test_get_field_utf8view_key() -> Result<()> { @@ -733,6 +780,160 @@ mod tests { Ok(()) } + #[test] + fn test_get_field_nested_struct_outer_nulls() -> Result<()> { + let inner_array = StructArray::new( + vec![Field::new("value", DataType::Int32, false)].into(), + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + None, + ); + + // Only the outer struct marks row 1 as null; its children are all valid. + let outer_array = StructArray::new( + vec![Field::new("inner", inner_array.data_type().clone(), false)].into(), + vec![Arc::new(inner_array)], + Some(NullBuffer::from(vec![true, false, true])), + ); + + let inner = extract_single_field( + ColumnarValue::Array(Arc::new(outer_array)), + ScalarValue::Utf8(Some("inner".to_string())), + )?; + let result = + extract_single_field(inner, ScalarValue::Utf8(Some("value".to_string())))? + .into_array(3)?; + + let expected = Int32Array::from(vec![Some(1), None, Some(3)]); + assert_eq!(result.as_ref(), &expected as &dyn Array); + + Ok(()) + } + + #[test] + fn test_get_field_map_parent_nulls() -> Result<()> { + use arrow::array::{FixedSizeListArray, MapArray}; + use arrow_buffer::OffsetBuffer; + + let keys = Arc::new(Int32Array::from(vec![7; 3])) as ArrayRef; + let nested_keys = Arc::new(FixedSizeListArray::new( + Arc::new(Field::new("item", DataType::Int32, false)), + 1, + Arc::clone(&keys), + None, + )) as ArrayRef; + + // Exercise both map lookup paths. The null map has a valid matching entry. + for keys in [keys, nested_keys] { + let key = ScalarValue::try_from_array(keys.as_ref(), 0)?; + let entries = StructArray::new( + vec![ + Field::new("key", keys.data_type().clone(), false), + Field::new("value", DataType::Int32, false), + ] + .into(), + vec![keys, Arc::new(Int32Array::from(vec![1, 2, 3]))], + None, + ); + let map = MapArray::new( + Arc::new(Field::new("entries", entries.data_type().clone(), false)), + OffsetBuffer::new(vec![0, 1, 2, 3].into()), + entries, + Some(NullBuffer::from(vec![true, false, true])), + false, + ); + let result = extract_single_field(ColumnarValue::Array(Arc::new(map)), key)? + .into_array(3)?; + let expected = Int32Array::from(vec![Some(1), None, Some(3)]); + assert_eq!(result.as_ref(), &expected as &dyn Array); + } + Ok(()) + } + + #[test] + fn test_get_field_null_typed_child() -> Result<()> { + use arrow::array::{DictionaryArray, NullArray, UInt32Array}; + use arrow::datatypes::UInt32Type; + + let values = Arc::new(StructArray::new( + vec![Field::new("value", DataType::Null, true)].into(), + vec![Arc::new(NullArray::new(2))], + Some(NullBuffer::from(vec![true, false])), + )) as ArrayRef; + let dictionary = Arc::new(DictionaryArray::::try_new( + UInt32Array::from(vec![0, 1]), + Arc::clone(&values), + )?) as ArrayRef; + + for input in [values, dictionary] { + let result = extract_single_field( + ColumnarValue::Array(input), + ScalarValue::Utf8(Some("value".to_string())), + )? + .into_array(2)?; + assert_eq!(result.logical_null_count(), 2); + } + Ok(()) + } + + #[test] + fn test_get_field_union_parent_nulls() -> Result<()> { + use arrow::array::{DictionaryArray, StringArray, UInt32Array}; + use arrow::datatypes::UInt32Type; + + let fields = UnionFields::try_new( + [3, 7], + [ + Field::new("int", DataType::Int32, true), + Field::new("string", DataType::Utf8, true), + ], + )?; + let ints = Int32Array::from(vec![Some(9), Some(1), Some(2), None, Some(4)]); + let strings = StringArray::from(vec!["x"; 5]); + + // Dense rows 1 and 2 share a value, but only row 2 has a null parent. + for offsets in [None, Some(vec![0, 1, 1, 3, 0].into())] { + let child = UnionArray::try_new( + fields.clone(), + vec![7, 3, 3, 3, 7].into(), + offsets, + vec![Arc::new(ints.clone()), Arc::new(strings.clone())], + )?; + let union_type = child.data_type().clone(); + let parent = StructArray::new( + vec![Field::new("u", union_type.clone(), true)].into(), + vec![Arc::new(child)], + Some(NullBuffer::from(vec![true, true, false, true, true])), + ); + let values = Arc::new(parent.slice(1, 4)) as ArrayRef; + let dictionary = Arc::new(DictionaryArray::::try_new( + UInt32Array::from(vec![0, 1, 2, 3]), + Arc::clone(&values), + )?) as ArrayRef; + + for input in [values, dictionary] { + let result = extract_single_field( + ColumnarValue::Array(input), + ScalarValue::Utf8(Some("u".to_string())), + )? + .into_array(4)?; + assert_eq!( + result.logical_nulls(), + Some(NullBuffer::from(vec![true, false, false, true])) + ); + let result = match result.data_type() { + DataType::Dictionary(_, _) => result.as_any_dictionary().values(), + _ => &result, + }; + assert_eq!(result.data_type(), &union_type); + result.to_data().validate_full()?; + let result = result.as_union(); + assert_eq!(result.value(0).as_ref(), &Int32Array::from(vec![1])); + assert_eq!(result.value(3).as_ref(), &StringArray::from(vec!["x"])); + } + } + Ok(()) + } + #[test] fn test_get_field_dict_encoded_struct() -> Result<()> { use arrow::array::{DictionaryArray, StringArray, UInt32Array}; diff --git a/datafusion/sqllogictest/src/test_context.rs b/datafusion/sqllogictest/src/test_context.rs index 87d39e34dc400..76dcb322a9b90 100644 --- a/datafusion/sqllogictest/src/test_context.rs +++ b/datafusion/sqllogictest/src/test_context.rs @@ -812,23 +812,25 @@ fn register_dictionary_struct_table(ctx: &SessionContext) { ctx.register_batch("dict_struct_table", batch).unwrap(); - // Second table: dictionary-encoded struct with nullable entries - let names_nullable = Arc::new(StringArray::from(vec!["X", "Y"])) as ArrayRef; - let ids_nullable = Arc::new(Int32Array::from(vec![10, 20])) as ArrayRef; + // Second table: null keys, null structs with valid children, and null children. + let names_nullable = + Arc::new(StringArray::from(vec!["X", "Y", "hidden"])) as ArrayRef; + let ids_nullable = + Arc::new(Int32Array::from(vec![Some(10), None, Some(30)])) as ArrayRef; let struct_fields_nullable: Fields = vec![ Field::new("name", DataType::Utf8, false), - Field::new("id", DataType::Int32, false), + Field::new("id", DataType::Int32, true), ] .into(); let values_struct_nullable = Arc::new( StructArray::try_new( struct_fields_nullable.clone(), vec![names_nullable, ids_nullable], - None, + Some(vec![true, true, false].into()), ) .unwrap(), ) as ArrayRef; - let keys_nullable = UInt32Array::from(vec![Some(0), None, Some(1), None]); + let keys_nullable = UInt32Array::from(vec![Some(0), None, Some(1), Some(2), Some(2)]); let dict_nullable = DictionaryArray::::try_new(keys_nullable, values_struct_nullable) .unwrap(); diff --git a/datafusion/sqllogictest/test_files/dictionary_struct.slt b/datafusion/sqllogictest/test_files/dictionary_struct.slt index d97cb5420e3d2..e2ddc8be1c88c 100644 --- a/datafusion/sqllogictest/test_files/dictionary_struct.slt +++ b/datafusion/sqllogictest/test_files/dictionary_struct.slt @@ -33,12 +33,14 @@ # {name: Bob, id: 2} # # dict_struct_nullable: -# ds Dictionary(UInt32, Struct(name: Utf8, id: Int32)) — 4 rows, keys [0, NULL, 1, NULL] +# ds Dictionary(UInt32, Struct(name: Utf8, id: Int32)) — 5 rows, keys [0, NULL, 1, 2, 2] +# Value 2 is a null struct with valid child values ("hidden", 30). # # Rows (logical values): # {name: X, id: 10} # NULL -# {name: Y, id: 20} +# {name: Y, id: NULL} +# NULL # NULL # Verify schema of dict_struct_table @@ -67,11 +69,11 @@ SELECT dict_struct['id'] FROM dict_struct_table; 1 2 -# Verify the extracted field preserves dictionary encoding -query T -SELECT arrow_typeof(dict_struct['name']) FROM dict_struct_table LIMIT 1; +# Verify the extracted field preserves dictionary encoding and nullability +query TBB +SELECT arrow_typeof(dict_struct['name']), arrow_field(dict_struct['name'])['nullable'], arrow_field(plain_struct['name'])['nullable'] FROM dict_struct_table LIMIT 1; ---- -Dictionary(UInt32, Utf8) +Dictionary(UInt32, Utf8) false false query T SELECT arrow_typeof(dict_struct['id']) FROM dict_struct_table LIMIT 1; @@ -107,21 +109,23 @@ Carol Alice Bob -# Field extraction from dict-encoded struct with NULLs -query T -SELECT ds['name'] FROM dict_struct_nullable; +# A nullable parent makes its non-nullable child nullable, preserving the dictionary type. +query TTB +SELECT ds['name'], arrow_typeof(ds['name']), arrow_field(ds['name'])['nullable'] FROM dict_struct_nullable; ---- -X -NULL -Y -NULL +X Dictionary(UInt32, Utf8) true +NULL Dictionary(UInt32, Utf8) true +Y Dictionary(UInt32, Utf8) true +NULL Dictionary(UInt32, Utf8) true +NULL Dictionary(UInt32, Utf8) true query ? SELECT ds['id'] FROM dict_struct_nullable; ---- 10 NULL -20 +NULL +NULL NULL # Filtering on extracted dict-encoded struct field diff --git a/datafusion/sqllogictest/test_files/map.slt b/datafusion/sqllogictest/test_files/map.slt index debaf7f8b344f..b69f04defc782 100644 --- a/datafusion/sqllogictest/test_files/map.slt +++ b/datafusion/sqllogictest/test_files/map.slt @@ -87,10 +87,11 @@ GET 27 PUT 25 DELETE 24 -query T -SELECT strings['not_found'] FROM data LIMIT 1; +# A missing key makes the result nullable even when the map and its values are not. +query TBB +SELECT strings['not_found'], arrow_field(strings)['nullable'], arrow_field(strings['not_found'])['nullable'] FROM data LIMIT 1; ---- -NULL +NULL false true # Select non existent key, expect NULL for each row query I diff --git a/datafusion/sqllogictest/test_files/struct.slt b/datafusion/sqllogictest/test_files/struct.slt index 183bfbd04e506..aad12353c2257 100644 --- a/datafusion/sqllogictest/test_files/struct.slt +++ b/datafusion/sqllogictest/test_files/struct.slt @@ -766,22 +766,23 @@ select get_field(CAST(NULL AS STRUCT(a STRUCT(b INT))), 'a', 'b'); ---- NULL -# Null handling: null in middle of chain +# Null parents at either level make a non-nullable child nullable. statement ok -create table null_mid_test (s STRUCT(a STRUCT(b INT))); - -statement ok -insert into null_mid_test values ({'a': NULL}); +create table null_mid_test as +select arrow_cast(NULL, 'Struct("a": non-null Struct("b": non-null Int32))') as null_base, + arrow_cast({'a': NULL}, 'Struct("a": Struct("b": non-null Int32))') as s; query I select s['a']['b'] from null_mid_test; ---- NULL -query I -select get_field(s, 'a', 'b') from null_mid_test; +query IBIB +select get_field(null_base, 'a', 'b'), arrow_field(null_base['a']['b'])['nullable'], + get_field(s, 'a', 'b'), arrow_field(s['a']['b'])['nullable'] +from null_mid_test; ---- -NULL +NULL true NULL true statement ok drop table null_mid_test;