comphead commented on code in PR #4802: URL: https://github.com/apache/datafusion-comet/pull/4802#discussion_r3970401803
########## 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: This returns a sketch where Spark returns NULL. `HllUnion` is a `TernaryExpression` using `nullSafeEval`, and `TernaryExpression.eval` returns NULL if **any** of the three inputs is NULL: ```scala // Expression.scala override def eval(input: InternalRow): Any = { val value1 = first.eval(input) if (value1 != null) { val value2 = second.eval(input) if (value2 != null) { val value3 = third.eval(input) if (value3 != null) { return nullSafeEval(value1, value2, value3) } } } null } ``` `CometHllUnion.getSupportLevel` only requires `expr.third.foldable`, and `Literal(null, BooleanType)` is foldable, so ```sql SELECT hll_union(s1, s2, cast(null as boolean)) ``` is planned natively and returns a non-null sketch. That is a categorical wrong answer rather than an approximation, so it is not something the `Incompatible` opt-in covers. Folding the third argument into the null branch a few lines up should be all it needs: ```rust if a.is_null(i) || b.is_null(i) || allow.is_null(i) { out.append_null(); continue; } ``` `CometHllUnionAgg` is safe here for a different reason: its `convert` falls back when `right.eval()` is not a `Boolean`, and Spark's `null.asInstanceOf[Boolean]` coerces to `false`, so the two agree. Worth a test covering both. ########## 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: This test passes even if Comet never plans the aggregate natively. `checkSparkAnswerMaybeThrows` runs the same query twice, once with `COMET_ENABLED=false` and once with it on, then returns `(expected.failed.toOption, actual.failed.toOption)`. If Comet fell back for the whole plan, both runs raise Spark's exception and both assertions below still hold. The native mismatch check in `hll_union_agg.rs` is never proven to run. The two messages do differ, so asserting on the Comet-only half distinguishes them. Spark raises `HLL_UNION_DIFFERENT_LG_K`: > Sketches have different `lgConfigK` values: 10 and 12. Set the `allowDifferentLgConfigK` parameter to true to call `hll_union_agg` with different `lgConfigK` values. while the native error ends "to enable unions of different lgConfigK". So: ```scala assert( cometErr.get.getMessage.contains("to enable unions of different lgConfigK"), s"expected Comet's native error, got: ${cometErr.get.getMessage}") ``` Worth noting separately that this is itself a user-visible difference: Comet raises a plain execution error where Spark raises `SparkRuntimeException` with condition `HLL_UNION_DIFFERENT_LG_K` and sqlState 22000. Same for the invalid-buffer case, where Spark raises `HLL_INVALID_INPUT_SKETCH_BUFFER`. Probably acceptable for an opt-in incompatible function, but it should be a listed incompatibility rather than an accident. ########## 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: HLL_4 input carrying a non-empty aux map decodes to silently wrong values here, and the compact workaround does not cover it. `Array4::deserialize` reads exactly `aux_count` coupons from the aux region regardless of the COMPACT flag: ```rust // datasketches 0.3.0, src/hll/array4.rs let aux_count = cursor.read_u32_le()...; ... if aux_count > 0 { let mut aux = AuxMap::new(lg_config_k); for i in 0..aux_count { let coupon = cursor.read_u32_le()...; aux.insert(slot, value); } } ``` That is the *compact* aux layout. DataSketches-Java's *updatable* HLL_4 form writes `1 << lgAuxArrInts` ints including the empty slots, so reading the first `aux_count` of those pulls empty slots in as `Coupon(0)` and drops real exceptions. Nothing errors. This one sits on the non-compact path, so `normalize_compact_hll_array` does not reach it. It is a different failure from the one you fixed, and it is the one with real exposure: Spark and Comet both fix `TgtHllType.HLL_8`, but `hll_sketch_estimate`, `hll_union`, and `hll_union_agg` accept **any** binary column, and DataSketches-Java's default target type is HLL_4. A user reading third-party HLL_4 sketch columns gets a correct answer from Spark and a silently wrong one from Comet. We already parse the mode byte right here, so the cheapest guard is to reject anything that is not HLL_8 (`extract_tgt_hll_type`, bits 2-3 of the mode byte, 2 == HLL_8) with an `Execution` error. Given the functions are opt-in anyway, a clear error beats a wrong estimate. If you would rather keep HLL_4 and HLL_6 readable, narrowing the reject to HLL_4 with `aux_count > 0` also works. -- 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]
