This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-5778-394ad88c9d0e1050a2d992dbecd5bb39b1998320 in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit d2f73d1da985bacd4907dc0d79cfa0923180538f Author: Liang-Chi Hsieh <[email protected]> AuthorDate: Sat Sep 12 18:51:51 2026 +0000 perf: batch the nested-element list hash for flat struct elements (#5778) * perf: batch the nested-element list hash for flat struct elements `hash_list_array!` handles a list whose elements are not primitives, which is the path `array<struct<..>>` takes. For every element it sliced a one-element Arrow array and re-entered the hash dispatch for it, so each element cost an array allocation plus a full type match. Spark chains the element hashes in order, so the elements of one row do have to be hashed in sequence, but the allocation and dispatch do not have to happen per element. This hashes one element per row at a time in a single batched call, seeding each slot with the running hash of its row, which is exactly what the per-element call did. Rows are independent, so the number of dispatches drops from the total element count to the length of the longest list. Only rows that still have an element at the current position take part, and a row never becomes alive again once exhausted, so the surviving rows are carried forward rather than the whole batch being rescanned each pass. Rescanning would cost rows x longest-list, which for one long list among short ones is almost all wasted. The same initialisation pass records whether every row is non-null with the same length and takes a simpler walk when it is, since no row drops out early there and the bookkeeping would be pure overhead. The gather is indexed by the list's own offset width, since a `LargeList` can hold more than `u32::MAX` elements and narrowing would silently wrap. Batching only pays when a pass covers several rows. With one row left there is nothing to gather across, and `take` would copy that row's remaining element payloads without reducing the dispatch count. That is checked at the top of every pass, not just before the loop, because the common case is the tail after the shorter rows finish: lengths [1, 1, 8] spend seven of eight passes on a single row. A lone row is hashed directly, element by element, as the previous implementation did throughout. Batching also needs several rows to batch across, so a column with one non-empty row takes the per-element path before any scheduling work happens -- no sliced element view, no survivor list, no per-pass buffers. Recursion makes that common rather than a corner case: a deep singleton list reaches a one-row batch at every level below the first, and without this check those levels allocated scheduler buffers for a batch of one, costing 46% more total allocation on a 5-deep fixture and doubling the peak on a single-row one. It also only pays when the gather itself is cheap, and how much `take` 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, so slicing away the row that held ten million ints leaves two cheap visible rows and a child `take` still 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 -- costed that way, a struct with an unreferenced 8 MiB dictionary value went from 70 us to 2.3 ms. And a nested child recurses, each level holding its gather while the level below builds its own. Rather than model all of that, only the shape whose cost is easy to bound is batched: a list of structs whose fields are all flat leaves, with the child's *retained* buffers under `GATHER_ELIGIBLE_CHILD_BYTES` (4 MiB), measured as retained size rather than an average per element so a large buffer kept alive by a slice disqualifies it even when few rows are visible. Everything else keeps the previous per-element path, and the decision is made before any scheduling work, so those shapes allocate nothing new: no survivor scan, no sliced element view, no per-pass buffers. 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. Timing, base 4c14448c2 against this commit, 8192 rows unless noted, Apple M4 Max: batched shapes base head list_of_struct_x10 10106 us -> 718 us 14.1x list_of_struct_half_null_x10 21067 us -> 716 us 29.4x list_of_struct_skewed_x1024 1208 us -> 203 us 6.0x batched, then a single-row tail base head list_of_struct_long_tail_x1024 136 us -> 123 us (3 rows) shapes left on the previous path base head list_of_int32_x10 149 us -> 147 us list_of_struct_1kb_string_x4 3693 us -> 3564 us list_of_list_x5x5 1684 us -> 1632 us struct_of_dict_unreferenced_8mb 2228 us -> 2210 us deep_singleton_list_5_deep 3146 us -> 3132 us width_skewed_8mb_first 5428 us -> 5326 us null_parent_struct_64kb_x4 180 us -> 185 us sliced_list_retaining_10m_ints 81.8 ns -> 82.6 ns Allocation is measured too, since a gather trades many small allocations for one large one and a timing figure cannot show that. `benches/hash_alloc.rs` counts bytes with a counting global allocator on the same generators `hash.rs` times. Each shape is hashed twice in the same build and the two results are compared, so a number is never reported for a kernel that hashes non-deterministically; agreement with Spark is what the Scala tests below cover: shape total base -> head peak base -> head list_of_struct_x10 35389440 -> 1502381 432 -> 386281 list_of_struct_skewed_x1024 3980880 -> 828465 432 -> 386281 list_of_struct_half_null_x10 49807360 -> 1842024 712 -> 419964 Every shape left on the previous path allocates byte-for-byte what it did before, which is what deciding before any scheduling work buys: `1kb_string_x4` 3538944/432, `list_of_list_x5x5` 5242880/128, `struct_of_dict_unreferenced_8mb` 6291456/768, `deep_singleton_list_5_deep` 1933312/944, `width_skewed_8mb_first` 884736/432, `sliced_list_retaining_10m_ints` 256/128, `null_parent_struct_64kb_x4` 401408/712. `int32` and `list_of_int32_x10` allocate nothing on either side: they do not reach this path. So the batched shapes trade peak live bytes -- about one element becomes one gathered pass -- for a 6x to 29x drop in time and a 4x to 27x drop in total allocation. The shapes whose gather cost cannot be bounded cheaply are not part of that trade at all. The hash values must not move, since they decide which partition a row lands in. Adds golden-value tests for `array<struct<..>>` covering empty lists, null lists, null struct elements, null fields and repeated values, plus element-order sensitivity, skewed lengths, `LargeList` against `List`, `array<array<int>>`, rows exhausting on different passes, a sliced list whose first offset is non-zero, a null list that still covers a non-empty element range, and descending lengths so the uniform-length path is not taken. Two more cover the eligibility rule: both sides of the size threshold on the same data must produce the same hashes, since the rule only decides which path runs, and a nested or dictionary child must not be eligible however small it is. Adds a test that drives the shared scheduler through xxhash64 rather than murmur3, on descending lengths with distinct per-row seeds and a following column, against a per-row reference. A Scala test hashes `array<struct<int, string>>` through both `hash()` and `xxhash64()` against Spark, over rows of differing length plus an empty array, a null array, a null struct element and null fields. `checkSparkAnswerAndOperator` also asserts Comet ran it natively, so a fallback on both sides cannot hide a mismatch. A second one covers element types the eligibility rule leaves alone, since the rule only picks an implementation. An off-by-one in the batched element index fails the first of these. `benches/hash_alloc.rs` checks its own counters before quoting them: a shrinking realloc releases its bytes, a pre-window free does not offset in-window live bytes, a pre-window block grown twice counts only its growth both times, and a window does not inherit the previous window's records. Each corresponds to a bug this allocator had. They allocate through `black_box`, since the optimiser is otherwise entitled to delete the allocations and the checks pass vacuously. Two more cover the paths this adds. One lowers the gather budget to zero so every pass takes the slice fallback, on uneven null-containing rows with a following column, and asserts the hashes match the batched path: without it an off-by-one in that path's element index passes the whole suite. The other drives the shared scheduler through xxhash64 rather than murmur3, on descending lengths with distinct per-row seeds and a following column, against a per-row reference. Co-authored-by: Claude Code <[email protected]> * test: clarify nested hash fallback coverage --------- Co-authored-by: Claude Code <[email protected]> Co-authored-by: Andy Grove <[email protected]> --- native/spark-expr/Cargo.toml | 4 + native/spark-expr/benches/common/hash_shapes.rs | 469 ++++++++++++++++++++ native/spark-expr/benches/hash.rs | 238 ++-------- native/spark-expr/benches/hash_alloc.rs | 485 ++++++++++++++++++++ native/spark-expr/src/hash_funcs/murmur3.rs | 492 +++++++++++++++++++++ native/spark-expr/src/hash_funcs/utils.rs | 263 ++++++++++- native/spark-expr/src/hash_funcs/xxhash64.rs | 62 +++ .../apache/comet/CometHashExpressionSuite.scala | 44 ++ 8 files changed, 1839 insertions(+), 218 deletions(-) diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml index f00f1a7c49..9019dffbb5 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 0000000000..77e3877ef4 --- /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::<Int32Builder>(0) + .unwrap() + .append_value(i as i32); + sb.field_builder::<StringBuilder>(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::<Vec<_>>())) +} + +/// `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::<Vec<_>>(), + )) +} + +/// `struct<a: int32, b: utf8>`: 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<int32>`: 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<struct<..>>`: 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<struct<..>>` 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<utf8, int32>`: 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: int32, m: map<utf8, int32>>`: 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<array<int32>>`: 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<utf8, struct<..>>`: 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<struct<a: int32, m: map<..>>>`: 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<i32> = (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<struct<..>>` 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::<Int32Builder>(0) + .unwrap() + .append_value((i * 31 + j) as i32); + s.field_builder::<StringBuilder>(1) + .unwrap() + .append_value(&big); + s.append(true); + } + lb.append(true); + } + Arc::new(lb.finish()) +} + +/// `array<struct<..>>` 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::<Int32Builder>(0) + .unwrap() + .append_value((i * 31 + j) as i32); + s.field_builder::<StringBuilder>(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::<Int32Builder>(0) + .unwrap() + .append_value(i as i32); + sb.field_builder::<StringBuilder>(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::<Int32Builder>(0) + .unwrap() + .append_value(i as i32); + sb.field_builder::<StringBuilder>(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::<Int32Type>::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::<Int32Builder>(0) + .unwrap() + .append_value(i as i32); + let s = if i == 0 { + "y".repeat(big_len) + } else { + "x".to_string() + }; + sb.field_builder::<StringBuilder>(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::<Int32Builder>(0) + .unwrap() + .append_value(i as i32); + sb.field_builder::<StringBuilder>(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::<Int32Builder>(0) + .unwrap() + .append_value(i as i32); + sb.field_builder::<StringBuilder>(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 4d8d4026cf..5cdfaa10f1 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::<Int32Builder>(0) - .unwrap() - .append_value(i as i32); - sb.field_builder::<StringBuilder>(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::<Vec<_>>())) -} - -/// `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::<Vec<_>>(), - )) -} - -/// `struct<a: int32, b: utf8>`: 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<int32>`: 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<struct<..>>`: 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<struct<..>>` 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<utf8, int32>`: 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: int32, m: map<utf8, int32>>`: 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<array<int32>>`: 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<utf8, struct<..>>`: 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<struct<a: int32, m: map<..>>>`: 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<i32> = (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 0000000000..0003ad4e4a --- /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<HashMap>` is enough: +/// the benchmark is single-threaded, and correctness here matters more than speed. +static TRACKED: Mutex<Option<HashMap<usize, Tracked>>> = 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<bool> = 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<R>(f: impl FnOnce(&mut HashMap<usize, Tracked>) -> R) -> Option<R> { + 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<Tracked> { + 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<u32>) { + 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 a9e5b67aa9..5bf7399b19 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<a: Int32, b: Utf8>` element: `None` is a null struct, and the fields are + /// independently nullable. + type StructElem = Option<(Option<i32>, Option<&'static str>)>; + /// One `array<struct<..>>` row: `None` is a null list. + type ListRow = Option<Vec<StructElem>>; + + /// Builds `array<struct<a: Int32, b: Utf8>>` from `rows`. + fn list_of_struct(rows: Vec<ListRow>) -> 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::<Int32Builder>(0).unwrap().append_null(); + sb.field_builder::<StringBuilder>(1).unwrap().append_null(); + sb.append(false); + } + Some((a, b)) => { + match a { + Some(v) => { + sb.field_builder::<Int32Builder>(0).unwrap().append_value(v) + } + None => { + sb.field_builder::<Int32Builder>(0).unwrap().append_null() + } + } + match b { + Some(v) => sb + .field_builder::<StringBuilder>(1) + .unwrap() + .append_value(v), + None => { + sb.field_builder::<StringBuilder>(1).unwrap().append_null() + } + } + sb.append(true); + } + } + } + lb.append(true); + } + } + } + Arc::new(lb.finish()) + } + + fn hash_of(array: ArrayRef, num_rows: usize) -> Vec<u32> { + let mut hashes = vec![42u32; num_rows]; + create_murmur3_hashes(&[array], &mut hashes).unwrap(); + hashes + } + + /// `array<struct<..>>` 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<ListRow> = 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::<Vec<_>>(), + )); + 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<u32> = 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::<Int32Builder>(0).unwrap().append_value(v); + s.field_builder::<StringBuilder>(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<array<int>>`: 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<Vec<i32>>]| -> 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<u32> = rows + .iter() + .map(|row| hash_of(build(std::slice::from_ref(row)), 1)[0]) + .collect(); + assert_eq!( + batched, per_row, + "array<array<int>> 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<ListRow> = 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<u32> = 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<ListRow> = 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::<Int32Builder>(0).unwrap().append_value(v); + sb.field_builder::<StringBuilder>(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<ListRow> = (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<u32> = 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::<i8, Int8Array>( @@ -460,4 +823,133 @@ mod tests { test_murmur3_hash::<String, StringArray>(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::<Int32Builder>(0) + .unwrap() + .append_value((row * 10 + i) as i32); + sb.field_builder::<StringBuilder>(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::<arrow::array::ListArray>() + .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::<arrow::array::ListArray>() + .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<list<int>>: 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<dictionary<string>>: `take` shares the values, so batching is not modelled here. + let dict: ArrayRef = Arc::new( + DictionaryArray::<Int32Type>::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 e6eedd31df..db9f95db82 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<usize> = 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<usize> = 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<usize> = Vec::with_capacity(active.len()); + let mut still_active: Vec<usize> = 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 93ac9304fc..45c273bb9f 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::<Int32Builder>(0) + .unwrap() + .append_value((row * 10 + i) as i32); + sb.field_builder::<StringBuilder>(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 c997da0012..68c4471e05 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<struct<..>>` 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<STRUCT<a: INT, b: STRING>>) 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<a: INT, b: STRING>))), + (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<STRUCT<l: ARRAY<INT>>>, + plain ARRAY<STRUCT<a: INT, b: STRING>>) + 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<int>))), 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) --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
