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..aa6822afe8956 --- /dev/null +++ b/datafusion/core/tests/sql/aggregates/nested_nullability.rs @@ -0,0 +1,179 @@ +// 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, SchemaRef}; +use datafusion::datasource::MemTable; +use datafusion::prelude::*; +use datafusion_common::Result; +use datafusion_execution::memory_pool::FairSpillPool; +use datafusion_execution::runtime_env::RuntimeEnvBuilder; + +/// Returns the fields of the struct column `b`: a single `colA Boolean`. +/// +/// `col_a_nullable` controls whether `colA` is declared nullable — the only +/// difference between the table's declared schema (`true`) and the actual +/// batches (`false`). +fn make_struct_fields(col_a_nullable: bool) -> Fields { + Fields::from(vec![Field::new("colA", DataType::Boolean, col_a_nullable)]) +} + +/// Returns the schema `(a UInt32 NOT NULL, b Struct("colA" Boolean) NOT NULL)` +/// with the nested field `b.colA` nullable per `col_a_nullable`. +/// +/// See [`make_struct_fields`]. +fn make_schema(col_a_nullable: bool) -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new( + "b", + DataType::Struct(make_struct_fields(col_a_nullable)), + false, + ), + ])) +} + +/// Runs a SQL aggregation over a table whose batches are stricter than its +/// declared schema. +/// +/// [`Self::run`] 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 the query +/// and returns the collected result. +struct AggregateBatchesTest { + /// Number of rows in the table. `a` is `0..num_rows` (so also the number + /// of groups for `GROUP BY a`) and `b.colA` alternates `true` / `false`. + num_rows: u32, + /// If set, the context uses a [`FairSpillPool`] of this size (and a small + /// batch size) so the aggregation is forced to spill. + memory_limit: Option, +} + +impl AggregateBatchesTest { + fn new() -> Self { + Self { + num_rows: 100, + memory_limit: None, + } + } + + fn with_num_rows(mut self, num_rows: u32) -> Self { + self.num_rows = num_rows; + self + } + + fn with_memory_limit(mut self, memory_limit: usize) -> Self { + self.memory_limit = Some(memory_limit); + self + } + + /// Runs `sql` against the table described above and asserts the result + /// has one output row per group (i.e. [`Self::num_rows`] rows in total). + async fn run(self, sql: &str) -> Result<()> { + // The table's declared schema: the nested field `b.colA` is + // nullable ... + let declared_schema = make_schema(true); + + // ... while the batches are stricter: `b.colA` is non-nullable. + // `MemTable::try_new` accepts this combination via + // `Schema::contains`. + let batch_struct_fields = make_struct_fields(false); + let batch = RecordBatch::try_new( + make_schema(false), + vec![ + Arc::new(UInt32Array::from_iter_values(0..self.num_rows)), + Arc::new(StructArray::new( + batch_struct_fields, + vec![Arc::new(BooleanArray::from_iter( + (0..self.num_rows).map(|i| Some(i % 2 == 0)), + ))], + None, + )), + ], + )?; + + let table = MemTable::try_new(declared_schema, vec![vec![batch]])?; + + let ctx = match self.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))?; + + let result = ctx.sql(sql).await?.collect().await?; + + let total_rows: usize = result.iter().map(|batch| batch.num_rows()).sum(); + assert_eq!(total_rows, self.num_rows as usize); + Ok(()) + } +} + +#[tokio::test] +async fn array_agg_struct_from_stricter_batches() -> Result<()> { + AggregateBatchesTest::new() + .run("SELECT a, array_agg(b) FROM t GROUP BY a") + .await +} + +#[tokio::test] +async fn array_agg_distinct_struct_from_stricter_batches() -> Result<()> { + AggregateBatchesTest::new() + .run("SELECT a, array_agg(DISTINCT b) FROM t GROUP BY a") + .await +} + +#[tokio::test] +async fn array_agg_struct_from_stricter_batches_with_spilling() -> Result<()> { + AggregateBatchesTest::new() + .with_num_rows(10_000) + .with_memory_limit(4_000_000) + .run("SELECT a, array_agg(b) FROM t GROUP BY a") + .await +} + +#[tokio::test] +async fn array_agg_distinct_struct_from_stricter_batches_with_spilling() -> Result<()> { + AggregateBatchesTest::new() + .with_num_rows(10_000) + .with_memory_limit(4_000_000) + .run("SELECT a, array_agg(DISTINCT b) FROM t GROUP BY a") + .await +}