diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml index f00f1a7c495..9019dffbb50 100644 --- a/native/spark-expr/Cargo.toml +++ b/native/spark-expr/Cargo.toml @@ -63,6 +63,10 @@ path = "src/lib.rs" name = "hash" harness = false +[[bench]] +name = "hash_alloc" +harness = false + [[bench]] name = "cast_from_string" harness = false diff --git a/native/spark-expr/benches/common/hash_shapes.rs b/native/spark-expr/benches/common/hash_shapes.rs new file mode 100644 index 00000000000..77e3877ef4d --- /dev/null +++ b/native/spark-expr/benches/common/hash_shapes.rs @@ -0,0 +1,469 @@ +// 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. + +//! Nested array shapes shared by the hash benchmarks, pulled in with +//! `#[path = "common/hash_shapes.rs"] mod hash_shapes;`. This lives in a subdirectory so that +//! Cargo's bench auto-discovery, which only looks at `benches/*.rs`, does not treat it as a bench +//! target. +//! +//! `hash.rs` measures time and `hash_alloc.rs` measures allocation, and both have to build the +//! same inputs for the two sets of numbers to describe the same work. +#![allow(dead_code)] + +use arrow::array::builder::{Int32Builder, ListBuilder, MapBuilder, StringBuilder, StructBuilder}; +use arrow::array::{Int32Array, ListArray, MapFieldNames, StringArray, StructArray}; +// Re-exported so a bench that uses these shapes does not have to import it separately. +pub use arrow::array::ArrayRef; +use arrow::buffer::OffsetBuffer; +use arrow::datatypes::{DataType, Field, Fields}; +use std::sync::Arc; + +pub const NUM_ROWS: usize = 8192; + +pub fn struct_fields() -> Fields { + vec![ + Arc::new(Field::new("a", DataType::Int32, true)), + Arc::new(Field::new("b", DataType::Utf8, true)), + ] + .into() +} + +pub fn struct_builder() -> StructBuilder { + StructBuilder::new( + struct_fields(), + vec![ + Box::new(Int32Builder::new()), + Box::new(StringBuilder::new()), + ], + ) +} + +pub fn append_struct(sb: &mut StructBuilder, i: usize) { + sb.field_builder::(0) + .unwrap() + .append_value(i as i32); + sb.field_builder::(1) + .unwrap() + .append_value(format!("v{}", i % 97)); + sb.append(true); +} + +/// `int32`, the cheapest leaf, as a reference point for the nested shapes. +pub fn primitive(num_rows: usize) -> ArrayRef { + Arc::new(Int32Array::from((0..num_rows as i32).collect::>())) +} + +/// `utf8`: variable-width, so the hash reads from the values buffer per row. +pub fn string(num_rows: usize) -> ArrayRef { + Arc::new(StringArray::from( + (0..num_rows) + .map(|i| format!("v{}", i % 97)) + .collect::>(), + )) +} + +/// `struct`: hashed field by field across the whole batch. +pub fn structs(num_rows: usize) -> ArrayRef { + let mut sb = struct_builder(); + for i in 0..num_rows { + append_struct(&mut sb, i); + } + Arc::new(sb.finish()) +} + +/// `array`: elements are primitives, so this takes the vectorized element path. +pub fn list_of_primitive(num_rows: usize, elems: usize) -> ArrayRef { + let mut lb = ListBuilder::new(Int32Builder::new()); + for i in 0..num_rows { + for j in 0..elems { + lb.values().append_value((i * 31 + j) as i32); + } + lb.append(true); + } + Arc::new(lb.finish()) +} + +/// `array>`: elements are nested, so this takes the per-element path. +pub fn list_of_struct(num_rows: usize, elems: usize) -> ArrayRef { + let mut lb = ListBuilder::new(struct_builder()); + for i in 0..num_rows { + for j in 0..elems { + append_struct(lb.values(), i * 31 + j); + } + lb.append(true); + } + Arc::new(lb.finish()) +} + +/// `array>` where one row is far longer than the rest, so the element count is spread +/// very unevenly across rows rather than uniformly. The per-element path slices and re-dispatches +/// once per element, so a batch dominated by a single long list has the same total work in a very +/// different distribution, which a uniform shape cannot show. +pub fn skewed_list_of_struct(num_rows: usize, long_len: usize) -> ArrayRef { + let mut lb = ListBuilder::new(struct_builder()); + for i in 0..num_rows { + let len = if i == 0 { long_len } else { 1 }; + for j in 0..len { + append_struct(lb.values(), i * 31 + j); + } + lb.append(true); + } + Arc::new(lb.finish()) +} + +/// `map`: keys and values are hashed entry by entry. +pub fn maps(num_rows: usize, entries: usize) -> ArrayRef { + let mut mb = MapBuilder::new( + Some(MapFieldNames { + entry: "entries".into(), + key: "key".into(), + value: "value".into(), + }), + StringBuilder::new(), + Int32Builder::new(), + ); + for i in 0..num_rows { + for j in 0..entries { + mb.keys().append_value(format!("k{}", (i + j) % 97)); + mb.values().append_value((i * 31 + j) as i32); + } + mb.append(true).unwrap(); + } + Arc::new(mb.finish()) +} + +/// `struct>`: a map inside a struct, so the struct branch recurses +/// into the map specialization rather than into a leaf. +pub fn struct_of_map(num_rows: usize, entries: usize) -> ArrayRef { + let fields: Fields = vec![ + Arc::new(Field::new("a", DataType::Int32, true)), + Arc::new(Field::new( + "m", + DataType::Map( + Arc::new(Field::new( + "entries", + DataType::Struct( + vec![ + Arc::new(Field::new("key", DataType::Utf8, false)), + Arc::new(Field::new("value", DataType::Int32, true)), + ] + .into(), + ), + false, + )), + false, + ), + true, + )), + ] + .into(); + let ints = primitive(num_rows); + let ms = maps(num_rows, entries); + Arc::new(StructArray::new(fields, vec![ints, ms], None)) +} + +/// `array>`: the element is a list, so the non-primitive element path recurses into +/// the vectorized leaf loop one level down. +pub fn list_of_list(num_rows: usize, outer: usize, inner: usize) -> ArrayRef { + let mut lb = ListBuilder::new(ListBuilder::new(Int32Builder::new())); + for i in 0..num_rows { + for j in 0..outer { + for k in 0..inner { + lb.values() + .values() + .append_value((i * 31 + j * 7 + k) as i32); + } + lb.values().append(true); + } + lb.append(true); + } + Arc::new(lb.finish()) +} + +/// `map>`: a struct as the map value, which the key/value specializations do not +/// cover, so the value array is hashed recursively instead. +pub fn map_of_struct(num_rows: usize, entries: usize) -> ArrayRef { + let mut mb = MapBuilder::new( + Some(MapFieldNames { + entry: "entries".into(), + key: "key".into(), + value: "value".into(), + }), + StringBuilder::new(), + struct_builder(), + ); + for i in 0..num_rows { + for j in 0..entries { + mb.keys().append_value(format!("k{}", (i + j) % 97)); + append_struct(mb.values(), i * 31 + j); + } + mb.append(true).unwrap(); + } + Arc::new(mb.finish()) +} + +/// `array>>`: three levels, so the per-element path recurses through a +/// struct into a map. +pub fn list_of_struct_of_map(num_rows: usize, elems: usize, entries: usize) -> ArrayRef { + let inner = struct_of_map(num_rows * elems, entries); + let offsets: Vec = (0..=num_rows).map(|i| (i * elems) as i32).collect(); + Arc::new(ListArray::new( + Arc::new(Field::new("item", inner.data_type().clone(), true)), + OffsetBuffer::new(offsets.into()), + inner, + None, + )) +} + +/// `array>` with a large string child, so the gather copies real payload rather than +/// the two- or three-byte strings the other shapes use. +pub fn list_of_struct_big_string(num_rows: usize, elems: usize, len: usize) -> ArrayRef { + let big = "x".repeat(len); + let mut lb = ListBuilder::new(struct_builder()); + for i in 0..num_rows { + for j in 0..elems { + let s = lb.values(); + s.field_builder::(0) + .unwrap() + .append_value((i * 31 + j) as i32); + s.field_builder::(1) + .unwrap() + .append_value(&big); + s.append(true); + } + lb.append(true); + } + Arc::new(lb.finish()) +} + +/// `array>` where half the elements are null structs whose children still hold values, +/// so the child data under a null parent is carried through the same path. +pub fn list_of_struct_half_null(num_rows: usize, elems: usize) -> ArrayRef { + let mut lb = ListBuilder::new(struct_builder()); + for i in 0..num_rows { + for j in 0..elems { + let s = lb.values(); + s.field_builder::(0) + .unwrap() + .append_value((i * 31 + j) as i32); + s.field_builder::(1) + .unwrap() + .append_value("payload"); + s.append((i + j) % 2 == 0); + } + lb.append(true); + } + Arc::new(lb.finish()) +} + +/// Short rows and then one much longer row, so all but the first pass has a single surviving row. +/// A shape that batches badly: there is nothing to gather across once the short rows finish. +pub fn list_of_struct_long_tail(num_rows: usize, tail: usize) -> ArrayRef { + let mut lb = ListBuilder::new(struct_builder()); + for i in 0..num_rows { + append_struct(lb.values(), i); + lb.append(true); + } + for j in 0..tail { + append_struct(lb.values(), j); + } + lb.append(true); + Arc::new(lb.finish()) +} + +/// Deeply nested singleton lists wrapping a struct with a wide string payload. Each recursion +/// level gathers while the level above it still holds its own gather, so this is the shape that +/// shows whether peak live bytes is one gather or the sum of the nesting depth. +pub fn deep_singleton_list_of_struct_big_string( + num_rows: usize, + depth: usize, + str_len: usize, +) -> ArrayRef { + let mut sb = StructBuilder::new( + struct_fields(), + vec![ + Box::new(Int32Builder::new()), + Box::new(StringBuilder::new()), + ], + ); + let payload = "x".repeat(str_len); + for i in 0..num_rows { + sb.field_builder::(0) + .unwrap() + .append_value(i as i32); + sb.field_builder::(1) + .unwrap() + .append_value(&payload); + sb.append(true); + } + let mut current: ArrayRef = Arc::new(sb.finish()); + // One element per row at every level, so every level gathers `num_rows` elements. + for _ in 0..depth { + let offsets = OffsetBuffer::from_lengths(std::iter::repeat_n(1usize, num_rows)); + let field = Arc::new(Field::new("item", current.data_type().clone(), true)); + current = Arc::new(ListArray::new(field, offsets, current, None)); + } + current +} + +/// A struct whose children hold live payload but whose parent row is null. Arrow's gather copies +/// the children anyway, so the copied bytes never reach the hash. +pub fn null_parent_struct_big_string(num_rows: usize, elems: usize, str_len: usize) -> ArrayRef { + let mut lb = ListBuilder::new(struct_builder()); + let payload = "x".repeat(str_len); + for _ in 0..num_rows { + for i in 0..elems { + let sb = lb.values(); + sb.field_builder::(0) + .unwrap() + .append_value(i as i32); + sb.field_builder::(1) + .unwrap() + .append_value(&payload); + // Null parent, live children. + sb.append(false); + } + lb.append(true); + } + Arc::new(lb.finish()) +} + +/// A struct whose child is a dictionary holding one huge, unreferenced value. `take` on a +/// dictionary copies the keys and shares the values, so the payload is never duplicated, but an +/// estimate based on the child's total memory size counts it anyway. +pub fn struct_of_dict_unreferenced_big_value(num_rows: usize, big_len: usize) -> ArrayRef { + use arrow::array::{DictionaryArray, Int32Array as I32}; + use arrow::datatypes::Int32Type; + + // Every key points at "x"; the 8 MiB value exists in the dictionary but is never referenced. + let values = StringArray::from(vec!["x".to_string(), "y".repeat(big_len)]); + let keys = I32::from(vec![0i32; num_rows]); + let dict: ArrayRef = Arc::new( + DictionaryArray::::try_new(keys, Arc::new(values)).expect("dictionary"), + ); + let ints: ArrayRef = Arc::new(I32::from_iter_values((0..num_rows).map(|i| i as i32))); + + let fields: Fields = vec![ + Arc::new(Field::new("a", DataType::Int32, true)), + Arc::new(Field::new( + "b", + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + true, + )), + ] + .into(); + let structs: ArrayRef = Arc::new(StructArray::new(fields, vec![ints, dict], None)); + + // One struct element per row, so the list path is taken. + let offsets = OffsetBuffer::from_lengths(std::iter::repeat_n(1usize, num_rows)); + let field = Arc::new(Field::new("item", structs.data_type().clone(), true)); + Arc::new(ListArray::new(field, offsets, structs, None)) +} + +/// Two rows of 1,024 structs where only the first element's string is huge. An average width over +/// the child is diluted by the short elements, so a cost estimate based on it under-reads what the +/// gather will copy. +pub fn width_skewed_list_of_struct(big_len: usize) -> ArrayRef { + let mut lb = ListBuilder::new(struct_builder()); + for _ in 0..2 { + for i in 0..1024usize { + let sb = lb.values(); + sb.field_builder::(0) + .unwrap() + .append_value(i as i32); + let s = if i == 0 { + "y".repeat(big_len) + } else { + "x".to_string() + }; + sb.field_builder::(1) + .unwrap() + .append_value(&s); + sb.append(true); + } + lb.append(true); + } + Arc::new(lb.finish()) +} + +/// A list sliced down to two rows whose child still backs ten million ints. `take` pre-sizes from +/// the retained child, so the gather is far larger than the two visible rows suggest. +pub fn sliced_list_retaining_big_child(kept_rows: usize, inner_elems: usize) -> ArrayRef { + let total_rows = 1000usize; + let values: ArrayRef = Arc::new(Int32Array::from_iter_values( + (0..inner_elems).map(|i| i as i32), + )); + // First row owns nearly all the elements; the rest are empty. + let mut lengths = vec![inner_elems - (total_rows - 1)]; + lengths.extend(std::iter::repeat_n(1usize, total_rows - 1)); + let offsets = OffsetBuffer::from_lengths(lengths); + let field = Arc::new(Field::new("item", DataType::Int32, true)); + let inner: ArrayRef = Arc::new(ListArray::new(field, offsets, values, None)); + + let outer_offsets = OffsetBuffer::from_lengths(std::iter::repeat_n(1usize, total_rows)); + let outer_field = Arc::new(Field::new("item", inner.data_type().clone(), true)); + let outer: ArrayRef = Arc::new(ListArray::new(outer_field, outer_offsets, inner, None)); + // Slice past the huge first row: only cheap rows are visible, the buffer stays. + outer.slice(1, kept_rows) +} + +/// A single row of flat structs. Eligible for batching, but with one row there is nothing to batch +/// across, so the scheduling buffers are pure overhead. +pub fn single_row_list_of_struct(elems: usize) -> ArrayRef { + let mut lb = ListBuilder::new(struct_builder()); + for i in 0..elems { + let sb = lb.values(); + sb.field_builder::(0) + .unwrap() + .append_value(i as i32); + sb.field_builder::(1) + .unwrap() + .append_value(format!("e{i}")); + sb.append(true); + } + lb.append(true); + Arc::new(lb.finish()) +} + +/// Deep singleton lists over a narrow struct. The outer levels are not eligible, but recursing +/// reaches an eligible single-row flat struct at the bottom, which is where scheduler buffers get +/// allocated for a batch of one. +pub fn deep_singleton_list_narrow(num_rows: usize, depth: usize) -> ArrayRef { + let mut sb = StructBuilder::new( + struct_fields(), + vec![ + Box::new(Int32Builder::new()), + Box::new(StringBuilder::new()), + ], + ); + for i in 0..num_rows { + sb.field_builder::(0) + .unwrap() + .append_value(i as i32); + sb.field_builder::(1) + .unwrap() + .append_value("abcd"); + sb.append(true); + } + let mut current: ArrayRef = Arc::new(sb.finish()); + for _ in 0..depth { + let offsets = OffsetBuffer::from_lengths(std::iter::repeat_n(1usize, num_rows)); + let field = Arc::new(Field::new("item", current.data_type().clone(), true)); + current = Arc::new(ListArray::new(field, offsets, current, None)); + } + current +} diff --git a/native/spark-expr/benches/hash.rs b/native/spark-expr/benches/hash.rs index 4d8d4026cfd..5cdfaa10f12 100644 --- a/native/spark-expr/benches/hash.rs +++ b/native/spark-expr/benches/hash.rs @@ -27,208 +27,18 @@ //! it. The two share `create_hashes_internal!`, so the shape of the work is the same and a change //! to that macro shows up here. -use arrow::array::builder::{Int32Builder, ListBuilder, MapBuilder, StringBuilder, StructBuilder}; -use arrow::array::{ArrayRef, Int32Array, ListArray, StringArray, StructArray}; -use arrow::buffer::OffsetBuffer; -use arrow::datatypes::{DataType, Field, Fields}; use criterion::{criterion_group, criterion_main, Criterion}; use datafusion_comet_spark_expr::murmur3::create_murmur3_hashes; use std::hint::black_box; -use std::sync::Arc; + +#[path = "common/hash_shapes.rs"] +mod hash_shapes; +use hash_shapes::*; mod common; #[path = "common/matched_maps.rs"] mod matched_maps; -const NUM_ROWS: usize = 8192; - -fn struct_fields() -> Fields { - vec![ - Arc::new(Field::new("a", DataType::Int32, true)), - Arc::new(Field::new("b", DataType::Utf8, true)), - ] - .into() -} - -fn struct_builder() -> StructBuilder { - StructBuilder::new( - struct_fields(), - vec![ - Box::new(Int32Builder::new()), - Box::new(StringBuilder::new()), - ], - ) -} - -fn append_struct(sb: &mut StructBuilder, i: usize) { - sb.field_builder::(0) - .unwrap() - .append_value(i as i32); - sb.field_builder::(1) - .unwrap() - .append_value(format!("v{}", i % 97)); - sb.append(true); -} - -/// `int32`, the cheapest leaf, as a reference point for the nested shapes. -fn primitive(num_rows: usize) -> ArrayRef { - Arc::new(Int32Array::from((0..num_rows as i32).collect::>())) -} - -/// `utf8`: variable-width, so the hash reads from the values buffer per row. -fn string(num_rows: usize) -> ArrayRef { - Arc::new(StringArray::from( - (0..num_rows) - .map(|i| format!("v{}", i % 97)) - .collect::>(), - )) -} - -/// `struct`: hashed field by field across the whole batch. -fn structs(num_rows: usize) -> ArrayRef { - let mut sb = struct_builder(); - for i in 0..num_rows { - append_struct(&mut sb, i); - } - Arc::new(sb.finish()) -} - -/// `array`: elements are primitives, so this takes the vectorized element path. -fn list_of_primitive(num_rows: usize, elems: usize) -> ArrayRef { - let mut lb = ListBuilder::new(Int32Builder::new()); - for i in 0..num_rows { - for j in 0..elems { - lb.values().append_value((i * 31 + j) as i32); - } - lb.append(true); - } - Arc::new(lb.finish()) -} - -/// `array>`: elements are nested, so this takes the per-element path. -fn list_of_struct(num_rows: usize, elems: usize) -> ArrayRef { - let mut lb = ListBuilder::new(struct_builder()); - for i in 0..num_rows { - for j in 0..elems { - append_struct(lb.values(), i * 31 + j); - } - lb.append(true); - } - Arc::new(lb.finish()) -} - -/// `array>` where one row is far longer than the rest, so the element count is spread -/// very unevenly across rows rather than uniformly. The per-element path slices and re-dispatches -/// once per element, so a batch dominated by a single long list has the same total work in a very -/// different distribution, which a uniform shape cannot show. -fn skewed_list_of_struct(num_rows: usize, long_len: usize) -> ArrayRef { - let mut lb = ListBuilder::new(struct_builder()); - for i in 0..num_rows { - let len = if i == 0 { long_len } else { 1 }; - for j in 0..len { - append_struct(lb.values(), i * 31 + j); - } - lb.append(true); - } - Arc::new(lb.finish()) -} - -/// `map`: keys and values are hashed entry by entry. -fn maps(num_rows: usize, entries: usize) -> ArrayRef { - let mut mb = MapBuilder::new( - Some(common::map_field_names()), - StringBuilder::new(), - Int32Builder::new(), - ); - for i in 0..num_rows { - for j in 0..entries { - mb.keys().append_value(format!("k{}", (i + j) % 97)); - mb.values().append_value((i * 31 + j) as i32); - } - mb.append(true).unwrap(); - } - Arc::new(mb.finish()) -} - -/// `struct>`: a map inside a struct, so the struct branch recurses -/// into the map specialization rather than into a leaf. -fn struct_of_map(num_rows: usize, entries: usize) -> ArrayRef { - let fields: Fields = vec![ - Arc::new(Field::new("a", DataType::Int32, true)), - Arc::new(Field::new( - "m", - DataType::Map( - Arc::new(Field::new( - "entries", - DataType::Struct( - vec![ - Arc::new(Field::new("key", DataType::Utf8, false)), - Arc::new(Field::new("value", DataType::Int32, true)), - ] - .into(), - ), - false, - )), - false, - ), - true, - )), - ] - .into(); - let ints = primitive(num_rows); - let ms = maps(num_rows, entries); - Arc::new(StructArray::new(fields, vec![ints, ms], None)) -} - -/// `array>`: the element is a list, so the non-primitive element path recurses into -/// the vectorized leaf loop one level down. -fn list_of_list(num_rows: usize, outer: usize, inner: usize) -> ArrayRef { - let mut lb = ListBuilder::new(ListBuilder::new(Int32Builder::new())); - for i in 0..num_rows { - for j in 0..outer { - for k in 0..inner { - lb.values() - .values() - .append_value((i * 31 + j * 7 + k) as i32); - } - lb.values().append(true); - } - lb.append(true); - } - Arc::new(lb.finish()) -} - -/// `map>`: a struct as the map value, which the key/value specializations do not -/// cover, so the value array is hashed recursively instead. -fn map_of_struct(num_rows: usize, entries: usize) -> ArrayRef { - let mut mb = MapBuilder::new( - Some(common::map_field_names()), - StringBuilder::new(), - struct_builder(), - ); - for i in 0..num_rows { - for j in 0..entries { - mb.keys().append_value(format!("k{}", (i + j) % 97)); - append_struct(mb.values(), i * 31 + j); - } - mb.append(true).unwrap(); - } - Arc::new(mb.finish()) -} - -/// `array>>`: three levels, so the per-element path recurses through a -/// struct into a map. -fn list_of_struct_of_map(num_rows: usize, elems: usize, entries: usize) -> ArrayRef { - let inner = struct_of_map(num_rows * elems, entries); - let offsets: Vec = (0..=num_rows).map(|i| (i * elems) as i32).collect(); - Arc::new(ListArray::new( - Arc::new(Field::new("item", inner.data_type().clone(), true)), - OffsetBuffer::new(offsets.into()), - inner, - None, - )) -} - fn bench(c: &mut Criterion) { let cases: Vec<(&str, ArrayRef)> = vec![ ("int32", primitive(NUM_ROWS)), @@ -249,13 +59,51 @@ fn bench(c: &mut Criterion) { "list_of_struct_of_map_x5x5", list_of_struct_of_map(NUM_ROWS, 5, 5), ), + // Shapes where gathering could cost more than the dispatches it saves. + ( + "list_of_struct_1kb_string_x4", + list_of_struct_big_string(2048, 4, 1024), + ), + ( + "list_of_struct_half_null_x10", + list_of_struct_half_null(NUM_ROWS, 10), + ), + ( + "list_of_struct_long_tail_x1024", + list_of_struct_long_tail(2, 1024), + ), + ( + "struct_of_dict_unreferenced_8mb", + struct_of_dict_unreferenced_big_value(8192, 8 * 1024 * 1024), + ), + // Shapes deliberately left on the per-element path, so a change to the eligibility rule + // shows up as a timing move here and not only in the allocation table. + ( + "null_parent_struct_64kb_x4", + null_parent_struct_big_string(128, 4, 65536), + ), + ( + "deep_singleton_list_5_deep", + deep_singleton_list_of_struct_big_string(2048, 5, 4096), + ), + ( + "width_skewed_8mb_first", + width_skewed_list_of_struct(8 * 1024 * 1024), + ), + ( + "sliced_list_retaining_10m_ints", + sliced_list_retaining_big_child(2, 10_000_000), + ), ]; let mut group = c.benchmark_group("murmur3"); for (name, array) in &cases { + // Size the buffer from the array rather than assuming `NUM_ROWS`, since not every shape + // uses that row count. + let rows = array.len(); group.bench_function(*name, |b| { b.iter(|| { - let mut hashes = vec![42u32; NUM_ROWS]; + let mut hashes = vec![42u32; rows]; create_murmur3_hashes(std::slice::from_ref(array), &mut hashes).unwrap(); black_box(&hashes); }) diff --git a/native/spark-expr/benches/hash_alloc.rs b/native/spark-expr/benches/hash_alloc.rs new file mode 100644 index 00000000000..0003ad4e4a4 --- /dev/null +++ b/native/spark-expr/benches/hash_alloc.rs @@ -0,0 +1,485 @@ +// 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. + +//! Allocation cost of the murmur3 hash kernel, on the same shapes `hash.rs` times. +//! +//! The nested-element list path gathers elements with `arrow::compute::take`, which allocates an +//! index array per pass and copies the selected element payloads. A timing benchmark cannot show +//! that cost, and it is the part that scales with payload width rather than with element count, so +//! it is reported separately here. +//! +//! Two numbers per shape: +//! +//! - **total** allocated-plus-growth bytes over the whole hash call: every `alloc` size, plus the +//! growth part of every `realloc` that grew. It counts a temporary even if it is freed +//! immediately, and it does not re-count the bytes a `realloc` carried over. This is the +//! throughput-relevant figure: what the allocator has to service. +//! - **peak** live bytes: the high-water mark of bytes *requested* and not yet freed, counting only +//! allocations made inside the measured window. It is a logical figure, not RSS, and it cannot see +//! a transient inside `System.realloc` that briefly holds the old and new blocks at once. This is the footprint-relevant figure. Successive +//! passes do not add up, because each frees its gather before the next allocates, but nesting +//! levels do: an outer gather stays live while the recursion below it builds its own, so a deep +//! element type can peak at the sum down the nesting. +//! +//! The two move in opposite directions for the batched gather, which is why both are reported. +//! Slicing one element at a time allocates a small array per element, so its total is large while +//! its peak stays near one element. Gathering a whole pass allocates once per pass for all +//! surviving rows, so its total is much smaller while its peak is a gathered pass. For a wide or +//! deeply nested child the gather copies more than the elements it hashes, so both figures can be +//! worse than the per-element path; `GATHER_ELIGIBLE_CHILD_BYTES` in the kernel is what keeps those shapes +//! on the slice path, and `null_parent_struct_64kb_x4` and `deep_singleton_list_5_deep` are here to +//! keep that guard measured. +//! +//! Run with `cargo bench --bench hash_alloc`. It prints a table rather than asserting a bound: the +//! numbers are for comparing two builds, and a threshold would either be too loose to catch a +//! regression or too tight to survive an Arrow upgrade. + +use datafusion_comet_spark_expr::murmur3::create_murmur3_hashes; +use std::alloc::{GlobalAlloc, Layout, System}; +use std::cell::Cell; +use std::collections::HashMap; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Mutex; + +#[path = "common/hash_shapes.rs"] +mod hash_shapes; +use hash_shapes::*; + +/// Counting allocator. +/// +/// Only allocations made inside the measured window are tracked, and they are tracked by pointer. +/// Without that, a buffer allocated before the window and freed inside it would be subtracted from +/// `live` and hide part of the peak, and a shrinking `realloc` would never give its bytes back. +/// Clamping the subtraction instead would only paper over both. +/// +/// The pointer set makes the allocator non-reentrant, so `IN_HOOK` guards against the set's own +/// allocations recursing back in. +struct Counting; + +static TOTAL: AtomicUsize = AtomicUsize::new(0); +static LIVE: AtomicUsize = AtomicUsize::new(0); +static PEAK: AtomicUsize = AtomicUsize::new(0); +static ON: AtomicUsize = AtomicUsize::new(0); + +/// What the window knows about one live allocation. +/// +/// The two fields are not the same number, and conflating them was a bug: a block allocated inside +/// the window is charged in full, while a block that existed *before* the window and was then grown +/// is charged only for the growth. A second `realloc` of either has to adjust by what this window +/// actually counted, not by the block's physical size. +#[derive(Clone, Copy)] +struct Tracked { + /// Current physical size of the block. + physical: usize, + /// How many bytes of it this window has counted toward `live`. + counted: usize, +} + +/// Live allocations the window knows about, keyed by address. A plain `Mutex` is enough: +/// the benchmark is single-threaded, and correctness here matters more than speed. +static TRACKED: Mutex>> = Mutex::new(None); + +thread_local! { + /// Set while inside the allocator hook, so the bookkeeping's own allocations are not counted + /// and cannot recurse. + static IN_HOOK: Cell = const { Cell::new(false) }; +} + +fn counting() -> bool { + ON.load(Ordering::Relaxed) == 1 && !IN_HOOK.with(|f| f.get()) +} + +/// Runs `f` with the hook flag set, so allocations it makes itself are ignored. +fn in_hook(f: impl FnOnce(&mut HashMap) -> R) -> Option { + IN_HOOK.with(|flag| { + if flag.get() { + return None; + } + flag.set(true); + let mut guard = TRACKED.lock().unwrap(); + let map = guard.get_or_insert_with(HashMap::new); + let out = f(map); + drop(guard); + flag.set(false); + Some(out) + }) +} + +/// Records a block allocated inside the window: physical size and counted size are the same. +fn record_alloc(ptr: *mut u8, size: usize) { + let entry = Tracked { + physical: size, + counted: size, + }; + if in_hook(|map| map.insert(ptr as usize, entry)).is_some() { + TOTAL.fetch_add(size, Ordering::Relaxed); + let live = LIVE.fetch_add(size, Ordering::Relaxed) + size; + PEAK.fetch_max(live, Ordering::Relaxed); + } +} + +/// Records a block that existed before the window and has now grown: only the growth is this +/// window's, so that is what `counted` holds even though `physical` is the whole block. +fn adopt_grown(ptr: *mut u8, physical: usize, grew: usize) { + let entry = Tracked { + physical, + counted: grew, + }; + if in_hook(|map| map.insert(ptr as usize, entry)).is_some() { + TOTAL.fetch_add(grew, Ordering::Relaxed); + let live = LIVE.fetch_add(grew, Ordering::Relaxed) + grew; + PEAK.fetch_max(live, Ordering::Relaxed); + } +} + +/// Drops `ptr` from the tracked set, returning what the window knew about it. `None` means the +/// window never counted it, so it must not touch `live`. +fn forget_alloc(ptr: *mut u8) -> Option { + in_hook(|map| map.remove(&(ptr as usize))).flatten() +} + +unsafe impl GlobalAlloc for Counting { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + let ptr = unsafe { System.alloc(layout) }; + if !ptr.is_null() && counting() { + record_alloc(ptr, layout.size()); + } + ptr + } + + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + if counting() { + if let Some(t) = forget_alloc(ptr) { + LIVE.fetch_sub(t.counted, Ordering::Relaxed); + } + } + unsafe { System.dealloc(ptr, layout) } + } + + unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { + let new_ptr = unsafe { System.realloc(ptr, layout, new_size) }; + if new_ptr.is_null() || !counting() { + return new_ptr; + } + match forget_alloc(ptr) { + // A block this window knows about. Release what the window counted before recording the + // new size: recording first would leave both live at once and push `peak` to old + new, + // which no moment of the program holds -- `realloc` either extends in place or frees as + // it copies. + Some(t) => { + LIVE.fetch_sub(t.counted, Ordering::Relaxed); + // Newly requested bytes are measured against the block's PHYSICAL size, since that + // is what already existed. The window's `counted` may be smaller -- a block adopted + // from before the window carries only its growth -- so using it here would charge + // the pre-existing bytes again on every subsequent realloc. + let newly_requested = new_size.saturating_sub(t.physical); + // Shrinking gives bytes back. Subtract the physical reduction rather than clamping + // to `new_size`: a block adopted from before the window counts less than its + // physical size, so clamping would leave counted bytes that no longer exist. An + // adopted 100-byte block grown to 200 and shrunk to 150 has counted 100 and must + // end at 50, not at `min(100, 150)`. + let released = t.physical.saturating_sub(new_size); + let entry = Tracked { + physical: new_size, + counted: t + .counted + .saturating_add(newly_requested) + .saturating_sub(released) + .min(new_size), + }; + if in_hook(|map| map.insert(new_ptr as usize, entry)).is_some() { + TOTAL.fetch_add(newly_requested, Ordering::Relaxed); + let live = LIVE.fetch_add(entry.counted, Ordering::Relaxed) + entry.counted; + PEAK.fetch_max(live, Ordering::Relaxed); + } + } + // A block from before the window. Count only the growth, and start tracking it so a + // later shrink or free has something to release -- but record its physical size too, so + // a *second* realloc adjusts against the right number. + None => { + let grew = new_size.saturating_sub(layout.size()); + if grew > 0 { + adopt_grown(new_ptr, new_size, grew); + } + } + } + new_ptr + } +} + +#[global_allocator] +static ALLOC: Counting = Counting; + +/// Hashes `array` once with counting on, returning `(total, peak, hashes)`. +/// +/// The hash buffer is allocated before counting starts, so the numbers describe the kernel's own +/// temporaries rather than the caller's output buffer. +/// Opens a measurement window: counters AND tracked records are reset together. +/// +/// Resetting only the counters leaves the previous window's live allocations in the map, and freeing +/// one of them inside this window subtracts from a `live` that never counted it, wrapping it toward +/// `usize::MAX` and poisoning `peak`. +fn begin_window() { + in_hook(|map| map.clear()); + TOTAL.store(0, Ordering::SeqCst); + LIVE.store(0, Ordering::SeqCst); + PEAK.store(0, Ordering::SeqCst); + ON.store(1, Ordering::SeqCst); +} + +/// Closes the window and returns `(total, peak)`. Anything still tracked is dropped with the window, +/// so it cannot affect a later one. +fn end_window() -> (usize, usize) { + ON.store(0, Ordering::SeqCst); + let total = TOTAL.load(Ordering::SeqCst); + let peak = PEAK.load(Ordering::SeqCst); + in_hook(|map| map.clear()); + (total, peak) +} + +fn measure(array: &arrow::array::ArrayRef) -> (usize, usize, Vec) { + let mut hashes = vec![42u32; array.len()]; + + begin_window(); + create_murmur3_hashes(std::slice::from_ref(array), &mut hashes).unwrap(); + let (total, peak) = end_window(); + + (total, peak, hashes) +} + +/// The counters are the instrument, so they are checked before the numbers they produce are quoted. +/// Every case below corresponds to a bug this allocator actually had: a clamped subtraction that +/// never released a shrink, a pre-window free that offset in-window bytes, a second `realloc` of an +/// adopted block adjusting by the wrong field, and a window that inherited the previous window's +/// records. +/// +/// Uses raw `alloc`/`realloc`/`dealloc` so each step is exact and nothing depends on how `Vec` +/// chooses capacity. Every pointer goes through `black_box`: these allocations are otherwise dead, +/// and the optimiser is entitled to delete them, which silently makes the whole check vacuous. +/// `alloc_zeroed` routes through the default `GlobalAlloc` implementation, which calls `alloc`, so +/// it would be counted too. +fn self_check() { + let l = |n: usize| Layout::from_size_align(n, 8).unwrap(); + + // A plain allocation is counted in full, and freeing it inside the window releases it. + unsafe { + begin_window(); + let p = std::hint::black_box(std::alloc::alloc(l(100))); + assert_eq!( + (TOTAL.load(Ordering::SeqCst), LIVE.load(Ordering::SeqCst)), + (100, 100), + "a plain allocation must be counted in full" + ); + std::alloc::dealloc(p, l(100)); + assert_eq!(LIVE.load(Ordering::SeqCst), 0, "a free must release it"); + let (total, peak) = end_window(); + assert_eq!((total, peak), (100, 100)); + } + + // A shrinking realloc has to give its bytes back, or peak keeps the pre-shrink size. + unsafe { + begin_window(); + let p = std::hint::black_box(std::alloc::alloc(l(100))); + let p = std::hint::black_box(std::alloc::realloc(p, l(100), 10)); + assert_eq!( + LIVE.load(Ordering::SeqCst), + 10, + "a shrink must release the difference" + ); + let q = std::hint::black_box(std::alloc::alloc(l(100))); + let (total, peak) = end_window(); + std::alloc::dealloc(p, l(10)); + std::alloc::dealloc(q, l(100)); + // 100 allocated, shrunk to 10 (no new bytes), then 100 more: peak is 10 + 100. + assert_eq!((total, peak), (200, 110), "shrink then grow accounting"); + } + + // Freeing a block from before the window must not offset the window's own live bytes. + unsafe { + let before = std::hint::black_box(std::alloc::alloc(l(100))); + begin_window(); + let a = std::hint::black_box(std::alloc::alloc(l(100))); + std::alloc::dealloc(before, l(100)); + let b = std::hint::black_box(std::alloc::alloc(l(100))); + let (_, peak) = end_window(); + std::alloc::dealloc(a, l(100)); + std::alloc::dealloc(b, l(100)); + assert_eq!( + peak, 200, + "a pre-window free must not offset in-window live bytes" + ); + } + + // A pre-window block grown twice inside the window: only growth counts, and the second grow + // must adjust against what the window counted, not the block's physical size. + unsafe { + let before = std::hint::black_box(std::alloc::alloc(l(100))); + begin_window(); + let p = std::hint::black_box(std::alloc::realloc(before, l(100), 200)); + assert_eq!( + (TOTAL.load(Ordering::SeqCst), LIVE.load(Ordering::SeqCst)), + (100, 100), + "growing a pre-window block counts only the growth" + ); + let p = std::hint::black_box(std::alloc::realloc(p, l(200), 300)); + let (total, live) = (TOTAL.load(Ordering::SeqCst), LIVE.load(Ordering::SeqCst)); + let (_, _) = end_window(); + std::alloc::dealloc(p, l(300)); + assert_eq!( + (total, live), + (200, 200), + "a second grow adds only its own growth" + ); + } + + // An adopted block grown then shrunk: the growth counts, the shrink gives back the physical + // reduction, not a clamp to the new size. + unsafe { + let before = std::hint::black_box(std::alloc::alloc(l(100))); + begin_window(); + let p = std::hint::black_box(std::alloc::realloc(before, l(100), 200)); + let p = std::hint::black_box(std::alloc::realloc(p, l(200), 150)); + let live_after_shrink = LIVE.load(Ordering::SeqCst); + let q = std::hint::black_box(std::alloc::alloc(l(100))); + let (total, peak) = end_window(); + std::alloc::dealloc(p, l(150)); + std::alloc::dealloc(q, l(100)); + assert_eq!( + live_after_shrink, 50, + "an adopted block grown to 200 then shrunk to 150 must count 50" + ); + assert_eq!( + (total, peak), + (200, 150), + "growth counts once; peak is 50 live plus the later 100" + ); + } + + // A window must not inherit the previous window's records, or freeing an older block wraps + // `live` toward usize::MAX. + unsafe { + begin_window(); + let stale = std::hint::black_box(std::alloc::alloc(l(100))); + let _ = end_window(); + + begin_window(); + std::alloc::dealloc(stale, l(100)); + let live = LIVE.load(Ordering::SeqCst); + let (_, peak) = end_window(); + assert_eq!( + live, 0, + "freeing a block tracked by an earlier window must not underflow live" + ); + assert_eq!(peak, 0, "and must not poison peak"); + } +} + +fn main() { + self_check(); + + // Same shapes and sizes as `hash.rs`, so the timing and allocation tables line up. The + // gather-heavy ones are the point; the primitive-element and non-list shapes are controls that + // must not move. + let cases: Vec<(&str, arrow::array::ArrayRef)> = vec![ + ("int32", primitive(NUM_ROWS)), + ("list_of_int32_x10", list_of_primitive(NUM_ROWS, 10)), + ("list_of_struct_x10", list_of_struct(NUM_ROWS, 10)), + ( + "list_of_struct_skewed_x1024", + skewed_list_of_struct(NUM_ROWS, 1024), + ), + ("list_of_list_x5x5", list_of_list(NUM_ROWS, 5, 5)), + ( + "list_of_struct_of_map_x5x5", + list_of_struct_of_map(NUM_ROWS, 5, 5), + ), + ( + "list_of_struct_1kb_string_x4", + list_of_struct_big_string(2048, 4, 1024), + ), + ( + "list_of_struct_half_null_x10", + list_of_struct_half_null(NUM_ROWS, 10), + ), + ( + "list_of_struct_long_tail_x1024", + list_of_struct_long_tail(2, 1024), + ), + // Adverse shapes: the gather copies bytes the hash never reads, and nested levels hold + // their gathers at the same time. + ( + "null_parent_struct_64kb_x4", + null_parent_struct_big_string(128, 4, 65536), + ), + ( + "deep_singleton_list_5_deep", + deep_singleton_list_of_struct_big_string(2048, 5, 4096), + ), + // The gather shares a dictionary's values instead of copying them, so a cost estimate + // taken from the child's total size can abandon batching for payload that never moves. + ( + "struct_of_dict_unreferenced_8mb", + struct_of_dict_unreferenced_big_value(8192, 8 * 1024 * 1024), + ), + // An average width is diluted by short elements, and a slice keeps its child's buffers, so + // these two are where a width-based estimate under-reads the gather. + ( + "width_skewed_8mb_first", + width_skewed_list_of_struct(8 * 1024 * 1024), + ), + ( + "sliced_list_retaining_10m_ints", + sliced_list_retaining_big_child(2, 10_000_000), + ), + // Eligible but with nothing to batch across, directly and via recursion: the scheduling + // buffers must not be allocated for a batch of one. + ( + "single_row_list_of_struct_x10", + single_row_list_of_struct(10), + ), + ( + "deep_singleton_narrow_5_deep", + deep_singleton_list_narrow(2048, 5), + ), + ]; + + println!( + "{:<32} {:>6} {:>14} {:>14} {:>12}", + "shape", "rows", "total bytes", "peak bytes", "hash" + ); + for (name, array) in &cases { + // Warm up first: the first call through a shape can allocate one-off caches that are not + // part of the per-call cost, and counting them would make the first row an outlier. + let (_, _, warm) = measure(array); + let (total, peak, hashes) = measure(array); + // The allocation numbers are only meaningful for a kernel that still produces the right + // hashes, so keep a check in the same run rather than trusting a separate one. Comparing + // two builds compares this digest too. + assert_eq!(warm, hashes, "{name}: hashing is not deterministic"); + let digest = hashes.iter().fold(0u64, |a, h| { + a.wrapping_mul(1_000_003).wrapping_add(*h as u64) + }); + println!( + "{:<32} {:>6} {:>14} {:>14} {:>12}", + name, + array.len(), + total, + peak, + digest % 1_000_000_007 + ); + } +} diff --git a/native/spark-expr/src/hash_funcs/murmur3.rs b/native/spark-expr/src/hash_funcs/murmur3.rs index a9e5b67aa9f..5bf7399b19b 100644 --- a/native/spark-expr/src/hash_funcs/murmur3.rs +++ b/native/spark-expr/src/hash_funcs/murmur3.rs @@ -194,6 +194,10 @@ pub fn create_murmur3_hashes<'a>( #[cfg(test)] mod tests { + /// Produced by the per-element `hash_list_array!` implementation. An empty list and a null + /// list both leave the seed untouched. + const EXPECTED_LIST_OF_STRUCT: [u32; 6] = + [262891156, 1206178823, 42, 42, 3798669693, 2032748937]; use arrow::array::{Float32Array, Float64Array}; use std::sync::Arc; @@ -387,6 +391,365 @@ mod tests { ); } + /// One `struct` element: `None` is a null struct, and the fields are + /// independently nullable. + type StructElem = Option<(Option, Option<&'static str>)>; + /// One `array>` row: `None` is a null list. + type ListRow = Option>; + + /// Builds `array>` from `rows`. + fn list_of_struct(rows: Vec) -> ArrayRef { + use arrow::array::{Int32Builder, ListBuilder, StringBuilder, StructBuilder}; + use arrow::datatypes::{DataType, Field, Fields}; + + let fields: Fields = vec![ + Arc::new(Field::new("a", DataType::Int32, true)), + Arc::new(Field::new("b", DataType::Utf8, true)), + ] + .into(); + let struct_builder = StructBuilder::new( + fields.clone(), + vec![ + Box::new(Int32Builder::new()), + Box::new(StringBuilder::new()), + ], + ); + let mut lb = ListBuilder::new(struct_builder); + for row in rows { + match row { + None => lb.append(false), + Some(elems) => { + for elem in elems { + let sb = lb.values(); + match elem { + None => { + sb.field_builder::(0).unwrap().append_null(); + sb.field_builder::(1).unwrap().append_null(); + sb.append(false); + } + Some((a, b)) => { + match a { + Some(v) => { + sb.field_builder::(0).unwrap().append_value(v) + } + None => { + sb.field_builder::(0).unwrap().append_null() + } + } + match b { + Some(v) => sb + .field_builder::(1) + .unwrap() + .append_value(v), + None => { + sb.field_builder::(1).unwrap().append_null() + } + } + sb.append(true); + } + } + } + lb.append(true); + } + } + } + Arc::new(lb.finish()) + } + + fn hash_of(array: ArrayRef, num_rows: usize) -> Vec { + let mut hashes = vec![42u32; num_rows]; + create_murmur3_hashes(&[array], &mut hashes).unwrap(); + hashes + } + + /// `array>` is the shape that goes through `hash_list_array!`, the per-element path. + /// These values were produced by that implementation and pin it: the hash decides which + /// partition a row lands in, so any rewrite of that path has to reproduce them exactly. + #[test] + fn test_list_of_struct_hashes_are_stable() { + let rows = vec![ + Some(vec![Some((Some(1), Some("x"))), Some((Some(2), Some("y")))]), + Some(vec![Some((Some(3), Some("z")))]), + // empty list: contributes nothing, so the seed survives + Some(vec![]), + // null list + None, + // null struct element, and elements with null fields + Some(vec![None, Some((None, Some("w"))), Some((Some(4), None))]), + // repeated element values, to catch an implementation that dedupes or reorders + Some(vec![Some((Some(5), Some("s"))), Some((Some(5), Some("s")))]), + ]; + let n = rows.len(); + assert_eq!(hash_of(list_of_struct(rows), n), EXPECTED_LIST_OF_STRUCT); + } + + /// Element order must matter: Spark chains the element hashes in sequence. + #[test] + fn test_list_of_struct_is_order_sensitive() { + let forward = list_of_struct(vec![Some(vec![ + Some((Some(1), Some("a"))), + Some((Some(2), Some("b"))), + ])]); + let reversed = list_of_struct(vec![Some(vec![ + Some((Some(2), Some("b"))), + Some((Some(1), Some("a"))), + ])]); + assert_ne!( + hash_of(forward, 1), + hash_of(reversed, 1), + "element order must change the hash" + ); + } + + /// The batched implementation makes one pass per element position, so a single long list forces + /// as many passes as its length while every other row is already finished. Check that a skewed + /// batch still agrees with hashing each row on its own, which is what the per-element + /// implementation effectively did. + #[test] + fn test_list_of_struct_skewed_lengths() { + let mut rows: Vec = vec![Some(vec![Some((Some(1), Some("a")))]); 8]; + // One row far longer than the rest. + rows.push(Some( + (0..64) + .map(|i| Some((Some(i), Some("long")))) + .collect::>(), + )); + rows.push(Some(vec![])); + + let batched = hash_of(list_of_struct(rows.clone()), rows.len()); + + // Hash each row as its own batch of one; the result must match position by position. + let per_row: Vec = rows + .into_iter() + .map(|row| hash_of(list_of_struct(vec![row]), 1)[0]) + .collect(); + assert_eq!(batched, per_row, "skewed batch must match per-row hashing"); + } + + /// `LargeList` reaches the same code path with 64-bit offsets. This pins that the two list + /// widths agree on the same data. It does not exercise the 64-bit gather itself: narrowing the + /// indices only goes wrong past `u32::MAX` elements, which is far larger than a test can build, + /// so the index width is chosen from the offset type rather than guarded by a test here. + #[test] + fn test_large_list_of_struct_matches_list() { + use arrow::array::{Int32Builder, LargeListBuilder, StringBuilder, StructBuilder}; + use arrow::datatypes::{DataType, Field, Fields}; + + let fields: Fields = vec![ + Arc::new(Field::new("a", DataType::Int32, true)), + Arc::new(Field::new("b", DataType::Utf8, true)), + ] + .into(); + let sb = StructBuilder::new( + fields.clone(), + vec![ + Box::new(Int32Builder::new()), + Box::new(StringBuilder::new()), + ], + ); + let mut lb = LargeListBuilder::new(sb); + for row in [vec![1, 2], vec![3], vec![]] { + for v in row { + let s = lb.values(); + s.field_builder::(0).unwrap().append_value(v); + s.field_builder::(1) + .unwrap() + .append_value(format!("s{v}")); + s.append(true); + } + lb.append(true); + } + let large: ArrayRef = Arc::new(lb.finish()); + + let small = list_of_struct(vec![ + Some(vec![ + Some((Some(1), Some("s1"))), + Some((Some(2), Some("s2"))), + ]), + Some(vec![Some((Some(3), Some("s3")))]), + Some(vec![]), + ]); + + assert_eq!( + hash_of(large, 3), + hash_of(small, 3), + "LargeList and List must hash identically" + ); + } + + /// `array>`: the element is a list rather than a struct, so it takes the same + /// non-primitive element path one level deeper. + #[test] + fn test_list_of_list_hashes_match_per_row() { + use arrow::array::{Int32Builder, ListBuilder}; + + let build = |rows: &[Vec>]| -> ArrayRef { + let mut lb = ListBuilder::new(ListBuilder::new(Int32Builder::new())); + for row in rows { + for inner in row { + for v in inner { + lb.values().values().append_value(*v); + } + lb.values().append(true); + } + lb.append(true); + } + Arc::new(lb.finish()) + }; + + let rows = vec![ + vec![vec![1, 2], vec![3]], + vec![vec![4]], + vec![], + vec![vec![5, 6, 7], vec![], vec![8]], + ]; + let batched = hash_of(build(&rows), rows.len()); + let per_row: Vec = rows + .iter() + .map(|row| hash_of(build(std::slice::from_ref(row)), 1)[0]) + .collect(); + assert_eq!( + batched, per_row, + "array> must match per-row hashing" + ); + } + + /// The cursor drops a row once its elements run out, so rows finishing on different passes, + /// including ones in the middle of the batch, exercise the survivor bookkeeping. Compared + /// against hashing each row as its own batch. + #[test] + fn test_list_of_struct_rows_exhaust_on_different_passes() { + let rows: Vec = vec![ + Some((0..5).map(|i| Some((Some(i), Some("a")))).collect()), + Some(vec![Some((Some(9), Some("b")))]), + Some((0..3).map(|i| Some((Some(i), Some("c")))).collect()), + Some(vec![]), + Some((0..7).map(|i| Some((Some(i), Some("d")))).collect()), + None, + Some(vec![Some((Some(1), Some("e"))), Some((Some(2), Some("f")))]), + ]; + let batched = hash_of(list_of_struct(rows.clone()), rows.len()); + let per_row: Vec = rows + .into_iter() + .map(|row| hash_of(list_of_struct(vec![row]), 1)[0]) + .collect(); + assert_eq!( + batched, per_row, + "rows exhausting on different passes must agree" + ); + } + + /// A sliced list has a non-zero first offset, and the cursor subtracts it when building gather + /// indices, so an off-by-one there would only show up on a slice. + #[test] + fn test_sliced_list_of_struct_matches_unsliced() { + let rows: Vec = vec![ + Some(vec![Some((Some(1), Some("x")))]), + Some(vec![Some((Some(2), Some("y"))), Some((Some(3), Some("z")))]), + Some((0..4).map(|i| Some((Some(i), Some("w")))).collect()), + Some(vec![Some((Some(8), Some("v")))]), + ]; + let full = list_of_struct(rows.clone()); + + // Hash rows 1..3 through a slice, and the same rows built on their own. + let sliced = full.slice(1, 2); + let mut from_slice = vec![42u32; 2]; + create_murmur3_hashes(&[sliced], &mut from_slice).unwrap(); + + let standalone = list_of_struct(rows[1..3].to_vec()); + let mut from_standalone = vec![42u32; 2]; + create_murmur3_hashes(&[standalone], &mut from_standalone).unwrap(); + + assert_eq!( + from_slice, from_standalone, + "a sliced list must hash like the same rows built unsliced" + ); + } + + /// A null list can still cover a non-empty range of elements. Those elements must not be + /// hashed, and the row must not join the cursor. + #[test] + fn test_null_list_with_populated_range_is_skipped() { + use arrow::array::{Int32Builder, ListArray, StringBuilder, StructBuilder}; + use arrow::buffer::{NullBuffer, OffsetBuffer}; + use arrow::datatypes::{DataType, Field, Fields}; + + let fields: Fields = vec![ + Arc::new(Field::new("a", DataType::Int32, true)), + Arc::new(Field::new("b", DataType::Utf8, true)), + ] + .into(); + let mut sb = StructBuilder::new( + fields.clone(), + vec![ + Box::new(Int32Builder::new()), + Box::new(StringBuilder::new()), + ], + ); + for v in 0..3 { + sb.field_builder::(0).unwrap().append_value(v); + sb.field_builder::(1) + .unwrap() + .append_value("hidden"); + sb.append(true); + } + let elements: ArrayRef = Arc::new(sb.finish()); + + // Row 0 covers elements 0..1, row 1 is null but still covers 1..3. + let list = ListArray::new( + Arc::new(Field::new("item", elements.data_type().clone(), true)), + OffsetBuffer::new(vec![0i32, 1, 3].into()), + elements, + Some(NullBuffer::from(vec![true, false])), + ); + let mut hashes = vec![42u32; 2]; + create_murmur3_hashes(&[Arc::new(list) as ArrayRef], &mut hashes).unwrap(); + assert_eq!(hashes[1], 42, "a null list must leave the seed untouched"); + + // And it must not disturb the visible row either. + let only_visible = list_of_struct(vec![Some(vec![Some((Some(0), Some("hidden")))])]); + let mut expected = vec![42u32; 1]; + create_murmur3_hashes(&[only_visible], &mut expected).unwrap(); + assert_eq!(hashes[0], expected[0]); + } + + /// Lengths that step down rather than being either all equal or one long outlier, so the + /// uniform-length gate is not taken and rows leave the cursor on consecutive passes. + #[test] + fn test_list_of_struct_descending_lengths() { + let rows: Vec = (1..=12) + .rev() + .map(|n| Some((0..n).map(|i| Some((Some(i), Some("s")))).collect())) + .collect(); + let batched = hash_of(list_of_struct(rows.clone()), rows.len()); + let per_row: Vec = rows + .into_iter() + .map(|row| hash_of(list_of_struct(vec![row]), 1)[0]) + .collect(); + assert_eq!(batched, per_row); + } + + /// A lone-row list must not stop the caller hashing the remaining columns. + #[test] + fn test_single_row_list_then_another_column() { + let l = list_of_struct(vec![Some(vec![ + Some((Some(1), Some("a"))), + Some((Some(2), Some("b"))), + ])]); + let other: ArrayRef = Arc::new(arrow::array::Int32Array::from(vec![7])); + + // both columns together + let mut both = vec![42u32; 1]; + create_murmur3_hashes(&[Arc::clone(&l), Arc::clone(&other)], &mut both).unwrap(); + + // chaining them by hand must agree + let mut step = vec![42u32; 1]; + create_murmur3_hashes(&[l], &mut step).unwrap(); + create_murmur3_hashes(&[other], &mut step).unwrap(); + + assert_eq!(both, step, "the second column must still be hashed"); + } + #[test] fn test_i8() { test_murmur3_hash::( @@ -460,4 +823,133 @@ mod tests { test_murmur3_hash::(input.clone(), expected); } + /// Both sides of the eligibility threshold must produce the same hashes, since the check only + /// decides which path runs. A struct of flat leaves under the size limit is batched; the same + /// data behind a child large enough to fail the limit is sliced. An off-by-one in either path's + /// element index shows up as a mismatch. + #[test] + fn eligible_and_ineligible_shapes_hash_alike() { + use crate::hash_funcs::utils::{gather_is_eligible, GATHER_ELIGIBLE_CHILD_BYTES}; + use arrow::array::builder::{Int32Builder, ListBuilder, StringBuilder, StructBuilder}; + use arrow::datatypes::{DataType, Field, Fields}; + + fn build(payload_len: usize) -> ArrayRef { + let fields: Fields = vec![ + Arc::new(Field::new("a", DataType::Int32, true)), + Arc::new(Field::new("b", DataType::Utf8, true)), + ] + .into(); + let mut lb = ListBuilder::new(StructBuilder::new( + fields, + vec![ + Box::new(Int32Builder::new()), + Box::new(StringBuilder::new()), + ], + )); + // Uneven lengths with an empty row and a null row, so rows drop out on different passes. + let payload = "x".repeat(payload_len); + for (row, len) in [3usize, 0, 4, 1, 2].iter().enumerate() { + for i in 0..*len { + let sb = lb.values(); + sb.field_builder::(0) + .unwrap() + .append_value((row * 10 + i) as i32); + sb.field_builder::(1) + .unwrap() + .append_value(&payload); + sb.append(i % 3 != 2); + } + lb.append(row != 1); + } + Arc::new(lb.finish()) + } + + let small = build(4); + let small_elements = small + .as_any() + .downcast_ref::() + .unwrap() + .values(); + assert!( + gather_is_eligible(small_elements.as_ref()), + "a small flat struct should be batched" + ); + + // One wide value pushes the retained child past the limit, so the same shape is sliced. + let big = build(GATHER_ELIGIBLE_CHILD_BYTES); + let big_elements = big + .as_any() + .downcast_ref::() + .unwrap() + .values(); + assert!( + !gather_is_eligible(big_elements.as_ref()), + "a child over the retained-size limit should not be batched" + ); + + // A following column, so leaving the pass loop must not skip it on either path. + let following: ArrayRef = Arc::new(Int32Array::from(vec![5, 6, 7, 8, 9])); + let seeds = [11u32, 22, 33, 44, 55]; + + for (name, array) in [("batched", &small), ("sliced", &big)] { + let mut got = seeds; + create_murmur3_hashes(&[Arc::clone(array), Arc::clone(&following)], &mut got).unwrap(); + + // Reference: hash each row alone, which cannot batch across rows at all. + let mut want = [0u32; 5]; + for row in 0..5 { + let mut one = [seeds[row]]; + create_murmur3_hashes(&[array.slice(row, 1), following.slice(row, 1)], &mut one) + .unwrap(); + want[row] = one[0]; + } + assert_eq!(got, want, "{name}: must agree with hashing each row alone"); + } + } + + /// A nested element type is not eligible however small it is, so the shape that used to + /// accumulate a gather per nesting level keeps the per-element path. + #[test] + fn nested_and_dictionary_elements_are_not_eligible() { + use crate::hash_funcs::utils::gather_is_eligible; + use arrow::array::{DictionaryArray, Int32Array as I32, StringArray, StructArray}; + use arrow::datatypes::{DataType, Field, Fields, Int32Type}; + + // struct>: the child recurses, so each level would hold its own gather. + let inner: ArrayRef = Arc::new(I32::from(vec![1, 2, 3, 4])); + let offsets = arrow::buffer::OffsetBuffer::from_lengths([2usize, 2]); + let list_child: ArrayRef = Arc::new(arrow::array::ListArray::new( + Arc::new(Field::new("item", DataType::Int32, true)), + offsets, + inner, + None, + )); + let nested_fields: Fields = vec![Arc::new(Field::new( + "l", + list_child.data_type().clone(), + true, + ))] + .into(); + let nested: ArrayRef = Arc::new(StructArray::new(nested_fields, vec![list_child], None)); + assert!( + !gather_is_eligible(nested.as_ref()), + "a nested child must keep the per-element path" + ); + + // struct>: `take` shares the values, so batching is not modelled here. + let dict: ArrayRef = Arc::new( + DictionaryArray::::try_new( + I32::from(vec![0, 0]), + Arc::new(StringArray::from(vec!["x"])), + ) + .unwrap(), + ); + let dict_fields: Fields = + vec![Arc::new(Field::new("d", dict.data_type().clone(), true))].into(); + let with_dict: ArrayRef = Arc::new(StructArray::new(dict_fields, vec![dict], None)); + assert!( + !gather_is_eligible(with_dict.as_ref()), + "a dictionary child must keep the per-element path" + ); + } } diff --git a/native/spark-expr/src/hash_funcs/utils.rs b/native/spark-expr/src/hash_funcs/utils.rs index e6eedd31df0..db9f95db824 100644 --- a/native/spark-expr/src/hash_funcs/utils.rs +++ b/native/spark-expr/src/hash_funcs/utils.rs @@ -17,6 +17,8 @@ //! This includes utilities for hashing and murmur3 hashing. +use arrow::array::Array; + #[macro_export] macro_rules! hash_array { ($array_type: ident, $column: ident, $hashes: ident, $hash_method: ident) => { @@ -509,6 +511,78 @@ macro_rules! hash_list_with_primitive_elements { }; } +/// Whether the batched gather is used for a list whose elements are `values`. +/// +/// Batching replaces a per-element slice and dispatch with one `arrow::compute::take` per element +/// position. That is a large win for a small flat struct, but `take` copies the selected payload, +/// and how much it copies is not something this kernel can predict cheaply: +/// +/// - A width-based estimate is diluted by short elements. Two rows of 1024 structs where only the +/// first string is 8 MiB average out to about 16 KB per element while the gather copies 16 MiB. +/// - A sliced list keeps its child's buffers. Slicing away the one row that held ten million ints +/// leaves two cheap visible rows and a child that `take` still pre-sizes from, turning a 128-byte +/// peak into 40 MB. +/// - A dictionary child is shared rather than copied, so charging for its payload abandons batching +/// on a shape that copies nothing. +/// - A nested child recurses, and each level's gather stays live while the level below builds its +/// own, so cost accumulates down the depth. +/// +/// Rather than model all of that, this admits only the shape whose cost is easy to bound -- a struct +/// of flat leaves -- and requires the child's *retained* buffers to fit a conservative limit, not an +/// average per element, so a large buffer kept alive by a slice disqualifies the gather even when +/// few rows are visible. Everything else keeps the previous per-element path. +/// +/// The limit is an eligibility condition for the optimization, not a bound on the peak memory of a +/// hash call: the index array, the row mapping and Arrow's own metadata are extra. +pub fn gather_is_eligible(values: &dyn Array) -> bool { + use arrow::datatypes::DataType; + + let DataType::Struct(fields) = values.data_type() else { + return false; + }; + if !fields.iter().all(|f| is_flat_leaf(f.data_type())) { + return false; + } + // Retained size, not per element: a slice that hides a huge child must not qualify. + values.get_array_memory_size() <= GATHER_ELIGIBLE_CHILD_BYTES +} + +/// Leaf types whose gather cost is proportional to the rows picked, with no shared or nested +/// payload behind them. Deliberately conservative: a type absent here keeps the previous path, and +/// adding one means measuring it. +fn is_flat_leaf(data_type: &arrow::datatypes::DataType) -> bool { + use arrow::datatypes::DataType; + + matches!( + data_type, + DataType::Boolean + | DataType::Int8 + | DataType::Int16 + | DataType::Int32 + | DataType::Int64 + | DataType::UInt8 + | DataType::UInt16 + | DataType::UInt32 + | DataType::UInt64 + | DataType::Float32 + | DataType::Float64 + | DataType::Date32 + | DataType::Date64 + | DataType::Timestamp(_, _) + | DataType::Decimal128(_, _) + | DataType::Utf8 + | DataType::Binary + ) +} + +/// Ceiling on a gathered child's retained buffers for the batched path to be used. +/// +/// Compared against the whole child rather than a per-pass estimate, which is what makes a large +/// buffer held alive by a slice fail the check. It is not a per-pass bound: when every row holds one +/// element, a single pass gathers nearly the whole child, so a shape just under the limit can gather +/// close to it in one pass. +pub const GATHER_ELIGIBLE_CHILD_BYTES: usize = 4 * 1024 * 1024; + #[macro_export] macro_rules! hash_list_array { ($array_type:ident, $offset_type:ty, $column: ident, $hashes: ident, $recursive_hash_method: ident) => { @@ -526,35 +600,178 @@ macro_rules! hash_list_array { let values = list_array.values(); let offsets = list_array.offsets(); - if list_array.null_count() == 0 { - // Fast path: no nulls, skip null checks - for (row_idx, hash) in $hashes.iter_mut().enumerate() { - let start = offsets[row_idx] as usize; - let end = offsets[row_idx + 1] as usize; - let len = end - start; - // Hash each element in sequence, chaining the hash values - for elem_idx in 0..len { - let elem_array = values.slice(start + elem_idx, 1); - let mut single_hash = [*hash]; - $recursive_hash_method(&[elem_array], &mut single_hash)?; - *hash = single_hash[0]; + // Spark chains the element hashes in order, so the elements of one row have to be hashed + // in sequence. What does not have to happen per element is the allocation and dispatch: + // slicing a one-element array and re-entering the hash dispatch for it costs an Arrow + // array plus a full type match every time, and for a struct element the dispatch also + // copies the field vector on every call. + // + // Instead, hash one element per row at a time in a single batched call, seeding each + // slot with the running hash of the row it belongs to. That is exactly what the + // per-element call did, so the result is bit-identical. + let total_elements = offsets[$hashes.len()] as usize - offsets[0] as usize; + if total_elements == 0 { + // Every list is empty or null; the seeds already hold the answer. + } else { + // Decide before any scheduling work. Eligibility depends only on the element type and + // the size of the child's retained buffers, not on which rows are active, so a shape + // that keeps the previous path allocates nothing new at all: no survivor scan, no + // sliced element view, no per-pass buffers. + // Batching needs several rows to batch across. One row, whether the whole column is one + // row or only one row is non-empty, has nothing to gather with, so take the per-element + // path before allocating a sliced element view, a survivor list or any per-pass buffer. + // Recursion makes this common rather than a corner case: a deep singleton list reaches a + // one-row batch at every level below the first. + let mut non_empty_rows = 0usize; + for row_idx in 0..$hashes.len() { + if !list_array.is_null(row_idx) + && offsets[row_idx + 1] > offsets[row_idx] + { + non_empty_rows += 1; + if non_empty_rows > 1 { + break; + } } } - } else { - // Slow path: array has nulls, check each row - for (row_idx, hash) in $hashes.iter_mut().enumerate() { - if !list_array.is_null(row_idx) { + if non_empty_rows <= 1 + || !$crate::hash_funcs::utils::gather_is_eligible(values.as_ref()) + { + for row_idx in 0..$hashes.len() { + if list_array.is_null(row_idx) { + continue; + } + let start = offsets[row_idx] as usize; + let end = offsets[row_idx + 1] as usize; + for elem_idx in start..end { + let elem = values.slice(elem_idx, 1); + let mut single = [$hashes[row_idx]]; + $recursive_hash_method(&[elem], &mut single)?; + $hashes[row_idx] = single[0]; + } + } + } else { + let first_offset = offsets[0] as usize; + let elements = values.slice(first_offset, total_elements); + + // Chaining means element k of a row can only be hashed once element k-1 is known, so + // batch by position: all the first elements together, then all the second, and so on. + // Rows are independent, so one pass per position is enough. + // + // Only rows that still have an element at the current position take part, and a row + // never becomes alive again once exhausted, so carry the surviving rows forward instead + // of rescanning all of them each pass. Rescanning would cost rows x longest-list, which + // for one long list among short ones is almost all wasted: 8192 rows with one list of + // 1024 scans 8.4M slots for 9215 elements. Carrying the survivors makes the scheduling + // work proportional to the elements actually hashed. + // + // Index the gather by the list's own offset width. A `LargeList` can hold more than + // `u32::MAX` elements, so narrowing the positions to `u32` would silently wrap and + // hash the wrong elements. + let mut active: Vec = Vec::with_capacity($hashes.len()); + // The same pass records whether every row is non-null with the same length. When it + // is, no row ever drops out early, so the survivor bookkeeping is pure overhead and + // the rows can simply be walked directly. + let mut uniform_len: Option = None; + let mut all_same = true; + for row_idx in 0..$hashes.len() { + if list_array.is_null(row_idx) { + all_same = false; + continue; + } + let len = offsets[row_idx + 1] as usize - offsets[row_idx] as usize; + if len > 0 { + active.push(row_idx); + } + match uniform_len { + None => uniform_len = Some(len), + Some(seen) if seen == len => {} + Some(_) => all_same = false, + } + } + let uniform = all_same && uniform_len.unwrap_or(0) > 0; + + // Only a batch that will actually gather needs this decision, and a single row always + // takes the direct path below, so skip the check for one row. Decided once for the + // column, never per pass. + + // Allocated only for the batched path, after the decision above. + let mut positions: Vec<$offset_type> = Vec::with_capacity(active.len()); + let mut rows_at_position: Vec = Vec::with_capacity(active.len()); + let mut still_active: Vec = Vec::with_capacity(active.len()); + let mut position_hashes = Vec::with_capacity(active.len()); + let mut position = 0usize; + let uniform_passes = if uniform { uniform_len.unwrap_or(0) } else { 0 }; + while (uniform && position < uniform_passes) || (!uniform && !active.is_empty()) { + // Batching pays only when a pass covers several rows. Once one row is left there is + // nothing to gather across: `take` would copy that row's remaining element payloads + // without saving a dispatch. That happens both for a batch that starts with a + // single non-empty row and, more often, for the tail after the shorter rows finish + // -- lengths [1, 1, 8] spend seven of eight passes on one row. Finish it by slicing, + // the way the previous implementation did throughout. + if active.len() == 1 { + let row_idx = active[0]; let start = offsets[row_idx] as usize; let end = offsets[row_idx + 1] as usize; - let len = end - start; - // Hash each element in sequence, chaining the hash values - for elem_idx in 0..len { - let elem_array = values.slice(start + elem_idx, 1); - let mut single_hash = [*hash]; - $recursive_hash_method(&[elem_array], &mut single_hash)?; - *hash = single_hash[0]; + for elem_idx in (start + position)..end { + let elem = values.slice(elem_idx, 1); + let mut single = [$hashes[row_idx]]; + $recursive_hash_method(&[elem], &mut single)?; + $hashes[row_idx] = single[0]; } + // `break`, not `return`: this macro runs inside the caller's loop over + // columns, so returning would skip every column after this one. Leaving + // `active` as it is costs nothing, since nothing reads it after the loop. + break; } + positions.clear(); + rows_at_position.clear(); + if uniform { + // Every row survives every pass, so skip the survivor bookkeeping. + for row_idx in active.iter().copied() { + let start = offsets[row_idx] as usize; + positions.push((start + position - first_offset) as $offset_type); + rows_at_position.push(row_idx); + } + } else { + still_active.clear(); + for row_idx in active.iter().copied() { + let start = offsets[row_idx] as usize; + let end = offsets[row_idx + 1] as usize; + positions.push((start + position - first_offset) as $offset_type); + rows_at_position.push(row_idx); + // Alive for the next pass only if it has an element beyond this one. + if start + position + 1 < end { + still_active.push(row_idx); + } + } + std::mem::swap(&mut active, &mut still_active); + } + position += 1; + // `take` accepts any integer index type, so index by the offset width: a + // `LargeList` can exceed `u32::MAX` elements. + let taken = if std::mem::size_of::<$offset_type>() > 4 { + let indices = arrow::array::Int64Array::from_iter_values( + positions.iter().map(|p| *p as i64), + ); + arrow::compute::take(&elements, &indices, None)? + } else { + let indices = arrow::array::Int32Array::from_iter_values( + positions.iter().map(|p| *p as i32), + ); + arrow::compute::take(&elements, &indices, None)? + }; + // The hash width differs per algorithm (u32 for murmur3, u64 for xxhash64), so + // let the element type come from the buffer rather than naming it here. Reused + // across passes so the gather does not reallocate each time. + position_hashes.clear(); + for row_idx in rows_at_position.iter() { + position_hashes.push($hashes[*row_idx]); + } + $recursive_hash_method(&[taken], &mut position_hashes)?; + for (slot, row_idx) in rows_at_position.iter().enumerate() { + $hashes[*row_idx] = position_hashes[slot]; + } + } } } }; diff --git a/native/spark-expr/src/hash_funcs/xxhash64.rs b/native/spark-expr/src/hash_funcs/xxhash64.rs index 93ac9304fc8..45c273bb9f5 100644 --- a/native/spark-expr/src/hash_funcs/xxhash64.rs +++ b/native/spark-expr/src/hash_funcs/xxhash64.rs @@ -385,4 +385,66 @@ mod tests { ], ) } + /// The nested-element list scheduler is shared between murmur3 and xxhash64 through + /// `create_hashes_internal!`, but its regressions live in the murmur3 tests. This covers the + /// combination that the scheduling is most likely to get wrong, on the other algorithm and the + /// wider hash width: rows whose lists have different lengths (so rows drop out on different + /// passes and the batch narrows to a single survivor), distinct incoming seeds per row (so a + /// mis-mapped slot shows up rather than being masked by a uniform seed), and a second column + /// hashed after the list (so breaking out of the pass loop must not skip it). + #[test] + fn xxhash64_list_of_struct_survivor_transition_with_seeds_and_following_column() { + use arrow::array::builder::{Int32Builder, ListBuilder, StringBuilder, StructBuilder}; + use arrow::datatypes::{DataType, Field, Fields}; + + let fields: Fields = vec![ + Arc::new(Field::new("a", DataType::Int32, true)), + Arc::new(Field::new("b", DataType::Utf8, true)), + ] + .into(); + let mut lb = ListBuilder::new(StructBuilder::new( + fields, + vec![ + Box::new(Int32Builder::new()), + Box::new(StringBuilder::new()), + ], + )); + // Descending lengths, so the uniform-length path is not taken and the batch narrows to one + // surviving row before the longest row is finished. + for (row, len) in [4usize, 2, 1, 3].iter().enumerate() { + for i in 0..*len { + let sb = lb.values(); + sb.field_builder::(0) + .unwrap() + .append_value((row * 10 + i) as i32); + sb.field_builder::(1) + .unwrap() + .append_value(format!("r{row}e{i}")); + sb.append(true); + } + lb.append(true); + } + let list: ArrayRef = Arc::new(lb.finish()); + let following: ArrayRef = Arc::new(Int32Array::from(vec![7, 8, 9, 10])); + + // Distinct seeds: a slot written back to the wrong row changes the result. + let seeds = [1u64, 2, 3, 4]; + + let mut batched = seeds; + create_xxhash64_hashes(&[Arc::clone(&list), Arc::clone(&following)], &mut batched).unwrap(); + + // Reference: hash each row on its own, which cannot batch across rows at all. + let mut per_row = [0u64; 4]; + for row in 0..4 { + let mut one = [seeds[row]]; + create_xxhash64_hashes(&[list.slice(row, 1), following.slice(row, 1)], &mut one) + .unwrap(); + per_row[row] = one[0]; + } + + assert_eq!( + batched, per_row, + "batched scheduling must agree with hashing each row alone" + ); + } } diff --git a/spark/src/test/scala/org/apache/comet/CometHashExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometHashExpressionSuite.scala index c997da0012f..68c4471e05d 100644 --- a/spark/src/test/scala/org/apache/comet/CometHashExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometHashExpressionSuite.scala @@ -524,6 +524,50 @@ class CometHashExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelpe } } + test("hash - array of struct") { + // `array>` is the shape the batched nested-element path handles, so this is the + // Spark oracle for it: the rows differ in length, and cover an empty array, a null array, null + // struct elements and null fields, which are the cases where the element scheduling could + // assign a seed to the wrong row without changing the row count. `checkSparkAnswerAndOperator` + // also asserts Comet actually ran it natively, so a fallback on both sides cannot hide a + // mismatch. + withTable("t") { + sql("CREATE TABLE t(id INT, c ARRAY>) USING parquet") + sql("""INSERT INTO t VALUES + (1, array(named_struct('a', 1, 'b', 'x'))), + (2, array(named_struct('a', 1, 'b', 'x'), named_struct('a', 2, 'b', 'yy'))), + (3, array(named_struct('a', -1, 'b', ''), named_struct('a', 0, 'b', 'z'), + named_struct('a', 7, 'b', 'w'))), + (4, array()), + (5, null), + (6, array(named_struct('a', null, 'b', 'nullfield'))), + (7, array(cast(null as struct))), + (8, array(named_struct('a', 9, 'b', null), named_struct('a', 9, 'b', null))), + (9, array(named_struct('a', 1, 'b', 'x'), named_struct('a', 1, 'b', 'x')))""") + checkSparkAnswerAndOperator("SELECT id, hash(c), xxhash64(c) FROM t ORDER BY id") + } + } + + test("hash - array of struct not eligible for batching") { + // Nested elements retain the per-element path. Compare with Spark and also cover chaining + // a flat-struct column with a nested column in one hash call. + withTable("t") { + sql("""CREATE TABLE t( + id INT, + nested ARRAY>>, + plain ARRAY>) + USING parquet""") + sql("""INSERT INTO t VALUES + (1, array(named_struct('l', array(1, 2))), array(named_struct('a', 1, 'b', 'x'))), + (2, array(named_struct('l', array()), named_struct('l', array(3))), + array(named_struct('a', 2, 'b', 'y'))), + (3, array(named_struct('l', cast(null as array))), null), + (4, null, array(named_struct('a', 3, 'b', 'z')))""") + checkSparkAnswerAndOperator( + "SELECT id, hash(nested), xxhash64(nested), hash(plain, nested) FROM t ORDER BY id") + } + } + test("hash - fuzz test") { val r = new Random(42) val options = SchemaGenOptions(generateStruct = true)