diff --git a/datafusion/core/tests/sql/aggregates/mod.rs b/datafusion/core/tests/sql/aggregates/mod.rs index b209e91cc81e7..0c7e152295400 100644 --- a/datafusion/core/tests/sql/aggregates/mod.rs +++ b/datafusion/core/tests/sql/aggregates/mod.rs @@ -1021,3 +1021,4 @@ pub fn split_fuzz_timestamp_data_into_batches( pub mod basic; pub mod dict_nulls; +pub mod nested_nullability; diff --git a/datafusion/core/tests/sql/aggregates/nested_nullability.rs b/datafusion/core/tests/sql/aggregates/nested_nullability.rs new file mode 100644 index 0000000000000..193e1f7b8e3b1 --- /dev/null +++ b/datafusion/core/tests/sql/aggregates/nested_nullability.rs @@ -0,0 +1,163 @@ +// 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. + +//! Regression tests for aggregating batches whose data types are *stricter* +//! than the table's declared schema. See +//! . +//! +//! A `RecordBatch` is a valid instance of a schema that is a superset of its +//! own (see [`Schema::contains`] / `Field::contains`): most commonly the +//! schema declares a (possibly nested) field as nullable while the batch's +//! arrays mark it non-nullable. `MemTable::try_new` accepts such batches via +//! exactly that check, and engines embedding DataFusion (e.g. Comet) feed +//! such batches over FFI. Aggregations must therefore not fail when the +//! runtime arrays are stricter than the planned schema. +//! +//! [`Schema::contains`]: arrow::datatypes::Schema::contains + +use std::sync::Arc; + +use arrow::array::{BooleanArray, RecordBatch, StructArray, UInt32Array}; +use arrow::datatypes::{DataType, Field, Fields, Schema}; +use datafusion::datasource::MemTable; +use datafusion::prelude::*; +use datafusion_common::Result; +use datafusion_execution::memory_pool::FairSpillPool; +use datafusion_execution::runtime_env::RuntimeEnvBuilder; + +/// Registers table `t(a UInt32, b Struct("colA" Boolean))` where the declared +/// schema marks the nested field `colA` as nullable, but the batches carry a +/// stricter, non-nullable `colA`, then runs `sql` and returns the collected +/// result. +/// +/// With `memory_limit`, the context uses a small [`FairSpillPool`] so the +/// aggregation is forced to spill. +async fn run_aggregate_over_stricter_batches( + sql: &str, + num_rows: u32, + memory_limit: Option, +) -> Result> { + // The table's declared schema: `b.colA` is nullable + let declared_schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new( + "b", + DataType::Struct(Fields::from(vec![Field::new( + "colA", + DataType::Boolean, + true, + )])), + false, + ), + ])); + + // The batches are stricter: `b.colA` is non-nullable. `MemTable::try_new` + // accepts this via `Schema::contains`. + let batch_struct_fields = + Fields::from(vec![Field::new("colA", DataType::Boolean, false)]); + let batch_schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new("b", DataType::Struct(batch_struct_fields.clone()), false), + ])); + let batch = RecordBatch::try_new( + batch_schema, + vec![ + Arc::new(UInt32Array::from_iter_values(0..num_rows)), + Arc::new(StructArray::new( + batch_struct_fields, + vec![Arc::new(BooleanArray::from_iter( + (0..num_rows).map(|i| Some(i % 2 == 0)), + ))], + None, + )), + ], + )?; + + let table = MemTable::try_new(declared_schema, vec![vec![batch]])?; + + let ctx = match memory_limit { + Some(limit) => { + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::new(FairSpillPool::new(limit))) + .build_arc()?; + SessionContext::new_with_config_rt( + SessionConfig::new().with_batch_size(100), + runtime, + ) + } + None => SessionContext::new(), + }; + ctx.register_table("t", Arc::new(table))?; + + ctx.sql(sql).await?.collect().await +} + +#[tokio::test] +async fn array_agg_struct_from_stricter_batches() -> Result<()> { + let num_rows = 100; + let result = run_aggregate_over_stricter_batches( + "SELECT a, array_agg(b) FROM t GROUP BY a", + num_rows, + None, + ) + .await?; + let total_rows: usize = result.iter().map(|batch| batch.num_rows()).sum(); + assert_eq!(total_rows, num_rows as usize); + Ok(()) +} + +#[tokio::test] +async fn array_agg_distinct_struct_from_stricter_batches() -> Result<()> { + let num_rows = 100; + let result = run_aggregate_over_stricter_batches( + "SELECT a, array_agg(DISTINCT b) FROM t GROUP BY a", + num_rows, + None, + ) + .await?; + let total_rows: usize = result.iter().map(|batch| batch.num_rows()).sum(); + assert_eq!(total_rows, num_rows as usize); + Ok(()) +} + +#[tokio::test] +async fn array_agg_struct_from_stricter_batches_with_spilling() -> Result<()> { + let num_rows = 10_000; + let result = run_aggregate_over_stricter_batches( + "SELECT a, array_agg(b) FROM t GROUP BY a", + num_rows, + Some(4_000_000), + ) + .await?; + let total_rows: usize = result.iter().map(|batch| batch.num_rows()).sum(); + assert_eq!(total_rows, num_rows as usize); + Ok(()) +} + +#[tokio::test] +async fn array_agg_distinct_struct_from_stricter_batches_with_spilling() -> Result<()> { + let num_rows = 10_000; + let result = run_aggregate_over_stricter_batches( + "SELECT a, array_agg(DISTINCT b) FROM t GROUP BY a", + num_rows, + Some(4_000_000), + ) + .await?; + let total_rows: usize = result.iter().map(|batch| batch.num_rows()).sum(); + assert_eq!(total_rows, num_rows as usize); + Ok(()) +}