From ff0599a507387987b7556840425f8709092e7323 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Wed, 9 Sep 2026 16:03:54 +0100 Subject: [PATCH 1/4] fix getfield thing only for structs Signed-off-by: Adam Gutglick --- datafusion/functions/src/core/getfield.rs | 42 +++++++++++++++++++++-- 1 file changed, 39 insertions(+), 3 deletions(-) diff --git a/datafusion/functions/src/core/getfield.rs b/datafusion/functions/src/core/getfield.rs index a0f024bbc7ea4..b51ecc93d5257 100644 --- a/datafusion/functions/src/core/getfield.rs +++ b/datafusion/functions/src/core/getfield.rs @@ -21,6 +21,7 @@ use arrow::array::{ Array, Capacities, MutableArrayData, Scalar, cast::AsArray, make_array, make_comparator, }; +use arrow::buffer::NullBuffer; use arrow::compute::SortOptions; use arrow::datatypes::{DataType, Field, FieldRef}; @@ -220,9 +221,15 @@ 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)) => { + let nulls = NullBuffer::union(col.nulls(), Some(parent_nulls)); + let data = col.to_data().into_builder().nulls(nulls).build()?; + Ok(ColumnarValue::Array(make_array(data))) + } } } (DataType::Struct(_), name, _) => exec_err!( @@ -733,6 +740,35 @@ 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_dict_encoded_struct() -> Result<()> { use arrow::array::{DictionaryArray, StringArray, UInt32Array}; From 5b43f042e7612fb3e36303b23a653cf03a8b0bbd Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Wed, 9 Sep 2026 16:16:27 +0100 Subject: [PATCH 2/4] Map nulls Signed-off-by: Adam Gutglick --- datafusion/functions/src/core/getfield.rs | 257 +++++++++++++++++++++- 1 file changed, 254 insertions(+), 3 deletions(-) diff --git a/datafusion/functions/src/core/getfield.rs b/datafusion/functions/src/core/getfield.rs index b51ecc93d5257..8f08b56774db9 100644 --- a/datafusion/functions/src/core/getfield.rs +++ b/datafusion/functions/src/core/getfield.rs @@ -126,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; @@ -161,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; @@ -205,9 +213,16 @@ fn extract_single_field(base: ColumnarValue, name: ScalarValue) -> Result { // The lookup key is a single scalar. `eq` does not support nested @@ -226,6 +241,10 @@ fn extract_single_field(base: ColumnarValue, name: ScalarValue) -> Result exec_err!("Field {k} not found in struct"), (Some(col), None) => Ok(ColumnarValue::Array(Arc::clone(col))), (Some(col), Some(parent_nulls)) => { + // NullArray is already entirely null and cannot have a validity bitmap. + if col.data_type().is_null() { + return Ok(ColumnarValue::Array(Arc::clone(col))); + } let nulls = NullBuffer::union(col.nulls(), Some(parent_nulls)); let data = col.to_data().into_builder().nulls(nulls).build()?; Ok(ColumnarValue::Array(make_array(data))) @@ -769,6 +788,238 @@ mod tests { Ok(()) } + #[test] + fn test_get_field_map_parent_nulls() -> Result<()> { + use arrow::array::{Int32Builder, MapBuilder, StringBuilder}; + + let mut builder = + MapBuilder::new(None, StringBuilder::new(), Int32Builder::new()); + for (key, value, valid) in [ + ("key", Some(1), true), + ("key", Some(2), false), + ("key", Some(3), true), + ("key", None, true), + ("other", Some(4), true), + ] { + // The null map still has a matching key and a valid child value. + builder.keys().append_value(key); + builder.values().append_option(value); + builder.append(valid)?; + } + builder.append(true)?; // An empty map. + let map = builder.finish(); + let expected = Int32Array::from(vec![Some(1), None, Some(3), None, None, None]); + + for offset in [0, 1] { + let len = map.len() - offset; + let result = extract_single_field( + ColumnarValue::Array(Arc::new(map.slice(offset, len))), + ScalarValue::Utf8(Some("key".to_string())), + )? + .into_array(len)?; + assert_eq!(result.as_ref(), &expected.slice(offset, len) as &dyn Array); + } + Ok(()) + } + + #[test] + fn test_get_field_map_nested_keys_parent_nulls() -> Result<()> { + use arrow::array::{FixedSizeListArray, ListArray, MapArray}; + use arrow::datatypes::Int32Type; + use arrow_buffer::OffsetBuffer; + + let keys: Vec = vec![ + Arc::new(ListArray::from_iter_primitive::([Some( + vec![Some(7)], + )])), + Arc::new(FixedSizeListArray::new( + Arc::new(Field::new("item", DataType::Int32, false)), + 1, + Arc::new(Int32Array::from(vec![7])), + None, + )), + Arc::new(StructArray::new( + vec![Field::new("x", DataType::Int32, false)].into(), + vec![Arc::new(Int32Array::from(vec![7]))], + None, + )), + ]; + + for keys in 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![42]))], + None, + ); + let map = MapArray::new( + Arc::new(Field::new("entries", entries.data_type().clone(), false)), + OffsetBuffer::new(vec![0, 1].into()), + entries, + Some(NullBuffer::from(vec![false])), + false, + ); + let result = + extract_single_field(ColumnarValue::Array(Arc::new(map)), key.clone())? + .into_array(1)?; + let expected = Int32Array::from(vec![None]); + assert_eq!(result.as_ref(), &expected as &dyn Array, "key: {key}"); + } + Ok(()) + } + + #[test] + fn test_get_field_nested_dict_struct_parent_nulls() -> Result<()> { + use arrow::array::{DictionaryArray, UInt32Array}; + use arrow::datatypes::UInt32Type; + + let inner = StructArray::new( + vec![Field::new("value", DataType::Int32, true)].into(), + vec![Arc::new(Int32Array::from(vec![Some(10), Some(20), None]))], + None, + ); + let values = StructArray::new( + vec![Field::new("inner", inner.data_type().clone(), false)].into(), + vec![Arc::new(inner)], + Some(NullBuffer::from(vec![true, false, true])), + ); + let keys = UInt32Array::from(vec![Some(0), Some(1), None, Some(2), Some(1)]); + let dict = + DictionaryArray::::try_new(keys.clone(), Arc::new(values))?; + + let inner = extract_single_field( + ColumnarValue::Array(Arc::new(dict)), + ScalarValue::Utf8(Some("inner".to_string())), + )?; + let result = + extract_single_field(inner, ScalarValue::Utf8(Some("value".to_string())))? + .into_array(keys.len())?; + let result_dict = result + .as_any() + .downcast_ref::>() + .unwrap(); + assert_eq!(result_dict.keys(), &keys); + assert_eq!(result_dict.values().len(), 3); + + let decoded = arrow::compute::cast(&result, &DataType::Int32)?; + let expected = Int32Array::from(vec![Some(10), None, None, None, None]); + assert_eq!(decoded.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_nested_return_field_nullability() -> Result<()> { + let dictionary = |value_type| { + DataType::Dictionary(Box::new(DataType::UInt32), Box::new(value_type)) + }; + let map = |value_type, nullable| { + let entries = DataType::Struct( + vec![ + Field::new("key", DataType::Utf8, false), + Field::new("value", value_type, nullable), + ] + .into(), + ); + DataType::Map(Arc::new(Field::new("entries", entries, false)), false) + }; + let names = [ + ScalarValue::Utf8(Some("inner".to_string())), + ScalarValue::Utf8(Some("value".to_string())), + ]; + + for (outer_nullable, inner_nullable, leaf_nullable, expected_nullable) in [ + (false, false, false, false), + (true, false, false, true), + (false, true, false, true), + (false, false, true, true), + (true, true, true, true), + ] { + let inner_type = DataType::Struct( + vec![Field::new("value", DataType::Int32, leaf_nullable)].into(), + ); + let outer_type = DataType::Struct( + vec![Field::new("inner", inner_type.clone(), inner_nullable)].into(), + ); + let struct_with_dict = DataType::Struct( + vec![Field::new( + "inner", + dictionary(inner_type.clone()), + inner_nullable, + )] + .into(), + ); + let struct_with_map = DataType::Struct( + vec![Field::new( + "inner", + map(DataType::Int32, leaf_nullable), + inner_nullable, + )] + .into(), + ); + + for (data_type, expected_type, nullable) in [ + (outer_type.clone(), DataType::Int32, expected_nullable), + ( + dictionary(outer_type), + dictionary(DataType::Int32), + expected_nullable, + ), + ( + struct_with_dict, + dictionary(DataType::Int32), + expected_nullable, + ), + // Map lookups can return null even when every field is non-nullable. + (map(inner_type, inner_nullable), DataType::Int32, true), + (struct_with_map, DataType::Int32, true), + ] { + let arg_fields = [ + Arc::new(Field::new("outer", data_type, outer_nullable)), + Arc::new(Field::new("key", DataType::Utf8, false)), + Arc::new(Field::new("key", DataType::Utf8, false)), + ]; + let result = + GetFieldFunc::new().return_field_from_args(ReturnFieldArgs { + arg_fields: &arg_fields, + scalar_arguments: &[None, Some(&names[0]), Some(&names[1])], + })?; + assert_eq!(result.data_type(), &expected_type, "{arg_fields:?}"); + assert_eq!(result.is_nullable(), nullable, "{arg_fields:?}"); + } + } + Ok(()) + } + #[test] fn test_get_field_dict_encoded_struct() -> Result<()> { use arrow::array::{DictionaryArray, StringArray, UInt32Array}; From 87a3975780318c83bfe4a12a5c416f2b293a2d7f Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Wed, 9 Sep 2026 16:44:14 +0100 Subject: [PATCH 3/4] different tests Signed-off-by: Adam Gutglick --- datafusion/functions/src/core/getfield.rs | 200 ++---------------- datafusion/sqllogictest/src/test_context.rs | 14 +- .../test_files/dictionary_struct.slt | 32 +-- datafusion/sqllogictest/test_files/map.slt | 7 +- datafusion/sqllogictest/test_files/struct.slt | 17 +- 5 files changed, 56 insertions(+), 214 deletions(-) diff --git a/datafusion/functions/src/core/getfield.rs b/datafusion/functions/src/core/getfield.rs index 8f08b56774db9..d505ec892174b 100644 --- a/datafusion/functions/src/core/getfield.rs +++ b/datafusion/functions/src/core/getfield.rs @@ -790,62 +790,19 @@ mod tests { #[test] fn test_get_field_map_parent_nulls() -> Result<()> { - use arrow::array::{Int32Builder, MapBuilder, StringBuilder}; - - let mut builder = - MapBuilder::new(None, StringBuilder::new(), Int32Builder::new()); - for (key, value, valid) in [ - ("key", Some(1), true), - ("key", Some(2), false), - ("key", Some(3), true), - ("key", None, true), - ("other", Some(4), true), - ] { - // The null map still has a matching key and a valid child value. - builder.keys().append_value(key); - builder.values().append_option(value); - builder.append(valid)?; - } - builder.append(true)?; // An empty map. - let map = builder.finish(); - let expected = Int32Array::from(vec![Some(1), None, Some(3), None, None, None]); - - for offset in [0, 1] { - let len = map.len() - offset; - let result = extract_single_field( - ColumnarValue::Array(Arc::new(map.slice(offset, len))), - ScalarValue::Utf8(Some("key".to_string())), - )? - .into_array(len)?; - assert_eq!(result.as_ref(), &expected.slice(offset, len) as &dyn Array); - } - Ok(()) - } - - #[test] - fn test_get_field_map_nested_keys_parent_nulls() -> Result<()> { - use arrow::array::{FixedSizeListArray, ListArray, MapArray}; - use arrow::datatypes::Int32Type; + use arrow::array::{FixedSizeListArray, MapArray}; use arrow_buffer::OffsetBuffer; - let keys: Vec = vec![ - Arc::new(ListArray::from_iter_primitive::([Some( - vec![Some(7)], - )])), - Arc::new(FixedSizeListArray::new( - Arc::new(Field::new("item", DataType::Int32, false)), - 1, - Arc::new(Int32Array::from(vec![7])), - None, - )), - Arc::new(StructArray::new( - vec![Field::new("x", DataType::Int32, false)].into(), - vec![Arc::new(Int32Array::from(vec![7]))], - None, - )), - ]; + 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; - for keys in keys { + // 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![ @@ -853,64 +810,24 @@ mod tests { Field::new("value", DataType::Int32, false), ] .into(), - vec![keys, Arc::new(Int32Array::from(vec![42]))], + 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].into()), + OffsetBuffer::new(vec![0, 1, 2, 3].into()), entries, - Some(NullBuffer::from(vec![false])), + Some(NullBuffer::from(vec![true, false, true])), false, ); - let result = - extract_single_field(ColumnarValue::Array(Arc::new(map)), key.clone())? - .into_array(1)?; - let expected = Int32Array::from(vec![None]); - assert_eq!(result.as_ref(), &expected as &dyn Array, "key: {key}"); + 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_nested_dict_struct_parent_nulls() -> Result<()> { - use arrow::array::{DictionaryArray, UInt32Array}; - use arrow::datatypes::UInt32Type; - - let inner = StructArray::new( - vec![Field::new("value", DataType::Int32, true)].into(), - vec![Arc::new(Int32Array::from(vec![Some(10), Some(20), None]))], - None, - ); - let values = StructArray::new( - vec![Field::new("inner", inner.data_type().clone(), false)].into(), - vec![Arc::new(inner)], - Some(NullBuffer::from(vec![true, false, true])), - ); - let keys = UInt32Array::from(vec![Some(0), Some(1), None, Some(2), Some(1)]); - let dict = - DictionaryArray::::try_new(keys.clone(), Arc::new(values))?; - - let inner = extract_single_field( - ColumnarValue::Array(Arc::new(dict)), - ScalarValue::Utf8(Some("inner".to_string())), - )?; - let result = - extract_single_field(inner, ScalarValue::Utf8(Some("value".to_string())))? - .into_array(keys.len())?; - let result_dict = result - .as_any() - .downcast_ref::>() - .unwrap(); - assert_eq!(result_dict.keys(), &keys); - assert_eq!(result_dict.values().len(), 3); - - let decoded = arrow::compute::cast(&result, &DataType::Int32)?; - let expected = Int32Array::from(vec![Some(10), None, None, None, None]); - assert_eq!(decoded.as_ref(), &expected as &dyn Array); - Ok(()) - } - #[test] fn test_get_field_null_typed_child() -> Result<()> { use arrow::array::{DictionaryArray, NullArray, UInt32Array}; @@ -937,89 +854,6 @@ mod tests { Ok(()) } - #[test] - fn test_get_field_nested_return_field_nullability() -> Result<()> { - let dictionary = |value_type| { - DataType::Dictionary(Box::new(DataType::UInt32), Box::new(value_type)) - }; - let map = |value_type, nullable| { - let entries = DataType::Struct( - vec![ - Field::new("key", DataType::Utf8, false), - Field::new("value", value_type, nullable), - ] - .into(), - ); - DataType::Map(Arc::new(Field::new("entries", entries, false)), false) - }; - let names = [ - ScalarValue::Utf8(Some("inner".to_string())), - ScalarValue::Utf8(Some("value".to_string())), - ]; - - for (outer_nullable, inner_nullable, leaf_nullable, expected_nullable) in [ - (false, false, false, false), - (true, false, false, true), - (false, true, false, true), - (false, false, true, true), - (true, true, true, true), - ] { - let inner_type = DataType::Struct( - vec![Field::new("value", DataType::Int32, leaf_nullable)].into(), - ); - let outer_type = DataType::Struct( - vec![Field::new("inner", inner_type.clone(), inner_nullable)].into(), - ); - let struct_with_dict = DataType::Struct( - vec![Field::new( - "inner", - dictionary(inner_type.clone()), - inner_nullable, - )] - .into(), - ); - let struct_with_map = DataType::Struct( - vec![Field::new( - "inner", - map(DataType::Int32, leaf_nullable), - inner_nullable, - )] - .into(), - ); - - for (data_type, expected_type, nullable) in [ - (outer_type.clone(), DataType::Int32, expected_nullable), - ( - dictionary(outer_type), - dictionary(DataType::Int32), - expected_nullable, - ), - ( - struct_with_dict, - dictionary(DataType::Int32), - expected_nullable, - ), - // Map lookups can return null even when every field is non-nullable. - (map(inner_type, inner_nullable), DataType::Int32, true), - (struct_with_map, DataType::Int32, true), - ] { - let arg_fields = [ - Arc::new(Field::new("outer", data_type, outer_nullable)), - Arc::new(Field::new("key", DataType::Utf8, false)), - Arc::new(Field::new("key", DataType::Utf8, false)), - ]; - let result = - GetFieldFunc::new().return_field_from_args(ReturnFieldArgs { - arg_fields: &arg_fields, - scalar_arguments: &[None, Some(&names[0]), Some(&names[1])], - })?; - assert_eq!(result.data_type(), &expected_type, "{arg_fields:?}"); - assert_eq!(result.is_nullable(), nullable, "{arg_fields:?}"); - } - } - 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; From 9591266d5573a2f9d1d8fbe30962003f9f0b8b52 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Fri, 11 Sep 2026 11:32:23 +0100 Subject: [PATCH 4/4] Handle union and refactor Signed-off-by: Adam Gutglick --- datafusion/functions/src/core/getfield.rs | 112 ++++++++++++++++++---- 1 file changed, 96 insertions(+), 16 deletions(-) diff --git a/datafusion/functions/src/core/getfield.rs b/datafusion/functions/src/core/getfield.rs index d505ec892174b..c7003a372eeaf 100644 --- a/datafusion/functions/src/core/getfield.rs +++ b/datafusion/functions/src/core/getfield.rs @@ -18,7 +18,7 @@ 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; @@ -191,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])?; @@ -213,12 +244,8 @@ fn extract_single_field(base: ColumnarValue, name: ScalarValue) -> Result Result exec_err!("Field {k} not found in struct"), (Some(col), None) => Ok(ColumnarValue::Array(Arc::clone(col))), (Some(col), Some(parent_nulls)) => { - // NullArray is already entirely null and cannot have a validity bitmap. - if col.data_type().is_null() { - return Ok(ColumnarValue::Array(Arc::clone(col))); - } - let nulls = NullBuffer::union(col.nulls(), Some(parent_nulls)); - let data = col.to_data().into_builder().nulls(nulls).build()?; - Ok(ColumnarValue::Array(make_array(data))) + Ok(ColumnarValue::Array(apply_parent_nulls(col, parent_nulls)?)) } } } @@ -684,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<()> { @@ -854,6 +875,65 @@ mod tests { 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};