Skip to content

perf: compare decimal bounds in the Arrow data buffer when gathering cache statistics - #7235

Open
andygrove wants to merge 19 commits into
apache:mainfrom
andygrove:perf/cache-stats-wide-types
Open

andygrove wants to merge 19 commits into
apache:mainfrom
andygrove:perf/cache-stats-wide-types

Conversation

@andygrove

@andygrove andygrove commented Oct 26, 2026 •

Copy link
Copy Markdown
Member

Which issue does this PR close?

Part of #6834.

This PR is stacked on #6988 and holds its commits as well. The change here is four commits (e30bc831b9, 8a9394fdaa, f20e6e123e, 158e105a71), and the registration of its test suite in the CI suite lists (4b2656ee46). The rest are #6988's. The diff shrinks to these once #6988 lands.

Rationale for this change

ArrowCachedBatchSerializer.gatherColumnStats records a lower and an upper bound for each decimal column of every cached batch. For each row, gatherDecimalStats called CometVector.getDecimal, which for a precision past 18 digits builds a BigInteger and a BigDecimal from the 16 bytes in Arrow's data buffer, and then a Decimal around them. Most of the statistics time of a decimal column was that allocation. #6834 lists it as the next step after the null counts of #6853.

What changes are included in this PR?

  • gatherDecimalStats reads a plain Arrow decimal column's 128-bit unscaled values straight from its data buffer and compares them without building anything per row. A value is a signed high word over an unsigned low word, little-endian. A narrow decimal is the sign extension of its low word, so one comparison covers every precision. All the values of a column share one scale, so their unscaled order is their numeric order. Only the rows of the two bounds become a Decimal, through the same getDecimal as before, so the bounds keep the type they had.
  • A decimal column that is not a plain Arrow vector (a dictionary-encoded one, for instance) takes the earlier by-value loop, which is kept as gatherDecimalStatsByValue.
  • CometInMemoryCacheBenchmark gets a decimal case in its statistics benchmark: precisions 10, 18 and 38 at scale 2, one row in ten null, and values on both sides of zero.

How are these changes tested?

