Skip to content

Commit 2347165

Browse files
viiryaviirya
committed
[SPARK-57268][SQL] Add Apache Arrow as a native cache format for in-memory Dataset caching
### What changes were proposed in this pull request? This PR adds Apache Arrow as a native cache format for Spark in-memory Dataset caching, available alongside the existing `DefaultCachedBatchSerializer`. It is one of the sub-tasks of [SPARK-56978](https://issues.apache.org/jira/browse/SPARK-56978) (SPIP: Faster queries in local laptop mode for Apache Spark), specifically the "Arrow-based `df.cache` reimplementation" item. The new `ArrowCachedBatchSerializer` stores cached data in Apache Arrow IPC streaming format. It is opt-in via `spark.sql.cache.serializer`: ```scala spark.conf.set("spark.sql.cache.serializer", "org.apache.spark.sql.execution.columnar.ArrowCachedBatchSerializer") ``` Main components: - **`ArrowCachedBatch`** -- a `SimpleMetricsCachedBatch` holding `numRows`, the serialized Arrow `RecordBatch` (IPC streaming format, optionally compressed), and per-column statistics for partition pruning. - **`ArrowCachedBatchSerializer`** -- the serializer: - Write paths for both `InternalRow` and `ColumnarBatch` input, with a zero-copy fast path when the input is already backed by `ArrowColumnVector`. - Read paths for both `ColumnarBatch` output (wrapping Arrow vectors directly) and `InternalRow` output. The row path uses pre-built typed `ArrowColumnReader`s that write directly into an `UnsafeRowWriter` to avoid per-row pattern matching, and falls back to a columnar-to-row path for complex types (Array/Struct/Map/UDT/Variant/etc.). - Optional background prefetch of the next batch (decompress/deserialize off the consumer thread), controlled by a new config (off by default). - Min/max statistics collection over Arrow vectors, kept consistent with the row-based `ColumnStats` path (NaN handling, collation-aware string comparison, null/decimal bounds). - **`ArrowUtils.isSupportedByArrow`** -- recursive type-support check used by `supportsColumnarInput`. - **`ObjectColumnStats`** -- now skips `getSizeInBytes` for columnar complex types (`ColumnarArray`/`ColumnarMap`/`ColumnarRow`), which are views into `ColumnVector`s and do not expose a size. - New config `spark.sql.execution.arrow.cache.prefetch.enabled` (default `false`), Kryo registration for the new classes, and documentation (`sql-arrow-cache-format.html`, linked from the SQL docs menu). ### Why are the changes needed? The default cache format is row/column-encoded specifically for Spark. Using Arrow as the cache format provides: - Zero-copy columnar reads when the cached data is already in Arrow form (e.g. re-caching Arrow-cached data with column projection). - Interoperability with the Arrow ecosystem and off-heap memory management via Arrow allocators. - Min/max statistics for partition pruning, consistent with the default path. In our benchmarks, the Arrow format is competitive with or faster than the default format on columnar/primitive workloads, with the largest gains on the zero-copy re-cache path. The default format can still be faster in some cases (for example, at higher compression levels), so this is offered as an opt-in alternative rather than a replacement. See the committed `sql/core/benchmarks/ArrowCacheBenchmark-jdk{17,21,25}-results.txt` files, generated by the `ArrowCacheBenchmark` suite via the GitHub Actions benchmark workflow. ### Does this PR introduce _any_ user-facing change? Yes, additively. A new opt-in cache serializer (`ArrowCachedBatchSerializer`) and a new config `spark.sql.execution.arrow.cache.prefetch.enabled` (default `false`) are added. The default cache behavior is unchanged: `spark.sql.cache.serializer` still defaults to `DefaultCachedBatchSerializer`. ### How was this patch tested? - New `ArrowCachedBatchSerializerSuite` covering primitive and complex/nested types, null handling, collation, NaN bounds, statistics correctness for both the row and columnar (Arrow-vector) paths, columnar input from Parquet, column projection, filter pushdown, and compression codecs (none/zstd/lz4), plus a check that the Arrow serializer is actually used. - `ArrowCachedBatchKryoRegistrationSuite` verifying Kryo registration. - Added `ArrowCacheBenchmark` for performance comparison against the default cache format. Result files for JDK 17/21/25 are generated in the consistent GitHub Actions environment via the benchmark workflow. Locally: `catalyst/compile` + `sql/Test/compile` pass; the two suites above run green (0 failures). ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code (Claude Opus 4.8) Closes #56334 from viirya/arrow-cache-format. Authored-by: Liang-Chi Hsieh <viirya@gmail.com> Signed-off-by: Liang-Chi Hsieh <viirya@gmail.com>
1 parent c369cf2 commit 2347165

