fix: Handle nulls correctly when extracting nested arrays from nullable struct/map

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<ArrayRef> {
+    // 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<ColumnarValue> {
     let arrays = ColumnarValue::values_to_arrays(&[base])?;
@@ -204,9 +244,12 @@ fn extract_single_field(base: ColumnarValue, name: ScalarValue) -> Result<Column
                         "Field {field_name} not found in dictionary struct"
                     )
                 })?;
-            Ok(ColumnarValue::Array(
-                dict.with_values(Arc::clone(field_col)),
-            ))
+            let field_col = if let Some(parent_nulls) = values_struct.nulls() {
+                apply_parent_nulls(field_col, parent_nulls)?
+            } else {
+                Arc::clone(field_col)
+            };
+            Ok(ColumnarValue::Array(dict.with_values(field_col)))
         }
         (DataType::Map(_, _), key, _) => {
             // 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<Column
         }
         (DataType::Struct(_), _, Some(k)) => {
             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::<UInt32Type>::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::<UInt32Type>::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};

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::<UInt32Type>::try_new(keys_nullable, values_struct_nullable)
             .unwrap();

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

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