perf: read cache statistics null counts from Comet's Arrow vectors

[andygrove]

Which issue does this PR close?

Part of #6834: the first step of its statistics-gathering item.

Rationale for this change

ArrowCachedBatchSerializer.gatherColumnStats visits every row of every column to count nulls, including binary, array, map and struct columns, which record a null count and no bounds. Every batch the cache writes is made of Comet's Arrow vectors, either as Comet operators and the row write path produce them or after a Spark vectorized reader's batch is converted, and each of those vectors already holds that count: CometDecodedVector takes it from Arrow when the vector is built. The pass over the rows only recomputes it.

What changes are included in this PR?

• encodeBatches gathers the statistics from the Arrow batch it serializes, after converting a batch that is not written directly, instead of from the input batch. Every column the collector reads is then a CometDecodedVector.
• gatherColumnStats takes each column's null count from its vector. A column without bounds records that count without visiting its rows, and a column of any type whose rows are all null skips its bounds loop.
• Spark's own vectors are no longer read for statistics, which matters for correctness as well as speed. Spark's vectorized Parquet reader marks a column missing from a file all-null without counting any nulls, so that vector's numNulls is 0 while every row isNullAt. The conversion writes a null for every row that isNullAt, so the converted vector's count is right.
• The statistics cases in CometInMemoryCacheBenchmark read Arrow vectors, the only kind the collector sees now, and add binary and nested columns.

CometInMemoryCacheBenchmark --stats-only, 5M row visits, best of three interleaved runs of each side (Apple M3 Max, JDK 17, Spark 4.1):

Case	main	This PR
3 bigint and 3 string columns, Arrow vectors	156 ms	155 ms
Struct, array and binary columns, 1 row in 10 null, Arrow vectors	27 ms	0 ms

The flat columns are unchanged, as expected. Their loops still visit every row for the bounds, and the null check is a small part of each visit.

Batches from Spark's vectorized Parquet reader used to be read through Spark's vectors and are now read after conversion, so I timed that path separately: gatherColumnStats column by column over every batch the reader returned for a 5M-row, 16-file table with Comet off, best of the last three of eight passes, in two interleaved runs of each side. Each cell is the range across the two runs:

Column, 1 row in 8 null	Spark's vectors (main)	Converted Arrow vectors (this PR)
bigint	5.3–5.7 ms	6.2–6.3 ms
bigint, dictionary-encoded	8.7–8.9 ms	6.0–6.1 ms
double	4.6–4.7 ms	5.2–5.4 ms
decimal(10,2)	22 ms	19–20 ms
decimal(38,2)	97–118 ms	80–83 ms
string	162–164 ms	129–140 ms
string, dictionary-encoded	187–193 ms	138–148 ms
binary, struct and array, each	0.9–1.0 ms	0.1 ms
All ten columns	493–516 ms	385–409 ms

Converting and serializing the batches cost the same on both sides. Strings gain the most, dictionary-encoded ones especially, since reading one from Spark's vector copies its bytes out of the Parquet dictionary every time. Most of the time that remains is in strings and wide decimals, where every row builds a UTF8String, or a BigInteger and a BigDecimal. Those are the next steps for this item.

How are these changes tested?

New CometInMemoryCacheSuite tests:

• "statistics count nulls in columns that record no bounds": binary, array, map and struct columns, and a bigint column with no values, cached from the row write path (with a void column too), from Comet's native scan and from Spark's vectorized reader. Each leg first checks that the relation is stored as CometCachedBatch, because Spark's DefaultCachedBatch records the same counts for these columns.
• "statistics count the nulls of a column missing from a file": the missing-column case above, through Spark's vectorized reader. It checks both the recorded count and the IS NULL answer.

Each test fails against a mutant of the code it guards. Gathering the statistics from the input batch and trusting each vector's numNulls makes the IS NULL query return 0 rows instead of 400. Zeroing the count for columns without bounds, or for all-null columns, fails the first test.

CometInMemoryCacheSuite, CometInMemoryCachePruningSuite and CometInMemoryCacheKryoSuite pass locally on Spark 3.4, 3.5, 4.0 and 4.1.

[sunchao]

Reviewed full head 320b384a8f3a0bddbd187250df65c1cbbe8f4605 against base dc880d6d265f8391b6202efa43574cfeea71e75a. The PR is not a draft. No introduced P1/P2 issues found within this review.

