Skip to content

Commit 1871883

Browse files
authored
fix: preserve nulls when unnesting structs (#25389)
## Which issue does this PR close? Closes #25293. ## Rationale for this change Unnesting a NULL struct can expose its stored child values. A non-nullable child field can also make a valid query fail. Both cases now produce NULL values. ## What changes are included in this PR? Output fields include the parent's nullability. UNNEST uses the parent-null helper shared with field extraction, including its null-free fast path. The helper retains value buffers where the array type supports a validity bitmap and keeps the union/run-end fallback. The operator benchmark covers integer and 128-byte string fields across batch sizes and null densities. ## What is the testing strategy for this PR? ### Testing Done SQL regressions cover schema nullability. One physical-operator regression checks values and buffer identity together across null-free, mixed-null and all-null parents, sliced and nested arrays, and encoded children. Run from the repository root: ```sh cargo run --locked --profile ci -p datafusion-cli -- --quiet --format json \ -c "CREATE TABLE base AS SELECT named_struct('a', 1, 'b', 'x') AS s; CREATE TABLE masked AS SELECT nullif(s, s) AS s FROM base; SELECT unnest(s) FROM masked;" cargo run --locked --profile ci -p datafusion-cli -- --quiet --format json \ -c "SELECT unnest(arrow_cast(NULL, 'Struct(\"a\": non-null Int32)'));" ``` The JSON writer omits NULL-valued fields. Before the fix, the first query exposes values and the second fails: ```text [{"masked.s.a":1,"masked.s.b":"x"}] Error: Arrow error: Invalid argument error: Column '__unnest_placeholder(arrow_cast(NULL,Utf8("Struct("a": non-null Int32)"))).a' is declared as non-nullable but contains null values ``` With the fix, both queries return: ```text [{}] [{}] ``` The SQL rerun used a cached CLI built from source matching this PR. The full local workspace run stopped during linking with `ld: write() failed, errno=28 (No space left on device)`, before tests started. #### Operator measurement The recorded comparison of [the shared-helper implementation](9fffb0a) against [the main control](d477a2b) used one ARM64 host and `release-nonlto`. For 65,536 rows of 128-byte strings with 50% null parents, the control took about 3.05 us and the helper took 4.15 us. Main omits the parent mask and returns incorrect null results, so this measures the added correctness cost, not a speedup. <details> <summary>Benchmark commands and raw output</summary> From this PR's checkout, prepare the main control with only the benchmark and its manifest entry: ```bash root=$(mktemp -d) before=d477a2b038730647bef258cb98e58978b0c74699 after=9fffb0a72273cef82b7c48091a257ac1bf77684c bench=datafusion/physical-plan/benches/unnest.rs manifest=datafusion/physical-plan/Cargo.toml git worktree add --detach "$root/before" "$before" git worktree add --detach "$root/after" "$after" git diff "$before" "$after" -- "$manifest" "$bench" | git -C "$root/before" apply export CARGO_BUILD_JOBS=2 CARGO_TARGET_DIR="$root/target" ``` The measured commands use identical inputs and options: ```bash (cd "$root/before" && cargo bench --locked --profile release-nonlto -p datafusion-physical-plan --bench unnest -- \ 'utf8_128_bytes_nulls_50/65536' --noplot --sample-size 30 \ --warm-up-time 1 --measurement-time 2 --save-baseline current-main) (cd "$root/after" && cargo bench --locked --profile release-nonlto -p datafusion-physical-plan --bench unnest -- \ 'utf8_128_bytes_nulls_50/65536' --noplot --sample-size 30 \ --warm-up-time 1 --measurement-time 2 --baseline current-main) ``` Main control: ```text unnest_struct/utf8_128_bytes_nulls_50/65536 time: [3.0444 µs 3.0483 µs 3.0536 µs] thrpt: [21.462 Gelem/s 21.499 Gelem/s 21.527 Gelem/s] ``` Candidate: ```text unnest_struct/utf8_128_bytes_nulls_50/65536 time: [4.1454 µs 4.1524 µs 4.1638 µs] thrpt: [15.740 Gelem/s 15.783 Gelem/s 15.809 Gelem/s] change: time: [+36.598% +37.187% +37.828%] (p = 0.00 < 0.05) thrpt: [−27.446% −27.107% −26.792%] Performance has regressed. ``` These are operator measurements, not end-to-end query timings. No null-free performance claim is made. </details> ## Are there any user-facing changes? NULL structs now unnest into NULL fields rather than stored values or an error. --------- Signed-off-by: 1fanwang <1fannnw@gmail.com>
1 parent 49c0aa3 commit 1871883

7 files changed

Lines changed: 527 additions & 57 deletions

File tree

‎datafusion/common/src/utils/mod.rs‎

Lines changed: 85 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -35,13 +35,13 @@ use arrow::array::{
3535
};
3636
use arrow::array::{
3737
ArrowPrimitiveType, BooleanArray, Datum, GenericListArray, Int32Array, Int64Array,
38-
MutableArrayData, PrimitiveArray, make_array,
38+
MutableArrayData, PrimitiveArray, layout, make_array,
3939
};
4040
use arrow::array::{LargeListViewArray, ListViewArray};
41-
use arrow::buffer::{OffsetBuffer, ScalarBuffer};
41+
use arrow::buffer::{NullBuffer, OffsetBuffer, ScalarBuffer};
4242
use arrow::compute::kernels::cmp::eq;
4343
use arrow::compute::kernels::length::length;
44-
use arrow::compute::{SortColumn, SortOptions, partition};
44+
use arrow::compute::{SortColumn, SortOptions, nullif, partition};
4545
use arrow::datatypes::{
4646
ArrowNativeType, DataType, Field, Int32Type, Int64Type, SchemaRef,
4747
};
@@ -1561,20 +1561,100 @@ pub fn normalize_float_zero_scalar(scalar: ScalarValue) -> ScalarValue {
15611561
}
15621562
}
15631563

