[rich7420]

Which issue does this PR close?

Closes #6093.

Rationale for this change

DataFrame.na.drop uses AtLeastNNonNulls, which Comet does not support, so the filter falls back to Spark.

What changes are included in this PR?

• Evaluate AtLeastNNonNulls natively. Like Spark, it counts values that are neither NULL nor NaN and stops evaluating a row once it reaches n. There is no tuning config.
• Fall back when a skippable child shares a subexpression that Spark evaluates on every row. The shared serde guard compares expressions after decimal promotion, preserving ANSI errors for repeated decimal arithmetic and the same protection for null short-circuiting. Plain column-based na.drop stays native.
• Add Rust and Spark benchmarks.

How are these changes tested?

• Rust and Scala tests cover NULL/NaN counting, batch boundaries, complex types, child errors and fallback. The two decimal regressions fail without the fix. With CSE, DIVIDE_BY_ZERO survives through fallback; without it, execution stays native.
• After merging current main (JDK 21): CometExpressionSuite, QueryPlanSerdeSuite, CometDecimalPromotionSuite and CometCodegenSuite passed on Spark 4.1.3 (323 tests). The AtLeastNNonNulls, serde and decimal promotion tests passed on Spark 3.5, 4.0 and 4.2 (15 each); Spark 3.4 needs JDK 17 and was not run locally. The Rust tests, workspace Clippy, Spotless and Scalastyle passed.
• PR CI passed on dae0c86f0 before merging main. Spark 4.1 SQL CI is pending a maintainer adding run-spark-4.1-tests; Actions is disabled on the fork.

On a narrow, tall table where every column is consumed, the native plan is slower than fallback because DataFusion’s FilterExec copies surviving rows. That shared filter cost is tracked in #6908.

[rich7420]

Outdated; see the PR description.

Status updated 2026-10-07: #6350, #6396 and #6397 are merged. Feature refresh c34480068 onto main c8e6553da is now published and mergeable; its CI is running. See the current delivery plan. The results below belong to historical candidate 4528a2a67, not the new feature head.

Historical integration progress and validation

#6350 has merged. #6396 remains open with passing PR CI and is the remaining implementation prerequisite. #6397 contributes independent IF semantic tests and does not block this feature.

The refreshed integration candidate 4528a2a67 is published against actual main 62ed9d16d. It preserves the previous fork candidate through a normal merge. CASE/IF and float normalization now match actual main, including the UTC fix and regression. Comparing against a prerequisite-only tree still leaves 13 AtLeastNNonNulls files.

Local validation of this exact tree passed: release native build; 24 Spark 4.1.3 tests; six Spark 3.5 tests with strict warnings (both on JDK 21); 16 conditional, seven AtLeastNNonNulls and 50 planner Rust tests; workspace/all-target release Clippy, cargo fmt, Spotless and RAT. Neither Spark run had failures, cancellations or skipped tests.

Two independent processes passed all full-query answers and native/dispatcher-path assertions on the M3 Pro/Spark 4.1.3 workload: 1,048,576 Parquet rows, 32 Double columns, no NULLs, na.drop(16), count and all 32 cleaned sums. Twelve rounds alternated mode order; medians exclude the first three. The second process reversed initial mode and NaN-rate order. At 5% NaN, native/fallback medians were 329.92/333.53 ms and 319.96/360.13 ms (1.1% and 11.2% lower native median time). At 0% NaN, they were 290.71/319.72 ms and 302.48/329.58 ms. The permanent focus-double benchmark also passed; its 5% NaN native/fallback averages were 325/354 ms, with standard deviations of 11/37 ms. Variance is material, so this is workload-specific combined-candidate evidence, not a fixed speedup. The original issue workload remains unverified.

The refreshed candidate is published in fork draft #42. Current-head fork CI passed at 4528a2a67: Required Checks, all Comet suites, Rust, strict Spark 3.5/JDK17 compilation, benchmark checks, TPC-H/TPC-DS and all nine Spark 4.1 SQL shards. The earlier same-head run cancelled during this push is superseded by the completed successful run. Other Spark runtime profiles, Iceberg and macOS suites were skipped by this run. This is the combined candidate's verdict; the final narrowed upstream head still needs independent verification after #6396 merges.

Upstream #6180 remains Draft at 22c04aaba. After #6396 merges, I will update this branch from actual main, apply the review follow-ups, narrow the feature diff and repeat validation on its resulting head. The delivery plan is also recorded in the original issue.

[hsiang-c]

Let's see if you could port this test to Comet SQL tests: https://datafusion.apache.org/comet/contributor-guide/sql-file-tests.html

[rich7420]

Moved the projection result regression to spark/src/test/resources/sql-tests/operators/filter_projection.sql in #6396, which is now merged. It covers count-only output, aliases/order/duplicates, full projections, empty results and a rejected row that must not reach an ANSI cast. The Scala test retains the native output-row metrics checks, which the SQL fixture cannot assert. These generic projection changes are no longer in this PR’s diff.