Reviewed all nine files in the base-to-head comparison. Five native-write/test differences reflect base-only fixes from #6808 and #6821. Those files are unchanged from this head’s parent, so these differences are not reversions introduced by the PR. The PR adds one commit affecting four cache-related files. The snapshot and live discussion checks contained no reviews, conversation comments, inline comments, or review threads.

Routed skills: review-comet-pr and review-comet-ffi-pr, plus review-comet-expression-pr and review-comet-iceberg-write-pr for the serde/write files in the full comparison.

Summary

• Prior state and problem: gatherColumnStats visited every row to count nulls, including binary and nested columns that record no bounds. Comet’s decoded Arrow vectors already retain that count.
• Design approach: Use the retained count when a CometDecodedVector contains exactly the batch’s rows. Skip bounds collection for all-null columns and preserve row-based counting for other vectors.
• Correctness: The type and row-count guards preserve the existing statistics contract. Checked empty batches, all-null columns, shortened batches, sliced vectors, nested vectors, and Spark’s missing-column behavior. Spark’s IsNull and IsNotNull pruning predicates continue to receive the expected counts. No introduced P1/P2 correctness issue was identified.
• Compatibility analysis: Compared relevant Spark sources at 3.4.3, 3.5.9, 4.0.4, 4.1.3, and 4.2.0. Their missing-column flags can represent all-null vectors without updating numNulls, which justifies excluding generic Spark vectors. Compared with branch-1.1 at e9efd9f764ee0a59b7898ff028d6985d4a7a28e1: its statistics-gathering methods match the supplied base, and this PR’s intended difference is reduced collection work.
• Key design decisions: Restricting the optimization to CometDecodedVector avoids relying on a broader contract Spark does not satisfy. Checking numValues() == numRows prevents nulls outside the visible batch from affecting pruning.
• Implementation sketch: Adds arrowNullCount, extracts the fallback countNulls loop, and introduces the all-null shortcut. Tests cover row, native-scan, Spark-columnar, missing-column, and shortened-batch inputs. Benchmarks add Arrow-backed and nested/binary cases.
• Performance: Eligible binary/nested columns avoid a row loop, and all-null columns avoid bounds scans. Mixed bounded columns retain their existing loops. The change adds constant work per column without new per-row conversion or allocation. The author’s reported nested-column improvement is consistent with that mechanism, but timings were not independently reproduced.
• Design: The optimization stays local to statistics collection and preserves the existing specialized bounds implementations. It does not alter cache encoding, pruning rules, or execution routing.
• Abstraction & complexity: The two small helpers make the trusted-count contract and fallback explicit. The test helper remains a thin accessor. No unnecessary abstraction or complexity meeting the P1/P2 bar was identified.
• Behavioral changes worth calling out: Eligible Arrow-backed batches require less statistics work. Recorded bounds and null counts remain equivalent to the prior collector. Spark vectors and batches exposing fewer rows than their vectors retain their previous behavior.
• Suggested improvements: No additional P1/P2 fix is supported by the reviewed evidence.

Exact-head CI: 20 checks succeeded, 8 were running, and 28 were skipped at the final inspection, with no failures reported. Native build, Java lint across all five Spark profiles, and strict Scala warnings passed. The Spark 4.1 execution suite and benchmark compile/lint check were still running. Spark SQL and Iceberg suites were skipped, so this is not a completed green CI verdict.

Validation: An isolated harness using source-extracted base/head statistics methods and the actual Comet vector classes passed 8,702 differential checks on Spark 4.1.3/JDK 17. It covered exhaustive small validity patterns, visible prefixes, slices, all-null and void columns, structs/lists, and Spark’s missing-column flag. The full serializer/native pipeline was not exercised. Root-reactor test-compile was blocked by resolution of org.scala-sbt:compiler-bridge_2.13:1.10.8:sources, so the repository cache suites were not independently run. Project source remained unchanged.

[andygrove]

@peterxcli you worked on the in-memory cache in #5869, so you know this area. If you have time, would you be able to review this one? Thanks.

[viirya]

Nice change, and the comment on arrowNullCount explaining why Spark's vectors are excluded is really helpful.

One question on the new no-bounds test. nullCountsWithoutBounds only reads the stats row, and Spark's DefaultCachedBatchSerializer records the same (null, null, nullCount, ...) layout for binary, nested and void columns. So if these relations ever end up as DefaultCachedBatch, for example through the serializer pinning that the comment at the top of this suite describes, or if NullType is later delegated to the fallback, the row-path and native-scan legs (here and at line 1203) would still pass without touching Comet's collector. Could you assert cachedBatchTypes(...) is CometCachedBatch in those two legs, the way withSparkColumnarCache does?

