andygrove commented on code in PR #4802: URL: https://github.com/apache/datafusion-comet/pull/4802#discussion_r3970632838
########## native/spark-expr/src/hll_scalar.rs: ########## @@ -0,0 +1,135 @@ +// 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::estimate_from_bytes; +use arrow::array::{Array, BinaryArray, Int64Array}; +use datafusion::common::{DataFusionError, Result}; +use datafusion::physical_plan::ColumnarValue; +use std::sync::Arc; + +/// Spark hll_sketch_estimate: Binary sketch -> Long distinct-count estimate. +pub fn spark_hll_sketch_estimate(args: &[ColumnarValue]) -> Result<ColumnarValue> { + let arrays = ColumnarValue::values_to_arrays(args)?; + let input = arrays[0].as_any().downcast_ref::<BinaryArray>().unwrap(); + let mut out = Int64Array::builder(input.len()); + for i in 0..input.len() { + if input.is_null(i) { + out.append_null(); + } else { + out.append_value(estimate_from_bytes(input.value(i))?); + } + } + Ok(ColumnarValue::Array(Arc::new(out.finish()))) +} + +// Spark's HllUnion is a TernaryExpression (first, second, third=allowDifferentLgConfigK). +// It builds `new Union(min(k1, k2))`, throws when the two sketches have different +// lgConfigK and the flag is false, and returns an HLL_8 sketch. +/// Spark hll_union(first, second, allowDifferentLgConfigK): union two sketch columns. +pub fn spark_hll_union(args: &[ColumnarValue]) -> Result<ColumnarValue> { + use crate::agg_funcs::{SparkHllSketch, SparkHllUnion}; + use arrow::array::BooleanArray; + let arrays = ColumnarValue::values_to_arrays(args)?; + let a = arrays[0].as_any().downcast_ref::<BinaryArray>().unwrap(); + let b = arrays[1].as_any().downcast_ref::<BinaryArray>().unwrap(); + let allow = arrays[2].as_any().downcast_ref::<BooleanArray>().unwrap(); + let mut out = arrow::array::BinaryBuilder::new(); + for i in 0..a.len() { + if a.is_null(i) || b.is_null(i) { + out.append_null(); + continue; + } + let sa = SparkHllSketch::from_bytes(a.value(i))?; + let sb = SparkHllSketch::from_bytes(b.value(i))?; + let allow_i = !allow.is_null(i) && allow.value(i); Review Comment: Fixed in 1fa3f260d, exactly as you wrote it — the third argument is now part of the null check: ```rust if a.is_null(i) || b.is_null(i) || allow.is_null(i) { out.append_null(); continue; } ``` `allow_i` is then just `allow.value(i)`, since the NULL case has already returned. The old `!allow.is_null(i) && allow.value(i)` was quietly doing Spark's `asInstanceOf[Boolean]` coercion on the wrong expression — that coercion is what makes `HllUnionAgg` correct, and I had applied it to the `TernaryExpression` where the null has to propagate instead. Agreed that `Incompatible` does not cover this. The opt-in says estimates may differ slightly; returning a sketch where Spark returns NULL is a different kind of claim. Test is `a_null_allow_flag_yields_null`, with two rows — a NULL flag and a real one — so it fails if the fix nulls the whole column rather than the affected row. I noted your point about `CometHllUnionAgg` being safe for the different reason in the comment on the null check, so the asymmetry between the two does not read as an oversight later. ########## native/spark-expr/src/agg_funcs/hll_sketch.rs: ########## @@ -0,0 +1,271 @@ +// 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, +} + +/// Byte offsets into the DataSketches HLL preamble, and the bits we need from it. +mod preamble { + /// Serialization flags. Bit 3 is COMPACT. + pub const FLAGS: usize = 5; + /// Mode byte. Low two bits are the current mode (0 LIST, 1 SET, 2 HLL). + pub const MODE: usize = 7; + pub const COMPACT_FLAG: u8 = 8; + pub const CUR_MODE_MASK: u8 = 0x3; + pub const CUR_MODE_HLL: u8 = 2; +} + +/// Work around a decoding bug in `datasketches` 0.3.0 for compact sketches in an HLL array mode. +/// +/// `Array4::deserialize` (and the `Array6` / `Array8` equivalents) skip the register block +/// entirely when the COMPACT flag is set, leaving every register zero: +/// +/// ```text +/// let mut data = vec![0u8; num_bytes]; +/// if !compact { +/// cursor.read_exact(&mut data)?; +/// } else { +/// cursor.advance(num_bytes as u64); +/// } +/// ``` +/// +/// The damage is quiet, which is what makes it worth guarding: the decoded sketch's own +/// `estimate()` still looks correct because it comes back from the HIP accumulator in the +/// preamble, but every union built from it is wrong. Two disjoint 1,000-value sketches union to +/// ~989 rather than ~1991. +/// +/// Clearing the flag is a correct parse rather than a guess. The register block is present in +/// both the compact and updatable forms, and the crate reads the HLL_4 auxiliary map as +/// `aux_count` coupons regardless of the flag - which is the compact layout. LIST and SET mode +/// compaction *is* a genuinely different layout, and the crate handles those correctly, so this +/// only touches HLL array mode. +/// +/// Returns `None` when the input needs no rewriting, so the common path does not copy. +/// +/// `compact_input_survives_a_union` pins the behaviour: if a future `datasketches` release fixes +/// the register read, that test is what tells us this can be deleted. +fn normalize_compact_hll_array(bytes: &[u8]) -> Option<Vec<u8>> { + if bytes.len() <= preamble::MODE + || bytes[preamble::MODE] & preamble::CUR_MODE_MASK != preamble::CUR_MODE_HLL + || bytes[preamble::FLAGS] & preamble::COMPACT_FLAG == 0 + { + return None; + } + let mut owned = bytes.to_vec(); + owned[preamble::FLAGS] &= !preamble::COMPACT_FLAG; + Some(owned) +} + +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> { Review Comment: You're right, and this is the more serious of the two. Fixed in 1fa3f260d by rejecting the shape rather than trying to decode it. I had actually noticed the aux-layout asymmetry while fixing the compact bug — the crate reading `aux_count` coupons regardless of the flag is what makes clearing the COMPACT flag a *correct* parse for compact input — and I did not follow it through to the updatable direction, where the same code is wrong. Your framing of the exposure is the part I had missed: Comet only ever writes HLL_8, but these functions accept any binary column and DataSketches-Java defaults to HLL_4, so third-party sketch columns are the realistic input. I took your narrower option rather than rejecting everything that is not HLL_8. Rejecting HLL_4 outright would refuse most third-party sketches, including the low-cardinality ones the crate reads perfectly well; the exception map only appears once a register exceeds `curMin + 15`. So `reject_undecodable_hll4` errors only on HLL array mode + HLL_4 + not compact + `aux_count > 0`, read straight from the preamble at offset 36 before the crate sees the bytes. Two tests, and I checked the premise in both rather than trusting the cardinality: - `updatable_hll4_with_aux_entries_is_rejected` — 100,000 values at lgK=12, which I measured as the point where `aux_count` becomes non-zero (1,000 gives 0). It asserts mode, target type, the absent COMPACT flag and `aux_count > 0` before asserting the rejection, so it cannot silently stop testing the shape it is named for. - `hll4_without_aux_entries_still_reads` — the same construction at 1,000 values, asserting `aux_count == 0` and that the estimate is right, which pins the guard as narrow. The error names the shape and tells the user to convert to HLL_8 or to the compact HLL_4 form. ########## spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala: ########## @@ -3549,4 +3549,96 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } + test("hll_sketch_agg and hll_sketch_estimate (incompatible, opt-in)") { + assume(isSpark40Plus) + // HLL is approximate: Comet's Rust DataSketches estimator differs slightly from + // Spark's after a merge, so these functions are Incompatible. Opt in, assert the + // query runs natively (no fallback), and that the estimate is within HLL error of + // the TRUE distinct count (700). Do NOT compare bit-exactly to Spark. + withSQLConf( + "spark.comet.expression.HllSketchAgg.allowIncompatible" -> "true", + "spark.comet.expression.HllSketchEstimate.allowIncompatible" -> "true") { + withParquetTable((0 until 1000).map(i => (i % 700, i)), "tbl") { + def checkEstimate(query: String): Unit = { + val df = sql(query) + checkCometOperators(stripAQEPlan(df.queryExecution.executedPlan)) + val est = df.collect().head.getLong(0) + assert( + math.abs(est - 700).toDouble / 700 <= 0.05, + s"estimate $est not within 5% of the true distinct count 700 for: $query") + } + checkEstimate("SELECT hll_sketch_estimate(hll_sketch_agg(_1)) FROM tbl") + checkEstimate("SELECT hll_sketch_estimate(hll_sketch_agg(_1, 14)) FROM tbl") + checkEstimate("SELECT hll_sketch_estimate(hll_sketch_agg(cast(_1 as string))) FROM tbl") + } + } + } + + test("hll_union_agg and hll_union (incompatible, opt-in)") { + assume(isSpark40Plus) + withSQLConf( + "spark.comet.expression.HllSketchAgg.allowIncompatible" -> "true", + "spark.comet.expression.HllSketchEstimate.allowIncompatible" -> "true", + "spark.comet.expression.HllUnionAgg.allowIncompatible" -> "true", + "spark.comet.expression.HllUnion.allowIncompatible" -> "true") { + withParquetTable((0 until 1000).map(i => (i % 3, i)), "tbl") { + // hll_union_agg: union the per-group sketches -> ~1000 distinct. + val aggDf = sql( + "SELECT hll_sketch_estimate(hll_union_agg(s)) FROM " + + "(SELECT _1 AS g, hll_sketch_agg(_2) AS s FROM tbl GROUP BY _1)") + checkCometOperators(stripAQEPlan(aggDf.queryExecution.executedPlan)) + val aggEst = aggDf.collect().head.getLong(0) + assert(math.abs(aggEst - 1000).toDouble / 1000 <= 0.05, s"union_agg estimate $aggEst") + + // hll_union: union two disjoint group sketches -> ~667 distinct. + val unionDf = sql( + "SELECT hll_sketch_estimate(hll_union(a.s, b.s)) FROM " + + "(SELECT hll_sketch_agg(_2) AS s FROM tbl WHERE _1 = 0) a, " + + "(SELECT hll_sketch_agg(_2) AS s FROM tbl WHERE _1 = 1) b") + checkCometOperators(stripAQEPlan(unionDf.queryExecution.executedPlan)) + val unionEst = unionDf.collect().head.getLong(0) + assert(math.abs(unionEst - 667).toDouble / 667 <= 0.05, s"union estimate $unionEst") + } + } + } + + test("hll_union_agg rejects different lgConfigK when not allowed") { + assume(isSpark40Plus) + withSQLConf( + "spark.comet.expression.HllSketchAgg.allowIncompatible" -> "true", + "spark.comet.expression.HllUnionAgg.allowIncompatible" -> "true") { + withParquetTable((0 until 100).map(i => Tuple1(i)), "tbl") { + // A lgConfigK=10 sketch unioned with a lgConfigK=12 sketch (allowDifferentLgConfigK + // defaults false) must throw in BOTH Spark and Comet. + val df = sql( + "SELECT hll_union_agg(s) FROM (" + + " SELECT hll_sketch_agg(_1, 10) AS s FROM tbl UNION ALL" + + " SELECT hll_sketch_agg(_1, 12) AS s FROM tbl)") + val (sparkErr, cometErr) = checkSparkAnswerMaybeThrows(df) Review Comment: Fixed in 1fa3f260d. You're right that it was fallback-blind, and it is the same trap as the `expect_error` sentinel problem — both arms raising proves nothing when one of them may be Spark twice. ```scala assert( cometErr.get.getMessage.contains("to enable unions of different lgConfigK"), s"expected Comet's native lgConfigK error, got: ${cometErr.get.getMessage}") ``` with a comment saying why the weaker pair of assertions above it is not sufficient on its own, so the next person does not simplify it back. **On the error-class divergence.** Agreed it should be listed rather than accidental, and it now is: `CometHllUnionAgg`, `CometHllUnion` and `CometHllSketchEstimate` each carry a second entry in `getIncompatibleReasons()` recording that failures surface as plain Comet execution errors rather than `SparkRuntimeException` with `HLL_UNION_DIFFERENT_LG_K` / `HLL_INVALID_INPUT_SKETCH_BUFFER` and sqlState 22000. That puts it in front of anyone opting in, which is the right place for it — and it is also what the strengthened assertion above now depends on, so the two are consistent. -- 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]
