rich7420 commented on code in PR #4802: URL: https://github.com/apache/datafusion-comet/pull/4802#discussion_r3950720196
########## native/spark-expr/src/agg_funcs/hll_sketch.rs: ########## @@ -0,0 +1,218 @@ +// 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. + +//! Thin wrapper over the `datasketches` crate's HLL sketch, isolating all +//! crate-specific API so Comet's aggregate/scalar code depends on a stable +//! surface. Every sketch uses `HllType::Hll8` and DataSketches' +//! `DEFAULT_UPDATE_SEED` (9001), matching Spark's `HllSketchAgg`. +//! +//! Input hashing goes through the crate's `hash_value` wrappers +//! (`raw_bytes` for strings/binary without Rust's length prefix, `sign_extend` +//! for narrow integers) so the MurmurHash3-x64-128 input bytes are identical to +//! DataSketches-Java. This makes the sketches mutually readable with Spark. +//! +//! Note: the crate serializes List/Set (low-cardinality) modes in DataSketches +//! *compact* form, whereas Spark emits the *updatable* form. The bytes are +//! therefore not byte-identical to Spark's output for small inputs, but +//! DataSketches `deserialize` reads both forms, so estimates round-trip in both +//! directions. Comet must own both Partial and Final aggregation +//! (`supportsMixedPartialFinal = false`) so this compact intermediate is only +//! ever read back by Comet. + +use datafusion::error::DataFusionError; +use datasketches::hash_value::{raw_bytes, sign_extend}; +use datasketches::hll::{HllSketch, HllType, HllUnion}; + +/// A DataSketches HLL_8 sketch configured to match Spark's `HllSketchAgg`. +#[derive(Debug)] +pub struct SparkHllSketch { + inner: HllSketch, +} + +impl SparkHllSketch { + /// Create an empty HLL_8 sketch with the given `lgConfigK`. + pub fn new(lg_config_k: u8) -> Self { + Self { + inner: HllSketch::new(lg_config_k, HllType::Hll8), + } + } + + /// Update with a 64-bit integer. Spark widens narrower integrals to `long` + /// before hashing; callers should pass the already-widened value here. + /// Rust's `Hash` for `i64` writes 8 little-endian bytes with no prefix, + /// matching DataSketches-Java `update(long)`. + pub fn update_i64(&mut self, v: i64) { + self.inner.update(v); + } + + /// Update with a narrow signed integer, sign-extending to 64 bits exactly as + /// Spark's `toLong` does before hashing. + pub fn update_i32(&mut self, v: i32) { + self.inner.update(sign_extend::from_i32(v)); + } + pub fn update_i16(&mut self, v: i16) { + self.inner.update(sign_extend::from_i16(v)); + } + pub fn update_i8(&mut self, v: i8) { + self.inner.update(sign_extend::from_i8(v)); + } + + /// Update with raw bytes (used for both StringType UTF-8 bytes and + /// BinaryType), hashing without Rust's slice length prefix. Empty inputs are + /// skipped, matching DataSketches (and Spark), which ignore empty values. + pub fn update_bytes(&mut self, v: &[u8]) { + if v.is_empty() { + return; + } + self.inner.update(raw_bytes::from_slice(v)); + } + + /// Serialize to DataSketches bytes (compact for List/Set modes, full for HLL + /// array modes). Readable by Spark's `hll_sketch_estimate` / `hll_union_agg`. + pub fn to_sketch_bytes(&self) -> Vec<u8> { + self.inner.serialize() + } + + /// Deserialize a DataSketches sketch (either compact or updatable form). + pub fn from_bytes(bytes: &[u8]) -> Result<Self, DataFusionError> { + HllSketch::deserialize(bytes) Review Comment: Compact `HLL_4` decoding skips the registers. With disjoint 1,000-value sketches, both native union paths return 989 while Spark SQL returns 1991. Please fix the decoder and add a compact-input union test. ########## native/spark-expr/src/agg_funcs/hll_union_agg.rs: ########## @@ -0,0 +1,224 @@ +// 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, SparkHllUnion}; +use arrow::array::{Array, ArrayRef, BinaryArray}; +use arrow::datatypes::{DataType, Field, FieldRef}; +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; + +// NOTE: matches bloom_filter_agg.rs for DataFusion 54.0.0 - no `as_any` method on +// AggregateUDFImpl, and PartialEq/Eq/Hash are required (DynEq/DynHash). +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct HllUnionAgg { + signature: Signature, + allow_different_lg_config_k: bool, +} + +impl HllUnionAgg { + pub fn new(allow_different_lg_config_k: bool) -> Self { + Self { + signature: Signature::uniform(1, vec![DataType::Binary], Volatility::Immutable), + allow_different_lg_config_k, + } + } +} + +impl AggregateUDFImpl for HllUnionAgg { + fn name(&self) -> &str { + "hll_union_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(HllUnionAccumulator::new( + self.allow_different_lg_config_k, + ))) + } + 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 + } +} + +/// Default `lgMaxK` used by Spark's `new Union()` when constructing the empty +/// union returned for a group that never absorbed any sketch. +const DEFAULT_LG_K: u8 = 12; + +#[derive(Debug)] +pub struct HllUnionAccumulator { + // Spark's HllUnionAgg defers creating the Union until the first sketch is seen, + // then builds `new Union(sketch.getLgConfigK)` - so lgMaxK is NOT a fixed 12. + union: Option<SparkHllUnion>, + allow_different_lg_config_k: bool, + seen_lg_config_k: Option<u8>, +} + +impl HllUnionAccumulator { + pub fn new(allow_different_lg_config_k: bool) -> Self { + Self { + union: None, + allow_different_lg_config_k, + seen_lg_config_k: None, + } + } + + fn absorb(&mut self, bytes: &[u8]) -> Result<()> { + let sketch = SparkHllSketch::from_bytes(bytes)?; + let k = sketch.lg_config_k(); + match self.seen_lg_config_k { + None => { + // Lazily instantiate the union from the first sketch's lgConfigK. + self.seen_lg_config_k = Some(k); + self.union = Some(SparkHllUnion::new(k)); + } + Some(prev) if prev != k && !self.allow_different_lg_config_k => { + return Err(DataFusionError::Execution(format!( + "Sketches have different lgConfigK values: {prev} and {k}. \ + Set allowDifferentLgConfigK to true to enable unions of different lgConfigK." + ))); + } + _ => {} + } + self.union.as_mut().unwrap().merge(&sketch); + Ok(()) + } +} + +impl Accumulator for HllUnionAccumulator { + fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { + if values.is_empty() { + return Ok(()); + } + let arr = downcast_value!(values[0], BinaryArray); + for i in 0..arr.len() { + if !arr.is_null(i) { + self.absorb(arr.value(i))?; + } + } + Ok(()) + } + fn evaluate(&mut self) -> Result<ScalarValue> { + // Spark's HllUnionAgg is declared non-nullable: an empty/all-null group + // still returns the serialized bytes of an empty `new Union()` (default + // lgMaxK), which estimates to 0, never NULL. + match &self.union { + Some(u) => Ok(ScalarValue::Binary(Some(u.to_sketch_bytes()))), + None => Ok(ScalarValue::Binary(Some( + SparkHllUnion::new(DEFAULT_LG_K).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) + + self + .seen_lg_config_k + .map(|k| 1usize << k as usize) + .unwrap_or(0) + } + fn state(&mut self) -> Result<Vec<ScalarValue>> { + match &self.union { + Some(u) => Ok(vec![ScalarValue::Binary(Some(u.to_sketch_bytes()))]), + None => Ok(vec![ScalarValue::Binary(Some( Review Comment: Please preserve NULL for an empty partial state. Creating an `lgConfigK=12` sketch here makes final merging fail when another partition contains only `lgConfigK=10` sketches. I reproduced this through DataFusion's Partial/Final execution. Please add a regression test. -- 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]
