fix: preserve nulls when unnesting structs

datafusion/common/src/utils/mod.rs

@@ -35,13 +35,13 @@ use arrow::array::{
 };
 use arrow::array::{
     ArrowPrimitiveType, BooleanArray, Datum, GenericListArray, Int32Array, Int64Array,
-    MutableArrayData, PrimitiveArray, make_array,
+    MutableArrayData, PrimitiveArray, layout, make_array,
 };
 use arrow::array::{LargeListViewArray, ListViewArray};
-use arrow::buffer::{OffsetBuffer, ScalarBuffer};
+use arrow::buffer::{NullBuffer, OffsetBuffer, ScalarBuffer};
 use arrow::compute::kernels::cmp::eq;
 use arrow::compute::kernels::length::length;
-use arrow::compute::{SortColumn, SortOptions, partition};
+use arrow::compute::{SortColumn, SortOptions, nullif, partition};
 use arrow::datatypes::{
     ArrowNativeType, DataType, Field, Int32Type, Int64Type, SchemaRef,
 };
@@ -1521,20 +1521,100 @@ pub fn normalize_float_zero_scalar(scalar: ScalarValue) -> ScalarValue {
     }
 }
 
+/// Apply a struct's nulls to one of its fields.
+///
+/// Arrays with validity bitmaps share their value buffers; unions and
+/// run-end encoded arrays are rebuilt.
+pub fn apply_parent_nulls(
+    col: &ArrayRef,
+    parent_nulls: Option<&NullBuffer>,
+) -> Result<ArrayRef> {
+    let Some(parent_nulls) = parent_nulls else {
+        // If there are no parent nulls to apply, we can just return
+        return Ok(Arc::clone(col));
+    };
+
+    // 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));
+    }
+
+    if layout(col.data_type()).can_contain_null_mask {
+        // `nullif` marks a row null where the mask is true and keeps the
+        // field's own nulls. Only the validity bitmap is rebuilt; the value
+        // buffers and child arrays are shared with `col`.
+        let null_parents = BooleanArray::new(!parent_nulls.inner(), None);
+        return Ok(nullif(col.as_ref(), &null_parents)?);
+    }
+
+    // Unions and run-end encoded arrays have no validity bitmap of their own
+    // and represent nulls in their children. Rebuild the array so null parents
+    // become null values in those children.
+    let data = col.to_data();
+    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()))
+}
+
 #[cfg(test)]
 mod tests {
     use std::sync::Arc;
 
     use super::*;
     use crate::ScalarValue::Null;
     use arrow::{
-        array::{Float64Array, Int32Array},
-        buffer::NullBuffer,
+        array::{Float64Array, Int32Array, NullArray},
         datatypes::Int32Type,
     };
     #[cfg(feature = "sql")]
     use sqlparser::ast::Ident;
 
+    #[test]
+    fn test_apply_parent_nulls_sliced() -> Result<()> {
+        let child = Arc::new(
+            Int32Array::from(vec![Some(1), None, Some(3), Some(4), Some(5)]).slice(1, 3),
+        ) as ArrayRef;
+        let parent_nulls =
+            NullBuffer::from(vec![true, true, true, false, true]).slice(1, 3);
+        let result = apply_parent_nulls(&child, Some(&parent_nulls))?;
+        assert_eq!(
+            result.as_ref(),
+            &Int32Array::from(vec![None, Some(3), None])
+        );
+        assert!(result.to_data().buffers()[0].ptr_eq(&child.to_data().buffers()[0]));
+        Ok(())
+    }
+
+    #[test]
+    fn test_apply_parent_nulls_noop() -> Result<()> {
+        let child = Arc::new(Int32Array::from(vec![Some(1), None, Some(3)])) as ArrayRef;
+        for parent_nulls in [None, Some(NullBuffer::new_valid(3))] {
+            assert!(Arc::ptr_eq(
+                &apply_parent_nulls(&child, parent_nulls.as_ref())?,
+                &child,
+            ));
+        }
+        let child = Arc::new(NullArray::new(3)) as ArrayRef;
+        assert!(Arc::ptr_eq(
+            &apply_parent_nulls(&child, Some(&NullBuffer::new_null(3)))?,
+            &child,
+        ));
+        let empty = child.slice(0, 0);
+        assert!(Arc::ptr_eq(
+            &apply_parent_nulls(&empty, Some(&NullBuffer::new_valid(0)))?,
+            &empty,
+        ));
+        Ok(())
+    }
+
     #[test]
     fn test_offset_span() {
         let offsets = OffsetBuffer::new(vec![0_i32, 5, 8, 8, 12].into());

datafusion/expr/src/logical_plan/plan.rs

@@ -5125,7 +5125,7 @@ impl Unnest {
                                 ));
                                 Ok(get_unnested_columns(
                                     &r.output_column.name,
-                                    original_field.data_type(),
+                                    original_field,
                                     r.depth,
                                 )?
                                 .into_iter()
@@ -5136,7 +5136,7 @@ impl Unnest {
                         if transformed_columns.is_empty() {
                             transformed_columns = get_unnested_columns(
                                 &column_to_unnest.name,
-                                original_field.data_type(),
+                                original_field,
                                 1,
                             )?;
                             match original_field.data_type() {
@@ -5216,9 +5216,10 @@ impl Unnest {
 // the recursion level
 fn get_unnested_columns(
     col_name: &String,
-    data_type: &DataType,
+    field: &Field,
     depth: usize,
 ) -> Result<Vec<(Column, Arc<Field>)>> {
+    let data_type = field.data_type();
     let mut qualified_columns = Vec::with_capacity(1);
 
     match data_type {
@@ -5242,7 +5243,11 @@ fn get_unnested_columns(
             qualified_columns.extend(fields.iter().map(|f| {
                 let new_name = format!("{}.{}", col_name, f.name());
                 let column = Column::from_name(&new_name);
-                let new_field = f.as_ref().clone().with_name(new_name);
+                let new_field = f
+                    .as_ref()
+                    .clone()
+                    .with_name(new_name)
+                    .with_nullable(field.is_nullable() || f.is_nullable());
                 // let column = Column::from((None, &f));
                 (column, Arc::new(new_field))
             }))

datafusion/functions/src/core/getfield.rs

@@ -17,14 +17,12 @@
 
 use std::sync::{Arc, OnceLock};
 
-use arrow::array::{
-    Array, ArrayRef, BooleanArray, MutableArrayData, cast::AsArray, layout, make_array,
-};
-use arrow::buffer::NullBuffer;
-use arrow::compute::{nullif, take};
+use arrow::array::{Array, cast::AsArray};
+use arrow::compute::take;
 use arrow::datatypes::{DataType, Field, FieldRef};
 
 use datafusion_common::cast::{as_map_array, as_struct_array};
+use datafusion_common::utils::apply_parent_nulls;
 use datafusion_common::{
     Result, ScalarValue, exec_datafusion_err, exec_err, internal_err, plan_datafusion_err,
 };
@@ -100,46 +98,6 @@ impl Default for GetFieldFunc {
     }
 }
 
-/// Apply a struct's nulls to one of its fields.
-fn apply_parent_nulls(
-    col: &ArrayRef,
-    parent_nulls: Option<&NullBuffer>,
-) -> Result<ArrayRef> {
-    let Some(parent_nulls) = parent_nulls else {
-        // If there are no parent nulls to apply, we can just return
-        return Ok(Arc::clone(col));
-    };
-
-    // 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));
-    }
-
-    if layout(col.data_type()).can_contain_null_mask {
-        // `nullif` marks a row null where the mask is true and keeps the
-        // field's own nulls. Only the validity bitmap is rebuilt; the value
-        // buffers and child arrays are shared with `col`.
-        let null_parents = BooleanArray::new(!parent_nulls.inner(), None);
-        return Ok(nullif(col.as_ref(), &null_parents)?);
-    }
-
-    // Unions and run-end encoded arrays have no validity bitmap of their own
-    // and represent nulls in their children. Rebuild the array so null parents
-    // become null values in those children.
-    let data = col.to_data();
-    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()))
-}
-
 /// 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])?;
@@ -617,6 +575,7 @@ mod tests {
         ArrayRef, Int32Array, Int32Builder, ListArray, ListBuilder, MapBuilder, RunArray,
         StructArray, UnionArray,
     };
+    use arrow::buffer::NullBuffer;
     use arrow::datatypes::{Fields, Int32Type, UnionFields};
 
     #[test]

datafusion/physical-plan/Cargo.toml

@@ -193,3 +193,7 @@ name = "window_filter"
 [[bench]]
 harness = false
 name = "range_repartition"
+
+[[bench]]
+harness = false
+name = "unnest"

datafusion/physical-plan/benches/unnest.rs

@@ -0,0 +1,108 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::hint::black_box;
+use std::sync::Arc;
+
+use arrow::array::{ArrayRef, Int64Array, StringArray, StructArray};
+use arrow::buffer::NullBuffer;
+use arrow::datatypes::{Field, Schema};
+use arrow::record_batch::RecordBatch;
+use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main};
+use datafusion_common::UnnestOptions;
+use datafusion_execution::{TaskContext, config::SessionConfig};
+use datafusion_physical_plan::test::TestMemoryExec;
+use datafusion_physical_plan::unnest::UnnestExec;
+use datafusion_physical_plan::{ExecutionPlan, collect};
+use rand::SeedableRng;
+use rand::rngs::StdRng;
+use rand::seq::SliceRandom;
+
+fn unnest_benchmark(c: &mut Criterion) {
+    let runtime = tokio::runtime::Builder::new_current_thread()
+        .enable_all()
+        .build()
+        .unwrap();
+    let mut group = c.benchmark_group("unnest_struct");
+    for num_rows in [1024, 8192, 65536] {
+        let integers = Arc::new(Int64Array::from_iter(
+            (0..num_rows).map(|i| (i % 10 != 0).then_some(i as i64)),
+        )) as ArrayRef;
+        let strings = Arc::new(StringArray::from_iter(
+            (0..num_rows).map(|i| (i % 10 != 0).then(|| format!("value-{i:0122}"))),
+        )) as ArrayRef;
+        let context = Arc::new(
+            TaskContext::default()
+                .with_session_config(SessionConfig::new().with_batch_size(num_rows)),
+        );
+        group.throughput(Throughput::Elements(num_rows as u64));
+        for (name, child) in [("int64", integers), ("utf8_128_bytes", strings)] {
+            for null_percent in [0, 1, 50, 100] {
+                let nulls = (null_percent != 0).then(|| {
+                    let mut validity = vec![true; num_rows];
+                    validity[..num_rows * null_percent / 100].fill(false);
+                    validity.shuffle(&mut StdRng::seed_from_u64(42));
+                    NullBuffer::from(validity)
+                });
+                let fields = vec![Field::new("value", child.data_type().clone(), true)];
+                let parent = StructArray::new(
+                    fields.clone().into(),
+                    vec![Arc::clone(&child)],
+                    nulls,
+                );
+                let batch =
+                    RecordBatch::try_from_iter(vec![("s", Arc::new(parent) as ArrayRef)])
+                        .unwrap();
+                let source = TestMemoryExec::try_new_exec(
+                    &[vec![batch.clone()]],
+                    batch.schema(),
+                    None,
+                )
+                .unwrap();
+                let plan: Arc<dyn ExecutionPlan> = Arc::new(
+                    UnnestExec::new(
+                        source,
+                        vec![],
+                        vec![0],
+                        Arc::new(Schema::new(fields)),
+                        UnnestOptions::default(),
+                    )
+                    .unwrap(),
+                );
+                group.bench_function(
+                    BenchmarkId::new(format!("{name}_nulls_{null_percent}"), num_rows),
+                    |b| {
+                        b.iter(|| {
+                            black_box(
+                                runtime
+                                    .block_on(collect(
+                                        Arc::clone(&plan),
+                                        Arc::clone(&context),
+                                    ))
+                                    .unwrap(),
+                            )
+                        });
+                    },
+                );
+            }
+        }
+    }
+    group.finish();
+}
+
+criterion_group!(benches, unnest_benchmark);
+criterion_main!(benches);