zhuqi-lucas commented on code in PR #24075:
URL: https://github.com/apache/datafusion/pull/24075#discussion_r3709976401
##########
datafusion/functions-aggregate/benches/first_last.rs:
##########
@@ -63,15 +75,64 @@ fn prepare_groups_accumulator(is_first: bool) -> Box<dyn
GroupsAccumulator> {
exprs: &[col("value", &schema).unwrap()],
};
- if is_first {
- FirstValue::new()
- .create_groups_accumulator(accumulator_args)
- .unwrap()
+ // Mirror the planner: use the native GroupsAccumulator when this value
type
+ // is supported, otherwise fall back to a GroupsAccumulatorAdapter wrapping
+ // one per-group Accumulator. Because the same benchmark case then runs the
+ // fallback on a build without native nested support and the native path on
+ // a build with it, the benchmark bot's before/after diff surfaces the win
+ // directly — no separate comparison case needed.
+ let result = if is_first {
+ FirstValue::new().create_groups_accumulator(accumulator_args)
} else {
- LastValue::new()
- .create_groups_accumulator(accumulator_args)
- .unwrap()
- }
+ LastValue::new().create_groups_accumulator(accumulator_args)
+ };
Review Comment:
Good point — switched to `groups_accumulator_supported()` to decide native
vs adapter, and only fall back when the type is genuinely unsupported. Now a
real `create_groups_accumulator` error `unwrap()`s loudly instead of being
silently swallowed into the fallback path.
##########
datafusion/functions-aggregate/benches/first_last.rs:
##########
@@ -270,6 +337,137 @@ fn trivial_update_bench(
});
}
+/// A 3-field struct value column `(Int64, Utf8, Float64)` — the shape produced
+/// by rewriting three peer `first_value(col ORDER BY o)` calls into a single
+/// `first_value(named_struct(..) ORDER BY o)` (the coalesce-peers
optimization).
+fn create_struct_array(n: usize, null_density: f32) -> ArrayRef {
+ let a = Arc::new(create_primitive_array::<Int64Type>(n, null_density)) as
ArrayRef;
+ let b =
+ Arc::new(create_string_array_with_len::<i32>(n, null_density, 16)) as
ArrayRef;
+ let d = Arc::new(create_primitive_array::<Float64Type>(n, null_density))
as ArrayRef;
+ let fields = Fields::from(vec![
+ Field::new("c0", DataType::Int64, true),
+ Field::new("c1", DataType::Utf8, true),
+ Field::new("c2", DataType::Float64, true),
+ ]);
+ // Struct-level nulls stay None: `named_struct` never produces a null
+ // struct, only null fields — match that shape here.
+ Arc::new(StructArray::new(fields, vec![a, b, d], None))
+}
+
+/// A `List<Int64>` value column with fixed-size lists of `list_len` elements.
+fn create_list_array(n: usize, list_len: usize, null_density: f32) -> ArrayRef
{
+ let child = Arc::new(create_primitive_array::<Int64Type>(
+ n * list_len,
+ null_density,
+ )) as ArrayRef;
+ let offsets = OffsetBuffer::from_lengths(std::iter::repeat_n(list_len, n));
+ let field = Arc::new(Field::new_list_field(DataType::Int64, true));
+ Arc::new(ListArray::new(field, offsets, child, None))
+}
+
+/// A `Map<Utf8, Int64>` value column with `entries_per_row` entries per row.
+/// Values carry `null_density` nulls (keys are never null), matching the null
+/// treatment of the struct / list generators.
+fn create_map_array(n: usize, entries_per_row: usize, null_density: f32) ->
ArrayRef {
+ let total = n * entries_per_row;
+ let values =
+ Arc::new(create_primitive_array::<Int64Type>(total, null_density)) as
ArrayRef;
+ let keys = Arc::new(StringArray::from_iter_values(
+ (0..total).map(|idx| format!("k{}", idx % entries_per_row)),
+ )) as ArrayRef;
+ let entry_fields = Fields::from(vec![
+ Field::new("keys", DataType::Utf8, false),
+ Field::new("values", DataType::Int64, true),
+ ]);
+ let entries = StructArray::new(entry_fields.clone(), vec![keys, values],
None);
+ let offsets =
OffsetBuffer::from_lengths(std::iter::repeat_n(entries_per_row, n));
+ let map_field =
+ Arc::new(Field::new("entries", DataType::Struct(entry_fields), false));
+ Arc::new(MapArray::new(map_field, offsets, entries, None, false))
+}
+
+/// A composite `List<Struct<a: Int64, b: Utf8>>` column — a list whose
+/// elements are structs (the "array of records" shape). Exercises the
+/// nested-within-nested case, which the generic value-state path must also
+/// handle.
+fn create_list_of_struct_array(n: usize, list_len: usize, null_density: f32)
-> ArrayRef {
+ let total = n * list_len;
+ let a = Arc::new(create_primitive_array::<Int64Type>(total, null_density))
as ArrayRef;
+ let b =
+ Arc::new(create_string_array_with_len::<i32>(total, null_density, 8))
as ArrayRef;
+ let struct_fields = Fields::from(vec![
+ Field::new("a", DataType::Int64, true),
+ Field::new("b", DataType::Utf8, true),
+ ]);
+ let child =
+ Arc::new(StructArray::new(struct_fields.clone(), vec![a, b], None)) as
ArrayRef;
+ let offsets = OffsetBuffer::from_lengths(std::iter::repeat_n(list_len, n));
+ let list_field =
+ Arc::new(Field::new_list_field(DataType::Struct(struct_fields), true));
+ Arc::new(ListArray::new(list_field, offsets, child, None))
+}
+
+fn first_last_nested_benchmark(c: &mut Criterion) {
+ const N: usize = 65536;
+ const NUM_GROUPS: usize = 1024;
+
+ let ord = Arc::new(create_primitive_array::<Int64Type>(N, 0.0)) as
ArrayRef;
+
+ for pct in [0, 90] {
+ let null_density = (pct as f32) / 100.0;
+
+ // One column per nested value type. Each type gets the same treatment
+ // as the primitive first_value / last_value benchmarks: update and
+ // merge (both first and last) plus evaluate, at 0% and 90% nulls. On a
+ // build without native nested support these run the fallback adapter;
+ // with this PR they run the native GroupsAccumulator, so the benchmark
+ // bot's before/after diff shows the win per type.
Review Comment:
Fixed by generating top-level nulls: a shared `top_level_nulls` helper
(derived from arrow's null generator) now feeds a parent validity buffer into
each nested array, so `nulls={pct}%` means the value itself is null at that
density — matching the primitive cases — not just inner fields/elements.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
