mizulun commented on code in PR #6073:
URL: https://github.com/apache/datafusion-comet/pull/6073#discussion_r4082484113
##########
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:
Thanks @sunchao, both points are addressed in the latest push.
**Equality path.** `=` / `<>` on arrays and structs with float leaves now go
through `spark_comparison`, which compares elements with Spark's equivalence
(`-0.0 == 0.0`, all NaNs equal) instead of Arrow's raw comparator. Other
operators and non-float nested types still build a plain `BinaryExpr`. The
regression from @rich7420 (singleton and duplicate candidates rewritten to
`EqualTo`, including `NOT IN`) passes and asserts the native `equalto` path.
**Dynamic membership.** When any candidate is a column, `spark_in_list`
compares element by element and stops at the first mismatch. It no longer
normalizes or copies the whole nested value buffers. Constant-only lists still
normalize the constants once and use DataFusion's static filter.
**Benchmark** (`native/spark-expr/benches/nested_comparison.rs`): 8,192 rows
of `ARRAY<DOUBLE>`, finite values only, so all three versions return identical
results. The bench asserts this before timing.
- `base`: DataFusion with no normalization. It is fast but wrong for signed
zero and NaN.
- `head`: the previous revision, which eagerly normalized every operand.
- `new`: this revision.
Each number is the minimum of two full runs, no nulls (i9-9880H, rustc
1.98.1):
| case | base | head | new |
|---|---:|---:|---:|
| `a IN (b)`, width 1024, mismatch at first element | 486 µs | 38.9 ms | 390
µs |
| `a IN (b)`, width 1024, all equal | 40.1 ms | 78.3 ms | 25.0 ms |
| `a IN (b)`, width 16, mismatch at first element | 86.3 µs | 262 µs | 82.9
µs |
| `a IN (b)`, width 1, all equal | 85.6 µs | 102 µs | 90.0 µs |
| `a IN (const, b)`, width 1024, mismatch at first element | 835 µs | 39.7
ms | 713 µs |
| `a = b`, width 1024, all equal | 40.2 ms | 39.6 ms | 25.0 ms |
| `a = b`, width 1024, mismatch at first element | 423 µs | 428 µs | 389 µs |
| `a = b`, width 16, mismatch at first element | 76.2 µs | 75.5 µs | 85.4 µs
|
| `a IN (const)`, width 1024, all equal | 85.4 ms | 105 ms | 105 ms |
| `a IN (const)`, width 16, mismatch at first element | 371 µs | 657 µs |
612 µs |
| static filter build, width 1024 | 8.3 µs | 10.0 µs | 11.9 µs |
- **Dynamic candidates.** Early exit is restored. Wide arrays with a
first-element mismatch return to base cost: the column-only case is 486 µs →
390 µs, where head took 38.9 ms. That is about 100x faster than head. Full
scans are about 1.5x faster than base on widths 16 and 1024, and single-element
arrays stay within about 20% of base.
- **Equality.** Wide or full-scan comparisons are 1.3–1.7x faster than base.
On short arrays with an early mismatch, the new path is up to about 25% slower
than base, roughly 10–25 µs per batch. That is the fixed per-row cost of the
Spark-equivalent comparator.
- **Constant lists.** These match head and are about 15–42% slower than
base, because the probe column is still normalized on each batch before the
static lookup. I kept that design because it keeps the lookup a hash probe. If
this overhead matters, I can follow up.
- **Nulls.** The sparse (every 16th row) and dense (every 2nd row) null
variants follow the same pattern. The full table is below.
Caveats: this is criterion with 10 samples × 0.5 s per case. Across the two
runs, the head/new ratio differed by 6% at the median. Only two width-1
constant cases changed direction, and they are within noise.
<details><summary>Full results (111 cases; head/new and base/new are
speedups, >1 means new is faster)</summary>
| case | base | head (eager normalize) | new | head/new | base/new |
|---|---:|---:|---:|---:|---:|
| `constant/1/equal/dense` | 129 µs | 142 µs | 207 µs | 0.69x | 0.62x |
| `constant/1/equal/no_nulls` | 160 µs | 246 µs | 214 µs | 1.15x | 0.75x |
| `constant/1/equal/sparse` | 176 µs | 222 µs | 306 µs | 0.73x | 0.58x |
| `constant/1/first/dense` | 123 µs | 139 µs | 171 µs | 0.81x | 0.72x |
| `constant/1/first/no_nulls` | 144 µs | 147 µs | 174 µs | 0.84x | 0.82x |
| `constant/1/first/sparse` | 182 µs | 222 µs | 239 µs | 0.93x | 0.76x |
| `constant/1/last/dense` | 116 µs | 207 µs | 195 µs | 1.06x | 0.59x |
| `constant/1/last/no_nulls` | 160 µs | 178 µs | 195 µs | 0.91x | 0.82x |
| `constant/1/last/sparse` | 181 µs | 258 µs | 219 µs | 1.17x | 0.83x |
| `constant/1024/equal/dense` | 58.6 ms | 77.4 ms | 77.4 ms | 1.00x | 0.76x |
| `constant/1024/equal/no_nulls` | 85.4 ms | 105 ms | 105 ms | 1.00x | 0.82x
|
| `constant/1024/equal/sparse` | 81.9 ms | 102 ms | 101 ms | 1.01x | 0.81x |
| `constant/1024/first/dense` | 38.0 ms | 57.4 ms | 57.2 ms | 1.00x | 0.66x |
| `constant/1024/first/no_nulls` | 45.4 ms | 64.6 ms | 64.1 ms | 1.01x |
0.71x |
| `constant/1024/first/sparse` | 44.6 ms | 63.6 ms | 63.4 ms | 1.00x | 0.70x
|
| `constant/1024/last/dense` | 58.2 ms | 78.0 ms | 76.6 ms | 1.02x | 0.76x |
| `constant/1024/last/no_nulls` | 85.0 ms | 104 ms | 105 ms | 0.99x | 0.81x |
| `constant/1024/last/sparse` | 82.1 ms | 101 ms | 101 ms | 1.00x | 0.81x |
| `constant/16/equal/dense` | 694 µs | 1.1 ms | 989 µs | 1.12x | 0.70x |
| `constant/16/equal/no_nulls` | 1.0 ms | 1.3 ms | 1.4 ms | 0.94x | 0.77x |
| `constant/16/equal/sparse` | 1.0 ms | 1.2 ms | 1.4 ms | 0.88x | 0.73x |
| `constant/16/first/dense` | 378 µs | 607 µs | 446 µs | 1.36x | 0.85x |
| `constant/16/first/no_nulls` | 371 µs | 657 µs | 612 µs | 1.07x | 0.61x |
| `constant/16/first/sparse` | 402 µs | 703 µs | 605 µs | 1.16x | 0.66x |
| `constant/16/last/dense` | 723 µs | 1.1 ms | 887 µs | 1.19x | 0.82x |
| `constant/16/last/no_nulls` | 1.0 ms | 1.3 ms | 1.4 ms | 0.99x | 0.77x |
| `constant/16/last/sparse` | 1.1 ms | 1.4 ms | 1.5 ms | 0.94x | 0.73x |
| `dynamic/1/equal/dense` | 72.5 µs | 122 µs | 76.4 µs | 1.59x | 0.95x |
| `dynamic/1/equal/no_nulls` | 85.6 µs | 102 µs | 90.0 µs | 1.13x | 0.95x |
| `dynamic/1/equal/sparse` | 107 µs | 172 µs | 123 µs | 1.40x | 0.87x |
| `dynamic/1/first/dense` | 69.0 µs | 85.5 µs | 79.4 µs | 1.08x | 0.87x |
| `dynamic/1/first/no_nulls` | 65.2 µs | 75.1 µs | 67.2 µs | 1.12x | 0.97x |
| `dynamic/1/first/sparse` | 97.0 µs | 149 µs | 122 µs | 1.22x | 0.79x |
| `dynamic/1/last/dense` | 64.9 µs | 132 µs | 70.8 µs | 1.87x | 0.92x |
| `dynamic/1/last/no_nulls` | 81.9 µs | 159 µs | 77.4 µs | 2.05x | 1.06x |
| `dynamic/1/last/sparse` | 95.4 µs | 122 µs | 108 µs | 1.13x | 0.88x |
| `dynamic/1024/equal/dense` | 20.3 ms | 59.3 ms | 12.5 ms | 4.75x | 1.63x |
| `dynamic/1024/equal/no_nulls` | 40.1 ms | 78.3 ms | 25.0 ms | 3.14x |
1.61x |
| `dynamic/1024/equal/sparse` | 35.8 ms | 76.5 ms | 22.6 ms | 3.38x | 1.58x |
| `dynamic/1024/first/dense` | 282 µs | 38.8 ms | 286 µs | 135.80x | 0.99x |
| `dynamic/1024/first/no_nulls` | 486 µs | 38.9 ms | 390 µs | 99.65x | 1.25x
|
| `dynamic/1024/first/sparse` | 507 µs | 38.9 ms | 521 µs | 74.55x | 0.97x |
| `dynamic/1024/last/dense` | 20.1 ms | 58.9 ms | 12.5 ms | 4.72x | 1.61x |
| `dynamic/1024/last/no_nulls` | 39.3 ms | 78.6 ms | 24.8 ms | 3.16x | 1.58x
|
| `dynamic/1024/last/sparse` | 37.3 ms | 77.2 ms | 23.7 ms | 3.25x | 1.57x |
| `dynamic/16/equal/dense` | 352 µs | 791 µs | 243 µs | 3.25x | 1.45x |
| `dynamic/16/equal/no_nulls` | 593 µs | 983 µs | 409 µs | 2.41x | 1.45x |
| `dynamic/16/equal/sparse` | 621 µs | 1.1 ms | 398 µs | 2.69x | 1.56x |
| `dynamic/16/first/dense` | 75.6 µs | 249 µs | 84.8 µs | 2.94x | 0.89x |
| `dynamic/16/first/no_nulls` | 86.3 µs | 262 µs | 82.9 µs | 3.16x | 1.04x |
| `dynamic/16/first/sparse` | 109 µs | 379 µs | 122 µs | 3.10x | 0.89x |
| `dynamic/16/last/dense` | 332 µs | 689 µs | 217 µs | 3.17x | 1.53x |
| `dynamic/16/last/no_nulls` | 615 µs | 1.0 ms | 407 µs | 2.53x | 1.51x |
| `dynamic/16/last/sparse` | 608 µs | 1.0 ms | 411 µs | 2.45x | 1.48x |
| `eq/1/equal/dense` | 67.7 µs | 58.7 µs | 85.4 µs | 0.69x | 0.79x |
| `eq/1/equal/no_nulls` | 76.7 µs | 86.9 µs | 82.4 µs | 1.05x | 0.93x |
| `eq/1/equal/sparse` | 104 µs | 103 µs | 121 µs | 0.85x | 0.86x |
| `eq/1/first/dense` | 59.3 µs | 58.7 µs | 77.8 µs | 0.76x | 0.76x |
| `eq/1/first/no_nulls` | 75.4 µs | 73.6 µs | 83.9 µs | 0.88x | 0.90x |
| `eq/1/first/sparse` | 94.2 µs | 99.0 µs | 120 µs | 0.82x | 0.78x |
| `eq/1/last/dense` | 57.3 µs | 56.9 µs | 70.7 µs | 0.81x | 0.81x |
| `eq/1/last/no_nulls` | 75.8 µs | 70.5 µs | 78.4 µs | 0.90x | 0.97x |
| `eq/1/last/sparse` | 86.2 µs | 89.8 µs | 107 µs | 0.84x | 0.81x |
| `eq/1024/equal/dense` | 20.3 ms | 19.5 ms | 12.2 ms | 1.60x | 1.67x |
| `eq/1024/equal/no_nulls` | 40.2 ms | 39.6 ms | 25.0 ms | 1.58x | 1.61x |
| `eq/1024/equal/sparse` | 37.1 ms | 37.3 ms | 22.9 ms | 1.63x | 1.62x |
| `eq/1024/first/dense` | 277 µs | 278 µs | 289 µs | 0.96x | 0.96x |
| `eq/1024/first/no_nulls` | 423 µs | 428 µs | 389 µs | 1.10x | 1.09x |
| `eq/1024/first/sparse` | 519 µs | 505 µs | 528 µs | 0.96x | 0.98x |
| `eq/1024/last/dense` | 17.5 ms | 20.3 ms | 12.6 ms | 1.61x | 1.39x |
| `eq/1024/last/no_nulls` | 40.2 ms | 40.3 ms | 25.2 ms | 1.60x | 1.60x |
| `eq/1024/last/sparse` | 36.3 ms | 34.3 ms | 23.4 ms | 1.46x | 1.55x |
| `eq/16/equal/dense` | 348 µs | 339 µs | 229 µs | 1.48x | 1.52x |
| `eq/16/equal/no_nulls` | 638 µs | 636 µs | 367 µs | 1.73x | 1.74x |
| `eq/16/equal/sparse` | 602 µs | 635 µs | 381 µs | 1.67x | 1.58x |
| `eq/16/first/dense` | 69.6 µs | 71.2 µs | 79.2 µs | 0.90x | 0.88x |
| `eq/16/first/no_nulls` | 76.2 µs | 75.5 µs | 85.4 µs | 0.88x | 0.89x |
| `eq/16/first/sparse` | 103 µs | 100 µs | 119 µs | 0.84x | 0.86x |
| `eq/16/last/dense` | 334 µs | 312 µs | 251 µs | 1.24x | 1.33x |
| `eq/16/last/no_nulls` | 578 µs | 628 µs | 368 µs | 1.71x | 1.57x |
| `eq/16/last/sparse` | 592 µs | 645 µs | 435 µs | 1.48x | 1.36x |
| `mixed/1/equal/dense` | 135 µs | 227 µs | 134 µs | 1.70x | 1.01x |
| `mixed/1/equal/no_nulls` | 88.3 µs | 99.7 µs | 81.6 µs | 1.22x | 1.08x |
| `mixed/1/equal/sparse` | 191 µs | 225 µs | 217 µs | 1.04x | 0.88x |
| `mixed/1/first/dense` | 122 µs | 190 µs | 135 µs | 1.40x | 0.90x |
| `mixed/1/first/no_nulls` | 155 µs | 213 µs | 165 µs | 1.29x | 0.94x |
| `mixed/1/first/sparse` | 189 µs | 251 µs | 195 µs | 1.28x | 0.97x |
| `mixed/1/last/dense` | 129 µs | 162 µs | 151 µs | 1.07x | 0.85x |
| `mixed/1/last/no_nulls` | 154 µs | 195 µs | 180 µs | 1.08x | 0.86x |
| `mixed/1/last/sparse` | 203 µs | 270 µs | 227 µs | 1.19x | 0.89x |
| `mixed/1024/equal/dense` | 40.2 ms | 79.0 ms | 24.8 ms | 3.19x | 1.63x |
| `mixed/1024/equal/no_nulls` | 40.1 ms | 58.7 ms | 23.7 ms | 2.48x | 1.69x |
| `mixed/1024/equal/sparse` | 75.0 ms | 115 ms | 46.2 ms | 2.48x | 1.63x |
| `mixed/1024/first/dense` | 484 µs | 39.2 ms | 493 µs | 79.58x | 0.98x |
| `mixed/1024/first/no_nulls` | 835 µs | 39.7 ms | 713 µs | 55.68x | 1.17x |
| `mixed/1024/first/sparse` | 905 µs | 39.3 ms | 925 µs | 42.49x | 0.98x |
| `mixed/1024/last/dense` | 40.4 ms | 79.9 ms | 24.8 ms | 3.22x | 1.63x |
| `mixed/1024/last/no_nulls` | 80.2 ms | 116 ms | 49.3 ms | 2.35x | 1.63x |
| `mixed/1024/last/sparse` | 73.6 ms | 113 ms | 46.3 ms | 2.44x | 1.59x |
| `mixed/16/equal/dense` | 702 µs | 1.0 ms | 480 µs | 2.18x | 1.46x |
| `mixed/16/equal/no_nulls` | 564 µs | 774 µs | 394 µs | 1.96x | 1.43x |
| `mixed/16/equal/sparse` | 1.3 ms | 1.7 ms | 858 µs | 2.01x | 1.51x |
| `mixed/16/first/dense` | 140 µs | 303 µs | 156 µs | 1.95x | 0.90x |
| `mixed/16/first/no_nulls` | 163 µs | 493 µs | 160 µs | 3.08x | 1.02x |
| `mixed/16/first/sparse` | 208 µs | 521 µs | 220 µs | 2.36x | 0.94x |
| `mixed/16/last/dense` | 603 µs | 1.1 ms | 480 µs | 2.30x | 1.26x |
| `mixed/16/last/no_nulls` | 1.2 ms | 1.8 ms | 843 µs | 2.17x | 1.44x |
| `mixed/16/last/sparse` | 1.3 ms | 1.6 ms | 752 µs | 2.09x | 1.68x |
| `nested_static_build/1` | 4.1 µs | 5.3 µs | 5.7 µs | 0.92x | 0.71x |
| `nested_static_build/1024` | 8.3 µs | 10.0 µs | 11.9 µs | 0.84x | 0.70x |
| `nested_static_build/16` | 4.3 µs | 5.3 µs | 6.3 µs | 0.84x | 0.69x |
</details>
Reproduce: `cd native && cargo bench -p datafusion-comet-spark-expr --bench
nested_comparison`
--
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]
