From 84dcc0b30e4edeff9850aa2f08dd1050002cb734 Mon Sep 17 00:00:00 2001 From: Huaijin Date: Tue, 9 Jun 2026 20:03:38 +0800 Subject: [PATCH 1/2] fix: approx_distinct over-counts for utf8view (#22815) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Which issue does this PR close? - Closes https://github.com/apache/datafusion/issues/22796 ## Rationale for this change `approx_distinct` over-counted distinct values for `Utf8View` columns when the same short string appeared across batches with different layouts. Arrow stores strings ≤ 12 bytes inline in the 128-bit view integer. The fast path (no data buffers) hashed these as raw `u128`. But when a batch also had a long string, it fell into a different branch that hashed **all** strings as `&str` — including the short inline ones. The same string hashed differently in different batches, so HyperLogLog counted it twice. ## What changes are included in this PR? - **`StringViewHLLAccumulator::update_batch`** and **`Utf8ViewHasher`**: in mixed batches (data buffers present), short strings (≤ 12 bytes) are still hashed as the raw `u128` view; only long strings hash as `&str`. This keeps hashing consistent regardless of batch layout. - **Two regression tests**: - `utf8view_acc_split_batches_match_single_mixed_batch` — scalar accumulator - `utf8view_groups_short_string_hashed_consistently_across_batches` — group accumulator ## Are these changes tested? Yes, two new regression tests cover the exact failure mode. ## Are there any user-facing changes? Yes. `approx_distinct` on `Utf8View` / `VARCHAR VIEW` columns now returns correct (lower) counts. Results may differ from the previously incorrect values. --------- Co-authored-by: Yongting You <2010youy01@gmail.com> --- .../src/approx_distinct.rs | 102 ++++++++++++++++-- 1 file changed, 93 insertions(+), 9 deletions(-) diff --git a/datafusion/functions-aggregate/src/approx_distinct.rs b/datafusion/functions-aggregate/src/approx_distinct.rs index 3550035635647..38b902964f546 100644 --- a/datafusion/functions-aggregate/src/approx_distinct.rs +++ b/datafusion/functions-aggregate/src/approx_distinct.rs @@ -251,16 +251,29 @@ impl Accumulator for StringViewHLLAccumulator { fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { let array: &StringViewArray = downcast_value!(values[0], StringViewArray); - // When all strings are stored inline in the StringView (≤ 12 bytes), - // hash the raw u128 view directly instead of materializing a &str. if array.data_buffers().is_empty() { + // Fast path: with no data buffers every value is inline, so they all + // take the u128 path — no need to check the length per row. for (i, &view) in array.views().iter().enumerate() { if !array.is_null(i) { self.hll.add_hashed(HLL_HASH_STATE.hash_one(view)); } } } else { - self.hll.extend(array.iter().flatten()); + // Mixed batch: decide per row by length. Short strings still use the + // u128 path so they match how they'd be hashed in an all-inline + // batch; only the genuinely out-of-line strings materialize a &str. + for (i, &view) in array.views().iter().enumerate() { + if array.is_null(i) { + continue; + } + // The low 32 bits of the u128 view encode the string length. + if (view as u32) <= 12 { + self.hll.add_hashed(HLL_HASH_STATE.hash_one(view)); + } else { + self.hll.add(array.value(i)); + } + } } Ok(()) @@ -567,9 +580,17 @@ impl HllValueHasher for Utf8ViewHasher { } } } else { + // Mixed batch: short strings (≤ 12 bytes) are still inline and must + // be hashed as the raw u128 view to match the all-inline fast path. + let views = array.views(); for i in 0..array.len() { if nulls.is_none_or(|n| n.is_valid(i)) { - f(i, HLL_HASH_STATE.hash_one(array.value(i))); + let view = views[i]; + if (view as u32) <= 12 { + f(i, HLL_HASH_STATE.hash_one(view)); + } else { + f(i, HLL_HASH_STATE.hash_one(array.value(i))); + } } } } @@ -1037,10 +1058,14 @@ fn is_hll_groups_type(data_type: &DataType) -> bool { } #[cfg(test)] -mod groups_tests { +mod tests { use super::*; + use arrow::array::{AsArray, Int64Array, StringViewArray}; + use std::sync::Arc; + + // A string longer than the 12-byte inline limit + const LONG: &str = "this string is definitely longer than twelve bytes"; - /// Hash a value the same way the accumulators do. fn h(v: u64) -> u64 { HLL_HASH_STATE.hash_one(v) } @@ -1061,6 +1086,13 @@ mod groups_tests { buf } + fn distinct_count(acc: &mut StringViewHLLAccumulator) -> u64 { + match acc.evaluate().unwrap() { + ScalarValue::UInt64(Some(v)) => v, + other => panic!("unexpected evaluate result: {other:?}"), + } + } + #[test] fn sparse_stays_sparse_for_small_groups() { let mut g = GroupHll::default(); @@ -1183,9 +1215,6 @@ mod groups_tests { /// must not be counted (null filter is treated the same as false). #[test] fn update_batch_nullable_filter_excludes_null_filter_rows() { - use arrow::array::Int64Array; - use std::sync::Arc; - let values: ArrayRef = Arc::new(Int64Array::from(vec![1i64, 2, 3, 4, 5])); // row 0: filter=true, row 1: filter=NULL, row 2: filter=false, // row 3: filter=NULL, row 4: filter=true @@ -1205,4 +1234,59 @@ mod groups_tests { let expected = reference_count(&[h(1), h(5)]); assert_eq!(counts.value(0), expected); } + + /// Regression: a short (≤ 12-byte) Utf8View string must hash identically + /// in an all-inline batch and in a mixed batch that also contains a long + /// string (which forces a data buffer). + #[test] + fn utf8view_groups_short_string_hashed_consistently_across_batches() { + // Batch 1: all-inline (no data buffers) — "aaa" is hashed as u128 view. + let batch1: ArrayRef = Arc::new(StringViewArray::from(vec!["aaa", "bbb"])); + assert!(batch1.as_string_view().data_buffers().is_empty()); + + // Batch 2: mixed — LONG forces a data buffer; "aaa" must still be + // hashed as u128 view so it matches its appearance in batch 1. + let batch2: ArrayRef = Arc::new(StringViewArray::from(vec!["aaa", LONG])); + assert!(!batch2.as_string_view().data_buffers().is_empty()); + + let group_indices = vec![0usize, 0]; + let mut acc = HllGroupsAccumulator::::new(); + acc.update_batch(&[batch1], &group_indices, None, 1) + .unwrap(); + acc.update_batch(&[batch2], &group_indices, None, 1) + .unwrap(); + + // True distinct values: {"aaa", "bbb", LONG} == 3. + let result = acc.evaluate(EmitTo::All).unwrap(); + let counts = result.as_any().downcast_ref::().unwrap(); + assert_eq!(counts.value(0), 3); + } + + /// Regression: a short (≤ 12-byte) Utf8View string must hash identically + /// regardless of which batch it appears in — all-inline or mixed. + #[test] + fn utf8view_acc_split_batches_match_single_mixed_batch() { + // Multiset: {"aaa" x2, "bbb", LONG}, so 3 distinct values. + let mixed: ArrayRef = + Arc::new(StringViewArray::from(vec!["aaa", "bbb", LONG, "aaa"])); + let mut acc_single = StringViewHLLAccumulator::new(); + acc_single.update_batch(&[mixed]).unwrap(); + + // Same multiset, but split so "aaa" lands in both an all-inline batch + // and a batch with a data buffer (forced by LONG). + let inline_only: ArrayRef = Arc::new(StringViewArray::from(vec!["aaa", "bbb"])); + let with_buffer: ArrayRef = Arc::new(StringViewArray::from(vec!["aaa", LONG])); + assert!(inline_only.as_string_view().data_buffers().is_empty()); + assert!(!with_buffer.as_string_view().data_buffers().is_empty()); + + let mut acc_split = StringViewHLLAccumulator::new(); + acc_split.update_batch(&[inline_only]).unwrap(); + acc_split.update_batch(&[with_buffer]).unwrap(); + + assert_eq!( + distinct_count(&mut acc_single), + distinct_count(&mut acc_split) + ); + assert_eq!(distinct_count(&mut acc_single), 3); + } } From 986a712ef61f5029ea184c7ce7b93cc2b33fc6de Mon Sep 17 00:00:00 2001 From: chakkk309 Date: Tue, 9 Jun 2026 20:10:49 +0800 Subject: [PATCH 2/2] refactor: Port CaseExpr proto serialization hooks (#22838) ## Which issue does this PR close? - Closes #22421. ## Rationale for this change CaseExpr should own its protobuf serialization and deserialization through the PhysicalExpr try_to_proto / try_from_proto hooks, matching the migration pattern used by other physical expressions and avoiding the central downcast chain. ## What changes are included in this PR? - Adds CaseExpr::try_to_proto in the PhysicalExpr impl. - Adds CaseExpr::try_from_proto as an inherent proto constructor. - Routes physical expression decoding for CaseExpr through the new hook. - Removes the old CaseExpr serialization arm and its helper from datafusion-proto. - Adds direct hook tests covering encode, decode, wrong node variant, missing when/then children, and child encode/decode error propagation. ## Are these changes tested? Yes: - cargo fmt --all - cargo test -p datafusion-physical-expr --features proto case::tests::try_ - cargo test -p datafusion-proto --test proto_integration - cargo clippy -p datafusion-physical-expr --features proto --tests -- -D warnings - cargo clippy -p datafusion-proto --test proto_integration -- -D warnings ## Are there any user-facing changes? No. This is an internal proto serialization refactor with the same wire format. --------- Co-authored-by: chakkk309 --- .../physical-expr/src/expressions/case.rs | 254 ++++++++++++++++++ .../proto/src/physical_plan/from_proto.rs | 36 +-- .../proto/src/physical_plan/to_proto.rs | 59 +--- 3 files changed, 257 insertions(+), 92 deletions(-) diff --git a/datafusion/physical-expr/src/expressions/case.rs b/datafusion/physical-expr/src/expressions/case.rs index 568ecb9cf336b..8a0f15467c47b 100644 --- a/datafusion/physical-expr/src/expressions/case.rs +++ b/datafusion/physical-expr/src/expressions/case.rs @@ -1410,6 +1410,86 @@ impl PhysicalExpr for CaseExpr { } write!(f, "END") } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + Ok(Some(protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(protobuf::physical_expr_node::ExprType::Case(Box::new( + protobuf::PhysicalCaseNode { + expr: self + .expr() + .map(|expr| ctx.encode_child(expr).map(Box::new)) + .transpose()?, + when_then_expr: self + .when_then_expr() + .iter() + .map(|(when_expr, then_expr)| { + Ok(protobuf::PhysicalWhenThen { + when_expr: Some(ctx.encode_child(when_expr)?), + then_expr: Some(ctx.encode_child(then_expr)?), + }) + }) + .collect::>>()?, + else_expr: self + .else_expr() + .map(|expr| ctx.encode_child(expr).map(Box::new)) + .transpose()?, + }, + ))), + })) + } +} + +#[cfg(feature = "proto")] +impl CaseExpr { + /// Reconstruct a [`CaseExpr`] from its protobuf representation. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalExprNode, + ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, + ) -> Result> { + use datafusion_physical_expr_common::expect_expr_variant; + use datafusion_proto_models::protobuf; + + let case = expect_expr_variant!( + node, + protobuf::physical_expr_node::ExprType::Case, + "CaseExpr", + ); + + Ok(Arc::new(CaseExpr::try_new( + case.expr + .as_deref() + .map(|expr| ctx.decode(expr)) + .transpose()?, + case.when_then_expr + .iter() + .map(|when_then| { + Ok(( + ctx.decode_required_expression( + when_then.when_expr.as_ref(), + "CaseExpr", + "when_expr", + )?, + ctx.decode_required_expression( + when_then.then_expr.as_ref(), + "CaseExpr", + "then_expr", + )?, + )) + }) + .collect::>>()?, + case.else_expr + .as_deref() + .map(|expr| ctx.decode(expr)) + .transpose()?, + )?)) + } } /// Attempts to const evaluate the given `predicate`. @@ -3193,3 +3273,177 @@ mod tests { Ok(()) } } + +#[cfg(all(test, feature = "proto"))] +mod proto_tests { + use super::*; + use crate::expressions::col; + use crate::proto_test_util::{ + StubDecoder, StubEncoder, UnreachableDecoder, column_node, + }; + use arrow::datatypes::Field; + use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; + use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx; + use datafusion_proto_models::protobuf; + use datafusion_proto_models::protobuf::{PhysicalExprNode, PhysicalWhenThen}; + + fn proto_case_fixture() -> CaseExpr { + let schema = Schema::new(vec![Field::new("a", DataType::Boolean, true)]); + CaseExpr::try_new( + Some(col("a", &schema).unwrap()), + vec![(lit(true), lit(1_i32))], + Some(lit(0_i32)), + ) + .unwrap() + } + + fn proto_when_then( + when_expr: Option, + then_expr: Option, + ) -> PhysicalWhenThen { + PhysicalWhenThen { + when_expr, + then_expr, + } + } + + fn proto_case_node( + expr: Option>, + when_then_expr: Vec, + else_expr: Option>, + ) -> PhysicalExprNode { + PhysicalExprNode { + expr_id: None, + expr_type: Some(protobuf::physical_expr_node::ExprType::Case(Box::new( + protobuf::PhysicalCaseNode { + expr, + when_then_expr, + else_expr, + }, + ))), + } + } + + #[test] + fn try_to_proto_encodes_case_expr() { + let case = proto_case_fixture(); + let encoder = StubEncoder::ok(); + let ctx = PhysicalExprEncodeCtx::new(&encoder); + + let node = case + .try_to_proto(&ctx) + .unwrap() + .expect("CaseExpr should encode to Some(node)"); + + assert!(node.expr_id.is_none()); + let case_node = match node.expr_type { + Some(protobuf::physical_expr_node::ExprType::Case(boxed)) => *boxed, + other => panic!("expected a CaseExpr node, got {other:?}"), + }; + assert!(case_node.expr.is_some()); + assert_eq!(case_node.when_then_expr.len(), 1); + assert!(case_node.when_then_expr[0].when_expr.is_some()); + assert!(case_node.when_then_expr[0].then_expr.is_some()); + assert!(case_node.else_expr.is_some()); + } + + #[test] + fn try_to_proto_propagates_child_encode_error() { + let case = proto_case_fixture(); + // Call 1 is the optional CASE expr, call 2 is the WHEN expr. + let encoder = StubEncoder::failing_on(2); + let ctx = PhysicalExprEncodeCtx::new(&encoder); + + let err = case.try_to_proto(&ctx).unwrap_err(); + assert!(matches!(err, DataFusionError::Internal(msg) if msg.contains("call 2"))); + } + + #[test] + fn try_from_proto_decodes_case_expr() { + let node = proto_case_node( + Some(Box::new(column_node("case"))), + vec![proto_when_then( + Some(column_node("when")), + Some(column_node("then")), + )], + Some(Box::new(column_node("else"))), + ); + let schema = Schema::empty(); + let decoder = StubDecoder::ok(); + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + + let decoded = CaseExpr::try_from_proto(&node, &ctx).unwrap(); + let case = decoded + .downcast_ref::() + .expect("decoded expr should be a CaseExpr"); + + assert!(case.expr().is_some()); + assert_eq!(case.when_then_expr().len(), 1); + assert!(case.else_expr().is_some()); + } + + #[test] + fn try_from_proto_rejects_non_case_node() { + let node = column_node("a"); + let schema = Schema::empty(); + let decoder = UnreachableDecoder; + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + + let err = CaseExpr::try_from_proto(&node, &ctx).unwrap_err(); + assert!( + matches!(err, DataFusionError::Internal(msg) if msg.contains("PhysicalExprNode is not a CaseExpr")) + ); + } + + #[test] + fn try_from_proto_rejects_missing_when_expr() { + let node = proto_case_node( + None, + vec![proto_when_then(None, Some(column_node("then")))], + None, + ); + let schema = Schema::empty(); + let decoder = UnreachableDecoder; + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + + let err = CaseExpr::try_from_proto(&node, &ctx).unwrap_err(); + assert!( + matches!(err, DataFusionError::Internal(msg) if msg.contains("CaseExpr is missing required field 'when_expr'")) + ); + } + + #[test] + fn try_from_proto_rejects_missing_then_expr() { + let node = proto_case_node( + None, + vec![proto_when_then(Some(column_node("when")), None)], + None, + ); + let schema = Schema::empty(); + let decoder = StubDecoder::ok(); + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + + let err = CaseExpr::try_from_proto(&node, &ctx).unwrap_err(); + assert!( + matches!(err, DataFusionError::Internal(msg) if msg.contains("CaseExpr is missing required field 'then_expr'")) + ); + } + + #[test] + fn try_from_proto_propagates_child_decode_error() { + let node = proto_case_node( + Some(Box::new(column_node("case"))), + vec![proto_when_then( + Some(column_node("when")), + Some(column_node("then")), + )], + Some(Box::new(column_node("else"))), + ); + let schema = Schema::empty(); + let decoder = StubDecoder::failing_on(2); + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + + let err = CaseExpr::try_from_proto(&node, &ctx).unwrap_err(); + assert!(matches!(err, DataFusionError::Internal(msg) if msg.contains("call 2"))); + } +} diff --git a/datafusion/proto/src/physical_plan/from_proto.rs b/datafusion/proto/src/physical_plan/from_proto.rs index c88663399908d..36751d8a61a3e 100644 --- a/datafusion/proto/src/physical_plan/from_proto.rs +++ b/datafusion/proto/src/physical_plan/from_proto.rs @@ -303,41 +303,7 @@ pub fn parse_physical_expr_with_converter( ExprType::NotExpr(_) => NotExpr::try_from_proto(proto, &decode_ctx)?, ExprType::Negative(_) => NegativeExpr::try_from_proto(proto, &decode_ctx)?, ExprType::InList(_) => InListExpr::try_from_proto(proto, &decode_ctx)?, - ExprType::Case(e) => Arc::new(CaseExpr::try_new( - e.expr - .as_ref() - .map(|e| { - proto_converter.proto_to_physical_expr(e.as_ref(), input_schema, ctx) - }) - .transpose()?, - e.when_then_expr - .iter() - .map(|e| { - Ok(( - parse_required_physical_expr( - e.when_expr.as_ref(), - ctx, - "when_expr", - input_schema, - proto_converter, - )?, - parse_required_physical_expr( - e.then_expr.as_ref(), - ctx, - "then_expr", - input_schema, - proto_converter, - )?, - )) - }) - .collect::>>()?, - e.else_expr - .as_ref() - .map(|e| { - proto_converter.proto_to_physical_expr(e.as_ref(), input_schema, ctx) - }) - .transpose()?, - )?), + ExprType::Case(_) => CaseExpr::try_from_proto(proto, &decode_ctx)?, ExprType::Cast(_) => CastExpr::try_from_proto(proto, &decode_ctx)?, ExprType::TryCast(_) => TryCastExpr::try_from_proto(proto, &decode_ctx)?, ExprType::ScalarUdf(e) => { diff --git a/datafusion/proto/src/physical_plan/to_proto.rs b/datafusion/proto/src/physical_plan/to_proto.rs index c45d432f9a6aa..d9315af431e22 100644 --- a/datafusion/proto/src/physical_plan/to_proto.rs +++ b/datafusion/proto/src/physical_plan/to_proto.rs @@ -35,7 +35,7 @@ use datafusion_physical_expr::ScalarFunctionExpr; use datafusion_physical_expr::scalar_subquery::ScalarSubqueryExpr; use datafusion_physical_expr::window::{SlidingAggregateWindowExpr, StandardWindowExpr}; use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; -use datafusion_physical_plan::expressions::{CaseExpr, DynamicFilterPhysicalExpr}; +use datafusion_physical_plan::expressions::DynamicFilterPhysicalExpr; use datafusion_physical_plan::udaf::AggregateFunctionExpr; use datafusion_physical_plan::windows::{PlainAggregateWindowExpr, WindowUDFExpr}; use datafusion_physical_plan::{ @@ -300,50 +300,7 @@ pub fn serialize_physical_expr_with_converter( return Ok(node); } - if let Some(expr) = expr.downcast_ref::() { - Ok(protobuf::PhysicalExprNode { - expr_id, - expr_type: Some( - protobuf::physical_expr_node::ExprType::Case( - Box::new( - protobuf::PhysicalCaseNode { - expr: expr - .expr() - .map(|exp| { - proto_converter - .physical_expr_to_proto(exp, codec) - .map(Box::new) - }) - .transpose()?, - when_then_expr: expr - .when_then_expr() - .iter() - .map(|(when_expr, then_expr)| { - serialize_when_then_expr( - when_expr, - then_expr, - codec, - proto_converter, - ) - }) - .collect::, - DataFusionError, - >>()?, - else_expr: expr - .else_expr() - .map(|a| { - proto_converter - .physical_expr_to_proto(a, codec) - .map(Box::new) - }) - .transpose()?, - }, - ), - ), - ), - }) - } else if let Some(expr) = expr.downcast_ref::() { + if let Some(expr) = expr.downcast_ref::() { let mut buf = Vec::new(); codec.try_encode_udf(expr.fun(), &mut buf)?; Ok(protobuf::PhysicalExprNode { @@ -500,18 +457,6 @@ fn serialize_range_split_point( }) } -fn serialize_when_then_expr( - when_expr: &Arc, - then_expr: &Arc, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, -) -> Result { - Ok(protobuf::PhysicalWhenThen { - when_expr: Some(proto_converter.physical_expr_to_proto(when_expr, codec)?), - then_expr: Some(proto_converter.physical_expr_to_proto(then_expr, codec)?), - }) -} - impl TryFromProto<&PartitionedFile> for protobuf::PartitionedFile { type Error = DataFusionError;