comphead commented on PR #5806:
URL:
https://github.com/apache/datafusion-comet/pull/5806#issuecomment-5605325013
Reviewed against the PR head `4f6a1cee80`.
Overall this is a well-built change. The approach is right, the Arrow usage
is correct, and the two hazards specific to this kernel (sliced `MapArray`
entry offsets, per-row lookup keys) are both handled and both tested at the
Rust and SQL levels. No blockers from me.
## Major
### 1. The map's validity buffer is never consulted, so a NULL map row is
answered from its entries
**Where:** `native/spark-expr/src/map_funcs/map_extract.rs:181-192`
The per-row scan derives NULL purely from "no matching entry":
```rust
let found = (start..end).find(|&i| matched.value(i));
if let Some(i) = found { indices[row] = (i + entries_start) as u32; }
nulls.append(found.is_some());
```
`map_array.nulls()` is never read. A row marked NULL in the `MapArray` whose
offset window is non-empty returns that entry's value instead of NULL.
Spark's `GetMapValue` and `ElementAt` are `BinaryExpression`s evaluated
through `nullSafeEval`, so a NULL map operand returns NULL unconditionally,
independent of what the underlying storage holds. Arrow does not require a null
list/map slot to have a zero-length offset window, only that offsets stay
monotonic, so the code relies on an invariant it does not own.
This is *not* a regression against the pinned DataFusion 55.0.0, which has
the same gap (`general_map_extract_inner` iterates `value_offsets().windows(2)`
with no validity test and passes `None` as the output `ListArray`'s null
buffer). But upstream closed exactly this gap in apache/datafusion#24999, which
added `if map_array.is_valid(row_index)` plus `map_array.nulls().cloned()`.
Since this PR forks the kernel into Comet, Comet will not inherit that fix when
the DataFusion pin moves.
I traced the producers that could deliver such an input (arrow-rs Parquet
list reader, arrow-java `ListVector.setValueCount`,
`MutableArrayData::extend_nulls` behind DataFusion's `CaseExpr` scatter,
`filter`, `concat`, Comet's own `spark_map_sort`) and all of them emit
degenerate offsets for null rows. So I could not construct a query that
triggers it today. Latent rather than live.
**Suggested direction:** one line at the end of the scan, before building
`indices`:
```rust
let nulls = NullBuffer::union(map_array.nulls(), nulls.finish().as_ref());
```
`map_array.nulls()` is already sliced to the visible rows, so no offset
arithmetic is needed. A Rust unit test building a `MapArray` via
`MapArray::try_new` with a null row over a non-empty offset window is the only
way to construct the input.
## Minor
### 2. A constant map argument is expanded to one copy per row before the
compare
**Where:** `map_extract.rs:120-123`
`ColumnarValue::Scalar(scalar) => scalar.to_array_of_size(number_rows)?`
materializes `number_rows x entries` map entries, and the vectorized `eq` then
runs over all of them. For an 8192-row batch against a folded map literal with
`E` entries that is `8192 * E` Arrow entries built and compared per batch.
This is reachable: `CometLiteral` rebuilds folded `MapType` literals
natively (gated on `MapKeySupport`), and `CometMapExpressionSuite` already
exercises `element_at(<folded map literal>, _1)`. It is not a regression (DF
55's `make_scalar_function` expands scalars identically), and the very large
case in the suite goes through `CometCreateMap`'s JVM dispatcher rather than a
literal, so the exposure is bounded by how large a map literal survives as a
`Literal`. Still, in a PR whose thesis is "hoist the constant out of the row
loop", the constant *map* is left un-hoisted while the constant *key* is
hoisted.
If you take it, the cheap version is a scalar-map branch that compares the
key array against the single map's `E` entry keys with `E` full-length `eq`
calls. Memory drops from `O(N * E)` to `O(N)`. If you would rather not add a
second path, a comment at line 122 recording that the expansion is deliberate
and bounded would keep the next reader from re-deriving this.
### 3. `spark_map_extract` is exported publicly with no consumer
**Where:** `native/spark-expr/src/lib.rs:64`, `map_funcs/mod.rs:19`
`SparkMapExtract` is used by `comet_scalar_funcs.rs` and the bench. The free
function `spark_map_extract` is only called from `invoke_with_args` and the
module's own tests, both of which see it through `super::*`. Compare
`spark_map_sort`, which is re-exported because `comet_scalar_funcs.rs` imports
it by name. Dropping `pub` on the function and the re-export keeps the crate's
public surface to the UDF.
### 4. `DataType::Null` in the map position is now an error rather than a
passthrough
**Where:** `map_extract.rs:104-112` and `124-129`
DF 55's `return_type`, `coerce_types`, and `map_extract_inner` all
special-case `map_type.is_null()`. The replacement does not, so a `Null`-typed
first argument produces `map_extract: the first argument must be a map, got
Null`. I believe this is unreachable from Comet (Spark's `ExtractValue.apply`
requires a `MapType` child, and `CometElementAt.getSupportLevel` returns
`Unsupported` for anything that is not an array or a map), so I would not add
the branch. Worth being deliberate about it, since the doc comment claims the
kernel "is never less capable than the one it replaces" and this is the one
place it is.
### 5. Test coverage does not span the key types the new compare primitive
accepts
The compare changed from `dyn Array` / `ArrayData` equality to
`arrow::compute::kernels::cmp::eq`, which is type-strict in ways `ArrayData`
equality is not (decimal precision and scale, timestamp time zone strings,
integer width). Everything that survives `MapKeySupport` reaches the new `eq`,
but the tests only exercise `Utf8`, `Int32`, and `Binary` keys
(`element_at_map.sql` covers string, int with coercion, and binary; the Rust
tests cover `Utf8` and `Int32`).
Not exercised anywhere: `decimal`, `date`, `timestamp` / `timestamp_ntz`,
`boolean`, and the narrower integer widths as map key types. Adding a handful
of rows to `element_at_map.sql` is cheap relative to the risk that one of these
takes the `exec_err!` path in `key_match_mask` at runtime instead of matching.
## Questions
### 6. The override is registered by name only, so `element_at` still
resolves to DataFusion's implementation
DataFusion's `SessionState::register_udf` registers a UDF under its aliases
as well as its name, and upstream `MapExtract` declares `aliases:
["element_at"]`. `SparkMapExtract` declares none, so after
`register_all_comet_functions` the registry holds `map_extract ->
SparkMapExtract` and `element_at -> MapExtract` (list-returning). Nothing in
Comet emits the name `element_at` (grepped `spark/src/main`, `native/core/src`,
`native/spark-expr/src`), so this is inert today, and it was equally inert
before the PR because the deleted `planner.rs` arm also matched only
`"map_extract"`. Worth either adding the alias or noting in the doc comment
that the override is name-scoped, so a future `element_at` emission does not
silently get a one-element list?
### 7. The PR description overstates one of the behavior changes
> a lookup key whose runtime type is not the map's key type is now rejected
rather than silently missing every row
DF 55's `map_extract_inner` already does this:
```rust
if key_type != key_arg.data_type() {
return exec_err!("The key type {} does not match the map key type {}",
...);
}
```
So that part is a no-op change, only the message text differs. Worth
correcting in the description so it is not weighed as new risk.
There *is* a real null-semantics improvement here that the description does
not claim. DF 55 compares a NULL lookup key against a NULL stored key with
`ArrayData` equality, which reports them equal, so a map carrying a NULL key
(reachable through the known `map_from_arrays` gap, #4680) would return that
key's value where Spark returns NULL. The new kernel returns NULL for a NULL
lookup key on both the scalar path (line 150-153) and the array path (the mask
is intersected with the comparison's null buffer at line 228). Upstream's
`make_comparator` fix does not close this, since `SortOptions::default()` also
treats two nulls as equal. That is a genuine correctness win worth mentioning
in the description.
## Scope
Appropriately scoped. Kernel, registration, the planner cleanup the kernel
enables, a bench, and tests. No drive-by refactoring, no unrelated formatting,
no new configuration, no API change beyond the one noted in finding 3.
The `elementwise_match_mask` backstop (lines 239-253) is the only arguably
speculative piece, since `MapKeySupport` declines every key type `eq` rejects,
so it is unreachable from Comet. It faithfully reproduces DF 55's comparison,
is 15 lines, and is unit-tested, so I would keep it.
`datafusion-comet-spark-expr` is a published crate with consumers outside
Comet's serde gate.
No duplicated abstraction. `spark_map_sort` handles the sliced-`MapArray`
trap by rebasing offsets rather than windowing them, which is right for its
output shape, so there is no helper to extract.
## Spark compatibility
Checked against `complexTypeExtractors.scala`. `GetMapValue.nullSafeEval`
delegates to `GetMapValueUtil.getValueEval`, which scans entries in order and
stops at the first `ordering.equiv` hit, returning null when not found or when
the matched value is null. `ElementAt`'s map overload shares that path. Both
are `nullSafeEval`, so a NULL map or a NULL key short-circuits to NULL. Since
SPARK-40066 (Spark 3.4) neither throws under ANSI for a missing key, so there
is no ANSI divergence to reproduce here.
Matches Spark: first-match-wins, missing key, NULL key, NULL map with a
degenerate window, empty map, NULL stored value. Improves on the replaced
kernel for a NULL lookup key against a NULL stored key.
Diverges: the NULL map row with a non-empty window (finding 1).
Pre-existing, not introduced here.
Remaining risk is carried entirely by `MapKeySupport`, unchanged by this PR.
The gate is correct for the new kernel for the same reasons it was correct for
the old one, since both compare raw Arrow values rather than Spark's normalized
keys.
One thing worth adding to the description: removing the `ListExtract`
wrapper also decouples Comet from apache/datafusion#24999, which changed the
absent-key result from a one-element NULL list to an empty list. `ListExtract`
with `fail_on_error=false` would have absorbed that, but not depending on it is
better. That answers my earlier question on this PR.
## Tests
Present and load-bearing:
- Sliced `MapArray` entry offsets, at both the Rust level
(`sliced_map_keeps_original_entry_offsets`, constant and per-row key) and the
SQL level (`element_at on a sliced map reads the visible entries`, via `ORDER
BY ... LIMIT ... OFFSET`). Reverting the `entries_start` arithmetic fails both.
- Per-row lookup key including NULL key and NULL map, Rust and SQL.
- Duplicate keys resolving to the first match, which pins the
`find`-stops-at-first-hit contract.
- Empty-window fast path, empty input, scalar map argument, non-string keys,
the nested-key backstop, and both rejection paths.
Missing:
- A NULL map row over a non-empty offset window (finding 1). Only
constructible in Rust.
- Key types beyond `Utf8` / `Int32` / `Binary` (finding 5).
- A scalar map combined with an array key. The code path exists at line 120
crossed with line 157 and is untested, though the risk is low.
The two new Scala tests would pass on the pre-PR implementation too, which
is correct for what they are. They are regression tests for the new kernel's
specific hazards, not for a bug the PR fixes, and the sliced-map one does fail
if the windowing is wrong.
Non-nullable map value types (`valueContainsNull = false`) are covered
indirectly by the existing `map_entries(element_at(map(1, map(1, 2)), _1))`
test, which routes a non-nullable-value map out of `element_at`.
## Performance
Evidence is required for this PR and it is supplied at both levels, which is
the right shape. The new criterion bench runs both implementations over the
same inputs so the comparison stays reproducible, and it is parameterized over
entries-per-map and constant versus per-row key.
The kernel-level claim (11x to 56x) is credible from the code. DF 55
allocates two `ArrayRef` slices per candidate entry and compares through `dyn
Array` equality, which the new path replaces with one `eq` over the batch plus
a bit scan plus one `take`. I confirmed the pinned 55.0.0 source is the slow
slice-based version, so the baseline is the one Comet actually ships.
By dimension:
- CPU: large reduction on the constant-key path, smaller but real on the
per-row path.
- Allocations: large reduction. The remaining per-call allocations are the
`gather` vector and the gathered per-entry key array on the per-row path, both
`O(total entries in the batch)`, plus `indices` at `O(rows)`.
- Memory: the per-row path materializes a full per-entry copy of the lookup
key at line 174, which for string keys copies the key bytes once per entry
rather than once per row. Bounded and worth the trade. The scalar-map expansion
at line 122 is the unbounded one, see finding 2.
- I/O, shuffle, latency: unaffected.
No regression risk I can see for the shapes that were already fast. The
`entries_start == entries_end` early return keeps the all-empty and all-NULL
case at one `new_null_array`.
The end-to-end table reads honestly, including the caveat that the Spark arm
drifted about 15% between runs and that the Comet columns should be read
against their own `size(attrs)` control. `+64 ms -> +1.5 ms` isolated lookup
cost is consistent with the kernel numbers.
--
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]
