andygrove commented on code in PR #4802: URL: https://github.com/apache/datafusion-comet/pull/4802#discussion_r4105237644
########## native/spark-expr/src/agg_funcs/hll_sketch_agg.rs: ########## @@ -0,0 +1,290 @@ +// 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. + +use crate::agg_funcs::hll_sketch::SparkHllSketch; +use arrow::array::Array; +use arrow::array::ArrayRef; +use arrow::array::BinaryArray; +use arrow::array::{as_primitive_array, GenericByteArray, PrimitiveArray, StringArray}; +use arrow::datatypes::{ + ArrowPrimitiveType, ByteArrayType, DataType, Field, FieldRef, Int16Type, Int32Type, Int64Type, + Int8Type, +}; +use datafusion::common::{downcast_value, ScalarValue}; +use datafusion::error::{DataFusionError, Result}; +use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs}; +use datafusion::logical_expr::{AggregateUDFImpl, Signature, Volatility}; +use datafusion::physical_plan::Accumulator; +use std::sync::Arc; + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct HllSketchAgg { + signature: Signature, + lg_config_k: i32, +} + +impl HllSketchAgg { + pub fn new(lg_config_k: i32) -> Self { + Self { + signature: Signature::uniform( + 1, + vec![ + DataType::Int8, + DataType::Int16, + DataType::Int32, + DataType::Int64, + DataType::Utf8, + DataType::Binary, + ], + Volatility::Immutable, + ), + lg_config_k, + } + } +} + +impl AggregateUDFImpl for HllSketchAgg { + fn name(&self) -> &str { + "hll_sketch_agg" + } + fn signature(&self) -> &Signature { + &self.signature + } + fn return_type(&self, _: &[DataType]) -> Result<DataType> { + Ok(DataType::Binary) + } + fn accumulator(&self, _: AccumulatorArgs) -> Result<Box<dyn Accumulator>> { + Ok(Box::new(HllSketchAccumulator::new(self.lg_config_k as u8))) + } + fn state_fields(&self, _: StateFieldsArgs) -> Result<Vec<FieldRef>> { + Ok(vec![Arc::new(Field::new("sketch", DataType::Binary, true))]) + } + fn groups_accumulator_supported(&self, _: AccumulatorArgs) -> bool { + false + } +} + +#[derive(Debug)] +pub struct HllSketchAccumulator { + sketch: SparkHllSketch, +} + +impl HllSketchAccumulator { + pub fn new(lg_config_k: u8) -> Self { + Self { + sketch: SparkHllSketch::new(lg_config_k), + } + } + + /// Spark widens every accepted integral to `long` before hashing, so all four widths + /// funnel through the same `i64` update. Nulls are ignored, matching `HllSketchAgg`. + fn update_ints<T>(&mut self, arr: &PrimitiveArray<T>) + where + T: ArrowPrimitiveType, + T::Native: Into<i64>, + { + for i in 0..arr.len() { + if !arr.is_null(i) { + self.sketch.update_i64(arr.value(i).into()); + } + } + } + + /// StringType hashes its UTF-8 bytes and BinaryType its bytes directly, so both share + /// this loop. + fn update_byte_slices<T>(&mut self, arr: &GenericByteArray<T>) + where + T: ByteArrayType, + for<'a> &'a T::Native: AsRef<[u8]>, + { + for i in 0..arr.len() { + if !arr.is_null(i) { + self.sketch.update_bytes(arr.value(i).as_ref()); + } + } + } +} + +impl Accumulator for HllSketchAccumulator { + fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { + if values.is_empty() { + return Ok(()); + } + let arr = &values[0]; + // Downcast once per batch rather than going through `ScalarValue::try_from_array` per + // row: for the string and binary cases that copies every value onto the heap only to + // hash it and drop it again. + match arr.data_type() { + DataType::Int8 => self.update_ints(as_primitive_array::<Int8Type>(arr)), + DataType::Int16 => self.update_ints(as_primitive_array::<Int16Type>(arr)), + DataType::Int32 => self.update_ints(as_primitive_array::<Int32Type>(arr)), + DataType::Int64 => self.update_ints(as_primitive_array::<Int64Type>(arr)), + DataType::Utf8 => self.update_byte_slices(downcast_value!(arr, StringArray)), + DataType::Binary => self.update_byte_slices(downcast_value!(arr, BinaryArray)), + other => { + return Err(DataFusionError::Internal(format!( + "hll_sketch_agg received an unsupported input type: {other:?}" + ))) + } + } + Ok(()) + } + + fn evaluate(&mut self) -> Result<ScalarValue> { + // Spark's HllSketchAgg is declared non-nullable: an empty/all-null group + // still returns a serialized empty sketch (which estimates to 0), never NULL. + Ok(ScalarValue::Binary(Some(self.sketch.to_sketch_bytes()))) + } + + fn size(&self) -> usize { + // An HLL_8 sketch at lgConfigK=k can heap-allocate up to 1 << k bytes; + // account for that so memory reservation reflects actual usage. + std::mem::size_of_val(self) + (1usize << self.sketch.lg_config_k() as usize) Review Comment: Fixed in 32b602559. The default lgK=12 had the same problem on a smaller scale: 4 KiB charged per group against 32 bytes held, so any high-cardinality GROUP BY with small groups over-reserved by about 128x. `size()` on both accumulators now reports what the sketch actually holds: an 8-slot coupon LIST (32 bytes), a SET of 4-byte slots that starts at 32 and doubles past 3/4 full, or the `2^lgConfigK` register array once the SET would outgrow `2^(lgConfigK - 3)` slots. `datasketches` keeps the mode private (`HllSketch::mode` is `pub(super)`), so the wrapper derives it from the estimate. In LIST and SET mode `estimate()` is `max(couponCount, interpolation)`, so it is never below the coupon count, and promotion seeds the HIP accumulator with it, so the estimate is past the promotion threshold as soon as the register array exists. The only other way into the array is a union with an input that already has one. That comes from the input's preamble and is kept as a sticky flag, since a union result estimates from its registers and can dip back under the threshold. Your reproduction is `singleton_groups_at_high_lg_config_k_fit_a_small_pool`, one per aggregate: 64 singleton lgK=21 groups through a real Partial/Final plan in a 16 MiB pool with two target partitions. Both fail against the old `size()` with exactly your `Failed to allocate additional 64.0 MB for FinalHashAggregateStream` error. `heap_size_follows_the_crates_layout` walks lgK 4, 7, 8, 11 and 14 through every SET resize and the promotion, duplicates included, and compares against the layout the crate itself serializes. The charge is never below it. From lgK 8 up it is exact except just below a resize, where the collision-corrected estimate runs ahead of the coupon count and the next size up is charged, and below lgK 8 it is a flat 128 bytes at most. `union_heap_size_follows_the_crates_layout` pins the preamble flag with a dense sketch that has a single register set, so every estimate involving it is tiny. -- 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]