16 files changed

Lines changed: 5644 additions & 5 deletions

File tree

‎core/src/main/scala/org/apache/spark/serializer/KryoSerializer.scala‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -620,6 +620,8 @@ private[serializer] object KryoSerializer {
620620
"org.apache.spark.sql.columnar.CachedBatchSerializer",
621621
"org.apache.spark.sql.columnar.SimpleMetricsCachedBatchSerializer",
622622
"org.apache.spark.sql.execution.columnar.DefaultCachedBatchSerializer",
623+
"org.apache.spark.sql.execution.columnar.ArrowCachedBatch",
624+
"org.apache.spark.sql.execution.columnar.ArrowCachedBatchSerializer",
623625

624626
"org.apache.spark.ml.attribute.Attribute",
625627
"org.apache.spark.ml.attribute.AttributeGroup",

‎docs/_data/menu-sql.yaml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,8 @@
8080
subitems:
8181
- text: Caching Data
8282
url: sql-performance-tuning.html#caching-data
83+
- text: Arrow Cache Format
84+
url: sql-arrow-cache-format.html
8385
- text: Tuning Partitions
8486
url: sql-performance-tuning.html#tuning-partitions
8587
- text: Leveraging Statistics

‎docs/sql-arrow-cache-format.md‎

Lines changed: 387 additions & 0 deletions
Large diffs are not rendered by default.

‎docs/sql-performance-tuning.md‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,10 @@ memory usage and GC pressure. You can call `spark.catalog.uncacheTable("tableNam
3232

3333
To list relations cached with an explicit name, use `spark.catalog.listCachedTables()`. Entries cached only via `Dataset.cache()` without a name are not included.
3434

35+
Spark supports two cache formats:
36+
- **Default cache format**: The standard in-memory columnar cache (used by default).
37+
- **Arrow cache format**: An Apache Arrow-based cache that can improve read performance for columnar workloads and enables Arrow ecosystem interoperability. See [Arrow Cache Format documentation](sql-arrow-cache-format.html) for details and configuration.
38+
3539
Configuration of in-memory caching can be done via `spark.conf.set` or by running
3640
`SET key=value` commands using SQL.
3741

‎sql/api/src/main/scala/org/apache/spark/sql/util/ArrowUtils.scala‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,51 @@ private[sql] object ArrowUtils {
3838

3939
// todo: support more types.
4040

41+
/**
42+
* Check if a Spark DataType is supported by Arrow. This recursively checks complex types
43+
* (Array, Struct, Map).
44+
*
45+
* Note: This checks compatibility with toArrowField(), not toArrowType(). Types like
46+
* GeometryType, GeographyType, and VariantType are not supported by toArrowType() (which only
47+
* handles primitive Arrow types), but ARE supported by toArrowField() which converts them to
48+
* Arrow Struct representations with metadata. Since Arrow cache uses toArrowField() via
49+
* toArrowSchema() to create the schema, these types are supported.
50+
*/
51+
def isSupportedByArrow(dt: DataType): Boolean = {
52+
dt match {
53+
// Primitive types
54+
case BooleanType | ByteType | ShortType | IntegerType | LongType | FloatType | DoubleType |
55+
_: StringType | BinaryType | NullType =>
56+
true
57+
58+
// Decimal
59+
case _: DecimalType => true
60+
61+
// Temporal types
62+
case DateType | TimestampType | TimestampNTZType | _: TimeType => true
63+
case _: TimestampNTZNanosType | _: TimestampLTZNanosType => true
64+
65+
// Interval types
66+
case _: YearMonthIntervalType | _: DayTimeIntervalType | CalendarIntervalType => true
67+
68+
// Complex types - recursively check element types
69+
case ArrayType(elementType, _) => isSupportedByArrow(elementType)
70+
case StructType(fields) => fields.forall(f => isSupportedByArrow(f.dataType))
71+
case MapType(keyType, valueType, _) =>
72+
isSupportedByArrow(keyType) && isSupportedByArrow(valueType)
73+
74+
// Special types
75+
// Note: These are not in toArrowType(), but are handled by toArrowField()
76+
case udt: UserDefinedType[_] => isSupportedByArrow(udt.sqlType)
77+
case _: GeometryType => true // Converted to Struct with srid + wkb fields
78+
case _: GeographyType => true // Converted to Struct with srid + wkb fields
79+
case _: VariantType => true // Converted to Struct with value + metadata fields
80+
81+
// Unsupported types
82+
case _ => false
83+
}
84+
}
85+
4186
/** Maps data type from Spark to Arrow. NOTE: timeZoneId required for TimestampTypes */
4287
def toArrowType(dt: DataType, timeZoneId: String, largeVarTypes: Boolean = false): ArrowType =
4388
TypeApiOps(dt)

‎sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4943,6 +4943,18 @@ object SQLConf {
49434943
.booleanConf
49444944
.createWithDefault(true)
49454945

4946+
val ARROW_CACHE_PREFETCH_ENABLED =
4947+
buildConf("spark.sql.execution.arrow.cache.prefetch.enabled")
4948+
.doc("When true, Arrow cache read path prefetches and decompresses the next batch " +
4949+
"in a background thread while the current batch is being consumed. This can " +
4950+
"significantly improve read performance for compressed Arrow caches (e.g., ZSTD) " +
4951+
"by overlapping decompression with consumption. Increases memory usage by up to " +
4952+
"one additional batch worth of Arrow vectors.")
4953+
.version("4.3.0")
4954+
.withBindingPolicy(ConfigBindingPolicy.SESSION)
4955+
.booleanConf
4956+
.createWithDefault(false)
4957+
49464958
val ARROW_TRANSFORM_WITH_STATE_IN_PYSPARK_MAX_STATE_RECORDS_PER_BATCH =
49474959
buildConf("spark.sql.execution.arrow.transformWithStateInPySpark.maxStateRecordsPerBatch")
49484960
.doc("When using TransformWithState in PySpark (both Python Row and Pandas), limit " +
@@ -8905,6 +8917,8 @@ class SQLConf extends Serializable with Logging with SqlApiConf {
89058917
def arrowPySparkUDFColumnarInputEnabled: Boolean =
89068918
getConf(ARROW_PYSPARK_UDF_COLUMNAR_INPUT_ENABLED)
89078919

8920+
def arrowCachePrefetchEnabled: Boolean = getConf(ARROW_CACHE_PREFETCH_ENABLED)
8921+
89088922
def arrowTransformWithStateInPySparkMaxStateRecordsPerBatch: Int =
89098923
getConf(ARROW_TRANSFORM_WITH_STATE_IN_PYSPARK_MAX_STATE_RECORDS_PER_BATCH)
89108924

‎sql/catalyst/src/main/scala/org/apache/spark/sql/internal/StaticSQLConf.scala‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -171,7 +171,10 @@ object StaticSQLConf {
171171
"org.apache.spark.sql.columnar.CachedBatchSerializer. It will be used to " +
172172
"translate SQL data into a format that can more efficiently be cached. The underlying " +
173173
"API is subject to change so use with caution. Multiple classes cannot be specified. " +
174-
"The class must have a no-arg constructor.")
174+
"The class must have a no-arg constructor. Available implementations include: " +
175+
"org.apache.spark.sql.execution.columnar.DefaultCachedBatchSerializer (default) and " +
176+
"org.apache.spark.sql.execution.columnar.ArrowCachedBatchSerializer (Arrow format with " +
177+
"zero-copy columnar reads and better Arrow ecosystem interoperability).")
175178
.version("3.1.0")
176179
.stringConf
177180
.createWithDefault("org.apache.spark.sql.execution.columnar.DefaultCachedBatchSerializer")

‎sql/core/benchmarks/ArrowCacheBenchmark-jdk21-results.txt‎

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
================================================================================================
2+
Arrow Cache vs Default Cache
3+
================================================================================================
4+
5+
================================================================================================
6+
Cache primitive types
7+
================================================================================================
8+
9+
OpenJDK 64-Bit Server VM 21.0.11+10-LTS on Linux 6.17.0-1018-azure
10+
AMD EPYC 7763 64-Core Processor
11+
Cache 5M rows with primitives: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
12+
---------------------------------------------------------------------------------------------------------------------------
13+
Default cache - write + read 1854 1922 97 2.7 370.8 1.0X
14+
Default cache - write + read (uncompressed) 1159 1165 8 4.3 231.8 1.6X
15+
Arrow cache - write + read 1300 1315 21 3.8 260.0 1.4X
16+
Arrow cache - write + read (zstd level -1) 1808 1811 4 2.8 361.6 1.0X
17+
Arrow cache - write + read (zstd level 1) 1814 1830 23 2.8 362.8 1.0X
18+
Arrow cache - write + read (zstd level 3) 1902 1929 39 2.6 380.4 1.0X
19+
20+
21+
================================================================================================
22+
Cache then filter
23+
================================================================================================
24+
25+
OpenJDK 64-Bit Server VM 21.0.11+10-LTS on Linux 6.17.0-1018-azure
26+
AMD EPYC 7763 64-Core Processor
27+
Cache 5M rows, then filter: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
28+
------------------------------------------------------------------------------------------------------------------------
29+
Default cache - filter 1662 1683 29 3.0 332.5 1.0X
30+
Default cache - filter (uncompressed) 1312 1312 0 3.8 262.4 1.3X
31+
Arrow cache - filter 1447 1462 21 3.5 289.4 1.1X
32+
Arrow cache - filter (zstd level -1) 1729 1757 40 2.9 345.8 1.0X
33+
Arrow cache - filter (zstd level 1) 1787 1799 17 2.8 357.3 0.9X
34+
Arrow cache - filter (zstd level 3) 1951 1955 5 2.6 390.3 0.9X
35+
36+
37+
================================================================================================
38+
Cache columnar input (Parquet)
39+
================================================================================================
40+
41+
OpenJDK 64-Bit Server VM 21.0.11+10-LTS on Linux 6.17.0-1018-azure
42+
AMD EPYC 7763 64-Core Processor
43+
Cache 2M rows from Parquet: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
44+
-----------------------------------------------------------------------------------------------------------------------------
45+
Default cache - columnar input 1545 1619 104 1.3 772.7 1.0X
46+
Default cache - columnar input (uncompressed) 1313 1336 33 1.5 656.4 1.2X
47+
Arrow cache - columnar input 1353 1378 35 1.5 676.7 1.1X
48+
Arrow cache - columnar input (zstd level -1) 1535 1573 54 1.3 767.6 1.0X
49+
Arrow cache - columnar input (zstd level 1) 1619 1622 5 1.2 809.6 1.0X
50+
Arrow cache - columnar input (zstd level 3) 1708 1709 2 1.2 853.8 0.9X
51+
52+
53+
================================================================================================
54+
Re-cache Arrow cached data (zero-copy test)
55+
================================================================================================
56+
57+
OpenJDK 64-Bit Server VM 21.0.11+10-LTS on Linux 6.17.0-1018-azure
58+
AMD EPYC 7763 64-Core Processor
59+
Re-cache 2M rows (zero-copy): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
60+
--------------------------------------------------------------------------------------------------------------------------------
61+
Default cache - cache a cached DF 411 428 20 4.9 205.7 1.0X
62+
Default cache - cache a cached DF (uncompressed) 191 210 26 10.5 95.7 2.2X
63+
Arrow cache - cache a cached DF (zero-copy) 137 156 24 14.6 68.4 3.0X
64+
Arrow cache - cache a cached DF (zstd level -1) 327 343 18 6.1 163.3 1.3X
65+
Arrow cache - cache a cached DF (zstd level 1) 338 341 3 5.9 168.8 1.2X
66+
Arrow cache - cache a cached DF (zstd level 3) 352 357 3 5.7 176.2 1.2X
67+
68+
69+
================================================================================================
70+
Cache with column pruning (select 1 of 20 columns)
71+
================================================================================================
72+
73+
OpenJDK 64-Bit Server VM 21.0.11+10-LTS on Linux 6.17.0-1018-azure
74+
AMD EPYC 7763 64-Core Processor
75+
Cache 5M rows, select 1 column: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
76+
-----------------------------------------------------------------------------------------------------------------------------
77+
Default cache - select 1 of 20 columns 10855 11142 406 0.5 2171.0 1.0X
78+
Default cache - select 1 of 20 (uncompressed) 4135 4149 20 1.2 827.0 2.6X
79+
Arrow cache - select 1 of 20 5179 5280 144 1.0 1035.8 2.1X
80+
Arrow cache - select 1 of 20 (zstd level -1) 9258 9283 35 0.5 1851.7 1.2X
81+
Arrow cache - select 1 of 20 (zstd level 1) 9437 9603 234 0.5 1887.4 1.2X
82+
Arrow cache - select 1 of 20 (zstd level 3) 9778 9794 23 0.5 1955.5 1.1X
83+
84+
85+

‎sql/core/benchmarks/ArrowCacheBenchmark-jdk25-results.txt‎

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
================================================================================================
2+
Arrow Cache vs Default Cache
3+
================================================================================================
4+
5+
================================================================================================
6+
Cache primitive types
7+
================================================================================================
8+
9+
OpenJDK 64-Bit Server VM 25.0.3+9-LTS on Linux 6.17.0-1018-azure
10+
AMD EPYC 7763 64-Core Processor
11+
Cache 5M rows with primitives: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
12+
---------------------------------------------------------------------------------------------------------------------------
13+
Default cache - write + read 1686 1723 53 3.0 337.2 1.0X
14+
Default cache - write + read (uncompressed) 1045 1065 27 4.8 209.1 1.6X
15+
Arrow cache - write + read 1268 1305 53 3.9 253.6 1.3X
16+
Arrow cache - write + read (zstd level -1) 1724 1725 1 2.9 344.8 1.0X
17+
Arrow cache - write + read (zstd level 1) 1770 1794 34 2.8 354.0 1.0X
18+
Arrow cache - write + read (zstd level 3) 1857 1893 50 2.7 371.4 0.9X
19+
20+
21+
================================================================================================
22+
Cache then filter
23+
================================================================================================
24+
25+
OpenJDK 64-Bit Server VM 25.0.3+9-LTS on Linux 6.17.0-1018-azure
26+
AMD EPYC 7763 64-Core Processor
27+
Cache 5M rows, then filter: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
28+
------------------------------------------------------------------------------------------------------------------------
29+
Default cache - filter 1426 1432 8 3.5 285.3 1.0X
30+
Default cache - filter (uncompressed) 1252 1274 31 4.0 250.4 1.1X
31+
Arrow cache - filter 1289 1295 8 3.9 257.8 1.1X
32+
Arrow cache - filter (zstd level -1) 1712 1716 7 2.9 342.4 0.8X
33+
Arrow cache - filter (zstd level 1) 1747 1759 16 2.9 349.5 0.8X
34+
Arrow cache - filter (zstd level 3) 1812 1848 50 2.8 362.4 0.8X
35+
36+
37+
================================================================================================
38+
Cache columnar input (Parquet)
39+
================================================================================================
40+
41+
OpenJDK 64-Bit Server VM 25.0.3+9-LTS on Linux 6.17.0-1018-azure
42+
AMD EPYC 7763 64-Core Processor
43+
Cache 2M rows from Parquet: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
44+
-----------------------------------------------------------------------------------------------------------------------------
45+
Default cache - columnar input 1461 1486 35 1.4 730.6 1.0X
46+
Default cache - columnar input (uncompressed) 1219 1227 12 1.6 609.3 1.2X
47+
Arrow cache - columnar input 1253 1273 27 1.6 626.7 1.2X
48+
Arrow cache - columnar input (zstd level -1) 1448 1460 17 1.4 723.8 1.0X
49+
Arrow cache - columnar input (zstd level 1) 1504 1504 0 1.3 752.0 1.0X
50+
Arrow cache - columnar input (zstd level 3) 1578 1587 13 1.3 788.9 0.9X
51+
52+
53+
================================================================================================
54+
Re-cache Arrow cached data (zero-copy test)
55+
================================================================================================
56+
57+
OpenJDK 64-Bit Server VM 25.0.3+9-LTS on Linux 6.17.0-1018-azure
58+
AMD EPYC 7763 64-Core Processor
59+
Re-cache 2M rows (zero-copy): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
60+
--------------------------------------------------------------------------------------------------------------------------------
61+
Default cache - cache a cached DF 386 409 28 5.2 193.1 1.0X
62+
Default cache - cache a cached DF (uncompressed) 194 217 26 10.3 96.8 2.0X
63+
Arrow cache - cache a cached DF (zero-copy) 132 144 10 15.2 65.9 2.9X
64+
Arrow cache - cache a cached DF (zstd level -1) 321 324 7 6.2 160.3 1.2X
65+
Arrow cache - cache a cached DF (zstd level 1) 333 341 7 6.0 166.7 1.2X
66+
Arrow cache - cache a cached DF (zstd level 3) 350 356 12 5.7 174.8 1.1X
67+
68+
69+
================================================================================================
70+
Cache with column pruning (select 1 of 20 columns)
71+
================================================================================================
72+
73+
OpenJDK 64-Bit Server VM 25.0.3+9-LTS on Linux 6.17.0-1018-azure
74+
AMD EPYC 7763 64-Core Processor
75+
Cache 5M rows, select 1 column: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
76+
-----------------------------------------------------------------------------------------------------------------------------
77+
Default cache - select 1 of 20 columns 9310 9426 164 0.5 1862.0 1.0X
78+
Default cache - select 1 of 20 (uncompressed) 3929 3994 92 1.3 785.7 2.4X
79+
Arrow cache - select 1 of 20 5150 5225 106 1.0 1030.0 1.8X
80+
Arrow cache - select 1 of 20 (zstd level -1) 9265 9376 156 0.5 1853.1 1.0X
81+
Arrow cache - select 1 of 20 (zstd level 1) 9296 9351 78 0.5 1859.3 1.0X
82+
Arrow cache - select 1 of 20 (zstd level 3) 9970 9982 18 0.5 1994.0 0.9X
83+
84+
85+

0 commit comments

Comments
 (0)