andygrove commented on code in PR #4802: URL: https://github.com/apache/datafusion-comet/pull/4802#discussion_r3969545545
########## 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: Fixed in 32793513d. You were right, and it is worse than HLL_4 — I traced it into the crate and all three array modes do it. `datasketches` 0.3.0, `Array4::deserialize` (and the `Array6` / `Array8` equivalents): ```rust let mut data = vec![0u8; num_bytes]; if !compact { cursor.read_exact(&mut data)?; } else { cursor.advance(num_bytes as u64); // registers dropped, `data` stays zeroed } ``` The register block is present in the compact form too, so skipping it is simply wrong. What makes it nasty is how quiet it is: the decoded sketch's own `estimate()` still reads correctly, because that comes back from the HIP accumulator in the preamble. Only a union exposes it. Reproducing your setup against the raw crate: ``` Hll4: compact-flagged a=996.2 b=989.8 union=989.8 Hll6: compact-flagged a=996.2 b=989.8 union=989.8 Hll8: compact-flagged a=996.2 b=989.8 union=0.0 ``` That 989 is your number exactly. `SparkHllSketch::from_bytes` now clears the COMPACT flag before handing HLL-array-mode bytes to the crate. I want to be clear that this is a correct parse and not a guess: the register block is identical in both forms, and the crate reads the HLL_4 auxiliary map as `aux_count` coupons *regardless* of the flag, which is already the compact layout. LIST and SET compaction genuinely is a different layout, and the crate gets those right, so the rewrite is scoped to `cur_mode == HLL` and returns `None` (no copy) otherwise. The compact-input union test you asked for is `compact_input_survives_a_union`. Against the unpatched decoder it fails with `union of two disjoint compact sketches estimated 0, expected ~2000`. This is an upstream bug rather than ours, so the workaround is written to be removable: the doc comment names the crate version and the exact code, and the test is what will tell us when a `datasketches` bump makes it unnecessary. I'll raise it upstream. ########## 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: Fixed in 32793513d. `state()` now returns `Binary(None)` for an accumulator that absorbed nothing; `evaluate()` is unchanged and still returns the empty-union bytes, since that is Spark's non-nullable contract for the user-visible result. The distinction is the whole bug: `state()` feeds the Final phase, and `merge_batch` already skips nulls, so an empty partial should contribute nothing at all. Emitting a concrete lgConfigK=12 sketch made the *empty* partition the first `lgConfigK` the Final accumulator latched onto, and every real sketch at a different k then failed the mismatch check. A partition with no input should not get a vote on the union's k. Two regression tests: - `empty_partial_state_is_null` — the direct assertion. - `empty_partial_does_not_fix_the_final_lg_config_k` — your scenario end to end: one empty partial, one partial holding only lgConfigK=10 sketches, merged in that order through Partial/Final states. Against the old code it fails with exactly the error you'd expect: ``` called `Result::unwrap()` on an `Err` value: Execution("Sketches have different lgConfigK values: 12 and 10. Set allowDifferentLgConfigK to true to enable unions of different lgConfigK.") ``` I checked both fail before the fix rather than only passing after. -- 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]