1564+
/// Apply a struct's nulls to one of its fields.
1565+
///
1566+
/// Arrays with validity bitmaps share their value buffers; unions and
1567+
/// run-end encoded arrays are rebuilt.
1568+
pub fn apply_parent_nulls(
1569+
col: &ArrayRef,
1570+
parent_nulls: Option<&NullBuffer>,
1571+
) -> Result<ArrayRef> {
1572+
let Some(parent_nulls) = parent_nulls else {
1573+
// If there are no parent nulls to apply, we can just return
1574+
return Ok(Arc::clone(col));
1575+
};
1576+
1577+
// NullArray is already entirely null and cannot have a validity bitmap.
1578+
// If we have 0 parent nulls, we can also avoid extra work.
1579+
if col.data_type().is_null() || parent_nulls.null_count() == 0 {
1580+
return Ok(Arc::clone(col));
1581+
}
1582+
1583+
if layout(col.data_type()).can_contain_null_mask {
1584+
// `nullif` marks a row null where the mask is true and keeps the
1585+
// field's own nulls. Only the validity bitmap is rebuilt; the value
1586+
// buffers and child arrays are shared with `col`.
1587+
let null_parents = BooleanArray::new(!parent_nulls.inner(), None);
1588+
return Ok(nullif(col.as_ref(), &null_parents)?);
1589+
}
1590+
1591+
// Unions and run-end encoded arrays have no validity bitmap of their own
1592+
// and represent nulls in their children. Rebuild the array so null parents
1593+
// become null values in those children.
1594+
let data = col.to_data();
1595+
let mut mutable = MutableArrayData::new(vec![&data], true, data.len());
1596+
let mut end = 0;
1597+
for (start, valid_end) in parent_nulls.valid_slices() {
1598+
mutable.try_extend_nulls(start - end)?;
1599+
mutable.try_extend(0, start, valid_end)?;
1600+
end = valid_end;
1601+
}
1602+
mutable.try_extend_nulls(data.len() - end)?;
1603+
1604+
Ok(make_array(mutable.freeze()))
1605+
}
1606+
15641607
#[cfg(test)]
15651608
mod tests {
15661609
use std::sync::Arc;
15671610

15681611
use super::*;
15691612
use crate::ScalarValue::Null;
15701613
use arrow::{
1571-
array::{Float64Array, Int32Array},
1572-
buffer::NullBuffer,
1614+
array::{Float64Array, Int32Array, NullArray},
15731615
datatypes::Int32Type,
15741616
};
15751617
#[cfg(feature = "sql")]
15761618
use sqlparser::ast::Ident;
15771619

1620+
#[test]
1621+
fn test_apply_parent_nulls_sliced() -> Result<()> {
1622+
let child = Arc::new(
1623+
Int32Array::from(vec![Some(1), None, Some(3), Some(4), Some(5)]).slice(1, 3),
1624+
) as ArrayRef;
1625+
let parent_nulls =
1626+
NullBuffer::from(vec![true, true, true, false, true]).slice(1, 3);
1627+
let result = apply_parent_nulls(&child, Some(&parent_nulls))?;
1628+
assert_eq!(
1629+
result.as_ref(),
1630+
&Int32Array::from(vec![None, Some(3), None])
1631+
);
1632+
assert!(result.to_data().buffers()[0].ptr_eq(&child.to_data().buffers()[0]));
1633+
Ok(())
1634+
}
1635+
1636+
#[test]
1637+
fn test_apply_parent_nulls_noop() -> Result<()> {
1638+
let child = Arc::new(Int32Array::from(vec![Some(1), None, Some(3)])) as ArrayRef;
1639+
for parent_nulls in [None, Some(NullBuffer::new_valid(3))] {
1640+
assert!(Arc::ptr_eq(
1641+
&apply_parent_nulls(&child, parent_nulls.as_ref())?,
1642+
&child,
1643+
));
1644+
}
1645+
let child = Arc::new(NullArray::new(3)) as ArrayRef;
1646+
assert!(Arc::ptr_eq(
1647+
&apply_parent_nulls(&child, Some(&NullBuffer::new_null(3)))?,
1648+
&child,
1649+
));
1650+
let empty = child.slice(0, 0);
1651+
assert!(Arc::ptr_eq(
1652+
&apply_parent_nulls(&empty, Some(&NullBuffer::new_valid(0)))?,
1653+
&empty,
1654+
));
1655+
Ok(())
1656+
}
1657+
15781658
#[test]
15791659
fn test_offset_span() {
15801660
let offsets = OffsetBuffer::new(vec![0_i32, 5, 8, 8, 12].into());

‎datafusion/expr/src/logical_plan/plan.rs‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5125,7 +5125,7 @@ impl Unnest {
51255125
));
51265126
Ok(get_unnested_columns(
51275127
&r.output_column.name,
5128-
original_field.data_type(),
5128+
original_field,
51295129
r.depth,
51305130
)?
51315131
.into_iter()
@@ -5136,7 +5136,7 @@ impl Unnest {
51365136
if transformed_columns.is_empty() {
51375137
transformed_columns = get_unnested_columns(
51385138
&column_to_unnest.name,
5139-
original_field.data_type(),
5139+
original_field,
51405140
1,
51415141
)?;
51425142
match original_field.data_type() {
@@ -5216,9 +5216,10 @@ impl Unnest {
52165216
// the recursion level
52175217
fn get_unnested_columns(
52185218
col_name: &String,
5219-
data_type: &DataType,
5219+
field: &Field,
52205220
depth: usize,
52215221
) -> Result<Vec<(Column, Arc<Field>)>> {
5222+
let data_type = field.data_type();
52225223
let mut qualified_columns = Vec::with_capacity(1);
52235224

52245225
match data_type {
@@ -5242,7 +5243,11 @@ fn get_unnested_columns(
52425243
qualified_columns.extend(fields.iter().map(|f| {
52435244
let new_name = format!("{}.{}", col_name, f.name());
52445245
let column = Column::from_name(&new_name);
5245-
let new_field = f.as_ref().clone().with_name(new_name);
5246+
let new_field = f
5247+
.as_ref()
5248+
.clone()
5249+
.with_name(new_name)
5250+
.with_nullable(field.is_nullable() || f.is_nullable());
52465251
// let column = Column::from((None, &f));
52475252
(column, Arc::new(new_field))
52485253
}))

‎datafusion/functions/src/core/getfield.rs‎

Lines changed: 4 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -17,14 +17,12 @@
1717

1818
use std::sync::{Arc, OnceLock};
1919

20-
use arrow::array::{
21-
Array, ArrayRef, BooleanArray, MutableArrayData, cast::AsArray, layout, make_array,
22-
};
23-
use arrow::buffer::NullBuffer;
24-
use arrow::compute::{nullif, take};
20+
use arrow::array::{Array, cast::AsArray};
21+
use arrow::compute::take;
2522
use arrow::datatypes::{DataType, Field, FieldRef};
2623

2724
use datafusion_common::cast::{as_map_array, as_struct_array};
25+
use datafusion_common::utils::apply_parent_nulls;
2826
use datafusion_common::{
2927
Result, ScalarValue, exec_datafusion_err, exec_err, internal_err, plan_datafusion_err,
3028
};
@@ -100,46 +98,6 @@ impl Default for GetFieldFunc {
10098
}
10199
}
102100

103-
/// Apply a struct's nulls to one of its fields.
104-
fn apply_parent_nulls(
105-
col: &ArrayRef,
106-
parent_nulls: Option<&NullBuffer>,
107-
) -> Result<ArrayRef> {
108-
let Some(parent_nulls) = parent_nulls else {
109-
// If there are no parent nulls to apply, we can just return
110-
return Ok(Arc::clone(col));
111-
};
112-
113-
// NullArray is already entirely null and cannot have a validity bitmap.
114-
// If we have 0 parent nulls, we can also avoid extra work.
115-
if col.data_type().is_null() || parent_nulls.null_count() == 0 {
116-
return Ok(Arc::clone(col));
117-
}
118-
119-
if layout(col.data_type()).can_contain_null_mask {
120-
// `nullif` marks a row null where the mask is true and keeps the
121-
// field's own nulls. Only the validity bitmap is rebuilt; the value
122-
// buffers and child arrays are shared with `col`.
123-
let null_parents = BooleanArray::new(!parent_nulls.inner(), None);
124-
return Ok(nullif(col.as_ref(), &null_parents)?);
125-
}
126-
127-
// Unions and run-end encoded arrays have no validity bitmap of their own
128-
// and represent nulls in their children. Rebuild the array so null parents
129-
// become null values in those children.
130-
let data = col.to_data();
131-
let mut mutable = MutableArrayData::new(vec![&data], true, data.len());
132-
let mut end = 0;
133-
for (start, valid_end) in parent_nulls.valid_slices() {
134-
mutable.try_extend_nulls(start - end)?;
135-
mutable.try_extend(0, start, valid_end)?;
136-
end = valid_end;
137-
}
138-
mutable.try_extend_nulls(data.len() - end)?;
139-
140-
Ok(make_array(mutable.freeze()))
141-
}
142-
143101
/// Extract a single field from a struct or map array
144102
fn extract_single_field(base: ColumnarValue, name: ScalarValue) -> Result<ColumnarValue> {
145103
let arrays = ColumnarValue::values_to_arrays(&[base])?;
@@ -617,6 +575,7 @@ mod tests {
617575
ArrayRef, Int32Array, Int32Builder, ListArray, ListBuilder, MapBuilder, RunArray,
618576
StructArray, UnionArray,
619577
};
578+
use arrow::buffer::NullBuffer;
620579
use arrow::datatypes::{Fields, Int32Type, UnionFields};
621580

622581
#[test]

‎datafusion/physical-plan/Cargo.toml‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -193,3 +193,7 @@ name = "window_filter"
193193
[[bench]]
194194
harness = false
195195
name = "range_repartition"
196+
197+
[[bench]]
198+
harness = false
199+
name = "unnest"
Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,108 @@
1+
// Licensed to the Apache Software Foundation (ASF) under one
2+
// or more contributor license agreements. See the NOTICE file
3+
// distributed with this work for additional information
4+
// regarding copyright ownership. The ASF licenses this file
5+
// to you under the Apache License, Version 2.0 (the
6+
// "License"); you may not use this file except in compliance
7+
// with the License. You may obtain a copy of the License at
8+
//
9+
// http://www.apache.org/licenses/LICENSE-2.0
10+
//
11+
// Unless required by applicable law or agreed to in writing,
12+
// software distributed under the License is distributed on an
13+
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
// KIND, either express or implied. See the License for the
15+
// specific language governing permissions and limitations
16+
// under the License.
17+
18+
use std::hint::black_box;
19+
use std::sync::Arc;
20+
21+
use arrow::array::{ArrayRef, Int64Array, StringArray, StructArray};
22+
use arrow::buffer::NullBuffer;
23+
use arrow::datatypes::{Field, Schema};
24+
use arrow::record_batch::RecordBatch;
25+
use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main};
26+
use datafusion_common::UnnestOptions;
27+
use datafusion_execution::{TaskContext, config::SessionConfig};
28+
use datafusion_physical_plan::test::TestMemoryExec;
29+
use datafusion_physical_plan::unnest::UnnestExec;
30+
use datafusion_physical_plan::{ExecutionPlan, collect};
31+
use rand::SeedableRng;
32+
use rand::rngs::StdRng;
33+
use rand::seq::SliceRandom;
34+
35+
fn unnest_benchmark(c: &mut Criterion) {
36+
let runtime = tokio::runtime::Builder::new_current_thread()
37+
.enable_all()
38+
.build()
39+
.unwrap();
40+
let mut group = c.benchmark_group("unnest_struct");
41+
for num_rows in [1024, 8192, 65536] {
42+
let integers = Arc::new(Int64Array::from_iter(
43+
(0..num_rows).map(|i| (i % 10 != 0).then_some(i as i64)),
44+
)) as ArrayRef;
45+
let strings = Arc::new(StringArray::from_iter(
46+
(0..num_rows).map(|i| (i % 10 != 0).then(|| format!("value-{i:0122}"))),
47+
)) as ArrayRef;
48+
let context = Arc::new(
49+
TaskContext::default()
50+
.with_session_config(SessionConfig::new().with_batch_size(num_rows)),
51+
);
52+
group.throughput(Throughput::Elements(num_rows as u64));
53+
for (name, child) in [("int64", integers), ("utf8_128_bytes", strings)] {
54+
for null_percent in [0, 1, 50, 100] {
55+
let nulls = (null_percent != 0).then(|| {
56+
let mut validity = vec![true; num_rows];
57+
validity[..num_rows * null_percent / 100].fill(false);
58+
validity.shuffle(&mut StdRng::seed_from_u64(42));
59+
NullBuffer::from(validity)
60+
});
61+
let fields = vec![Field::new("value", child.data_type().clone(), true)];
62+
let parent = StructArray::new(
63+
fields.clone().into(),
64+
vec![Arc::clone(&child)],
65+
nulls,
66+
);
67+
let batch =
68+
RecordBatch::try_from_iter(vec![("s", Arc::new(parent) as ArrayRef)])
69+
.unwrap();
70+
let source = TestMemoryExec::try_new_exec(
71+
&[vec![batch.clone()]],
72+
batch.schema(),
73+
None,
74+
)
75+
.unwrap();
76+
let plan: Arc<dyn ExecutionPlan> = Arc::new(
77+
UnnestExec::new(
78+
source,
79+
vec![],
80+
vec![0],
81+
Arc::new(Schema::new(fields)),
82+
UnnestOptions::default(),
83+
)
84+
.unwrap(),
85+
);
86+
group.bench_function(
87+
BenchmarkId::new(format!("{name}_nulls_{null_percent}"), num_rows),
88+
|b| {
89+
b.iter(|| {
90+
black_box(
91+
runtime
92+
.block_on(collect(
93+
Arc::clone(&plan),
94+
Arc::clone(&context),
95+
))
96+
.unwrap(),
97+
)
98+
});
99+
},
100+
);
101+
}
102+
}
103+
}
104+
group.finish();
105+
}
106+
107+
criterion_group!(benches, unnest_benchmark);
108+
criterion_main!(benches);

0 commit comments

Comments
 (0)