[andygrove]

You're right that the counts alone don't show which collector produced them. In 1b3990d the check is in nullCountsWithoutBounds, so every leg asserts the relation is a CometCachedBatch before it reads the statistics, and a failure names the format it found.

While checking this I found that the no_values column already catches a fallback by accident. Spark's LongColumnStats records Long.MaxValue and Long.MinValue as the bounds of an all-null bigint column, so with NullType delegated to Spark's format the row leg failed on column 4 rather than passing. That depends on the test keeping a primitive column, so the explicit check is still better. With the same mutant the test now fails with no_bounds_rows was cached as org.apache.spark.sql.execution.columnar.DefaultCachedBatch.

[peterxcli]

the comparison - numValues() == numRows here maybe is too defensive, looks like we dont need this guard

[andygrove]

You're right that it isn't needed once the statistics come from the written batch, so b979e1b drops it along with the test that pinned it. A converted batch holds exactly the input's rows. A batch written directly has to hold exactly its rows anyway, because serialize takes the payload's row count from the first vector, so a batch reporting fewer rows than its vectors would fail when read back.

[viirya]

I agree no producer hands the cache such a batch today, but I'm not sure the read-back check covers it. Stats pruning runs before the read. Say a direct batch reports 2 rows over vectors of 4, with the nulls in rows 2 and 3. numNulls is then 2, which equals numRows, so the batch is recorded all null and IS NOT NULL skips it without ever decoding it. The query then silently loses rows where it used to fail. Could the direct path require that each vector holds exactly batch.numRows() values, and convert otherwise? columnarBatchToArrowBatch copies exactly numRows rows, so that case would be written correctly instead of relying on the reader to catch it, which also matches how matchesReaderLayout converts rather than trusts.

[peterxcli]

Any chance we could completely get rid of this?

suggestion is: Gather the stats after the conversion in encodeBatches, from arrowBatch in the convert branch and from batch in the direct branch. Two things make this work:

1. Every CometVector subclass in the repo extends CometDecodedVector: Plain, Dictionary, Struct, List and Map. The only exception is the private IntervalChildVector, which never appears as a top-level column.
2. isArrowBacked lets only CometVectors through to the direct path.

After that change, every batch the collector sees has an Arrow null count, and ArrowWriter has already made the row pass while building it. countNulls and the comment about Spark’s Parquet reader could go,

[andygrove]

Thanks, that's a better structure. Done in b979e1b: encodeBatches converts a batch that isn't written directly first and gathers the statistics from the batch it serializes, so the collector only ever sees CometDecodedVectors. countNulls is gone, and the collector now throws if it is handed any other vector. I kept one sentence about Spark's Parquet reader in encodeBatches, since it's the reason the input's numNulls can't be used, and the missing-column test still guards it.

I measured it first, because Arrow's getters had looked slower than Spark's for strings and wide decimals in the synthetic benchmark. On real batches from Spark's vectorized Parquet reader the converted batch is faster to read: over 5M rows and ten columns, the statistics took 385 to 409 ms, against 493 to 516 ms from Spark's vectors. Strings and decimal(38,2) gained the most, plain bigint and double moved by under a millisecond, and the conversion costs the same either way. The description has the per-column numbers.

[viirya]

Thanks for adding the format check to nullCountsWithoutBounds. Every leg now fails with a clear message if the relation is not stored as CometCachedBatch, which covers what I was worried about. The note about no_values catching a fallback through LongColumnStats bounds was a nice find too.

I went through the paths that build a CometDecodedVector again: native import, rootAsBatch for the row write path and the Arrow sources, and the Python runner. Each wraps the Arrow vector after it is filled, and the native side zeroes offsets before export. So numNulls and isNullAt agree on every input the collector can see. The statistics come out the same as before, just with less work. LGTM.

[viirya]

Thanks for taking this further. Gathering the statistics from the batch that is actually written is a cleaner structure, and the numbers on real Parquet batches settle the concern I had about reading the converted copy, including the plain-encoded strings. I checked that every CometVector the direct path can see is a CometDecodedVector, and that the conversion writes validity through isNullAt, so a column Spark marks all-null gets the right count. One question inline about the row-count check.