[hsiang-c]

Is it possible to fix this issue in a separate PR?

[rich7420]

Split out as #6396 and merged. This branch now takes the projection implementation from main, so this PR’s diff is limited to AtLeastNNonNulls and its configuration, tests, benchmarks and documentation. The shared CASE/IF implementation (#6350) and independent IF tests (#6397) are also merged.

[hsiang-c]

Why 64?

[rich7420]

64 was a one-word starting point for the general counting path: bitmap counters process 64 rows per word, while small batches may not amortize their bookkeeping. It is a heuristic, not a required or universal crossover. I made it configurable as spark.comet.exec.atLeastNNonNulls.smallBatchThreshold, default 64; any positive value is accepted, independently of the bitmap word size and the na.drop threshold. Tests cover both strategies and non-word-aligned thresholds. Criterion now compares row, bitmap and default strategies on identical inputs around the boundary. Current-head timings are still pending, so I am not claiming that 64 is optimal.

[rich7420]

Threshold removed; 64 is now just the bitmap word size.

[andygrove]

I ran this PR against main (fallback), the codegen dispatcher and Spark on a release build (Spark 4.1, 64M values per dataset, interleaved runs, wall and process CPU time). The native kernel holds up well. For the default na.drop() it used 0.7-1.0x Spark's CPU on numeric columns and 0.5-0.75x on strings, and it keeps a following group-by native. Dispatch was the slowest option everywhere, roughly 2-16x the native predicate cost and up to 64% more CPU end to end, which matches what #6093 reported. So I think native is the right direction.

My main concern is the smallBatchThreshold config. The general-n counting step cost 18-65 ms per 64M values in my runs, against roughly 400 ms for the scan, and batches are normally 8192 rows, so I don't see anything for a user to tune there. Could we drop the config, its proto field and the tuning docs, and keep a single counting strategy? If the simpler row-counter path is close at 8192-row batches, I'd lean towards keeping only that one.

One thing to be aware of, though I don't think it should block this PR. On a narrow, tall table where every column is consumed (4 double columns, 16M rows, na.drop() followed by a sum of each column), the native plan was about 1.7x slower in wall time than the fallback on main. That cost is DataFusion's FilterExec copying the surviving rows, which Spark's fused codegen avoids, so it applies to any native filter rather than to this kernel. I've opened #6908 to track it.

[rich7420]

@andygrove Removed the config; kept bitmap, 5x faster at 8192.

[sunchao]

Reviewed all 13 changed files and all three commits against base d0f367dc02520fe68f524db22083cc0af6e58bca, at head dd8ce9f789fa71473d0317d675944266744f4747. The PR remains open and non-draft. Found one new P2 correctness issue.

Routed skills: .ai/skills/review-comet-pr/SKILL.md and .ai/skills/review-comet-expression-pr/SKILL.md. Read the existing reviews, conversation comments, inline comments, and threads before forming findings.

Summary

• Prior state and problem: DataFrame.na.drop uses AtLeastNNonNulls, which previously forced its filter back to Spark. This PR adds native execution so supported queries can retain a native pipeline.
• Design approach: A dedicated physical expression counts non-null, non-NaN children while preserving per-row short-circuit evaluation. Scala attempts to retain Spark fallback when subexpression elimination requires otherwise-skippable work.
• Correctness: Nulls, NaNs, outer complex-type validity, threshold boundaries, sliced bitmaps, and required child evaluation are handled carefully. However, the new fallback guard misses shared expressions changed by decimal promotion. A reproduced native projection returns [true, true] where Spark raises DIVIDE_BY_ZERO.
• Compatibility analysis: Read Spark sources for 3.4.3, 3.5.9, 4.0.4, 4.1.3, and 4.2.0. Their AtLeastNNonNulls implementations are identical. Using Spark's own EquivalentExpressions accommodates version-specific sharing rules, but comparing original shared nodes with promoted children breaks compatibility in the reported case.
• Key design decisions: The n=1 and n=children.length cases use bitmap OR and AND. General thresholds use saturating bit-plane counters. Impossible thresholds still evaluate required children, matching Spark's error behavior.
• Implementation sketch: The change adds serde registration, a protobuf variant, native planner construction, the physical expression, Rust and Scala tests, benchmarks, and documentation. canSkipEvaluation generalizes the existing shared-subexpression guard without changing its other callers' intended behavior.
• Performance: Columns and literals avoid selection filtering, and bitmap counters avoid per-row integer counters. The existing discussion and #6908 document a narrow-table regression: 1,203 ms native versus 695 ms fallback for four double columns and 16M rows followed by full-column aggregation. The relevant AND path and native filter construction remain unchanged. This substantiated existing concern remains unresolved and is not duplicated as a new inline finding. I did not rerun those timings.
• Design: A physical expression is appropriate because an eager scalar UDF would evaluate children that Spark skips. The main design is coherent. The correctness boundary between Spark's sharing analysis and Comet's expression rewriting needs repair.
• Abstraction & complexity: Reusing the shared-subexpression context avoids another independent analyzer. Removing the tuning configuration and alternate small-batch implementation simplifies the final change. The bit-plane implementation is localized and has useful boundary coverage.
• Behavioral changes worth calling out: Compared with the relevant paths on branch-1.1 at 89c729a9697467e8f8d89e3eadf1797a75e2328c, native na.drop execution is an intended change. The standard spark.comet.expression.AtLeastNNonNulls.enabled=false switch restores fallback. Successful execution instead of Spark's required ANSI error is unintended.
• Suggested improvements: Compare shared and skippable expressions in the same representation, and retain a regression covering repeated decimal expressions. Resolve or gate the already-reported performance regression. Obtain the repository-required Spark SQL verdict before queueing this serde/planner change.

Exact-head CI: Required Checks, Rust tests, all four Linux Spark 4.1 Comet shards, supported-profile compilation checks, and TPC-H/TPC-DS verification passed in run 39165588736. Logs confirm all five new Scala tests and eight Rust tests passed. Spark SQL suites, macOS, and benchmark checks were skipped. Alternate Spark profiles had no runtime verdict.

Local validation: all eight focused Rust tests passed. A disposable Spark 4.1.3 regression, using current JVM sources and the native artifact from the exact-head CI run, failed on the reported semantic mismatch. Disabling the new expression restored DIVIDE_BY_ZERO in both engines. Disabling subexpression elimination made both engines return matching native results. The disposable test was removed and the checkout is clean. Full local Spark SQL, cross-version runtime suites, and benchmarks were not run.

[sunchao]

[P2] [P2] Preserve shared-expression matching across decimal promotion

With ANSI mode and subexpression elimination enabled, project AtLeastNNonNulls(2, Seq(Literal(1), marker, r, r)), where r = Cast(Divide(a, b), IntegerType) and a and b are decimal columns. withSharedSubexpressions records the original shared cast, but exprToProto runs DecimalPrecision.promote before this check. The skippable child becomes Cast(CheckOverflow(Divide(...))), so it no longer matches the original shared node and this fallback is bypassed. For (marker=1, a=1, b=0), Spark evaluates the shared expression and raises DIVIDE_BY_ZERO, while Comet reaches the threshold before the division and returns true. Please compare both sides in the same representation, or perform the guard against the original expression tree, and add a repeated-decimal regression.

Evidence: Reproduced on the pinned head with Spark 4.1.3 and Temurin JDK 21, using native-lib-linux from exact-head CI run 39165588736. A disposable CometTestBase suite wrote Parquet rows (Some(1), BigDecimal(1), BigDecimal(0)) and (None, BigDecimal(7), BigDecimal(1)), then projected the expression above. The plan contained CometProject over CometNativeScan. Output: spark_error=Some([DIVIDE_BY_ZERO] ...), comet_error=None, native_rows=ArraySeq([true], [true]). Setting spark.comet.expression.AtLeastNNonNulls.enabled=false produced DIVIDE_BY_ZERO in both engines. Setting spark.sql.subexpressionElimination.enabled=false produced matching native results. The regression failed its error-parity assertion. Saved reproduction: /tmp/comet-6180-dd8ce9-review-source/Review6180DecimalCSESuite.scala. Log: /tmp/comet-6180-dd8ce9-repro-controls.log.

[andygrove]

Thanks @rich7420 for dropping the config. The single bitmap strategy is much simpler. sunchao's decimal-promotion finding needs fixing before this merges, though. Comet returns true where Spark raises DIVIDE_BY_ZERO, because the guard compares the original shared Cast against the promoted Cast(CheckOverflow(...)) child. Since canSkipEvaluation is the same helper canShortCircuitNulls uses, could the fix go there, comparing both sides in the same representation, so the existing null short-circuit callers are covered too? A regression test with a repeated decimal subexpression under ANSI would pin it. Once that's in, I'll apply run-spark-4.1-tests.

[andygrove]

Thanks for dropping the config. The decimal-promotion finding from sunchao's review is still open at dd8ce9f789.

exprToProto runs DecimalPrecision.promote before CometAtLeastNNonNulls.convert, so the skippable children are Cast(CheckOverflow(...)). sharedSubexpressions.common holds the unpromoted Cast. canSkipEvaluation never sees a match for those children, so the fallback is skipped and Comet returns true where Spark raises DIVIDE_BY_ZERO.

Could the fix go in canSkipEvaluation so the canShortCircuitNulls callers are covered too? A regression test with a repeated decimal Cast(Divide(a, b), IntegerType) under ANSI would pin it. The current test uses a string cast, which promote doesn't change, so it can't catch this.

[andygrove]

Thanks @rich7420. Promoting both sides in canSkipEvaluation is the fix I was hoping for, since it covers the canShortCircuitNulls callers too, and the repeated Cast(Divide(a, b), IntegerType) test pins the case that slipped through. Looks good to me.