From d9fe07d48c10331d97a279b901f5852869d93515 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 29 Sep 2026 03:04:39 +0000 Subject: [PATCH 1/2] fix: zero sliced boolean offsets at every level before exporting to the JVM (#6339) 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. --- 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 | 55 +++- native/spark-expr/src/jvm_udf/mod.rs | 10 +- .../org/apache/comet/CometCodegenSuite.scala | 54 ++++ .../apache/comet/exec/CometExecSuite.scala | 59 ++++ 8 files changed, 445 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 f45fb9beafd..029f62f0a08 100644 --- a/docs/source/contributor-guide/ffi.md +++ b/docs/source/contributor-guide/ffi.md @@ -250,6 +250,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..c0db0e6adc3 --- /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.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 c4153606a1f..978f0b45ac5 100644 --- a/native/common/src/lib.rs +++ b/native/common/src/lib.rs @@ -16,12 +16,14 @@ // under the License. mod error; +mod ffi_offsets; mod query_context; mod schema; pub mod tracing; 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 utils::{bytes_to_i128, decode_utf8_spark_lossy}; diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index d82ca3ed46c..154e462ef76 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, ScalarValue}; use datafusion::execution::disk_manager::DiskManagerMode; @@ -709,27 +708,9 @@ fn prepare_output( while i < results.len() { let array_ref = results.get(i).ok_or(CometError::IndexOutOfBounds(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(array_addrs[i], schema_addrs[i])?; - } else { - array_ref - .to_data() - .move_to_spark(array_addrs[i], schema_addrs[i])?; - } + array_ref + .to_data() + .move_to_spark(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 6195e3f0aea..cf908b2e32c 100644 --- a/native/core/src/execution/utils.rs +++ b/native/core/src/execution/utils.rs @@ -21,6 +21,7 @@ use arrow::{ array::ArrayData, ffi::{FFI_ArrowArray, FFI_ArrowSchema}, }; +use datafusion_comet_common::zero_offsets; pub trait SparkArrowConvert { /// Move Arrow Arrays to C data interface. @@ -35,12 +36,16 @@ impl SparkArrowConvert for ArrayData { let array_align = std::mem::align_of::(); let schema_align = std::mem::align_of::(); + // 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); + let ffi_schema = FFI_ArrowSchema::try_from(self.data_type())?; // Check if the pointer alignment is correct. if array_ptr.align_offset(array_align) != 0 || schema_ptr.align_offset(schema_align) != 0 { unsafe { - std::ptr::write_unaligned(array_ptr, FFI_ArrowArray::new(self)); - std::ptr::write_unaligned(schema_ptr, FFI_ArrowSchema::try_from(self.data_type())?); + std::ptr::write_unaligned(array_ptr, ffi_array); + std::ptr::write_unaligned(schema_ptr, ffi_schema); } } else { // SAFETY: `array_ptr` and `schema_ptr` are aligned correctly. @@ -55,8 +60,8 @@ impl SparkArrowConvert for ArrayData { "move_to_spark: schema_ptr not aligned" ); unsafe { - std::ptr::write(array_ptr, FFI_ArrowArray::new(self)); - std::ptr::write(schema_ptr, FFI_ArrowSchema::try_from(self.data_type())?); + std::ptr::write(array_ptr, ffi_array); + std::ptr::write(schema_ptr, ffi_schema); } } @@ -65,3 +70,45 @@ impl SparkArrowConvert for ArrayData { } pub use datafusion_comet_common::bytes_to_i128; + +#[cfg(test)] +mod tests { + use super::*; + use arrow::{ + array::{Array, ArrayRef, BooleanArray, StructArray}, + datatypes::{DataType, Field}, + ffi::from_ffi, + }; + use std::{mem::MaybeUninit, sync::Arc}; + + /// 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 mut ffi_array = MaybeUninit::::uninit(); + let mut ffi_schema = MaybeUninit::::uninit(); + + data.move_to_spark( + 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 8e148ef40ae..b1fb97c79d9 100644 --- a/native/spark-expr/src/jvm_udf/mod.rs +++ b/native/spark-expr/src/jvm_udf/mod.rs @@ -28,6 +28,7 @@ use datafusion::common::Result as DFResult; use datafusion::logical_expr::ColumnarValue; use datafusion::physical_expr::PhysicalExpr; +use datafusion_comet_common::zero_offsets; use datafusion_comet_jni_bridge::errors::{CometError, ExecutionError}; use datafusion_comet_jni_bridge::JVMClasses; use jni::objects::{Global, JObject, JValue}; @@ -140,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 5806cb35015..af317f14e35 100644 --- a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala @@ -538,6 +538,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 6000d7190e1..dab5c29d750 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala @@ -3676,6 +3676,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 6647d0f6c762c21281b3815d6b254f56074033fc Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 3 Oct 2026 15:41:55 -0600 Subject: [PATCH 2/2] fix: support Rust 1.99 lints --- native/Cargo.lock | 34 +++++++++---------- .../execution/memory_pools/spark_memory.rs | 2 +- .../execution/memory_pools/unified_pool.rs | 6 ++-- .../src/conversion_funcs/numeric.rs | 1 - 4 files changed, 21 insertions(+), 22 deletions(-) diff --git a/native/Cargo.lock b/native/Cargo.lock index 0a8edfbf708..6093206d3cd 100644 --- a/native/Cargo.lock +++ b/native/Cargo.lock @@ -198,7 +198,7 @@ dependencies = [ "snap", "strum", "strum_macros", - "thiserror 2.0.19", + "thiserror 2.0.20", "uuid", "zstd", ] @@ -598,9 +598,9 @@ checksum = "8b75356056920673b02621b35afd0f7dda9306d03c79a30f5c56c44cf256e3de" [[package]] name = "async-trait" -version = "0.1.91" +version = "0.1.92" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ae36dc4177970ef04fde5178d3e2429882def40e57a451f919c098f72baa6cec" +checksum = "82f6aeea286b8eb4dd3431a1be1b59d290ace00f5bfd8e2a159bc2a05e2c1667" dependencies = [ "proc-macro2", "quote", @@ -2002,7 +2002,7 @@ dependencies = [ "datafusion", "serde", "serde_json", - "thiserror 2.0.19", + "thiserror 2.0.20", ] [[package]] @@ -2020,7 +2020,7 @@ dependencies = [ "paste", "prost", "regex", - "thiserror 2.0.19", + "thiserror 2.0.20", ] [[package]] @@ -2703,7 +2703,7 @@ version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "10d60334b3b2e7c9d91ef8150abfb6fa4c1c39ebbcf4a81c2e346aad939fee3e" dependencies = [ - "thiserror 2.0.19", + "thiserror 2.0.20", ] [[package]] @@ -3899,7 +3899,7 @@ dependencies = [ "libloading", "log", "simd_cesu8", - "thiserror 2.0.19", + "thiserror 2.0.20", "walkdir", "windows-link", ] @@ -4169,7 +4169,7 @@ dependencies = [ "serde-value", "serde_json", "serde_yaml", - "thiserror 2.0.19", + "thiserror 2.0.20", "thread-id", "typemap-ors", "unicode-segmentation", @@ -4469,7 +4469,7 @@ dependencies = [ "serde", "serde_json", "serde_urlencoded", - "thiserror 2.0.19", + "thiserror 2.0.20", "tokio", "tracing", "url", @@ -5145,7 +5145,7 @@ dependencies = [ "spin 0.10.1", "symbolic-demangle", "tempfile", - "thiserror 2.0.19", + "thiserror 2.0.20", ] [[package]] @@ -5300,7 +5300,7 @@ dependencies = [ "rustc-hash", "rustls", "socket2", - "thiserror 2.0.19", + "thiserror 2.0.20", "tokio", "tracing", "web-time", @@ -5323,7 +5323,7 @@ dependencies = [ "rustls", "rustls-pki-types", "slab", - "thiserror 2.0.19", + "thiserror 2.0.20", "tinyvec", "tracing", "web-time", @@ -6467,11 +6467,11 @@ dependencies = [ [[package]] name = "thiserror" -version = "2.0.19" +version = "2.0.20" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "09a43598840e33d5b0331f38c5e30d13bb11c11210a4b58f0d9b18a5a5eefcd9" +checksum = "ec86235f5fcc2a73650310756d2ac5b138a5780bbbdfae3eeccec992c435ba4f" dependencies = [ - "thiserror-impl 2.0.19", + "thiserror-impl 2.0.20", ] [[package]] @@ -6487,9 +6487,9 @@ dependencies = [ [[package]] name = "thiserror-impl" -version = "2.0.19" +version = "2.0.20" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "43cbfe0cf76104d42a574802844187e84a305e531ed54455f11fbde0f10541cd" +checksum = "bc04cd3e1236dd4a98afca4569f2deb3f120e5422a4023be2cb683f8486292af" dependencies = [ "proc-macro2", "quote", diff --git a/native/core/src/execution/memory_pools/spark_memory.rs b/native/core/src/execution/memory_pools/spark_memory.rs index 1db196253e1..4cf89a31604 100644 --- a/native/core/src/execution/memory_pools/spark_memory.rs +++ b/native/core/src/execution/memory_pools/spark_memory.rs @@ -167,7 +167,7 @@ impl SparkMemory { fn repay(&self, size: usize) -> usize { let debt = self .overcommit - .fetch_update(Relaxed, Relaxed, |debt| Some(debt.saturating_sub(size))) + .try_update(Relaxed, Relaxed, |debt| Some(debt.saturating_sub(size))) .unwrap(); debt.min(size) } diff --git a/native/core/src/execution/memory_pools/unified_pool.rs b/native/core/src/execution/memory_pools/unified_pool.rs index f023d51af23..95c3561ef95 100644 --- a/native/core/src/execution/memory_pools/unified_pool.rs +++ b/native/core/src/execution/memory_pools/unified_pool.rs @@ -102,7 +102,7 @@ impl MemoryPool for CometUnifiedMemoryPool { } self.spark.acquire(additional); self.used - .fetch_update(Relaxed, Relaxed, |old| Some(old.saturating_add(additional))) + .try_update(Relaxed, Relaxed, |old| Some(old.saturating_add(additional))) .unwrap(); } @@ -115,7 +115,7 @@ impl MemoryPool for CometUnifiedMemoryPool { } if let Err(prev) = self .used - .fetch_update(Relaxed, Relaxed, |old| old.checked_sub(size)) + .try_update(Relaxed, Relaxed, |old| old.checked_sub(size)) { panic!( "Task {} overflow when releasing {size} of {prev} bytes", @@ -139,7 +139,7 @@ impl MemoryPool for CometUnifiedMemoryPool { } if let Err(prev) = self .used - .fetch_update(Relaxed, Relaxed, |old| old.checked_add(additional)) + .try_update(Relaxed, Relaxed, |old| old.checked_add(additional)) { return Err(resources_datafusion_err!( "Task {} failed to acquire {} bytes due to overflow. Reserved: {}", diff --git a/native/spark-expr/src/conversion_funcs/numeric.rs b/native/spark-expr/src/conversion_funcs/numeric.rs index 90fc5ffc5bf..346190ab987 100644 --- a/native/spark-expr/src/conversion_funcs/numeric.rs +++ b/native/spark-expr/src/conversion_funcs/numeric.rs @@ -1245,7 +1245,6 @@ mod tests { use super::*; use arrow::array::AsArray; use arrow::datatypes::TimestampMicrosecondType; - use core::f64; #[test] fn test_spark_cast_int_to_int_overflow() {