Skip to content

perf: compare string bounds in the Arrow data buffer when gathering cache statistics - #7237

Open
andygrove wants to merge 22 commits into
apache:mainfrom
andygrove:perf/cache-stats-strings
Open

andygrove wants to merge 22 commits into
apache:mainfrom
andygrove:perf/cache-stats-strings

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 #7235, which is stacked on #6988, and holds all of their commits. The change here is two commits (011718272e, 0f01c13981). The rest are #7235's and #6988's. The diff shrinks to those two once the others land.

Rationale for this change

ArrowCachedBatchSerializer.gatherColumnStats records a lower and an upper bound for each string column of every cached batch. For each row, gatherStringStats called getUTF8String, which allocates a UTF8String, and compared it through a boxed Ordering[Any]. That left strings as the most expensive column type in the statistics loop after the decimals fixed in #7235, and #6834 lists them as the next step.

What changes are included in this PR?

  • For a plain 32-bit Arrow string vector under the default StringType, gatherStringStats compares the bytes in the vector's data buffer and builds no UTF8String per row. The comparison is that of UTF8String.compareTo: eight bytes at a time as big-endian unsigned words, then byte by byte as unsigned values, and a shorter string sorts before a longer one it prefixes. A row is compared with the lower bound first, and with the upper bound only when it is not below it. Only the rows of the two bounds become UTF8Strings, as owned copies.
  • A collated string, a large-offset vector and a column that is not a plain Arrow vector keep the earlier loop, as gatherStringStatsByValue, and so keep comparing through the ordering boundsOrderings resolved for the column. The default collation is recognised by dataType == StringType, which a collated StringType is not equal to, so no version shim is needed.

How are these changes tested?

Four new tests in ArrowCachedBatchStatsSuite compare the bounds with the minimum and maximum of the strings' UTF-8 bytes under an unsigned byte-wise order:

  • lengths on both sides of the 8-byte words, a shared prefix of every length, a value that is a prefix of another, and one- to four-byte characters whose bytes are at or above 0x80, in three orders, with duplicates and nulls at the start, between and at the end, and single-value columns;
  • the values that sort after and before a long shared prefix, so that the byte after it decides the bound;
  • random strings over a small alphabet of one- and two-byte characters, so that many share long prefixes, with nulls;
  • a large-offset vector, which takes the fallback loop.

Each of these fails if the comparison is changed: treating bytes as signed, comparing the words in native byte order, or dropping the tie-break on length.

CometInMemoryCacheSuite, CacheBatchCoalescerSuite and ArrowCachedBatchStatsSuite pass on Spark 4.1 (108 tests), and the stats and cache suites pass on 3.5. Spotless and the strict-warnings build pass, and dev/ci/check-suites.py finds every suite in both workflow files.

Benchmark (CometInMemoryCacheBenchmark --stats-only, Apple M3 Max, local[1], 5M rows, Arrow vectors), run before, after, before in one session so that drift would show:

Case before after before, again
3 long and 3 string columns 161 ms 63 ms 157 ms
3 decimal columns (#7235) 21 ms 21 ms 20 ms

The decimal case did not move, so the machine was steady between the runs. The string columns' share of the first case is not measured separately.

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 0f01c13981bb126541b6c3a7881dc9d4baf88a97 across correctness, compatibility, performance, design and test coverage, including its prerequisites. I found no new runtime defect in the string optimization. Unsigned byte ordering, prefix handling, collation dispatch and ownership of the final bounds match the existing contracts. I did not independently reproduce the benchmark timings.

Approval is pending validation. This head inherits the missing ArrowCachedBatchStatsSuite registration from #7235, where I left the registration comment. I reproduced dev/ci/check-suites.py exiting with code 255 here, and the current-head Preflight job reports the same missing suite. The normal build/test jobs were skipped. Please carry the registration fix into this stack and get those checks running.

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
@andygrove

Copy link
Copy Markdown
Member Author

The registration of ArrowCachedBatchStatsSuite is in the suite lists of both pr_build_linux.yml and pr_build_macos.yml now. It comes in through #7235 (4b2656e), and dev/ci/check-suites.py exits 0 on this head (d8ea75a). I also merged apache/main and the updated #6988 in, which cleared the conflict left by #6853 being squash-merged.

@comphead comphead left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Comparing 8-byte words byte-reversed and unsigned, then the tail bytes and the length, gives the same order as UTF8String.compareTo, absolute offsets keep sliced vectors correct, and the dt == StringType guard sends collated and char/varchar types to the by-value path on 4.x. One question on the evidence, outside the diff. The 161 to 63 ms figure comes from the existing mixed case in CometInMemoryCacheBenchmark, and the description says the strings' share is not measured separately. Those strings are 9 to 12 bytes with a 6-byte shared prefix (str_a_ plus digits), so compareBytes never gets past its first word, and the longer shared prefixes the word loop is meant for are not exercised. #7235 added a dedicated decimal case to that benchmark. Could a string-only case with long shared prefixes, such as URL or path keys, and one with short distinct keys back the claim for those shapes?

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.

3 participants