From 5618dec7160cc94094e24bbafbcdc5bec4cca1f7 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 30 Sep 2026 13:40:42 -0600 Subject: [PATCH 1/2] fix: zero sliced boolean offsets at every level before exporting to the JVM (#6339) (#6449) Arrow Java's C Data import ignores ArrowArray.offset at every level (apache/arrow-java#88). arrow-rs folds a slice into the buffers for every Spark type except boolean, which keeps its bit offset, so a sliced boolean reached the JVM read from bit 0. The fix for #2051 only checked the top level of each output column, so a boolean inside a sliced struct, list or map came back wrong, and inputs to the JVM UDF bridge were never checked at all. Add zero_offsets to datafusion-comet-common. It returns the array unchanged when every offset is already zero, and otherwise re-slices boolean bitmaps to start at bit 0 at every level, sharing every other buffer. move_to_spark now applies it to every executed batch and decoded shuffle block, replacing the top-level take in prepare_output, and JvmScalarUdfExpr applies it to the bridge inputs. Closes #6288. (cherry picked from commit aee5e06133ca8e2fc1b1337ca2ac8cf3cb0514f3) (cherry picked from commit 66800830ea30cc75fddeed1b24e50d6d7ec61c57) --- docs/source/contributor-guide/ffi.md | 12 + native/common/src/ffi_offsets.rs | 255 ++++++++++++++++++ native/common/src/lib.rs | 2 + native/core/src/execution/jni_api.rs | 27 +- native/core/src/execution/utils.rs | 43 ++- native/spark-expr/src/jvm_udf/mod.rs | 11 +- .../org/apache/comet/CometCodegenSuite.scala | 54 ++++ .../apache/comet/exec/CometExecSuite.scala | 59 ++++ 8 files changed, 434 insertions(+), 29 deletions(-) create mode 100644 native/common/src/ffi_offsets.rs diff --git a/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md index e56e5ac0f55..ecb812af41b 100644 --- a/docs/source/contributor-guide/ffi.md +++ b/docs/source/contributor-guide/ffi.md @@ -256,6 +256,18 @@ pub extern "system" fn Java_..._exportVector( } ``` +### Array Offsets + +Arrow Java's C Data import ignores `ArrowArray.offset` at every level +([apache/arrow-java#88](https://github.com/apache/arrow-java/issues/88)) and reads each buffer from its start. arrow-rs +folds a slice into the buffers for almost every type, but a sliced `BooleanArray` keeps its bit offset, and a struct +exports offset 0 even when its children are sliced. So every array native exports to the JVM first goes through +`zero_offsets` (in `native/common/src/ffi_offsets.rs`), which re-slices boolean bitmaps to start at bit 0 at every +level and shares every other buffer. `move_to_spark` applies it to executed batches and decoded shuffle blocks, and +`JvmScalarUdfExpr` applies it to the inputs of the JVM UDF bridge. A new native to JVM export path has to call it too, +or sliced booleans reach the JVM misaligned +([#6288](https://github.com/apache/datafusion-comet/issues/6288)). + ### Wrapper Object Lifecycle (Native → JVM) ``` diff --git a/native/common/src/ffi_offsets.rs b/native/common/src/ffi_offsets.rs new file mode 100644 index 00000000000..6c1aac80404 --- /dev/null +++ b/native/common/src/ffi_offsets.rs @@ -0,0 +1,255 @@ +// 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. + +//! Zeroing array offsets before an array is exported to the JVM. +//! +//! Arrow Java's C Data importer ignores `ArrowArray.offset` at every level +//! () and reads each buffer from its start. arrow-rs +//! folds a slice into the buffers for almost every type, so a sliced `Int64Array`, `StringArray` or +//! `StructArray` exports offset 0 and nothing is lost. `BooleanArray` is the exception: it keeps its +//! bit offset in `ArrayData::offset`, and the export keeps the validity bitmap at that same offset, +//! so the JVM reads both the values and the nulls of a sliced boolean from bit 0. A struct exports +//! offset 0 even when its children are sliced, so checking only the top level misses a boolean +//! child (). + +use arrow::array::{ArrayData, MutableArrayData}; +use arrow::datatypes::DataType; +use arrow::error::ArrowError; +use std::borrow::Cow; + +/// `data` with a zero offset at every level, children and dictionary values included, so that a +/// consumer that ignores `ArrowArray.offset` reads the right rows. Apply it to every array exported +/// to the JVM. +/// +/// Returns `data` itself when every offset is already zero, which is the common case. A sliced +/// boolean gets its values bitmap re-sliced to start at bit 0, which shares the buffer when the +/// offset is a whole number of bytes and copies the bitmap otherwise. Every other buffer is shared. +/// The validity bitmap needs nothing here, because `FFI_ArrowArray::new` realigns it to the +/// exported offset. +pub fn zero_offsets(data: &ArrayData) -> Result, ArrowError> { + if !has_nonzero_offset(data) { + return Ok(Cow::Borrowed(data)); + } + let zeroed = rewrite(data)?; + debug_assert!(!has_nonzero_offset(&zeroed)); + Ok(Cow::Owned(zeroed)) +} + +/// Whether `data` or anything below it has a non-zero offset. A dictionary's values are its only +/// child in `ArrayData`, so they are covered too. +fn has_nonzero_offset(data: &ArrayData) -> bool { + data.offset() != 0 || data.child_data().iter().any(has_nonzero_offset) +} + +fn rewrite(data: &ArrayData) -> Result { + let data = match (data.offset(), data.data_type()) { + (0, _) => data.clone(), + (offset, DataType::Boolean) => { + let values = data.buffers()[0].bit_slice(offset, data.len()); + // SAFETY: the same bits, now starting at bit 0. The length and nulls are unchanged. + unsafe { + data.clone() + .into_builder() + .offset(0) + .buffers(vec![values]) + .build_unchecked() + } + } + // arrow-rs keeps an offset only for booleans and run-end encoded arrays, and Spark has no + // run-end encoded type, so this is not expected to run. A copy starts at offset 0. + _ => { + let mut copy = MutableArrayData::new(vec![data], false, data.len()); + copy.try_extend(0, 0, data.len())?; + copy.freeze() + } + }; + if !data.child_data().iter().any(has_nonzero_offset) { + return Ok(data); + } + let children = data + .child_data() + .iter() + .map(|child| zero_offsets(child).map(Cow::into_owned)) + .collect::, _>>()?; + // SAFETY: each child holds the same values as before, with its offsets folded into its buffers. + Ok(unsafe { data.into_builder().child_data(children).build_unchecked() }) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{ + Array, ArrayRef, BooleanArray, DictionaryArray, Int32Array, Int64Array, Int8Array, + ListArray, MapArray, StringArray, StructArray, + }; + use arrow::buffer::OffsetBuffer; + use arrow::datatypes::{Field, Fields, Int8Type}; + use arrow::ffi::{from_ffi, FFI_ArrowArray, FFI_ArrowSchema}; + use std::sync::Arc; + + /// Nullable booleans with nulls and values that differ from their neighbours, so that reading + /// from the wrong bit shows up in both. + fn booleans(len: usize) -> BooleanArray { + (0..len) + .map(|i| (i % 5 != 0).then_some(i % 3 == 0)) + .collect() + } + + fn assert_exported_offsets_are_zero(array: &FFI_ArrowArray) { + assert_eq!(array.offset(), 0); + (0..array.num_children()).for_each(|i| assert_exported_offsets_are_zero(array.child(i))); + if let Some(dictionary) = array.dictionary() { + assert_exported_offsets_are_zero(dictionary); + } + } + + /// Exports `data` the way Comet hands it to the JVM, checks that the export has offset 0 at + /// every level, and checks that importing it back gives the same values. + fn assert_exports_aligned(data: &ArrayData) { + let zeroed = zero_offsets(data).unwrap(); + let schema = FFI_ArrowSchema::try_from(data.data_type()).unwrap(); + let exported = FFI_ArrowArray::new(&zeroed); + assert_exported_offsets_are_zero(&exported); + let imported = unsafe { from_ffi(exported, &schema) }.unwrap(); + assert_eq!(imported, *data); + } + + #[test] + fn unsliced_arrays_are_returned_as_is() { + let columns: Vec = vec![ + Arc::new(booleans(20)), + // Slicing folds into the buffers for every type but boolean. + Arc::new(Int64Array::from_iter_values(0..20).slice(3, 10)), + Arc::new(StringArray::from_iter_values((0..20).map(|i| i.to_string())).slice(3, 10)), + ]; + for column in columns { + let data = column.to_data(); + assert!(matches!(zero_offsets(&data).unwrap(), Cow::Borrowed(_))); + } + } + + #[test] + fn sliced_booleans_start_at_bit_zero() { + let array = booleans(100); + // A whole number of bytes in, and part of one. + for offset in [8, 17] { + let data = array.slice(offset, 40).to_data(); + assert_exports_aligned(&data); + + let zeroed = zero_offsets(&data).unwrap(); + assert_eq!(zeroed.offset(), 0); + if offset % 8 == 0 { + // Re-slicing on a byte boundary shares the bitmap. + assert_eq!( + zeroed.buffers()[0].as_ptr(), + data.buffers()[0].as_ptr().wrapping_add(offset / 8) + ); + } + } + } + + #[test] + fn sliced_boolean_struct_children_start_at_bit_zero() { + let fields = Fields::from(vec![ + Field::new("b", DataType::Boolean, true), + Field::new("k", DataType::Int64, false), + ]); + let array = StructArray::new( + fields, + vec![ + Arc::new(booleans(100)), + Arc::new(Int64Array::from_iter_values(0..100)), + ], + None, + ); + let data = array.slice(17, 40).to_data(); + assert_eq!(data.offset(), 0); + assert_eq!(data.child_data()[0].offset(), 17); + assert_exports_aligned(&data); + + // The Int64 child is passed through, not copied. + let zeroed = zero_offsets(&data).unwrap(); + assert_eq!( + zeroed.child_data()[1].buffers()[0].as_ptr(), + data.child_data()[1].buffers()[0].as_ptr() + ); + } + + #[test] + fn sliced_booleans_in_a_list_of_structs_start_at_bit_zero() { + let fields = Fields::from(vec![Field::new("b", DataType::Boolean, true)]); + let structs = StructArray::new(fields.clone(), vec![Arc::new(booleans(100))], None); + // A list slice moves only the offsets, so the list has to be built over a sliced struct, + // which is what collect_list returns for a group with a single run of rows. + let list = ListArray::new( + Arc::new(Field::new_struct("item", fields, true)), + OffsetBuffer::from_lengths([10, 0, 20]), + Arc::new(structs.slice(5, 30)), + None, + ); + assert_exports_aligned(&list.to_data()); + } + + #[test] + fn sliced_boolean_map_values_start_at_bit_zero() { + let entries = StructArray::new( + Fields::from(vec![ + Field::new("key", DataType::Utf8, false), + Field::new("value", DataType::Boolean, true), + ]), + vec![ + Arc::new(StringArray::from_iter_values( + (0..100).map(|i| i.to_string()), + )), + Arc::new(booleans(100)), + ], + None, + ); + let map = MapArray::try_new( + Arc::new(Field::new_struct( + "entries", + entries.fields().clone(), + false, + )), + OffsetBuffer::from_lengths([3, 7, 0, 10]), + entries.slice(33, 20), + None, + false, + ) + .unwrap(); + assert_exports_aligned(&map.to_data()); + } + + #[test] + fn sliced_boolean_dictionary_values_start_at_bit_zero() { + let dictionary = DictionaryArray::::new( + Int8Array::from(vec![Some(2), None, Some(0), Some(1), Some(2)]), + Arc::new(booleans(20).slice(3, 3)), + ); + assert_exports_aligned(&dictionary.to_data()); + } + + #[test] + fn other_types_with_an_offset_are_copied() { + // arrow-rs never builds one of these, but an ArrayData can carry an offset for any type. + let data = Int32Array::from_iter((0..20).map(|i| (i % 4 != 0).then_some(i))) + .to_data() + .slice(5, 10); + assert_eq!(data.offset(), 5); + assert_exports_aligned(&data); + } +} diff --git a/native/common/src/lib.rs b/native/common/src/lib.rs index ffa9668bb92..31fa20ab704 100644 --- a/native/common/src/lib.rs +++ b/native/common/src/lib.rs @@ -17,6 +17,7 @@ pub mod cancellation; mod error; +mod ffi_offsets; pub mod offset_extents; mod query_context; mod schema; @@ -26,6 +27,7 @@ mod utf8; mod utils; pub use error::{decimal_overflow_error, SparkError, SparkErrorWithContext, SparkResult}; +pub use ffi_offsets::zero_offsets; pub use query_context::{create_query_context_map, QueryContext, QueryContextMap}; pub use schema::{cast_and_stamp_schema, widen_nested_nullability}; pub use struct_nulls::{child_with_parent_nulls, children_with_parent_nulls}; diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index a7f5f89aa2c..19a2d0dd75d 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -28,8 +28,7 @@ use crate::{ }; use std::collections::HashSet; -use arrow::array::{Array, RecordBatch, UInt32Array}; -use arrow::compute::{take, TakeOptions}; +use arrow::array::{Array, RecordBatch}; use arrow::datatypes::DataType as ArrowDataType; use datafusion::common::{DataFusionError, Result as DataFusionResult}; use datafusion::execution::disk_manager::DiskManagerMode; @@ -1075,27 +1074,9 @@ fn prepare_output( let array_ref = results.get(i).ok_or(CometError::IndexOutOfBounds(i))?; let field = output_schema.field(i); - if array_ref.offset() != 0 { - // https://github.com/apache/datafusion-comet/issues/2051 - // Bug with non-zero offset FFI, so take to a new array which will have an offset of 0. - // We expect this to be a cold code path, hence the check_bounds: true and assert_eq. - let indices = UInt32Array::from((0..num_rows as u32).collect::>()); - let new_array = take( - array_ref, - &indices, - Some(TakeOptions { check_bounds: true }), - )?; - - assert_eq!(new_array.offset(), 0); - - new_array - .to_data() - .move_to_spark(field, array_addrs[i], schema_addrs[i])?; - } else { - array_ref - .to_data() - .move_to_spark(field, array_addrs[i], schema_addrs[i])?; - } + array_ref + .to_data() + .move_to_spark(field, array_addrs[i], schema_addrs[i])?; i += 1; } } diff --git a/native/core/src/execution/utils.rs b/native/core/src/execution/utils.rs index 4102bfab345..d32359843a3 100644 --- a/native/core/src/execution/utils.rs +++ b/native/core/src/execution/utils.rs @@ -22,6 +22,7 @@ use arrow::{ datatypes::Field, ffi::{FFI_ArrowArray, FFI_ArrowSchema}, }; +use datafusion_comet_common::zero_offsets; pub trait SparkArrowConvert { /// Move Arrow Arrays to C data interface. @@ -36,7 +37,9 @@ impl SparkArrowConvert for ArrayData { let array_align = std::mem::align_of::(); let schema_align = std::mem::align_of::(); - let ffi_array = FFI_ArrowArray::new(self); + // Arrow Java ignores `ArrowArray.offset` on import, so every level has to start at 0. + let data = zero_offsets(self)?; + let ffi_array = FFI_ArrowArray::new(&data); // Spark owns the top-level name and nullability. Preserve the existing anonymous schema // shape while carrying logical extension metadata from the RecordBatch field. let ffi_schema = @@ -76,10 +79,11 @@ pub use datafusion_comet_common::bytes_to_i128; mod tests { use super::*; use arrow::{ - array::{Array, Int32Array}, + array::{Array, ArrayRef, BooleanArray, Int32Array, StructArray}, datatypes::DataType, + ffi::from_ffi, }; - use std::{collections::HashMap, mem::MaybeUninit}; + use std::{collections::HashMap, mem::MaybeUninit, sync::Arc}; #[test] fn test_move_to_spark_preserves_field_metadata() { @@ -110,4 +114,37 @@ mod tests { drop(ffi_array); drop(ffi_schema); } + + /// Arrow Java ignores `ArrowArray.offset`, so a boolean child sliced along with its struct + /// has to be exported at offset 0 (https://github.com/apache/datafusion-comet/issues/6288). + #[test] + fn test_move_to_spark_zeroes_nested_boolean_offsets() { + let booleans: BooleanArray = (0..64) + .map(|i| (i % 5 != 0).then_some(i % 3 == 0)) + .collect(); + let array = StructArray::from(vec![( + Arc::new(Field::new("b", DataType::Boolean, true)), + Arc::new(booleans) as ArrayRef, + )]) + .slice(17, 20); + let data = array.to_data(); + assert_eq!(data.child_data()[0].offset(), 17); + let field = Field::new("s", data.data_type().clone(), true); + let mut ffi_array = MaybeUninit::::uninit(); + let mut ffi_schema = MaybeUninit::::uninit(); + + data.move_to_spark( + &field, + ffi_array.as_mut_ptr() as i64, + ffi_schema.as_mut_ptr() as i64, + ) + .unwrap(); + + let ffi_array = unsafe { ffi_array.assume_init() }; + let ffi_schema = unsafe { ffi_schema.assume_init() }; + assert_eq!(ffi_array.offset(), 0); + assert_eq!(ffi_array.child(0).offset(), 0); + let imported = unsafe { from_ffi(ffi_array, &ffi_schema) }.unwrap(); + assert_eq!(imported, data); + } } diff --git a/native/spark-expr/src/jvm_udf/mod.rs b/native/spark-expr/src/jvm_udf/mod.rs index 75e75351fab..232926744d0 100644 --- a/native/spark-expr/src/jvm_udf/mod.rs +++ b/native/spark-expr/src/jvm_udf/mod.rs @@ -28,7 +28,7 @@ use datafusion::common::Result as DFResult; use datafusion::logical_expr::ColumnarValue; use datafusion::physical_expr::PhysicalExpr; -use datafusion_comet_common::decode_string_arrays; +use datafusion_comet_common::{decode_string_arrays, zero_offsets}; use datafusion_comet_jni_bridge::errors::{CometError, ExecutionError}; use datafusion_comet_jni_bridge::JVMClasses; use jni::objects::{Global, JObject, JValue}; @@ -141,10 +141,15 @@ impl PhysicalExpr for JvmScalarUdfExpr { .collect::>()?; // The JVM writes into the out_array/out_schema slots and reads from the in_ slots. + // Arrow Java ignores `ArrowArray.offset` on import, so every level has to start at 0. let in_ffi_arrays: Vec> = arrays .iter() - .map(|arr| Box::new(FFI_ArrowArray::new(&arr.to_data()))) - .collect(); + .map(|arr| { + let data = arr.to_data(); + let data = zero_offsets(&data).map_err(|e| CometError::Arrow { source: e })?; + Ok(Box::new(FFI_ArrowArray::new(&data))) + }) + .collect::>()?; let in_ffi_schemas: Vec> = arrays .iter() .map(|arr| { diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala index e3570e77331..0b9a19864ea 100644 --- a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala @@ -966,6 +966,60 @@ class CometCodegenSuite } } + // Arrow Java ignores ArrowArray.offset on import, so the bridge has to zero a sliced boolean's + // offset before handing it over, at the top level and inside a struct. + // https://github.com/apache/datafusion-comet/issues/6288 + private def withSlicedGroups(f: => Unit): Unit = { + withTempPath { dir => + spark + .range(0, 4000) + .selectExpr( + "id % 2000 AS k", + "id % 2000 % 3 = 0 AS b", + "CAST(id % 2000 AS STRING) AS s", + "named_struct('x', IF(id % 5 = 0, NULL, id % 2000 % 3 = 0)) AS st") + .write + .parquet(dir.getCanonicalPath) + // The hash aggregate slices its emitted groups into batch-size chunks, so with one + // partition every output batch after the first is a slice. + withSQLConf( + SQLConf.SHUFFLE_PARTITIONS.key -> "1", + CometConf.COMET_BATCH_SIZE.key -> "100") { + withParquetTable(dir.getCanonicalPath, "g")(f) + } + } + } + + test("boolean ScalaUDF argument sliced by an aggregate keeps its values") { + spark.udf.register("flip", (x: Boolean) => !x) + withSlicedGroups { + assertCodegenRan { + checkSparkAnswerAndOperator( + sql("SELECT k, flip(b) FROM (SELECT k, b, count(*) FROM g GROUP BY k, b)")) + } + } + } + + test("dispatched regexp_replace reads a sliced boolean with its own values") { + withSlicedGroups { + assertCodegenRan { + checkSparkAnswerAndOperator(sql("""SELECT k, regexp_replace(IF(b, s, 'zz'), '1', 'y') + |FROM (SELECT k, b, s, count(*) FROM g GROUP BY k, b, s)""".stripMargin)) + } + } + } + + test("struct ScalaUDF input with a sliced boolean child keeps its values") { + spark.udf.register("flip", (x: Boolean) => !x) + withSlicedGroups { + // `st.x` is dispatched along with the UDF, so the kernel's input is the whole struct. + assertCodegenRan { + checkSparkAnswerAndOperator( + sql("SELECT k, flip(st.x) FROM (SELECT k, st, count(*) FROM g GROUP BY k, st)")) + } + } + } + test("multi-arg ScalaUDF over string + literal routes through dispatcher") { spark.udf.register( "prepend", diff --git a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala index ffa2325302c..9eabbd1e179 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala @@ -3807,6 +3807,65 @@ class CometExecSuite extends CometTestBase { }) } + // Arrow Java ignores ArrowArray.offset on import, and a sliced boolean is the one array arrow-rs + // exports with a non-zero offset, so native has to zero it at every level, not only the top. + // https://github.com/apache/datafusion-comet/issues/6288 + test("boolean struct children sliced by an OFFSET keep their values") { + withTempPath { dir => + spark + .range(0, 1000) + .selectExpr( + "id AS v", + "named_struct('x', id % 7 = 0, 'y', id) AS st", + "named_struct('x', IF(id % 3 = 0, NULL, id % 7 = 0)) AS stn") + .write + .parquet(dir.getCanonicalPath) + withParquetTable(dir.getCanonicalPath, "t") { + // The OFFSET slices the sorted batch 17 rows in, which is not a whole byte of bits. + checkSparkAnswerAndOperator("SELECT st FROM t ORDER BY v LIMIT 40 OFFSET 17") + checkSparkAnswerAndOperator("SELECT stn FROM t ORDER BY v LIMIT 40 OFFSET 17") + } + } + } + + test("boolean struct children built over sliced aggregate output keep their values") { + withTempPath { dir => + spark + .range(0, 4000) + .selectExpr("id % 2000 AS k", "IF(id % 5 = 0, NULL, id % 2000 % 3 = 0) AS b") + .write + .parquet(dir.getCanonicalPath) + // The hash aggregate slices its emitted groups into batch-size chunks, so with one + // partition every output batch after the first is a slice, at offsets both on and off a + // byte boundary. + withSQLConf( + SQLConf.SHUFFLE_PARTITIONS.key -> "1", + CometConf.COMET_BATCH_SIZE.key -> "100") { + withParquetTable(dir.getCanonicalPath, "g") { + checkSparkAnswerAndOperator( + "SELECT named_struct('b', b, 'k', k) FROM (SELECT k, b, count(*) FROM g GROUP BY k, b)") + } + } + } + } + + test("collect_list of structs with sliced boolean children keeps their values") { + withTempPath { dir => + // One file, so the partial and final aggregates share a native plan and see one batch. + spark + .range(0, 100, 1, 1) + .selectExpr("0 AS k", "IF(id < 5, NULL, named_struct('b', id % 3 = 0)) AS st") + .write + .parquet(dir.getCanonicalPath) + withParquetTable(dir.getCanonicalPath, "t") { + // Dropping the leading null structs starts the group's only run past row 0, and + // collect_list returns a single run as a slice rather than a copy. + checkSparkAnswerAndOperator( + spark.table("t").coalesce(1).groupBy("k").agg(collect_list("st"))) + } + } + } + test("collect limit") { Seq("true", "false").foreach(aqe => { withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe) { From 64e096c91a33dd72b7cbfc0a91b8db6ef519f4bf Mon Sep 17 00:00:00 2001 From: msaf Date: Thu, 8 Oct 2026 22:44:05 +0100 Subject: [PATCH 2/2] test: nested boolean two struct levels deep sliced by an OFFSET Reproduces the export corruption seen in mongo.finance_order_costs_daily_snapshot: a boolean inside a nullable struct inside a struct comes back with the values of neighbouring rows on joom-1.1-4 and is fixed by the previous commit. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../apache/comet/exec/CometExecSuite.scala | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala index 9eabbd1e179..b79b412cc03 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala @@ -3866,6 +3866,25 @@ class CometExecSuite extends CometTestBase { } } + test("boolean two struct levels deep sliced by an OFFSET keeps its values") { + withTempPath { dir => + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + spark + .range(0, 1000, 1, 1) + .selectExpr( + "id AS v", + "named_struct('d', IF(id % 4 = 0, NULL, named_struct('lost', " + + "IF(id % 13 = 0, NULL, id % 11 = 0), 't', id)), 'p', named_struct('t', id)) AS ops") + .write + .parquet(dir.getCanonicalPath) + } + withParquetTable(dir.getCanonicalPath, "t") { + checkSparkAnswerAndOperator("SELECT ops FROM t ORDER BY v LIMIT 40 OFFSET 17") + checkSparkAnswerAndOperator("SELECT v, ops FROM t ORDER BY v LIMIT 500 OFFSET 333") + } + } + } + test("collect limit") { Seq("true", "false").foreach(aqe => { withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe) {