sunchao commented on code in PR #6073:
URL: https://github.com/apache/datafusion-comet/pull/6073#discussion_r4064607264
##########
native/spark-expr/src/array_funcs/nested_float_normalize.rs:
##########
@@ -20,9 +20,105 @@ use arrow::array::{
Array, ArrayRef, AsArray, FixedSizeListArray, Float32Array, Float64Array,
LargeListArray,
ListArray, StructArray,
};
-use arrow::datatypes::{DataType, Float32Type, Float64Type};
+use arrow::datatypes::{DataType, Float32Type, Float64Type, Schema};
+use arrow::record_batch::RecordBatch;
+use datafusion::common::{Result, ScalarValue};
+use datafusion::logical_expr::ColumnarValue;
+use datafusion::physical_expr::PhysicalExpr;
+use std::fmt::{Display, Formatter};
+use std::hash::{Hash, Hasher};
use std::sync::Arc;
+/// Normalizes nested IN operands, preserving constants for static membership
lookup.
+#[derive(Debug, Eq)]
+pub struct NormalizeNestedFloats {
+ child: Arc<dyn PhysicalExpr>,
+}
+
+impl NormalizeNestedFloats {
+ /// Wrap nested floating-point operands only; scalar floats keep their
existing semantics.
+ pub fn wrap_if_needed(
+ child: Arc<dyn PhysicalExpr>,
+ schema: &Schema,
+ ) -> Result<Arc<dyn PhysicalExpr>> {
+ let dt = child.data_type(schema)?;
+ if matches!(
+ dt,
+ DataType::List(_)
+ | DataType::LargeList(_)
+ | DataType::FixedSizeList(_, _)
+ | DataType::Struct(_)
+ ) && has_float_leaf(&dt)
+ {
+ Ok(Arc::new(Self { child }))
+ } else {
+ Ok(child)
+ }
+ }
+}
+
+impl PartialEq for NormalizeNestedFloats {
+ fn eq(&self, other: &Self) -> bool {
+ self.child.eq(&other.child)
+ }
+}
+
+impl Hash for NormalizeNestedFloats {
+ fn hash<H: Hasher>(&self, state: &mut H) {
+ self.child.hash(state);
+ }
+}
+
+impl Display for NormalizeNestedFloats {
+ fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+ write!(f, "NormalizeNestedFloats({})", self.child)
+ }
+}
+
+impl PhysicalExpr for NormalizeNestedFloats {
+ fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+ Display::fmt(self, f)
+ }
+
+ fn data_type(&self, schema: &Schema) -> Result<DataType> {
+ self.child.data_type(schema)
+ }
+
+ fn nullable(&self, schema: &Schema) -> Result<bool> {
+ self.child.nullable(schema)
+ }
+
+ fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
+ match self.child.evaluate(batch)? {
+ ColumnarValue::Array(array) => {
+ Ok(ColumnarValue::Array(normalize_nested_floats(&array)))
Review Comment:
### Performance
[P2] Could you keep the dynamic-candidate path from normalizing entire
nested columns before it can short-circuit, and add a focused microbenchmark
for that case? For `a IN (b)` with two `ARRAY<DOUBLE>` columns whose first
elements differ, DataFusion's dynamic list comparator currently stops after
that first element in each row. This wrapper runs `normalize_nested_floats`
over both complete value buffers first, and Arrow's `unary` allocates a new
buffer even when every value is an ordinary finite number. For 8,192 rows with
1,024 doubles per array, that adds 128 MiB of float-value buffers and
16,777,216 normalization operations to a comparison that can inspect one pair
per row. This is a source-derived cost, not a measured runtime claim. Please
compare base and head for short and wide arrays, early mismatches, nulls, and
constant versus column candidates. A comparison that applies Spark's
floating-point equivalence to the elements it visits could preserve the
early-exit behavior f
or dynamic candidates while retaining canonicalization for the static hash
path.
--
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]