New ArrowCachedBatchStatsSuite builds Arrow DecimalVectors at precisions 5, 10, 18, 19, 20, 30 and 38 with a range of scales, and compares the bounds and null count with a BigDecimal minimum and maximum:

  • values on either side of each word boundary of a 128-bit value (the low word's sign bit, a carry into the high word, both signs), in three orders;
  • the same values without the largest magnitudes, in positive-only and negative-only sets, because at the full range the high word alone decides the bounds and the low word's comparison never matters;
  • random values of random widths, with and without nulls, and single-value columns.

Comparing the low word as signed instead of unsigned fails the suite at 19, 20, 30 and 38 digits.

CometInMemoryCacheSuite, CacheBatchCoalescerSuite and the new suite pass on Spark 4.1 (104 tests) and on 3.5. dev/ci/check-suites.py finds the new suite in both workflow files.

Benchmark (CometInMemoryCacheBenchmark --stats-only, Apple M3 Max, local[1], 5M rows of the three decimal columns above, Arrow vectors), run old, new, old in one session so that drift would show:

Collector Time for 5M rows
before (by value) 153 ms
after (data buffer) 21 ms
before, again 154 ms

The long and string cases of the same benchmark did not move (160, 160 and 161 ms), so the machine was steady between the runs.

gatherColumnStats visited every row to count nulls, including for binary,
array, map and struct columns, which record a null count and no bounds.
A CometDecodedVector already holds that count, taken from Arrow when the
vector is built. Use it when the vector holds exactly the batch's rows: a
column without bounds records it without a pass over its rows, and an
all-null column of any type skips its bounds loop.

Spark's own vectors are still counted row by row. Its vectorized Parquet
reader marks a column missing from a file all-null without counting any
nulls, so numNulls is 0 there while every row isNullAt.

The statistics benchmark now measures Comet's Arrow vectors as well as
Spark's, and binary and nested columns.

Part of apache#6834.
Spark's DefaultCachedBatch records the same null counts and no bounds
for binary, nested and void columns. nullCountsWithoutBounds now checks
that the relation is stored as CometCachedBatch before reading its
statistics, so the row write path and native scan legs check the format
as the Spark vectorized reader leg already did.
encodeBatches now converts a batch that is not written directly before it
gathers the statistics, so every column the collector reads is a
CometDecodedVector holding the null count Arrow took when it was built. The
row-by-row null count for Spark's vectors and the check that a vector holds
exactly the batch's rows go, and so do the test of that check and its helper.
The statistics benchmark measures only Arrow vectors, which is all the
collector reads now.
The columnar cache write stored each input batch as one cached batch,
and nothing bounds how small those batches are. A shuffle read returns
at least one batch per map output block, so a relation cached straight
from a shuffle over 200 map tasks into 50 partitions held 10000 batches
of about 500 rows, each with its own message metadata, compression
framing and statistics row.

CacheBatchCoalescer now sits in front of the columnar write. An
Arrow-backed batch laid out the way the reader reads it is passed
through when it holds at least half of
spark.sql.inMemoryColumnarStorage.batchSize, so full native batches are
still written without a copy. Every other batch is copied through
ArrowWriter.writeColumns into batches of batchSize rows, the size the
row path cuts at: a small batch is worth the copy, and a batch of Spark
column vectors was already copied into Arrow. Row order is kept, and
each batch is copied before the input advances, because a native
producer releases the batch it returned last when asked for the next.

Closes apache#6950.
The cache statistics now come from the Arrow batch that is written, so a
vector longer than its batch would count nulls past the batch's last row
and could record the batch all null. Require each vector to hold exactly
the batch's rows to take the direct path, and convert otherwise.

@sunchao sunchao left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reviewed 158e105a719a4bb2c08e2335a8a915c1e1ae0349 across correctness, compatibility, performance, design and test coverage, including the inherited cache changes from #6853 and #6988.

The signed-high/unsigned-low decimal comparison and final getDecimal materialization look correct. I found no additional runtime defect in the reviewed paths. There is one validation issue to address before approval: the new statistics suite is missing from both CI suite lists. I reproduced the suite-registration failure locally and verified the same failure in this head's Preflight job. Normal build/test jobs were consequently skipped. Details are inline.

I did not independently rerun the JVM/native suites or benchmark timings.

import org.apache.comet.vector.CometVector

/** The bounds and null counts `ArrowCachedBatchSerializer` records for decimal columns. */
class ArrowCachedBatchStatsSuite extends AnyFunSuite {

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you register org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchStatsSuite in the suite lists in both .github/workflows/pr_build_linux.yml and pr_build_macos.yml? Neither currently includes this new suite. I reproduced dev/ci/check-suites.py exiting with code 255, and the current-head Preflight job fails with the same missing-suite message. This stops the normal build/test jobs before these decimal-boundary tests can run. Adding the suite to both lists is needed to restore CI and keep this regression coverage active.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Registered ArrowCachedBatchStatsSuite in the suite lists of both pr_build_linux.yml and pr_build_macos.yml, and dev/ci/check-suites.py exits 0 again. The branch also had a conflict left by #6853 being squash-merged, so I merged apache/main and the updated #6988 into it. The diff against main is now #6988 plus this PR's four commits. Fixed in 4b2656e.

The coalescer queued every fragment of a large Spark-vector batch before returning the first, so a small spark.sql.inMemoryColumnarStorage.batchSize multiplied the Arrow allocation of one input. It now keeps the batch it is copying and a row offset, and copies only as far as one output batch each time it is asked for one.
# Conflicts:
#	spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala
#	spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala
#	spark/src/test/scala/org/apache/spark/sql/benchmark/CometInMemoryCacheBenchmark.scala
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request performance

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants