sunchao commented on code in PR #4802: URL: https://github.com/apache/datafusion-comet/pull/4802#discussion_r4104650157
########## 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: [P2] Account for the current sketch allocation in `size()`. A singleton sketch still uses LIST mode with eight 4-byte coupons, but this expression charges `2^lgConfigK` bytes immediately. At `lgConfigK=21`, that means 2 MiB per tiny group. With 64 singleton groups and a 16 MiB native memory pool, both `hll_sketch_agg` and `hll_union_agg` fail during Final aggregation even with spilling enabled, whereas Spark returns all 64 groups. This makes supported high-precision grouped queries exhaust their reservation despite using little actual memory. Please track the allocated LIST/SET/HLL capacity and apply the same correction to `HllUnionAccumulator::size()`. Evidence: Reproduced using the exact-head UDAFs with `RuntimeEnvBuilder::with_memory_limit(16 * 1024 * 1024, 1.0)`, normal temporary-file spilling, and two target partitions. Register `HllSketchAgg::new(21)`, provide `id=0..63`, and collect `SELECT id, hll_sketch_agg(id) FROM t GROUP BY id`. The Partial/Final plan reports `Failed to allocate additional 64.0 MB for FinalHashAggregateStream ... 16.0 MB remain available`. Grouping 64 singleton lgK=21 sketch inputs through `HllUnionAgg::new(false)` produces the same failure. Both lgK=12 controls succeed. Spark 4.1.3 executes `SELECT id, hll_sketch_agg(id, 21) FROM range(64) GROUP BY id` successfully, producing 64 groups with 40-byte updatable sketches. Reproduction source was preserved at `/tmp/comet-4802-review_hll_memory.rs`. -- 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]
