From ec92ee5130d77b62626be2f48bf1994b7c3954ac Mon Sep 17 00:00:00 2001 From: peterxcli Date: Fri, 11 Sep 2026 20:27:42 +0800 Subject: [PATCH 01/12] perf: reduce Spark cache row conversion overhead --- .../cache-consumer/2026-09-11-row-reader.txt | 261 ++++++++++++++++++ .../arrow/ArrowCachedBatchSerializer.scala | 9 +- .../arrow/CachedBatchRowIterator.scala | 128 +++++++++ .../comet/exec/CometInMemoryCacheSuite.scala | 77 +++++- .../CometCacheRowReaderBenchmark.scala | 198 +++++++++++++ .../arrow/CachedBatchRowIteratorSuite.scala | 137 +++++++++ 6 files changed, 801 insertions(+), 9 deletions(-) create mode 100644 benchmarks/results/cache-consumer/2026-09-11-row-reader.txt create mode 100644 spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchRowIterator.scala create mode 100644 spark/src/test/scala/org/apache/spark/sql/benchmark/CometCacheRowReaderBenchmark.scala create mode 100644 spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchRowIteratorSuite.scala diff --git a/benchmarks/results/cache-consumer/2026-09-11-row-reader.txt b/benchmarks/results/cache-consumer/2026-09-11-row-reader.txt new file mode 100644 index 00000000000..56614da1c9f --- /dev/null +++ b/benchmarks/results/cache-consumer/2026-09-11-row-reader.txt @@ -0,0 +1,261 @@ +Comet cache row reader: main versus patch, 2026-09-11 + +Main: 8320ae481b4eb916beebc9ed611473abb79c3e71 (upstream/main at the start of the run). +Spark 4.1.3 / Scala 2.13, Zulu JDK 21.0.6, Apple M4 (10 cores, 24 GiB), macOS 26.6.2. +Heap: -Xms1g -Xmx6g. One local worker, 16 partitions, 10,000 rows per cached batch. +5M rows, six non-null columns: three longs/three strings (mixed), or six longs (numeric). +Warm cache in MEMORY_ONLY; cache creation, planning, warm-up and validation outside timing. +Each query has five warm-ups and 15 measured collect() actions per fresh JVM. +Two rounds: main/comet, patch/comet, main/spark, then reversed order. No concurrent builds. +Normal desktop background load. All 810 measured actions match their uncached answers. +The benchmark forces Spark row readers and CODEGEN_ONLY; no Comet consumer operators. +It checks payload types, residency, projection width and absence of columnar transitions. +Both Comet versions retain the same per-column Arrow streams and Spark LZ4 compression. +Spark uses DefaultCachedBatch with its own encodings. Cache creation is not measured. + +Medians across 30 actions per cell. Reduction is 1 - patch/main. + +| Schema | Read | Main ms | Patch ms | Reduction | Spark ms | +| --- | --- | ---: | ---: | ---: | ---: | +| mixed | count | 148.32 | 54.95 | 62.9% | 53.95 | +| mixed | long | 191.13 | 85.03 | 55.5% | 63.76 | +| mixed | string | 311.83 | 187.79 | 39.8% | 162.22 | +| mixed | three | 429.20 | 344.96 | 19.6% | 315.29 | +| mixed | all | 635.37 | 543.76 | 14.4% | 457.38 | +| numeric | count | 138.13 | 40.10 | 71.0% | 38.34 | +| numeric | long | 183.27 | 69.40 | 62.1% | 58.33 | +| numeric | three | 242.78 | 134.26 | 44.7% | 95.17 | +| numeric | all | 310.17 | 223.97 | 27.8% | 152.64 | + +These are cached aggregate read timings on one machine, not application speedups. +The patch still writes owned UnsafeRow values. It does not avoid Arrow decoding or the +necessary conversion of variable-width values. The remaining gap varies by projection. +Spark reference mixed-string cases varied noticeably between JVMs; per-JVM medians below +preserve that variation. Compare main and patch directly when assessing this change. +Nulls, nested values, collection, sorting, joins and early limits are correctness-tested; +this timing workload does not measure those cases or distributed/cache-eviction behavior. + +Reproduce from the repository root with the same benchmark source on main and the patch: + + SPARK_LOCAL_IP=127.0.0.1 BENCH_HEAP=6g PROFILES=-Pspark-4.1 \ + make benchmark-org.apache.spark.sql.benchmark.CometCacheRowReaderBenchmark \ + -- comet 5000000 15 all + +Use spark instead of comet for DefaultCachedBatch. Run each case in a separate JVM. +Arguments: [spark|comet|all] [rows] [iterations] [all|mixed|numeric]. + +Source SHA-256 (patch measured before commit; source hashes identify the tested code): +6182f794fe0540c8f1cb9a2c7edc45a42901875dc8c186a58d1f7e8efeb96459 spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala +7e2a211f107a55a485a651c5b5b73daaa83cd0044120d87fe8fba15a18d481d3 spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchRowIterator.scala +d03f3917f0bb8530f98c65ff3f801b1d2b220d18cf48b9b310be490712601322 spark/src/test/scala/org/apache/spark/sql/benchmark/CometCacheRowReaderBenchmark.scala + +Per-JVM medians and pooled interquartile ranges (ms): +{"revision":"main","format":"comet","schema":"mixed","query":"all","median_ms":635.3704375,"q1_ms":622.1365005,"q3_ms":639.452177,"run_medians_ms":{"0":622.089709,"1":636.8295}} +{"revision":"main","format":"comet","schema":"mixed","query":"count","median_ms":148.32439599999998,"q1_ms":135.3494585,"q3_ms":152.87797899999998,"run_medians_ms":{"0":151.064042,"1":132.294667}} +{"revision":"main","format":"comet","schema":"mixed","query":"long","median_ms":191.128417,"q1_ms":187.18945825,"q3_ms":194.69022875000002,"run_medians_ms":{"0":192.464041,"1":190.718916}} +{"revision":"main","format":"comet","schema":"mixed","query":"string","median_ms":311.8330625,"q1_ms":308.76664575,"q3_ms":314.50710425,"run_medians_ms":{"0":308.630958,"1":314.596625}} +{"revision":"main","format":"comet","schema":"mixed","query":"three","median_ms":429.202208,"q1_ms":425.5822505,"q3_ms":432.03943799999996,"run_medians_ms":{"0":425.506084,"1":432.090834}} +{"revision":"main","format":"comet","schema":"numeric","query":"all","median_ms":310.1708125,"q1_ms":308.65911425,"q3_ms":313.00083324999997,"run_medians_ms":{"0":310.2665,"1":309.892291}} +{"revision":"main","format":"comet","schema":"numeric","query":"count","median_ms":138.12770849999998,"q1_ms":136.36403099999998,"q3_ms":139.53172899999998,"run_medians_ms":{"0":138.10675,"1":138.148667}} +{"revision":"main","format":"comet","schema":"numeric","query":"long","median_ms":183.2734585,"q1_ms":182.447979,"q3_ms":184.56236474999997,"run_medians_ms":{"0":184.455708,"1":182.601333}} +{"revision":"main","format":"comet","schema":"numeric","query":"three","median_ms":242.7832295,"q1_ms":241.1498645,"q3_ms":245.72063550000001,"run_medians_ms":{"0":241.121708,"1":245.177625}} +{"revision":"main","format":"spark","schema":"mixed","query":"all","median_ms":457.375313,"q1_ms":395.95653125,"q3_ms":492.63080225,"run_medians_ms":{"0":492.720375,"1":395.84225}} +{"revision":"main","format":"spark","schema":"mixed","query":"count","median_ms":53.946812,"q1_ms":50.17786425,"q3_ms":58.00539575,"run_medians_ms":{"0":53.892083,"1":54.001541}} +{"revision":"main","format":"spark","schema":"mixed","query":"long","median_ms":63.7625625,"q1_ms":62.68343725,"q3_ms":65.3089895,"run_medians_ms":{"0":63.536125,"1":64.538708}} +{"revision":"main","format":"spark","schema":"mixed","query":"string","median_ms":162.2188745,"q1_ms":146.113125,"q3_ms":176.23776049999998,"run_medians_ms":{"0":176.437292,"1":145.93425}} +{"revision":"main","format":"spark","schema":"mixed","query":"three","median_ms":315.2902915,"q1_ms":266.62325999999996,"q3_ms":325.89978125,"run_medians_ms":{"0":325.907375,"1":266.612916}} +{"revision":"main","format":"spark","schema":"numeric","query":"all","median_ms":152.63875000000002,"q1_ms":149.42490650000002,"q3_ms":154.1288125,"run_medians_ms":{"0":149.380417,"1":154.073125}} +{"revision":"main","format":"spark","schema":"numeric","query":"count","median_ms":38.341916499999996,"q1_ms":36.89433325,"q3_ms":40.09940575,"run_medians_ms":{"0":36.865458,"1":39.495166}} +{"revision":"main","format":"spark","schema":"numeric","query":"long","median_ms":58.331770500000005,"q1_ms":56.25829175,"q3_ms":61.178666750000005,"run_medians_ms":{"0":57.005709,"1":58.92875}} +{"revision":"main","format":"spark","schema":"numeric","query":"three","median_ms":95.171271,"q1_ms":93.4965205,"q3_ms":97.22080199999999,"run_medians_ms":{"0":94.719333,"1":95.984042}} +{"revision":"patch","format":"comet","schema":"mixed","query":"all","median_ms":543.759125,"q1_ms":530.9317920000001,"q3_ms":546.5296352500001,"run_medians_ms":{"0":545.543958,"1":530.617417}} +{"revision":"patch","format":"comet","schema":"mixed","query":"count","median_ms":54.9545,"q1_ms":50.54442775,"q3_ms":58.91093725,"run_medians_ms":{"0":55.737291,"1":54.683083}} +{"revision":"patch","format":"comet","schema":"mixed","query":"long","median_ms":85.0292295,"q1_ms":81.33712474999999,"q3_ms":93.53949999999999,"run_medians_ms":{"0":85.575959,"1":84.4825}} +{"revision":"patch","format":"comet","schema":"mixed","query":"string","median_ms":187.7871665,"q1_ms":179.7058855,"q3_ms":191.64538550000003,"run_medians_ms":{"0":191.704167,"1":179.688}} +{"revision":"patch","format":"comet","schema":"mixed","query":"three","median_ms":344.960604,"q1_ms":327.21871875,"q3_ms":346.24670825,"run_medians_ms":{"0":346.015958,"1":327.184125}} +{"revision":"patch","format":"comet","schema":"numeric","query":"all","median_ms":223.975,"q1_ms":221.97183325,"q3_ms":226.91169799999997,"run_medians_ms":{"0":224.666459,"1":222.452125}} +{"revision":"patch","format":"comet","schema":"numeric","query":"count","median_ms":40.1013125,"q1_ms":38.992948,"q3_ms":42.20344825,"run_medians_ms":{"0":39.626,"1":40.237292}} +{"revision":"patch","format":"comet","schema":"numeric","query":"long","median_ms":69.3984375,"q1_ms":67.67601025,"q3_ms":72.21907325,"run_medians_ms":{"0":71.474375,"1":68.322209}} +{"revision":"patch","format":"comet","schema":"numeric","query":"three","median_ms":134.2613955,"q1_ms":132.09339625,"q3_ms":136.37547949999998,"run_medians_ms":{"0":135.490458,"1":133.329833}} + +Raw samples (nanoseconds, original iteration order) and result/storage records: +patch-timing-main-0-comet +CACHE_STORAGE,comet,mixed,512,168804118,169754384 +CACHE_PLAN,comet,mixed,count,columns=0 +CACHE_RESULT,comet,mixed,count,[5000000] +CACHE_PLAN,comet,mixed,long,columns=1 +CACHE_RESULT,comet,mixed,long,[12499997500000] +CACHE_PLAN,comet,mixed,string,columns=1 +CACHE_RESULT,comet,mixed,string,[54444500] +CACHE_PLAN,comet,mixed,three,columns=3 +CACHE_RESULT,comet,mixed,three,[12499997500000,54444500,49298480] +CACHE_PLAN,comet,mixed,all,columns=6 +CACHE_RESULT,comet,mixed,all,[12499997500000,2497500000,12500002500000,54444500,49298480,63888890] +CACHE_STORAGE,comet,numeric,512,107700478,108468272 +CACHE_PLAN,comet,numeric,count,columns=0 +CACHE_RESULT,comet,numeric,count,[5000000] +CACHE_PLAN,comet,numeric,long,columns=1 +CACHE_RESULT,comet,numeric,long,[12499997500000] +CACHE_PLAN,comet,numeric,three,columns=3 +CACHE_RESULT,comet,numeric,three,[12499997500000,2497500000,12500002500000] +CACHE_PLAN,comet,numeric,all,columns=6 +CACHE_RESULT,comet,numeric,all,[12499997500000,2497500000,12500002500000,249997500000,19787521156,37499992500000] +mixed,count,[156275458,156354916,159006750,153516625,151064042,153741000,151687250,148591750,147690250,149423500,149615500,145786917,147726375,200938958,144888167] +mixed,long,[218584542,193985500,189673417,192464041,233971708,199725958,193182042,194866083,195643375,190918959,184838333,186109500,187441333,188796166,188259875] +mixed,string,[311542625,309173709,307336334,307839708,310446125,311094125,314238542,306203958,307310458,306693416,306676375,312506250,307541583,310115209,308630958] +mixed,three,[429648875,424204500,425035834,430615792,428876375,425506084,426293250,424188250,427042333,423257958,425810750,423611750,426879292,425246833,424468625] +mixed,all,[621667000,621355458,619942250,620184959,622016083,623577750,637758833,624661292,620917500,622089709,621071833,622276875,626528375,653803459,654305167] +numeric,count,[138692209,138221792,147840000,136329333,139490791,135603459,139344333,136484959,140639625,136270458,138106750,135796667,138714292,135293083,136468125] +numeric,long,[185256916,183080625,183161209,183279792,182529666,184455708,182193959,186464500,183157000,186669417,209552625,225998542,187401625,183690542,184793208] +numeric,three,[240703792,240628667,241669125,241878042,250955125,239991792,245657667,241718791,241234334,241121708,239762083,240272416,240437708,237762042,242509334] +numeric,all,[312815083,314692042,310280417,310545042,308268917,307266208,308160542,310266500,309230209,318005209,307760125,308640833,313062750,310075125,310716583] +patch-timing-main-0-spark +CACHE_STORAGE,spark,mixed,512,242706686,243299704 +CACHE_PLAN,spark,mixed,count,columns=0 +CACHE_RESULT,spark,mixed,count,[5000000] +CACHE_PLAN,spark,mixed,long,columns=1 +CACHE_RESULT,spark,mixed,long,[12499997500000] +CACHE_PLAN,spark,mixed,string,columns=1 +CACHE_RESULT,spark,mixed,string,[54444500] +CACHE_PLAN,spark,mixed,three,columns=3 +CACHE_RESULT,spark,mixed,three,[12499997500000,54444500,49298480] +CACHE_PLAN,spark,mixed,all,columns=6 +CACHE_RESULT,spark,mixed,all,[12499997500000,2497500000,12500002500000,54444500,49298480,63888890] +CACHE_STORAGE,spark,numeric,512,30092440,30500288 +CACHE_PLAN,spark,numeric,count,columns=0 +CACHE_RESULT,spark,numeric,count,[5000000] +CACHE_PLAN,spark,numeric,long,columns=1 +CACHE_RESULT,spark,numeric,long,[12499997500000] +CACHE_PLAN,spark,numeric,three,columns=3 +CACHE_RESULT,spark,numeric,three,[12499997500000,2497500000,12500002500000] +CACHE_PLAN,spark,numeric,all,columns=6 +CACHE_RESULT,spark,numeric,all,[12499997500000,2497500000,12500002500000,249997500000,19787521156,37499992500000] +mixed,count,[60122458,58763584,64781250,57733083,60336000,53491125,51893667,50396208,49797750,57104333,52548834,58005625,53892083,48885750,45805125] +mixed,long,[68629250,64295750,62363917,63536125,63408792,67167542,68428292,63685792,62541167,62608583,63894250,62370375,61135167,63920084,62908000] +mixed,string,[175639166,174771208,191523459,173436959,173579167,172032792,173196916,171999166,182126292,179218417,186934333,178227292,177398416,183371500,176437292] +mixed,three,[324713334,325907375,326343083,326163917,327724459,325487458,324853083,325207417,326750042,326410750,326311834,328773875,325877000,323402917,320861458] +mixed,all,[494696875,492067459,492720375,494211667,517829542,507144750,489879167,502179250,487612625,488592625,487424459,492362084,491221708,511834083,522218875] +numeric,count,[36556333,38222000,36865458,37317208,37113583,35841958,36333125,46646708,36437750,37088791,35848709,36183834,40102916,39218792,35185333] +numeric,long,[56195750,59335792,54800083,53821583,54192666,58564416,59747000,62119625,56546875,79152125,57005709,57796167,56445917,55756708,58180500] +numeric,three,[93292167,95226792,94719333,95328500,94266584,93141458,98593417,99678666,91984250,101237916,92021875,93002208,101805333,91773792,98124666] +numeric,all,[148008625,152498791,149654250,151202708,147113750,148592875,152934666,148202458,147809875,148520041,149558375,150207583,164515167,148237833,149380417] +patch-timing-main-1-comet +CACHE_STORAGE,comet,mixed,512,168804118,169754384 +CACHE_PLAN,comet,mixed,count,columns=0 +CACHE_RESULT,comet,mixed,count,[5000000] +CACHE_PLAN,comet,mixed,long,columns=1 +CACHE_RESULT,comet,mixed,long,[12499997500000] +CACHE_PLAN,comet,mixed,string,columns=1 +CACHE_RESULT,comet,mixed,string,[54444500] +CACHE_PLAN,comet,mixed,three,columns=3 +CACHE_RESULT,comet,mixed,three,[12499997500000,54444500,49298480] +CACHE_PLAN,comet,mixed,all,columns=6 +CACHE_RESULT,comet,mixed,all,[12499997500000,2497500000,12500002500000,54444500,49298480,63888890] +CACHE_STORAGE,comet,numeric,512,107700478,108468272 +CACHE_PLAN,comet,numeric,count,columns=0 +CACHE_RESULT,comet,numeric,count,[5000000] +CACHE_PLAN,comet,numeric,long,columns=1 +CACHE_RESULT,comet,numeric,long,[12499997500000] +CACHE_PLAN,comet,numeric,three,columns=3 +CACHE_RESULT,comet,numeric,three,[12499997500000,2497500000,12500002500000] +CACHE_PLAN,comet,numeric,all,columns=6 +CACHE_RESULT,comet,numeric,all,[12499997500000,2497500000,12500002500000,249997500000,19787521156,37499992500000] +mixed,count,[154190417,152400916,151556041,153037000,148057042,129450416,130290875,129583000,125314083,128601333,132294667,123816791,126060292,144513833,147887000] +mixed,long,[195783667,192151625,190718916,194162666,193313792,198741583,211980875,191337875,189801333,184391000,186905041,187072250,183634625,187105500,184347083] +mixed,string,[312995792,312123500,313970959,317437792,317939458,323753667,337865875,349354959,316337375,311129042,315006500,313764541,313513542,310059417,314596625] +mixed,three,[431867084,442551250,431833500,435027333,433383208,428544792,433603083,434283834,433689834,432753583,431885250,430490000,428812916,432090834,429528041] +mixed,all,[636730875,635343583,641854875,643184292,636829500,634420334,635997709,640016625,641488625,635003542,635853792,662478125,637731209,641186291,635397292] +numeric,count,[142760625,136149000,138774334,137586167,139995000,136235500,136847250,140961125,139545375,145747542,138148667,140620417,136974500,135497666,137843750] +numeric,long,[183267125,182182250,182601333,182420750,181305458,182034584,182920042,184597917,183546125,184360500,182162000,184208333,181852375,184438208,181256667] +numeric,three,[245741625,243931500,243121292,245839250,248457209,246520667,242955416,244956833,245886542,246375042,242799125,248458917,245177625,242553291,242767334] +numeric,all,[343389084,315352541,345362834,314666708,309892291,355066834,309275333,307111250,308713958,306240542,307532458,309249208,310719625,308733125,312619083] +patch-timing-main-1-spark +CACHE_STORAGE,spark,mixed,512,242706686,243299704 +CACHE_PLAN,spark,mixed,count,columns=0 +CACHE_RESULT,spark,mixed,count,[5000000] +CACHE_PLAN,spark,mixed,long,columns=1 +CACHE_RESULT,spark,mixed,long,[12499997500000] +CACHE_PLAN,spark,mixed,string,columns=1 +CACHE_RESULT,spark,mixed,string,[54444500] +CACHE_PLAN,spark,mixed,three,columns=3 +CACHE_RESULT,spark,mixed,three,[12499997500000,54444500,49298480] +CACHE_PLAN,spark,mixed,all,columns=6 +CACHE_RESULT,spark,mixed,all,[12499997500000,2497500000,12500002500000,54444500,49298480,63888890] +CACHE_STORAGE,spark,numeric,512,30092440,30500288 +CACHE_PLAN,spark,numeric,count,columns=0 +CACHE_RESULT,spark,numeric,count,[5000000] +CACHE_PLAN,spark,numeric,long,columns=1 +CACHE_RESULT,spark,numeric,long,[12499997500000] +CACHE_PLAN,spark,numeric,three,columns=3 +CACHE_RESULT,spark,numeric,three,[12499997500000,2497500000,12500002500000] +CACHE_PLAN,spark,numeric,all,columns=6 +CACHE_RESULT,spark,numeric,all,[12499997500000,2497500000,12500002500000,249997500000,19787521156,37499992500000] +mixed,count,[57487083,59363583,54143834,58940458,52327167,58004708,54001541,53217583,50105083,47961959,59443417,46851667,55576958,46156166,49164541] +mixed,long,[72350583,64538708,63330750,70352959,74904000,64644917,63511042,65512125,66159875,62451416,62099125,63839333,63680792,64699583,61988875] +mixed,string,[152438583,145464083,148502500,147122041,148288542,146654625,145934250,144692417,145511667,145852208,146716667,145594500,146649750,145668167,144365416] +mixed,three,[309719125,268227166,298957209,266183125,270068583,300752834,266464208,266654292,265936708,265905833,266612916,266348917,266377666,266655500,266074292] +mixed,all,[392305916,395842250,396338000,396299375,427326167,401013625,391817500,406162167,391513166,401631375,394721875,397171875,392184375,390213708,394317084] +numeric,count,[39495166,37036000,40412917,39494500,38963792,40985375,38569000,40689625,40838500,45950750,47363000,37006208,36980959,40088875,38461833] +numeric,long,[63018792,75463750,73122750,60735542,57100917,61424125,56713833,62480875,58928750,53959541,58728584,55394500,58483041,54589291,61326375] +numeric,three,[98728834,93677667,97250000,99086792,93611083,94749708,96481792,97028375,96409583,95984042,93458333,94811542,95115750,92376916,97133208] +numeric,all,[156200042,160211417,151995625,153991125,154568542,153277959,156859209,154147375,152535708,157814750,155393208,153101875,153394750,154073125,152741792] +patch-timing-patch-0-comet +CACHE_STORAGE,comet,mixed,512,168804118,169754384 +CACHE_PLAN,comet,mixed,count,columns=0 +CACHE_RESULT,comet,mixed,count,[5000000] +CACHE_PLAN,comet,mixed,long,columns=1 +CACHE_RESULT,comet,mixed,long,[12499997500000] +CACHE_PLAN,comet,mixed,string,columns=1 +CACHE_RESULT,comet,mixed,string,[54444500] +CACHE_PLAN,comet,mixed,three,columns=3 +CACHE_RESULT,comet,mixed,three,[12499997500000,54444500,49298480] +CACHE_PLAN,comet,mixed,all,columns=6 +CACHE_RESULT,comet,mixed,all,[12499997500000,2497500000,12500002500000,54444500,49298480,63888890] +CACHE_STORAGE,comet,numeric,512,107700478,108468272 +CACHE_PLAN,comet,numeric,count,columns=0 +CACHE_RESULT,comet,numeric,count,[5000000] +CACHE_PLAN,comet,numeric,long,columns=1 +CACHE_RESULT,comet,numeric,long,[12499997500000] +CACHE_PLAN,comet,numeric,three,columns=3 +CACHE_RESULT,comet,numeric,three,[12499997500000,2497500000,12500002500000] +CACHE_PLAN,comet,numeric,all,columns=6 +CACHE_RESULT,comet,numeric,all,[12499997500000,2497500000,12500002500000,249997500000,19787521156,37499992500000] +mixed,count,[62337875,61327791,56817292,65798208,58202500,53627000,51277667,61614667,47866292,50476334,60204750,50157792,55737291,53457500,47255708] +mixed,long,[89675834,93536500,89428875,83268750,85575959,100159208,92236041,84258042,95596083,128454458,77327125,80336875,81303333,82338750,73700541] +mixed,string,[199106833,236384083,222001708,190832084,190944625,192168083,191469041,192773667,192143375,193418917,189756667,189414084,189339000,191704167,188841708] +mixed,three,[345616291,344921875,345123375,352998875,355030417,345967583,346323625,349292583,348524459,346015958,345910250,348346584,358505834,344999333,345473833] +mixed,all,[546806791,545775250,544199417,548145833,545403667,544608125,543799708,546250042,543718542,546622833,548913167,541548625,543101667,546973667,545543958] +numeric,count,[42205917,42656458,39518042,38891083,39330750,43295792,39038416,38956584,40436708,42119875,39161958,40589375,39626000,42196042,38676834] +numeric,long,[67438666,71493875,70207250,67902292,69103584,74414459,71598167,76024334,71474375,74995750,69474125,69322750,69228167,74330625,72708875] +numeric,three,[133468834,136875750,132058417,136080750,132691667,135932417,133862458,137188667,135100791,135490458,131474666,131887750,137495584,136610750,136438042] +numeric,all,[226937750,260211458,224666459,254753292,221156500,223965125,220584375,258870708,222454959,223984875,222399708,226644667,226833542,221015250,225697500] +patch-timing-patch-1-comet +CACHE_STORAGE,comet,mixed,512,168804118,169754384 +CACHE_PLAN,comet,mixed,count,columns=0 +CACHE_RESULT,comet,mixed,count,[5000000] +CACHE_PLAN,comet,mixed,long,columns=1 +CACHE_RESULT,comet,mixed,long,[12499997500000] +CACHE_PLAN,comet,mixed,string,columns=1 +CACHE_RESULT,comet,mixed,string,[54444500] +CACHE_PLAN,comet,mixed,three,columns=3 +CACHE_RESULT,comet,mixed,three,[12499997500000,54444500,49298480] +CACHE_PLAN,comet,mixed,all,columns=6 +CACHE_RESULT,comet,mixed,all,[12499997500000,2497500000,12500002500000,54444500,49298480,63888890] +CACHE_STORAGE,comet,numeric,512,107700478,108468272 +CACHE_PLAN,comet,numeric,count,columns=0 +CACHE_RESULT,comet,numeric,count,[5000000] +CACHE_PLAN,comet,numeric,long,columns=1 +CACHE_RESULT,comet,numeric,long,[12499997500000] +CACHE_PLAN,comet,numeric,three,columns=3 +CACHE_RESULT,comet,numeric,three,[12499997500000,2497500000,12500002500000] +CACHE_PLAN,comet,numeric,all,columns=6 +CACHE_RESULT,comet,numeric,all,[12499997500000,2497500000,12500002500000,249997500000,19787521156,37499992500000] +mixed,count,[59758583,54683083,57307625,60145791,54191000,55225917,55481709,56982959,49442208,50748709,53317417,49717916,59147083,47952167,49751375] +mixed,long,[83909916,95533000,125708084,87679958,86288917,103240000,93540500,84482500,119125167,74684375,76612875,82295500,77504375,81438500,77675500] +mixed,string,[180124375,182181458,179421625,184985833,186732625,179759542,179762042,178609083,180764625,178184708,176803583,177248583,179657292,178242125,179688000] +mixed,three,[330406625,326012875,328227666,328417042,327184125,325927750,332122541,397798875,326347708,325478459,324366333,327322500,329175334,325517750,326888709] +mixed,all,[530396042,535261667,524119916,528182709,521769250,528619709,542558458,521230459,559584375,524968916,530617417,536841750,531874917,565121375,595912375] +numeric,count,[46135291,59249000,51261500,40242667,43565833,39965333,38665375,39931083,38977792,37054375,37599166,43196542,40237292,38089916,41203750] +numeric,long,[67600583,67079833,67131125,66134250,66619458,68740917,72604292,68322209,72426042,71186208,70206167,67425833,68510708,66768000,75425083] +numeric,three,[131630959,133329833,133825500,135334500,136680708,131918334,132369084,138030375,136187792,132198334,131299708,130418542,130962500,134660333,139471667] +numeric,all,[222216583,222856334,255083917,222452125,224083209,252175917,220925375,224026334,222322458,258053208,220046250,219516042,227483208,221431459,221890250] diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala index 821f84e0c2b..c9f05289c30 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala @@ -21,13 +21,12 @@ package org.apache.spark.sql.comet.execution.arrow import java.lang.{Boolean => JBoolean, Byte => JByte, Double => JDouble, Float => JFloat, Integer => JInteger, Long => JLong, Short => JShort} -import scala.collection.JavaConverters._ import scala.util.control.NonFatal import org.apache.spark.TaskContext import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, GenericInternalRow, IsNotNull, IsNull, UnsafeProjection} +import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, GenericInternalRow, IsNotNull, IsNull} import org.apache.spark.sql.catalyst.util.TypeUtils import org.apache.spark.sql.columnar.{CachedBatch, SimpleMetricsCachedBatch, SimpleMetricsCachedBatchSerializer} import org.apache.spark.sql.comet.util.Utils @@ -646,11 +645,7 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { convertCachedBatchToColumnarBatch(input, cacheAttributes, selectedAttributes, conf) .mapPartitions { batches => - val toUnsafe = UnsafeProjection.create(selectedAttributes, selectedAttributes) - - batches.flatMap { batch => - batch.rowIterator().asScala.map(row => toUnsafe(row).copy()) - } + new CachedBatchRowIterator(selectedAttributes).createObject(batches) } } } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchRowIterator.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchRowIterator.scala new file mode 100644 index 00000000000..907236999aa --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchRowIterator.scala @@ -0,0 +1,128 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.spark.sql.comet.execution.arrow + +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Attribute, BoundReference, CodeGeneratorWithInterpretedFallback, InterpretedUnsafeProjection} +import org.apache.spark.sql.catalyst.expressions.codegen._ +import org.apache.spark.sql.catalyst.expressions.codegen.Block._ +import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} + +/** + * Reads vectors directly into Spark's reusable UnsafeRow buffer. The input iterator owns the + * batches and releases them on advancement or task completion. As with Spark's cache reader, + * callers must copy rows they retain across next(), but the returned row owns its variable-width + * values and remains valid when hasNext() releases the batch that supplied them. + */ +private[arrow] class CachedBatchRowIterator(attributes: Seq[Attribute]) + extends CodeGeneratorWithInterpretedFallback[Iterator[ColumnarBatch], Iterator[InternalRow]] { + + private def fields: Seq[BoundReference] = attributes.zipWithIndex.map { case (attr, i) => + BoundReference(i, attr.dataType, attr.nullable) + } + + override protected def createCodeGeneratedObject( + batches: Iterator[ColumnarBatch]): Iterator[InternalRow] = { + val ctx = new CodegenContext + val columns = attributes.indices.map { i => + ctx.addMutableState(classOf[ColumnVector].getName, s"column$i") + } + ctx.currentVars = attributes.zip(columns).map { case (attr, column) => + val value = JavaCode.variable(ctx.freshName("value"), attr.dataType) + val getter = CodeGenerator.getValueFromVector(column, attr.dataType, "rowId") + val javaType = CodeGenerator.javaType(attr.dataType) + if (attr.nullable) { + val isNull = JavaCode.isNullVariable(ctx.freshName("isNull")) + ExprCode( + code""" + boolean $isNull = $column.isNullAt(rowId); + $javaType $value = $isNull ? ${CodeGenerator.defaultValue(attr.dataType)} : ($getter); + """, + isNull, + value) + } else { + ExprCode(code"$javaType $value = $getter;", FalseLiteral, value) + } + } + val projection = GenerateUnsafeProjection.createCode(ctx, fields) + val bindColumns = columns.zipWithIndex + .map { case (column, i) => + s"$column = batch.column($i);" + } + .mkString("\n") + val code = s""" + public Object generate(Object[] references) { + return new SpecificCachedBatchRowIterator((scala.collection.Iterator) references[0]); + } + + class SpecificCachedBatchRowIterator extends scala.collection.AbstractIterator { + private final scala.collection.Iterator batches; + private int rowId = 0; + private int numRows = 0; + ${ctx.declareMutableStates()} + + public SpecificCachedBatchRowIterator(scala.collection.Iterator batches) { + this.batches = batches; + ${ctx.initMutableStates()} + } + + public boolean hasNext() { + while (rowId >= numRows && batches.hasNext()) { + ${classOf[ColumnarBatch].getName} batch = + (${classOf[ColumnarBatch].getName}) batches.next(); + numRows = batch.numRows(); + rowId = 0; + $bindColumns + } + return rowId < numRows; + } + + public InternalRow next() { + if (!hasNext()) throw new java.util.NoSuchElementException(); + ${projection.code} + rowId++; + return ${projection.value}; + } + + ${ctx.declareAddedFunctions()} + } + """ + val (compiled, _) = + CodeGenerator.compile(new CodeAndComment(code, ctx.getPlaceHolderToComments())) + compiled.generate(Array[Any](batches)).asInstanceOf[Iterator[InternalRow]] + } + + override protected def createInterpretedObject( + batches: Iterator[ColumnarBatch]): Iterator[InternalRow] = { + val toUnsafe = InterpretedUnsafeProjection.createProjection(fields) + batches.flatMap { batch => + new Iterator[InternalRow] { + private var rowId = 0 + override def hasNext: Boolean = rowId < batch.numRows() + override def next(): InternalRow = { + if (!hasNext) throw new NoSuchElementException + val row = toUnsafe(batch.getRow(rowId)) + rowId += 1 + row + } + } + } + } +} diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala index e72f4d10f74..f3472dbb6d5 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -24,14 +24,14 @@ import java.{util => ju} import org.apache.arrow.vector.types.pojo.ArrowType import org.apache.spark.CometDriverPlugin import org.apache.spark.SparkConf -import org.apache.spark.sql.{CometTestBase, Row} +import org.apache.spark.sql.{CometTestBase, DataFrame, Row} import org.apache.spark.sql.catalyst.expressions.{And, Attribute, Expression, GreaterThanOrEqual, LessThan, Literal} import org.apache.spark.sql.columnar.{CachedBatch, SimpleMetricsCachedBatch} import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometInMemoryTableScanExec, CometSortExec, CometSortMergeJoinExec} import org.apache.spark.sql.comet.execution.arrow.CometCachedBatchHelper import org.apache.spark.sql.execution.SortExec import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, QueryStageExec, ShuffleQueryStageExec} -import org.apache.spark.sql.execution.columnar.{CometInMemoryRelationHelper, InMemoryRelation} +import org.apache.spark.sql.execution.columnar.{CometInMemoryRelationHelper, InMemoryRelation, InMemoryTableScanExec} import org.apache.spark.sql.execution.exchange.{Exchange, ReusedExchangeExec, ShuffleExchangeLike} import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, SortMergeJoinExec} import org.apache.spark.sql.functions.max @@ -295,6 +295,79 @@ class CometInMemoryCacheSuite extends CometTestBase { } } + test("Spark row consumers of Comet cache preserve values across batches") { + Seq("CODEGEN_ONLY", "NO_CODEGEN").foreach { mode => + withSQLConf( + CometConf.COMET_ENABLED.key -> "false", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "false", + SQLConf.COLUMN_BATCH_SIZE.key -> "7", + SQLConf.CODEGEN_FACTORY_MODE.key -> mode, + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> (mode == "CODEGEN_ONLY").toString, + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.SHUFFLE_PARTITIONS.key -> "2") { + val scalars = Seq( + "boolean", + "tinyint", + "smallint", + "int", + "bigint", + "float", + "double", + "decimal(10,2)", + "decimal(38,2)", + "date", + "timestamp", + "timestamp_ntz").zipWithIndex.map { case (dt, i) => + val value = dt match { + case "date" | "timestamp" | "timestamp_ntz" => + s"cast(date_add(DATE '2000-01-01', cast(id AS INT)) AS $dt)" + case _ => s"cast(id AS $dt)" + } + s"if(id % 3 = 0, null, $value) AS c$i" + } + val source = spark + .range(0, 41, 1, 2) + .selectExpr((Seq("id AS key") ++ scalars ++ Seq( + "if(id % 3 = 0, null, repeat(concat('字', id), cast(id + 1 AS INT))) AS s", + "if(id % 3 = 0, null, cast(concat('binary', id) AS BINARY)) AS b", + "if(id % 3 = 0, null, array(cast(id AS STRING), null)) AS a", + "if(id % 3 = 0, null, named_struct('x', id, 'a', array(cast(id AS STRING)))) AS st", + "if(id % 3 = 0, null, map('k', array(cast(id AS STRING), null))) AS m", + "null AS n")): _*) + + def queries(df: DataFrame): Seq[DataFrame] = Seq( + df.select("*"), + df.selectExpr("s AS renamed", "key", "b", "a", "st", "m"), + df.orderBy($"s".desc, $"key"), + df.join(spark.range(41).toDF("join_key"), $"key" === $"join_key").select(df("*")), + df.selectExpr("count(*)"), + df.limit(1)) + + val expected = queries(source).map(_.collect().toSeq) + source.cache() + try { + assert(source.count() == 41) + val relation = + spark.sharedState.cacheManager.lookupCachedData(source).get.cachedRepresentation + val buffers = relation.cacheBuilder.cachedColumnBuffers.collect() + assert(buffers.length > 2) + assert(buffers.forall(_.getClass.getSimpleName == "CometCachedBatch")) + queries(source).zip(expected).foreach { case (df, answer) => + val scans = + df.queryExecution.executedPlan.collect { case scan: InMemoryTableScanExec => + scan + } + assert( + scans.nonEmpty && scans.forall(!_.supportsColumnar), + df.queryExecution.executedPlan.toString) + checkAnswer(df, answer) + } + } finally source.unpersist(blocking = true) + } + } + } + test("Comet cache serializer delegates unsupported types to Spark's cache format") { withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometCacheRowReaderBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometCacheRowReaderBenchmark.scala new file mode 100644 index 00000000000..7b254f2c50a --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometCacheRowReaderBenchmark.scala @@ -0,0 +1,198 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.spark.sql.benchmark + +import java.nio.charset.StandardCharsets + +import org.apache.spark.benchmark.BenchmarkBase +import org.apache.spark.sql.{DataFrame, Row, SparkSession} +import org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer +import org.apache.spark.sql.execution.ColumnarToRowExec +import org.apache.spark.sql.execution.columnar.{CometInMemoryRelationHelper, DefaultCachedBatch, DefaultCachedBatchSerializer, InMemoryRelation, InMemoryTableScanExec} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.storage.StorageLevel + +import org.apache.comet.{CometConf, CometSparkSessionExtensions} + +/** + * Compare Spark consumers of Comet and Spark caches (issue #5485). + * + * Arguments: [spark|comet|all] [rows] [iterations] [all|mixed|numeric]. Run one format per JVM in + * alternating order on main and the patch. Cache creation and validation are outside timing. + */ +object CometCacheRowReaderBenchmark extends BenchmarkBase { + private val warmups = 5 + + override def runBenchmarkSuite(args: Array[String]): Unit = { + require(args.length <= 4, "Expected format, rows, iterations, schema") + val format = args.headOption.getOrElse("all") + val rows = args.lift(1).map(_.toLong).getOrElse(5000000L) + val iterations = args.lift(2).map(_.toInt).getOrElse(15) + val schema = args.lift(3).getOrElse("all") + require(Set("all", "spark", "comet").contains(format)) + require(Set("all", "mixed", "numeric").contains(schema)) + require(rows > 0 && iterations > 0) + + emit("CACHE_SAMPLE,format,schema,query,rows,iteration,elapsed_ns") + val formats = + if (format == "all") Seq("spark", "comet") else Seq(format) + val schemas = if (schema == "all") Seq("mixed", "numeric") else Seq(schema) + formats.foreach { name => + CometInMemoryRelationHelper.clearSerializer() + SparkSession.clearActiveSession() + SparkSession.clearDefaultSession() + val serializer = if (name == "spark") { + classOf[DefaultCachedBatchSerializer].getName + } else { + classOf[ArrowCachedBatchSerializer].getName + } + val spark = SparkSession + .builder() + .master("local[1]") + .appName(getClass.getSimpleName) + .config("spark.ui.enabled", "false") + .config("spark.sql.cache.serializer", serializer) + .config("spark.sql.shuffle.partitions", "1") + .config("spark.sql.inMemoryColumnarStorage.batchSize", "10000") + .config("spark.sql.inMemoryColumnarStorage.compressed", "true") + .config("spark.io.compression.codec", "lz4") + .config(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key, "false") + .config(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key, "true") + .config(SQLConf.CACHE_VECTORIZED_READER_ENABLED.key, "false") + .config(SQLConf.CODEGEN_FACTORY_MODE.key, "CODEGEN_ONLY") + .config(CometConf.COMET_ENABLED.key, "false") + .config(CometConf.COMET_EXEC_ENABLED.key, "false") + .withExtensions(new CometSparkSessionExtensions) + .getOrCreate() + spark.sparkContext.setLogLevel("WARN") + try { + emit(s"CACHE_ENV,$name,Spark=${spark.version},Java=${System.getProperty("java.version")}") + schemas.foreach(runSchema(spark, name, _, rows, iterations, serializer)) + } finally { + spark.stop() + SparkSession.clearActiveSession() + SparkSession.clearDefaultSession() + CometInMemoryRelationHelper.clearSerializer() + } + } + } + + private def runSchema( + spark: SparkSession, + format: String, + schema: String, + rows: Long, + iterations: Int, + serializer: String): Unit = { + val mixed = schema == "mixed" + val first = Seq("id", "id % 1000 AS k", "id + 1 AS v") + val rest = if (mixed) { + Seq( + "concat('str_a_', cast(id % 100000 as string)) AS s1", + "concat('str_b_', cast(id % 7919 as string)) AS s2", + "concat('str_c_', cast(id as string)) AS s3") + } else { + Seq("id % 100000 AS n1", "id % 7919 AS n2", "id * 3 AS n3") + } + val source = spark.range(0, rows, 1, 16).selectExpr((first ++ rest): _*) + val columns = source.columns.toSeq + val three = if (mixed) Seq("id", "s1", "s2") else columns.take(3) + val projections = Seq("count" -> Seq.empty[String], "long" -> Seq("id")) ++ + (if (mixed) Seq("string" -> Seq("s1")) else Seq.empty) ++ + Seq("three" -> three, "all" -> columns) + def expressions(selected: Seq[String]): Seq[String] = { + if (selected.isEmpty) Seq("count(*)") + else + selected.map { name => + if (name.startsWith("s")) s"sum(length($name))" else s"sum($name)" + } + } + // Obtain the expected values before the relation is cached, using Spark's ordinary row plan. + val expected = projections.map { case (_, selected) => + source.selectExpr(expressions(selected): _*).collect() + } + val cached = source.persist(StorageLevel.MEMORY_ONLY) + try { + assert(cached.count() == rows) + val relation = cached.queryExecution.withCachedData.collectFirst { + case relation: InMemoryRelation => relation + }.get + val builder = relation.cacheBuilder + assert(builder.serializer.getClass.getName == serializer) + val batches = builder.cachedColumnBuffers + val batchSummary = batches + .map { batch => + // Spark's sizeInBytes comes from statistics; measure its encoded column buffers. + val bytes = batch match { + case b: DefaultCachedBatch => b.buffers.map(_.length.toLong).sum + case _ => batch.sizeInBytes + } + (batch.getClass.getSimpleName, batch.numRows.toLong, bytes) + } + .collect() + val expectedClass = if (format == "spark") "DefaultCachedBatch" else "CometCachedBatch" + assert(batchSummary.forall(_._1 == expectedClass), "Wrong cached payload format") + assert(batchSummary.map(_._2).sum == rows) + val storage = spark.sparkContext.getRDDStorageInfo.find(_.id == batches.id).get + assert(storage.numCachedPartitions == batches.getNumPartitions && storage.diskSize == 0) + emit( + s"CACHE_STORAGE,$format,$schema,${batchSummary.length}," + + s"${batchSummary.map(_._3).sum},${storage.memSize}") + + projections.zip(expected).foreach { case ((name, selected), answer) => + val query = cached.selectExpr(expressions(selected): _*) + val plan = query.queryExecution.executedPlan + val scans = plan.collect { case scan: InMemoryTableScanExec => scan } + assert(scans.size == 1, s"Expected one Spark cache scan:\n$plan") + val scan = scans.head + assert(scan.attributes.map(_.name).toSet == selected.toSet, s"Wrong projection:\n$plan") + assert(!scan.supportsColumnar, s"Expected the cache row reader:\n$plan") + assert(!plan.exists(_.isInstanceOf[ColumnarToRowExec]), s"Unexpected transition:\n$plan") + assert(!plan.exists(_.getClass.getName.startsWith("org.apache.spark.sql.comet."))) + emit(s"CACHE_PLAN,$format,$schema,$name,columns=${selected.size}\n$plan") + runQuery(query, answer, format, schema, name, rows, iterations) + } + } finally cached.unpersist(blocking = true) + } + + private def runQuery( + query: DataFrame, + expected: Array[Row], + format: String, + schema: String, + name: String, + rows: Long, + iterations: Int): Unit = { + (0 until warmups).foreach { _ => assert(query.collect().sameElements(expected)) } + (0 until iterations).foreach { i => + val start = System.nanoTime() + val actual = query.collect() + val elapsed = System.nanoTime() - start + assert(actual.sameElements(expected), s"Wrong result for $format/$schema/$name") + emit(s"CACHE_SAMPLE,$format,$schema,$name,$rows,$i,$elapsed") + } + emit(s"CACHE_RESULT,$format,$schema,$name,${expected.mkString(";")}") + } + + private def emit(line: String): Unit = { + println(line) + output.foreach(_.write((line + "\n").getBytes(StandardCharsets.UTF_8))) + } +} diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchRowIteratorSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchRowIteratorSuite.scala new file mode 100644 index 00000000000..1bea5c5e3e1 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchRowIteratorSuite.scala @@ -0,0 +1,137 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.spark.sql.comet.execution.arrow + +import java.nio.charset.StandardCharsets.UTF_8 + +import org.scalatest.funsuite.AnyFunSuite + +import org.apache.arrow.memory.RootAllocator +import org.apache.arrow.vector.VarCharVector +import org.apache.spark.sql.catalyst.expressions.{AttributeReference, UnsafeRow} +import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{IntegerType, StringType} +import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} + +import org.apache.comet.vector.CometPlainVector + +class CachedBatchRowIteratorSuite extends AnyFunSuite { + Seq("CODEGEN_ONLY", "NO_CODEGEN").foreach { mode => + def withMode(f: => Unit): Unit = { + val conf = new SQLConf + conf.setConfString(SQLConf.CODEGEN_FACTORY_MODE.key, mode) + SQLConf.withExistingConf(conf)(f) + } + + test(s"$mode: rows own Arrow values across batch release and reuse the output buffer") { + withMode { + val allocator = new RootAllocator(Long.MaxValue) + val vectors = Seq(Seq("first", null), Seq("字" * 1000, "last")).map { values => + val vector = new VarCharVector("s", allocator) + values.zipWithIndex.foreach { case (value, i) => + if (value == null) vector.setNull(i) else vector.setSafe(i, value.getBytes(UTF_8)) + } + vector.setValueCount(values.size) + vector + } + try { + // Match the cache decoder: hasNext releases a consumed batch before the next is read. + val batches = vectors.iterator.flatMap { vector => + new Iterator[ColumnarBatch] { + private var emitted = false + override def hasNext: Boolean = { + if (emitted) vector.close() + !emitted + } + override def next(): ColumnarBatch = { + emitted = true + new ColumnarBatch(Array(new CometPlainVector(vector, false)), 2) + } + } + } + val attributes = Seq(AttributeReference("s", StringType, nullable = true)()) + val rows = new CachedBatchRowIterator(attributes).createObject(batches) + assert(rows.hasNext && rows.hasNext) + val first = rows.next().asInstanceOf[UnsafeRow] + val saved = first.copy() + assert(first.getUTF8String(0).toString == "first") + assert(rows.next() eq first) + assert(first.isNullAt(0)) + assert(rows.hasNext && rows.hasNext) + assert(first.isNullAt(0)) + assert(rows.next().getUTF8String(0).toString == "字" * 1000) + val last = rows.next() + assert(!rows.hasNext && !rows.hasNext) + assert(allocator.getAllocatedMemory == 0) + assert(last.getUTF8String(0).toString == "last") + assert(saved.getUTF8String(0).toString == "first") + intercept[NoSuchElementException](rows.next()) + } finally { + vectors.foreach(_.close()) + allocator.close() + } + } + } + + test(s"$mode: empty input, empty batches, and zero-column rows") { + withMode { + val factory = new CachedBatchRowIterator(Seq.empty) + val empty = factory.createObject(Iterator.empty) + assert(!empty.hasNext) + intercept[NoSuchElementException](empty.next()) + val batches = Seq(0, 2, 0, 3, 0).map { n => + new ColumnarBatch(Array.empty[ColumnVector], n) + } + val rows = factory.createObject(batches.iterator) + assert(rows.map { row => + assert(row.isInstanceOf[UnsafeRow] && row.numFields == 0) + 1 + }.sum == 5) + intercept[NoSuchElementException](rows.next()) + } + } + + test(s"$mode: wide projections preserve nullable and required columns") { + withMode { + val attributes = (0 until 150).map { i => + AttributeReference(s"c$i", IntegerType, nullable = i % 2 == 0)() + } + val columns = attributes.indices.map { i => + val column = new OnHeapColumnVector(2, IntegerType) + column.putInt(0, i) + if (i % 2 == 0) column.putNull(1) else column.putInt(1, -i) + column + } + val batch = new ColumnarBatch(columns.toArray[ColumnVector], 2) + try { + val rows = new CachedBatchRowIterator(attributes).createObject(Iterator.single(batch)) + val first = rows.next().copy() + val second = rows.next() + attributes.indices.foreach { i => + assert(first.getInt(i) == i) + if (i % 2 == 0) assert(second.isNullAt(i)) else assert(second.getInt(i) == -i) + } + assert(!rows.hasNext) + } finally batch.close() + } + } + } +} From 17dcdc63d1858db9a8caa3aeb1b8b9f9c48d121d Mon Sep 17 00:00:00 2001 From: peterxcli Date: Fri, 11 Sep 2026 20:34:06 +0800 Subject: [PATCH 02/12] chore: remove cache row reader benchmark results --- .../cache-consumer/2026-09-11-row-reader.txt | 261 ------------------ 1 file changed, 261 deletions(-) delete mode 100644 benchmarks/results/cache-consumer/2026-09-11-row-reader.txt diff --git a/benchmarks/results/cache-consumer/2026-09-11-row-reader.txt b/benchmarks/results/cache-consumer/2026-09-11-row-reader.txt deleted file mode 100644 index 56614da1c9f..00000000000 --- a/benchmarks/results/cache-consumer/2026-09-11-row-reader.txt +++ /dev/null @@ -1,261 +0,0 @@ -Comet cache row reader: main versus patch, 2026-09-11 - -Main: 8320ae481b4eb916beebc9ed611473abb79c3e71 (upstream/main at the start of the run). -Spark 4.1.3 / Scala 2.13, Zulu JDK 21.0.6, Apple M4 (10 cores, 24 GiB), macOS 26.6.2. -Heap: -Xms1g -Xmx6g. One local worker, 16 partitions, 10,000 rows per cached batch. -5M rows, six non-null columns: three longs/three strings (mixed), or six longs (numeric). -Warm cache in MEMORY_ONLY; cache creation, planning, warm-up and validation outside timing. -Each query has five warm-ups and 15 measured collect() actions per fresh JVM. -Two rounds: main/comet, patch/comet, main/spark, then reversed order. No concurrent builds. -Normal desktop background load. All 810 measured actions match their uncached answers. -The benchmark forces Spark row readers and CODEGEN_ONLY; no Comet consumer operators. -It checks payload types, residency, projection width and absence of columnar transitions. -Both Comet versions retain the same per-column Arrow streams and Spark LZ4 compression. -Spark uses DefaultCachedBatch with its own encodings. Cache creation is not measured. - -Medians across 30 actions per cell. Reduction is 1 - patch/main. - -| Schema | Read | Main ms | Patch ms | Reduction | Spark ms | -| --- | --- | ---: | ---: | ---: | ---: | -| mixed | count | 148.32 | 54.95 | 62.9% | 53.95 | -| mixed | long | 191.13 | 85.03 | 55.5% | 63.76 | -| mixed | string | 311.83 | 187.79 | 39.8% | 162.22 | -| mixed | three | 429.20 | 344.96 | 19.6% | 315.29 | -| mixed | all | 635.37 | 543.76 | 14.4% | 457.38 | -| numeric | count | 138.13 | 40.10 | 71.0% | 38.34 | -| numeric | long | 183.27 | 69.40 | 62.1% | 58.33 | -| numeric | three | 242.78 | 134.26 | 44.7% | 95.17 | -| numeric | all | 310.17 | 223.97 | 27.8% | 152.64 | - -These are cached aggregate read timings on one machine, not application speedups. -The patch still writes owned UnsafeRow values. It does not avoid Arrow decoding or the -necessary conversion of variable-width values. The remaining gap varies by projection. -Spark reference mixed-string cases varied noticeably between JVMs; per-JVM medians below -preserve that variation. Compare main and patch directly when assessing this change. -Nulls, nested values, collection, sorting, joins and early limits are correctness-tested; -this timing workload does not measure those cases or distributed/cache-eviction behavior. - -Reproduce from the repository root with the same benchmark source on main and the patch: - - SPARK_LOCAL_IP=127.0.0.1 BENCH_HEAP=6g PROFILES=-Pspark-4.1 \ - make benchmark-org.apache.spark.sql.benchmark.CometCacheRowReaderBenchmark \ - -- comet 5000000 15 all - -Use spark instead of comet for DefaultCachedBatch. Run each case in a separate JVM. -Arguments: [spark|comet|all] [rows] [iterations] [all|mixed|numeric]. - -Source SHA-256 (patch measured before commit; source hashes identify the tested code): -6182f794fe0540c8f1cb9a2c7edc45a42901875dc8c186a58d1f7e8efeb96459 spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala -7e2a211f107a55a485a651c5b5b73daaa83cd0044120d87fe8fba15a18d481d3 spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchRowIterator.scala -d03f3917f0bb8530f98c65ff3f801b1d2b220d18cf48b9b310be490712601322 spark/src/test/scala/org/apache/spark/sql/benchmark/CometCacheRowReaderBenchmark.scala - -Per-JVM medians and pooled interquartile ranges (ms): -{"revision":"main","format":"comet","schema":"mixed","query":"all","median_ms":635.3704375,"q1_ms":622.1365005,"q3_ms":639.452177,"run_medians_ms":{"0":622.089709,"1":636.8295}} -{"revision":"main","format":"comet","schema":"mixed","query":"count","median_ms":148.32439599999998,"q1_ms":135.3494585,"q3_ms":152.87797899999998,"run_medians_ms":{"0":151.064042,"1":132.294667}} -{"revision":"main","format":"comet","schema":"mixed","query":"long","median_ms":191.128417,"q1_ms":187.18945825,"q3_ms":194.69022875000002,"run_medians_ms":{"0":192.464041,"1":190.718916}} -{"revision":"main","format":"comet","schema":"mixed","query":"string","median_ms":311.8330625,"q1_ms":308.76664575,"q3_ms":314.50710425,"run_medians_ms":{"0":308.630958,"1":314.596625}} -{"revision":"main","format":"comet","schema":"mixed","query":"three","median_ms":429.202208,"q1_ms":425.5822505,"q3_ms":432.03943799999996,"run_medians_ms":{"0":425.506084,"1":432.090834}} -{"revision":"main","format":"comet","schema":"numeric","query":"all","median_ms":310.1708125,"q1_ms":308.65911425,"q3_ms":313.00083324999997,"run_medians_ms":{"0":310.2665,"1":309.892291}} -{"revision":"main","format":"comet","schema":"numeric","query":"count","median_ms":138.12770849999998,"q1_ms":136.36403099999998,"q3_ms":139.53172899999998,"run_medians_ms":{"0":138.10675,"1":138.148667}} -{"revision":"main","format":"comet","schema":"numeric","query":"long","median_ms":183.2734585,"q1_ms":182.447979,"q3_ms":184.56236474999997,"run_medians_ms":{"0":184.455708,"1":182.601333}} -{"revision":"main","format":"comet","schema":"numeric","query":"three","median_ms":242.7832295,"q1_ms":241.1498645,"q3_ms":245.72063550000001,"run_medians_ms":{"0":241.121708,"1":245.177625}} -{"revision":"main","format":"spark","schema":"mixed","query":"all","median_ms":457.375313,"q1_ms":395.95653125,"q3_ms":492.63080225,"run_medians_ms":{"0":492.720375,"1":395.84225}} -{"revision":"main","format":"spark","schema":"mixed","query":"count","median_ms":53.946812,"q1_ms":50.17786425,"q3_ms":58.00539575,"run_medians_ms":{"0":53.892083,"1":54.001541}} -{"revision":"main","format":"spark","schema":"mixed","query":"long","median_ms":63.7625625,"q1_ms":62.68343725,"q3_ms":65.3089895,"run_medians_ms":{"0":63.536125,"1":64.538708}} -{"revision":"main","format":"spark","schema":"mixed","query":"string","median_ms":162.2188745,"q1_ms":146.113125,"q3_ms":176.23776049999998,"run_medians_ms":{"0":176.437292,"1":145.93425}} -{"revision":"main","format":"spark","schema":"mixed","query":"three","median_ms":315.2902915,"q1_ms":266.62325999999996,"q3_ms":325.89978125,"run_medians_ms":{"0":325.907375,"1":266.612916}} -{"revision":"main","format":"spark","schema":"numeric","query":"all","median_ms":152.63875000000002,"q1_ms":149.42490650000002,"q3_ms":154.1288125,"run_medians_ms":{"0":149.380417,"1":154.073125}} -{"revision":"main","format":"spark","schema":"numeric","query":"count","median_ms":38.341916499999996,"q1_ms":36.89433325,"q3_ms":40.09940575,"run_medians_ms":{"0":36.865458,"1":39.495166}} -{"revision":"main","format":"spark","schema":"numeric","query":"long","median_ms":58.331770500000005,"q1_ms":56.25829175,"q3_ms":61.178666750000005,"run_medians_ms":{"0":57.005709,"1":58.92875}} -{"revision":"main","format":"spark","schema":"numeric","query":"three","median_ms":95.171271,"q1_ms":93.4965205,"q3_ms":97.22080199999999,"run_medians_ms":{"0":94.719333,"1":95.984042}} -{"revision":"patch","format":"comet","schema":"mixed","query":"all","median_ms":543.759125,"q1_ms":530.9317920000001,"q3_ms":546.5296352500001,"run_medians_ms":{"0":545.543958,"1":530.617417}} -{"revision":"patch","format":"comet","schema":"mixed","query":"count","median_ms":54.9545,"q1_ms":50.54442775,"q3_ms":58.91093725,"run_medians_ms":{"0":55.737291,"1":54.683083}} -{"revision":"patch","format":"comet","schema":"mixed","query":"long","median_ms":85.0292295,"q1_ms":81.33712474999999,"q3_ms":93.53949999999999,"run_medians_ms":{"0":85.575959,"1":84.4825}} -{"revision":"patch","format":"comet","schema":"mixed","query":"string","median_ms":187.7871665,"q1_ms":179.7058855,"q3_ms":191.64538550000003,"run_medians_ms":{"0":191.704167,"1":179.688}} -{"revision":"patch","format":"comet","schema":"mixed","query":"three","median_ms":344.960604,"q1_ms":327.21871875,"q3_ms":346.24670825,"run_medians_ms":{"0":346.015958,"1":327.184125}} -{"revision":"patch","format":"comet","schema":"numeric","query":"all","median_ms":223.975,"q1_ms":221.97183325,"q3_ms":226.91169799999997,"run_medians_ms":{"0":224.666459,"1":222.452125}} -{"revision":"patch","format":"comet","schema":"numeric","query":"count","median_ms":40.1013125,"q1_ms":38.992948,"q3_ms":42.20344825,"run_medians_ms":{"0":39.626,"1":40.237292}} -{"revision":"patch","format":"comet","schema":"numeric","query":"long","median_ms":69.3984375,"q1_ms":67.67601025,"q3_ms":72.21907325,"run_medians_ms":{"0":71.474375,"1":68.322209}} -{"revision":"patch","format":"comet","schema":"numeric","query":"three","median_ms":134.2613955,"q1_ms":132.09339625,"q3_ms":136.37547949999998,"run_medians_ms":{"0":135.490458,"1":133.329833}} - -Raw samples (nanoseconds, original iteration order) and result/storage records: -patch-timing-main-0-comet -CACHE_STORAGE,comet,mixed,512,168804118,169754384 -CACHE_PLAN,comet,mixed,count,columns=0 -CACHE_RESULT,comet,mixed,count,[5000000] -CACHE_PLAN,comet,mixed,long,columns=1 -CACHE_RESULT,comet,mixed,long,[12499997500000] -CACHE_PLAN,comet,mixed,string,columns=1 -CACHE_RESULT,comet,mixed,string,[54444500] -CACHE_PLAN,comet,mixed,three,columns=3 -CACHE_RESULT,comet,mixed,three,[12499997500000,54444500,49298480] -CACHE_PLAN,comet,mixed,all,columns=6 -CACHE_RESULT,comet,mixed,all,[12499997500000,2497500000,12500002500000,54444500,49298480,63888890] -CACHE_STORAGE,comet,numeric,512,107700478,108468272 -CACHE_PLAN,comet,numeric,count,columns=0 -CACHE_RESULT,comet,numeric,count,[5000000] -CACHE_PLAN,comet,numeric,long,columns=1 -CACHE_RESULT,comet,numeric,long,[12499997500000] -CACHE_PLAN,comet,numeric,three,columns=3 -CACHE_RESULT,comet,numeric,three,[12499997500000,2497500000,12500002500000] -CACHE_PLAN,comet,numeric,all,columns=6 -CACHE_RESULT,comet,numeric,all,[12499997500000,2497500000,12500002500000,249997500000,19787521156,37499992500000] -mixed,count,[156275458,156354916,159006750,153516625,151064042,153741000,151687250,148591750,147690250,149423500,149615500,145786917,147726375,200938958,144888167] -mixed,long,[218584542,193985500,189673417,192464041,233971708,199725958,193182042,194866083,195643375,190918959,184838333,186109500,187441333,188796166,188259875] -mixed,string,[311542625,309173709,307336334,307839708,310446125,311094125,314238542,306203958,307310458,306693416,306676375,312506250,307541583,310115209,308630958] -mixed,three,[429648875,424204500,425035834,430615792,428876375,425506084,426293250,424188250,427042333,423257958,425810750,423611750,426879292,425246833,424468625] -mixed,all,[621667000,621355458,619942250,620184959,622016083,623577750,637758833,624661292,620917500,622089709,621071833,622276875,626528375,653803459,654305167] -numeric,count,[138692209,138221792,147840000,136329333,139490791,135603459,139344333,136484959,140639625,136270458,138106750,135796667,138714292,135293083,136468125] -numeric,long,[185256916,183080625,183161209,183279792,182529666,184455708,182193959,186464500,183157000,186669417,209552625,225998542,187401625,183690542,184793208] -numeric,three,[240703792,240628667,241669125,241878042,250955125,239991792,245657667,241718791,241234334,241121708,239762083,240272416,240437708,237762042,242509334] -numeric,all,[312815083,314692042,310280417,310545042,308268917,307266208,308160542,310266500,309230209,318005209,307760125,308640833,313062750,310075125,310716583] -patch-timing-main-0-spark -CACHE_STORAGE,spark,mixed,512,242706686,243299704 -CACHE_PLAN,spark,mixed,count,columns=0 -CACHE_RESULT,spark,mixed,count,[5000000] -CACHE_PLAN,spark,mixed,long,columns=1 -CACHE_RESULT,spark,mixed,long,[12499997500000] -CACHE_PLAN,spark,mixed,string,columns=1 -CACHE_RESULT,spark,mixed,string,[54444500] -CACHE_PLAN,spark,mixed,three,columns=3 -CACHE_RESULT,spark,mixed,three,[12499997500000,54444500,49298480] -CACHE_PLAN,spark,mixed,all,columns=6 -CACHE_RESULT,spark,mixed,all,[12499997500000,2497500000,12500002500000,54444500,49298480,63888890] -CACHE_STORAGE,spark,numeric,512,30092440,30500288 -CACHE_PLAN,spark,numeric,count,columns=0 -CACHE_RESULT,spark,numeric,count,[5000000] -CACHE_PLAN,spark,numeric,long,columns=1 -CACHE_RESULT,spark,numeric,long,[12499997500000] -CACHE_PLAN,spark,numeric,three,columns=3 -CACHE_RESULT,spark,numeric,three,[12499997500000,2497500000,12500002500000] -CACHE_PLAN,spark,numeric,all,columns=6 -CACHE_RESULT,spark,numeric,all,[12499997500000,2497500000,12500002500000,249997500000,19787521156,37499992500000] -mixed,count,[60122458,58763584,64781250,57733083,60336000,53491125,51893667,50396208,49797750,57104333,52548834,58005625,53892083,48885750,45805125] -mixed,long,[68629250,64295750,62363917,63536125,63408792,67167542,68428292,63685792,62541167,62608583,63894250,62370375,61135167,63920084,62908000] -mixed,string,[175639166,174771208,191523459,173436959,173579167,172032792,173196916,171999166,182126292,179218417,186934333,178227292,177398416,183371500,176437292] -mixed,three,[324713334,325907375,326343083,326163917,327724459,325487458,324853083,325207417,326750042,326410750,326311834,328773875,325877000,323402917,320861458] -mixed,all,[494696875,492067459,492720375,494211667,517829542,507144750,489879167,502179250,487612625,488592625,487424459,492362084,491221708,511834083,522218875] -numeric,count,[36556333,38222000,36865458,37317208,37113583,35841958,36333125,46646708,36437750,37088791,35848709,36183834,40102916,39218792,35185333] -numeric,long,[56195750,59335792,54800083,53821583,54192666,58564416,59747000,62119625,56546875,79152125,57005709,57796167,56445917,55756708,58180500] -numeric,three,[93292167,95226792,94719333,95328500,94266584,93141458,98593417,99678666,91984250,101237916,92021875,93002208,101805333,91773792,98124666] -numeric,all,[148008625,152498791,149654250,151202708,147113750,148592875,152934666,148202458,147809875,148520041,149558375,150207583,164515167,148237833,149380417] -patch-timing-main-1-comet -CACHE_STORAGE,comet,mixed,512,168804118,169754384 -CACHE_PLAN,comet,mixed,count,columns=0 -CACHE_RESULT,comet,mixed,count,[5000000] -CACHE_PLAN,comet,mixed,long,columns=1 -CACHE_RESULT,comet,mixed,long,[12499997500000] -CACHE_PLAN,comet,mixed,string,columns=1 -CACHE_RESULT,comet,mixed,string,[54444500] -CACHE_PLAN,comet,mixed,three,columns=3 -CACHE_RESULT,comet,mixed,three,[12499997500000,54444500,49298480] -CACHE_PLAN,comet,mixed,all,columns=6 -CACHE_RESULT,comet,mixed,all,[12499997500000,2497500000,12500002500000,54444500,49298480,63888890] -CACHE_STORAGE,comet,numeric,512,107700478,108468272 -CACHE_PLAN,comet,numeric,count,columns=0 -CACHE_RESULT,comet,numeric,count,[5000000] -CACHE_PLAN,comet,numeric,long,columns=1 -CACHE_RESULT,comet,numeric,long,[12499997500000] -CACHE_PLAN,comet,numeric,three,columns=3 -CACHE_RESULT,comet,numeric,three,[12499997500000,2497500000,12500002500000] -CACHE_PLAN,comet,numeric,all,columns=6 -CACHE_RESULT,comet,numeric,all,[12499997500000,2497500000,12500002500000,249997500000,19787521156,37499992500000] -mixed,count,[154190417,152400916,151556041,153037000,148057042,129450416,130290875,129583000,125314083,128601333,132294667,123816791,126060292,144513833,147887000] -mixed,long,[195783667,192151625,190718916,194162666,193313792,198741583,211980875,191337875,189801333,184391000,186905041,187072250,183634625,187105500,184347083] -mixed,string,[312995792,312123500,313970959,317437792,317939458,323753667,337865875,349354959,316337375,311129042,315006500,313764541,313513542,310059417,314596625] -mixed,three,[431867084,442551250,431833500,435027333,433383208,428544792,433603083,434283834,433689834,432753583,431885250,430490000,428812916,432090834,429528041] -mixed,all,[636730875,635343583,641854875,643184292,636829500,634420334,635997709,640016625,641488625,635003542,635853792,662478125,637731209,641186291,635397292] -numeric,count,[142760625,136149000,138774334,137586167,139995000,136235500,136847250,140961125,139545375,145747542,138148667,140620417,136974500,135497666,137843750] -numeric,long,[183267125,182182250,182601333,182420750,181305458,182034584,182920042,184597917,183546125,184360500,182162000,184208333,181852375,184438208,181256667] -numeric,three,[245741625,243931500,243121292,245839250,248457209,246520667,242955416,244956833,245886542,246375042,242799125,248458917,245177625,242553291,242767334] -numeric,all,[343389084,315352541,345362834,314666708,309892291,355066834,309275333,307111250,308713958,306240542,307532458,309249208,310719625,308733125,312619083] -patch-timing-main-1-spark -CACHE_STORAGE,spark,mixed,512,242706686,243299704 -CACHE_PLAN,spark,mixed,count,columns=0 -CACHE_RESULT,spark,mixed,count,[5000000] -CACHE_PLAN,spark,mixed,long,columns=1 -CACHE_RESULT,spark,mixed,long,[12499997500000] -CACHE_PLAN,spark,mixed,string,columns=1 -CACHE_RESULT,spark,mixed,string,[54444500] -CACHE_PLAN,spark,mixed,three,columns=3 -CACHE_RESULT,spark,mixed,three,[12499997500000,54444500,49298480] -CACHE_PLAN,spark,mixed,all,columns=6 -CACHE_RESULT,spark,mixed,all,[12499997500000,2497500000,12500002500000,54444500,49298480,63888890] -CACHE_STORAGE,spark,numeric,512,30092440,30500288 -CACHE_PLAN,spark,numeric,count,columns=0 -CACHE_RESULT,spark,numeric,count,[5000000] -CACHE_PLAN,spark,numeric,long,columns=1 -CACHE_RESULT,spark,numeric,long,[12499997500000] -CACHE_PLAN,spark,numeric,three,columns=3 -CACHE_RESULT,spark,numeric,three,[12499997500000,2497500000,12500002500000] -CACHE_PLAN,spark,numeric,all,columns=6 -CACHE_RESULT,spark,numeric,all,[12499997500000,2497500000,12500002500000,249997500000,19787521156,37499992500000] -mixed,count,[57487083,59363583,54143834,58940458,52327167,58004708,54001541,53217583,50105083,47961959,59443417,46851667,55576958,46156166,49164541] -mixed,long,[72350583,64538708,63330750,70352959,74904000,64644917,63511042,65512125,66159875,62451416,62099125,63839333,63680792,64699583,61988875] -mixed,string,[152438583,145464083,148502500,147122041,148288542,146654625,145934250,144692417,145511667,145852208,146716667,145594500,146649750,145668167,144365416] -mixed,three,[309719125,268227166,298957209,266183125,270068583,300752834,266464208,266654292,265936708,265905833,266612916,266348917,266377666,266655500,266074292] -mixed,all,[392305916,395842250,396338000,396299375,427326167,401013625,391817500,406162167,391513166,401631375,394721875,397171875,392184375,390213708,394317084] -numeric,count,[39495166,37036000,40412917,39494500,38963792,40985375,38569000,40689625,40838500,45950750,47363000,37006208,36980959,40088875,38461833] -numeric,long,[63018792,75463750,73122750,60735542,57100917,61424125,56713833,62480875,58928750,53959541,58728584,55394500,58483041,54589291,61326375] -numeric,three,[98728834,93677667,97250000,99086792,93611083,94749708,96481792,97028375,96409583,95984042,93458333,94811542,95115750,92376916,97133208] -numeric,all,[156200042,160211417,151995625,153991125,154568542,153277959,156859209,154147375,152535708,157814750,155393208,153101875,153394750,154073125,152741792] -patch-timing-patch-0-comet -CACHE_STORAGE,comet,mixed,512,168804118,169754384 -CACHE_PLAN,comet,mixed,count,columns=0 -CACHE_RESULT,comet,mixed,count,[5000000] -CACHE_PLAN,comet,mixed,long,columns=1 -CACHE_RESULT,comet,mixed,long,[12499997500000] -CACHE_PLAN,comet,mixed,string,columns=1 -CACHE_RESULT,comet,mixed,string,[54444500] -CACHE_PLAN,comet,mixed,three,columns=3 -CACHE_RESULT,comet,mixed,three,[12499997500000,54444500,49298480] -CACHE_PLAN,comet,mixed,all,columns=6 -CACHE_RESULT,comet,mixed,all,[12499997500000,2497500000,12500002500000,54444500,49298480,63888890] -CACHE_STORAGE,comet,numeric,512,107700478,108468272 -CACHE_PLAN,comet,numeric,count,columns=0 -CACHE_RESULT,comet,numeric,count,[5000000] -CACHE_PLAN,comet,numeric,long,columns=1 -CACHE_RESULT,comet,numeric,long,[12499997500000] -CACHE_PLAN,comet,numeric,three,columns=3 -CACHE_RESULT,comet,numeric,three,[12499997500000,2497500000,12500002500000] -CACHE_PLAN,comet,numeric,all,columns=6 -CACHE_RESULT,comet,numeric,all,[12499997500000,2497500000,12500002500000,249997500000,19787521156,37499992500000] -mixed,count,[62337875,61327791,56817292,65798208,58202500,53627000,51277667,61614667,47866292,50476334,60204750,50157792,55737291,53457500,47255708] -mixed,long,[89675834,93536500,89428875,83268750,85575959,100159208,92236041,84258042,95596083,128454458,77327125,80336875,81303333,82338750,73700541] -mixed,string,[199106833,236384083,222001708,190832084,190944625,192168083,191469041,192773667,192143375,193418917,189756667,189414084,189339000,191704167,188841708] -mixed,three,[345616291,344921875,345123375,352998875,355030417,345967583,346323625,349292583,348524459,346015958,345910250,348346584,358505834,344999333,345473833] -mixed,all,[546806791,545775250,544199417,548145833,545403667,544608125,543799708,546250042,543718542,546622833,548913167,541548625,543101667,546973667,545543958] -numeric,count,[42205917,42656458,39518042,38891083,39330750,43295792,39038416,38956584,40436708,42119875,39161958,40589375,39626000,42196042,38676834] -numeric,long,[67438666,71493875,70207250,67902292,69103584,74414459,71598167,76024334,71474375,74995750,69474125,69322750,69228167,74330625,72708875] -numeric,three,[133468834,136875750,132058417,136080750,132691667,135932417,133862458,137188667,135100791,135490458,131474666,131887750,137495584,136610750,136438042] -numeric,all,[226937750,260211458,224666459,254753292,221156500,223965125,220584375,258870708,222454959,223984875,222399708,226644667,226833542,221015250,225697500] -patch-timing-patch-1-comet -CACHE_STORAGE,comet,mixed,512,168804118,169754384 -CACHE_PLAN,comet,mixed,count,columns=0 -CACHE_RESULT,comet,mixed,count,[5000000] -CACHE_PLAN,comet,mixed,long,columns=1 -CACHE_RESULT,comet,mixed,long,[12499997500000] -CACHE_PLAN,comet,mixed,string,columns=1 -CACHE_RESULT,comet,mixed,string,[54444500] -CACHE_PLAN,comet,mixed,three,columns=3 -CACHE_RESULT,comet,mixed,three,[12499997500000,54444500,49298480] -CACHE_PLAN,comet,mixed,all,columns=6 -CACHE_RESULT,comet,mixed,all,[12499997500000,2497500000,12500002500000,54444500,49298480,63888890] -CACHE_STORAGE,comet,numeric,512,107700478,108468272 -CACHE_PLAN,comet,numeric,count,columns=0 -CACHE_RESULT,comet,numeric,count,[5000000] -CACHE_PLAN,comet,numeric,long,columns=1 -CACHE_RESULT,comet,numeric,long,[12499997500000] -CACHE_PLAN,comet,numeric,three,columns=3 -CACHE_RESULT,comet,numeric,three,[12499997500000,2497500000,12500002500000] -CACHE_PLAN,comet,numeric,all,columns=6 -CACHE_RESULT,comet,numeric,all,[12499997500000,2497500000,12500002500000,249997500000,19787521156,37499992500000] -mixed,count,[59758583,54683083,57307625,60145791,54191000,55225917,55481709,56982959,49442208,50748709,53317417,49717916,59147083,47952167,49751375] -mixed,long,[83909916,95533000,125708084,87679958,86288917,103240000,93540500,84482500,119125167,74684375,76612875,82295500,77504375,81438500,77675500] -mixed,string,[180124375,182181458,179421625,184985833,186732625,179759542,179762042,178609083,180764625,178184708,176803583,177248583,179657292,178242125,179688000] -mixed,three,[330406625,326012875,328227666,328417042,327184125,325927750,332122541,397798875,326347708,325478459,324366333,327322500,329175334,325517750,326888709] -mixed,all,[530396042,535261667,524119916,528182709,521769250,528619709,542558458,521230459,559584375,524968916,530617417,536841750,531874917,565121375,595912375] -numeric,count,[46135291,59249000,51261500,40242667,43565833,39965333,38665375,39931083,38977792,37054375,37599166,43196542,40237292,38089916,41203750] -numeric,long,[67600583,67079833,67131125,66134250,66619458,68740917,72604292,68322209,72426042,71186208,70206167,67425833,68510708,66768000,75425083] -numeric,three,[131630959,133329833,133825500,135334500,136680708,131918334,132369084,138030375,136187792,132198334,131299708,130418542,130962500,134660333,139471667] -numeric,all,[222216583,222856334,255083917,222452125,224083209,252175917,220925375,224026334,222322458,258053208,220046250,219516042,227483208,221431459,221890250] From 8dc61adaf6a7699e9870c85d50cd791f5593f438 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Fri, 11 Sep 2026 22:17:26 +0800 Subject: [PATCH 03/12] perf: feed cached Arrow columns into Spark codegen --- .github/workflows/pr_build_linux.yml | 1 + .github/workflows/pr_build_macos.yml | 1 + .../comet/CometSparkSessionExtensions.scala | 11 ++- .../comet/rules/CometCacheColumnarRule.scala | 69 ++++++++++++++++ .../comet/exec/CometInMemoryCacheSuite.scala | 79 ++++++++++++++++++- .../CometCacheRowReaderBenchmark.scala | 20 +++-- 6 files changed, 165 insertions(+), 16 deletions(-) create mode 100644 spark/src/main/scala/org/apache/comet/rules/CometCacheColumnarRule.scala diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 58c04f406a7..f30fedc5cab 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -460,6 +460,7 @@ jobs: org.apache.spark.sql.CometCollationSuite org.apache.comet.CometFuzzAggregateSuite org.apache.spark.sql.comet.execution.arrow.CometArrowStreamSuite + org.apache.spark.sql.comet.execution.arrow.CachedBatchRowIteratorSuite org.apache.spark.sql.CometSparkInternalFunctionsSuite - name: "expressions" value: | diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index 5731b2f5be7..28d768270cf 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -214,6 +214,7 @@ jobs: org.apache.spark.sql.CometCollationSuite org.apache.comet.CometFuzzAggregateSuite org.apache.spark.sql.comet.execution.arrow.CometArrowStreamSuite + org.apache.spark.sql.comet.execution.arrow.CachedBatchRowIteratorSuite org.apache.spark.sql.CometSparkInternalFunctionsSuite - name: "expressions" value: | diff --git a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala index 5a67d5eef80..e7a8f950dcf 100644 --- a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala +++ b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala @@ -34,7 +34,7 @@ import org.apache.spark.sql.internal.SQLConf import org.apache.comet.CometConf._ import org.apache.comet.iceberg.IcebergWriteStrategy -import org.apache.comet.rules.{CometExecRule, CometPlanAdaptiveDynamicPruningFilters, CometReuseSubquery, CometScanRule, CometSpark34AqeDppFallbackRule, EliminateRedundantTransitions, RevertNativeForTransitionHeavyStages} +import org.apache.comet.rules.{CometCacheColumnarRule, CometExecRule, CometPlanAdaptiveDynamicPruningFilters, CometReuseSubquery, CometScanRule, CometSpark34AqeDppFallbackRule, EliminateRedundantTransitions, RevertNativeForTransitionHeavyStages} import org.apache.comet.shims.ShimCometSparkSessionExtensions /** @@ -54,7 +54,7 @@ import org.apache.comet.shims.ShimCometSparkSessionExtensions * CometSubqueryBroadcastExec for exchange reuse with Comet broadcasts * b. insertTransitions: ColumnarToRow/RowToColumnar added * c. postColumnarTransitions: RevertNativeForTransitionHeavyStages, - * EliminateRedundantTransitions + * EliminateRedundantTransitions, CometCacheColumnarRule * 5. ReuseExchangeAndSubquery -- Spark deduplicates subqueries (sees Comet nodes) * }}} * @@ -78,7 +78,7 @@ import org.apache.comet.shims.ShimCometSparkSessionExtensions * a. preColumnarTransitions: CometScanRule, CometExecRule (no-ops, already converted) * b. insertTransitions * c. postColumnarTransitions: RevertNativeForTransitionHeavyStages, - * EliminateRedundantTransitions + * EliminateRedundantTransitions, CometCacheColumnarRule * }}} * * On Spark 3.4, injectQueryStageOptimizerRule is unavailable. CometExecRule does not wrap SABs, @@ -113,7 +113,10 @@ class CometSparkSessionExtensions override def postColumnarTransitions: Rule[SparkPlan] = { val rules = - Seq(RevertNativeForTransitionHeavyStages(session), EliminateRedundantTransitions(session)) + Seq( + RevertNativeForTransitionHeavyStages(session), + EliminateRedundantTransitions(session), + CometCacheColumnarRule) plan => rules.foldLeft(plan) { case (p, rule) => rule(p) } } } diff --git a/spark/src/main/scala/org/apache/comet/rules/CometCacheColumnarRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometCacheColumnarRule.scala new file mode 100644 index 00000000000..50b8a840a19 --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/rules/CometCacheColumnarRule.scala @@ -0,0 +1,69 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.comet.rules + +import org.apache.spark.sql.catalyst.expressions.LeafExpression +import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback +import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer +import org.apache.spark.sql.execution.{CodegenSupport, ColumnarToRowExec, ColumnarToRowTransition, SparkPlan, WholeStageCodegenExec} +import org.apache.spark.sql.execution.adaptive.QueryStageExec +import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec + +/** + * Lets Spark's generated consumers read cached Arrow vectors without an intermediate UnsafeRow. + */ +object CometCacheColumnarRule extends Rule[SparkPlan] { + override def apply(plan: SparkPlan): SparkPlan = { + if (!conf.wholeStageEnabled) return plan + + plan.transformUp { + case parent: CodegenSupport + if parent.supportCodegen && !parent.supportsColumnar && + !parent.isInstanceOf[ColumnarToRowTransition] && + !WholeStageCodegenExec.isTooManyFields(conf, parent.schema) && + !parent.children.exists(p => WholeStageCodegenExec.isTooManyFields(conf, p.schema)) && + !parent.expressions.exists(_.exists { + case _: LeafExpression => false + case _: CodegenFallback => true + case _ => false + }) => + // Match the consuming edge rather than every scan: an existing columnar consumer (or a + // cache stage being materialized by AQE) must keep receiving batches. Spark inserts an + // InputAdapter around the scan later, while this transition fuses with the row consumer. + parent.withNewChildren(parent.children.map { + case child if isColumnarCometCache(child) => ColumnarToRowExec(child) + case child => child + }) + } + } + + private def isColumnarCometCache(plan: SparkPlan): Boolean = { + plan.supportsColumnar && (plan match { + case scan: InMemoryTableScanExec => + // The materialized format is fixed even when Comet execution is later disabled. The + // serializer delegates unsupported schemas to Spark, whose cache keeps its own reader. + scan.relation.cacheBuilder.serializer.isInstanceOf[ArrowCachedBatchSerializer] && + ArrowCachedBatchSerializer.supportsSchema(scan.relation.output) + case stage: QueryStageExec => isColumnarCometCache(stage.plan) + case _ => false + }) + } +} diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala index f3472dbb6d5..fa45e255365 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -29,7 +29,7 @@ import org.apache.spark.sql.catalyst.expressions.{And, Attribute, Expression, Gr import org.apache.spark.sql.columnar.{CachedBatch, SimpleMetricsCachedBatch} import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometInMemoryTableScanExec, CometSortExec, CometSortMergeJoinExec} import org.apache.spark.sql.comet.execution.arrow.CometCachedBatchHelper -import org.apache.spark.sql.execution.SortExec +import org.apache.spark.sql.execution.{ColumnarToRowExec, RowToColumnarExec, SortExec} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, QueryStageExec, ShuffleQueryStageExec} import org.apache.spark.sql.execution.columnar.{CometInMemoryRelationHelper, InMemoryRelation, InMemoryTableScanExec} import org.apache.spark.sql.execution.exchange.{Exchange, ReusedExchangeExec, ShuffleExchangeLike} @@ -41,6 +41,7 @@ import org.apache.spark.storage.StorageLevel import org.apache.comet.{CometArrowAllocator, CometConf} import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus} +import org.apache.comet.rules.CometCacheColumnarRule import org.apache.comet.vector.CometVector class CometInMemoryCacheSuite extends CometTestBase { @@ -296,11 +297,14 @@ class CometInMemoryCacheSuite extends CometTestBase { } test("Spark row consumers of Comet cache preserve values across batches") { - Seq("CODEGEN_ONLY", "NO_CODEGEN").foreach { mode => + for { + mode <- Seq("CODEGEN_ONLY", "NO_CODEGEN") + vectorized <- Seq(false, true) + } { withSQLConf( CometConf.COMET_ENABLED.key -> "false", SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", - SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "false", + SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> vectorized.toString, SQLConf.COLUMN_BATCH_SIZE.key -> "7", SQLConf.CODEGEN_FACTORY_MODE.key -> mode, SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> (mode == "CODEGEN_ONLY").toString, @@ -359,8 +363,13 @@ class CometInMemoryCacheSuite extends CometTestBase { scan } assert( - scans.nonEmpty && scans.forall(!_.supportsColumnar), + scans.nonEmpty && scans.forall(_.supportsColumnar == vectorized), df.queryExecution.executedPlan.toString) + if (!vectorized || mode == "NO_CODEGEN") { + assert( + !df.queryExecution.executedPlan.exists(_.isInstanceOf[ColumnarToRowExec]), + df.queryExecution.executedPlan.toString) + } checkAnswer(df, answer) } } finally source.unpersist(blocking = true) @@ -368,6 +377,68 @@ class CometInMemoryCacheSuite extends CometTestBase { } } + test("Spark generated consumers read cold and warm Comet caches columnarly") { + for { + adaptive <- Seq(false, true) + cometEnabled <- Seq(false, true) + } { + withSQLConf( + CometConf.COMET_ENABLED.key -> cometEnabled.toString, + CometConf.COMET_EXEC_ENABLED.key -> "false", + CometConf.COMET_SHUFFLE_ENABLED.key -> "false", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "true", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true", + SQLConf.COLUMN_BATCH_SIZE.key -> "7", + SQLConf.SHUFFLE_PARTITIONS.key -> "2") { + val source = spark + .range(0, 41, 1, 2) + .selectExpr("id AS key", "if(id % 3 = 0, null, concat('字', id)) AS s") + def query = source + .filter("key >= 7") + .selectExpr("sum(key)", "sum(length(s))", "count(*)") + val expected = query.collect().toSeq + source.cache() + try { + val builder = spark.sharedState.cacheManager + .lookupCachedData(source) + .get + .cachedRepresentation + .cacheBuilder + Seq(true, false).foreach { cold => + val df = query + val plan = df.queryExecution.executedPlan + // Planning must not materialize the cache or replace AQE's cache-stage metadata. + assert(builder.isCachedColumnBuffersLoaded != cold, plan.toString) + checkAnswer(df, expected) + assert(builder.isCachedColumnBuffersLoaded) + val transitions = collect(plan) { + case c: ColumnarToRowExec if collect(c.child) { case s: InMemoryTableScanExec => + s + }.nonEmpty => + c + } + assert(transitions.size == 1, plan.toString) + assert(collect(plan) { case s: CometInMemoryTableScanExec => s }.isEmpty) + if (adaptive && isSpark35Plus) { + assert(collect(plan) { + case s: QueryStageExec + if s.getClass.getSimpleName == "TableCacheQueryStageExec" => + s + }.size == 1) + } + val scan = collect(plan) { case s: InMemoryTableScanExec => s }.head + // A cache scan can also be the root of a columnar request or already have a + // transition. Applying the rule again must preserve those input/output contracts. + Seq(scan, ColumnarToRowExec(scan), RowToColumnarExec(scan)).foreach { boundary => + assert(CometCacheColumnarRule(boundary).fastEquals(boundary)) + } + } + } finally source.unpersist(blocking = true) + } + } + } + test("Comet cache serializer delegates unsupported types to Spark's cache format") { withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometCacheRowReaderBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometCacheRowReaderBenchmark.scala index 7b254f2c50a..1e89afbd84e 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometCacheRowReaderBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometCacheRowReaderBenchmark.scala @@ -34,8 +34,9 @@ import org.apache.comet.{CometConf, CometSparkSessionExtensions} /** * Compare Spark consumers of Comet and Spark caches (issue #5485). * - * Arguments: [spark|comet|all] [rows] [iterations] [all|mixed|numeric]. Run one format per JVM in - * alternating order on main and the patch. Cache creation and validation are outside timing. + * Arguments: [spark|comet|comet-row|all] [rows] [iterations] [all|mixed|numeric]. Run one format + * per JVM in alternating order on main and the patch. comet-row disables vectorized cache reading + * to isolate the row iterator. Cache creation and validation are outside timing. */ object CometCacheRowReaderBenchmark extends BenchmarkBase { private val warmups = 5 @@ -46,13 +47,13 @@ object CometCacheRowReaderBenchmark extends BenchmarkBase { val rows = args.lift(1).map(_.toLong).getOrElse(5000000L) val iterations = args.lift(2).map(_.toInt).getOrElse(15) val schema = args.lift(3).getOrElse("all") - require(Set("all", "spark", "comet").contains(format)) + require(Set("all", "spark", "comet", "comet-row").contains(format)) require(Set("all", "mixed", "numeric").contains(schema)) require(rows > 0 && iterations > 0) emit("CACHE_SAMPLE,format,schema,query,rows,iteration,elapsed_ns") val formats = - if (format == "all") Seq("spark", "comet") else Seq(format) + if (format == "all") Seq("spark", "comet", "comet-row") else Seq(format) val schemas = if (schema == "all") Seq("mixed", "numeric") else Seq(schema) formats.foreach { name => CometInMemoryRelationHelper.clearSerializer() @@ -75,7 +76,7 @@ object CometCacheRowReaderBenchmark extends BenchmarkBase { .config("spark.io.compression.codec", "lz4") .config(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key, "false") .config(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key, "true") - .config(SQLConf.CACHE_VECTORIZED_READER_ENABLED.key, "false") + .config(SQLConf.CACHE_VECTORIZED_READER_ENABLED.key, (name != "comet-row").toString) .config(SQLConf.CODEGEN_FACTORY_MODE.key, "CODEGEN_ONLY") .config(CometConf.COMET_ENABLED.key, "false") .config(CometConf.COMET_EXEC_ENABLED.key, "false") @@ -163,10 +164,13 @@ object CometCacheRowReaderBenchmark extends BenchmarkBase { assert(scans.size == 1, s"Expected one Spark cache scan:\n$plan") val scan = scans.head assert(scan.attributes.map(_.name).toSet == selected.toSet, s"Wrong projection:\n$plan") - assert(!scan.supportsColumnar, s"Expected the cache row reader:\n$plan") - assert(!plan.exists(_.isInstanceOf[ColumnarToRowExec]), s"Unexpected transition:\n$plan") + val columnar = plan.exists(_.isInstanceOf[ColumnarToRowExec]) + if (format != "comet") { + assert(!columnar, s"Expected the cache row reader:\n$plan") + } assert(!plan.exists(_.getClass.getName.startsWith("org.apache.spark.sql.comet."))) - emit(s"CACHE_PLAN,$format,$schema,$name,columns=${selected.size}\n$plan") + val reader = if (columnar) "columnar" else "row" + emit(s"CACHE_PLAN,$format,$schema,$name,columns=${selected.size},reader=$reader\n$plan") runQuery(query, answer, format, schema, name, rows, iterations) } } finally cached.unpersist(blocking = true) From 091eb00207f633eda299ecfa97972e4a5e67800f Mon Sep 17 00:00:00 2001 From: peterxcli Date: Fri, 11 Sep 2026 22:55:40 +0800 Subject: [PATCH 04/12] docs: illustrate Comet cache columnar rewrite --- .../comet/rules/CometCacheColumnarRule.scala | 21 +++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/spark/src/main/scala/org/apache/comet/rules/CometCacheColumnarRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometCacheColumnarRule.scala index 50b8a840a19..7d8ba194aed 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometCacheColumnarRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometCacheColumnarRule.scala @@ -29,6 +29,27 @@ import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec /** * Lets Spark's generated consumers read cached Arrow vectors without an intermediate UnsafeRow. + * + * Data flows upward. Spark's InputAdapter/whole-stage wrappers and an optional AQE cache stage + * are omitted: + * {{{ + * Before After + * +------------------------+ +------------------------+ + * | Spark codegen consumer | | Spark codegen consumer | + * +------------------------+ +------------------------+ + * ^ ^ + * | UnsafeRow | column values + * +------------------------+ +------------------------+ + * | InMemoryTableScanExec | | ColumnarToRowExec | + * | row iterator | | fused with consumer | + * +------------------------+ +------------------------+ + * ^ + * | ColumnarBatch + * +------------------------+ + * | InMemoryTableScanExec | + * | Arrow vectors | + * +------------------------+ + * }}} */ object CometCacheColumnarRule extends Rule[SparkPlan] { override def apply(plan: SparkPlan): SparkPlan = { From eafdba5e661bbfcf6cf3dd85809e1d06263c0c87 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Sun, 13 Sep 2026 22:11:49 +0800 Subject: [PATCH 05/12] fix: honor Comet disable switches for fused cache reads --- .../scala/org/apache/comet/CometConf.scala | 20 +++--- .../comet/rules/CometCacheColumnarRule.scala | 7 +- .../comet/exec/CometInMemoryCacheSuite.scala | 66 +++++++++++-------- .../CometCacheRowReaderBenchmark.scala | 4 +- 4 files changed, 57 insertions(+), 40 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 43f030a7d9d..7e0e644ab78 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -262,16 +262,20 @@ object CometConf extends ShimCometConf { conf("spark.comet.exec.inMemoryCache.enabled") .category(CATEGORY_EXEC) .doc( - "Whether to enable Comet native execution for in-memory cached tables. Its value at " + - "startup also decides whether CometDriverPlugin installs Comet's cache serializer, " + - "which stores cached data in Arrow format. Because spark.sql.cache.serializer is a " + + "Whether to enable Comet native scans and fused Spark reads of in-memory cached tables. " + + "Requires spark.comet.enabled=true. At startup, this setting also decides whether " + + "CometDriverPlugin installs Comet's cache serializer, which stores cached data in " + + "Arrow format. Because spark.sql.cache.serializer is a " + "static config, the cached format is fixed for the application, and disabling this " + - "at runtime only sends cached scans back to Spark's execution path. Relations whose " + - "schema Comet's Arrow writer does not support are always cached in Spark's default " + + "or spark.comet.enabled at runtime sends cached scans back to Spark's execution path " + + "without the fused reader. Relations whose schema Comet's Arrow writer does not " + + "support are always cached in Spark's default " + "format. Each cached column is stored as its own compressed Arrow IPC stream, so a " + - "scan decodes only the columns it projected. Reads that feed Spark operators rather " + - "than Comet ones still pay a row conversion the default format avoids, and can be " + - "slower than Spark's cache. With spark.kryo.registrationRequired=true, also set " + + "scan decodes only the columns it projected. Eligible Spark whole-stage codegen " + + "consumers read cached vectors directly when vectorized cache reading is enabled; " + + "other Spark row consumers use a reusable row buffer. Decoding costs can still make " + + "wide numeric reads slower than Spark's default cache. With " + + "spark.kryo.registrationRequired=true, also set " + "spark.kryo.registrator=org.apache.comet.CometKryoRegistrator before creating the " + "SparkContext, otherwise caching fails as soon as a block is serialized, including " + "the disk half of the default MEMORY_AND_DISK storage level.") diff --git a/spark/src/main/scala/org/apache/comet/rules/CometCacheColumnarRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometCacheColumnarRule.scala index 7d8ba194aed..fe9383869c7 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometCacheColumnarRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometCacheColumnarRule.scala @@ -27,6 +27,9 @@ import org.apache.spark.sql.execution.{CodegenSupport, ColumnarToRowExec, Column import org.apache.spark.sql.execution.adaptive.QueryStageExec import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec +import org.apache.comet.CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED +import org.apache.comet.CometSparkSessionExtensions.isCometLoaded + /** * Lets Spark's generated consumers read cached Arrow vectors without an intermediate UnsafeRow. * @@ -53,6 +56,7 @@ import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec */ object CometCacheColumnarRule extends Rule[SparkPlan] { override def apply(plan: SparkPlan): SparkPlan = { + if (!isCometLoaded(conf) || !COMET_EXEC_IN_MEMORY_CACHE_ENABLED.get(conf)) return plan if (!conf.wholeStageEnabled) return plan plan.transformUp { @@ -79,8 +83,7 @@ object CometCacheColumnarRule extends Rule[SparkPlan] { private def isColumnarCometCache(plan: SparkPlan): Boolean = { plan.supportsColumnar && (plan match { case scan: InMemoryTableScanExec => - // The materialized format is fixed even when Comet execution is later disabled. The - // serializer delegates unsupported schemas to Spark, whose cache keeps its own reader. + // The serializer delegates unsupported schemas to Spark, whose cache keeps its own reader. scan.relation.cacheBuilder.serializer.isInstanceOf[ArrowCachedBatchSerializer] && ArrowCachedBatchSerializer.supportsSchema(scan.relation.output) case stage: QueryStageExec => isColumnarCometCache(stage.plan) diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala index fa45e255365..2944ee508ca 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -377,13 +377,16 @@ class CometInMemoryCacheSuite extends CometTestBase { } } - test("Spark generated consumers read cold and warm Comet caches columnarly") { + test("Spark generated cache consumers respect runtime Comet disable switches") { for { adaptive <- Seq(false, true) - cometEnabled <- Seq(false, true) + enabledKey <- Seq( + CometConf.COMET_ENABLED.key, + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key) } { withSQLConf( - CometConf.COMET_ENABLED.key -> cometEnabled.toString, + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", CometConf.COMET_EXEC_ENABLED.key -> "false", CometConf.COMET_SHUFFLE_ENABLED.key -> "false", SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, @@ -405,33 +408,38 @@ class CometInMemoryCacheSuite extends CometTestBase { .get .cachedRepresentation .cacheBuilder - Seq(true, false).foreach { cold => - val df = query - val plan = df.queryExecution.executedPlan - // Planning must not materialize the cache or replace AQE's cache-stage metadata. - assert(builder.isCachedColumnBuffersLoaded != cold, plan.toString) - checkAnswer(df, expected) - assert(builder.isCachedColumnBuffersLoaded) - val transitions = collect(plan) { - case c: ColumnarToRowExec if collect(c.child) { case s: InMemoryTableScanExec => + // Materialize with fusion enabled, then disable and re-enable it on the same cache. + Seq(true, false, true).zipWithIndex.foreach { case (enabled, index) => + withSQLConf(enabledKey -> enabled.toString) { + val cold = index == 0 + val df = query + val plan = df.queryExecution.executedPlan + // Planning must not materialize the cache or replace AQE's cache-stage metadata. + assert(builder.isCachedColumnBuffersLoaded != cold, plan.toString) + checkAnswer(df, expected) + assert(builder.isCachedColumnBuffersLoaded) + val transitions = collect(plan) { + case c: ColumnarToRowExec if collect(c.child) { case s: InMemoryTableScanExec => + s + }.nonEmpty => + c + } + assert(transitions.size == (if (enabled) 1 else 0), plan.toString) + assert(collect(plan) { case s: CometInMemoryTableScanExec => s }.isEmpty) + if (adaptive && isSpark35Plus) { + assert(collect(plan) { + case s: QueryStageExec + if s.getClass.getSimpleName == "TableCacheQueryStageExec" => s - }.nonEmpty => - c - } - assert(transitions.size == 1, plan.toString) - assert(collect(plan) { case s: CometInMemoryTableScanExec => s }.isEmpty) - if (adaptive && isSpark35Plus) { - assert(collect(plan) { - case s: QueryStageExec - if s.getClass.getSimpleName == "TableCacheQueryStageExec" => - s - }.size == 1) - } - val scan = collect(plan) { case s: InMemoryTableScanExec => s }.head - // A cache scan can also be the root of a columnar request or already have a - // transition. Applying the rule again must preserve those input/output contracts. - Seq(scan, ColumnarToRowExec(scan), RowToColumnarExec(scan)).foreach { boundary => - assert(CometCacheColumnarRule(boundary).fastEquals(boundary)) + }.size == 1) + } + val scan = collect(plan) { case s: InMemoryTableScanExec => s }.head + assert(scan.supportsColumnar) + // A cache scan can also be the root of a columnar request or already have a + // transition. Applying the rule again must preserve those input/output contracts. + Seq(scan, ColumnarToRowExec(scan), RowToColumnarExec(scan)).foreach { boundary => + assert(CometCacheColumnarRule(boundary).fastEquals(boundary)) + } } } } finally source.unpersist(blocking = true) diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometCacheRowReaderBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometCacheRowReaderBenchmark.scala index 1e89afbd84e..3197f061a89 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometCacheRowReaderBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometCacheRowReaderBenchmark.scala @@ -78,8 +78,10 @@ object CometCacheRowReaderBenchmark extends BenchmarkBase { .config(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key, "true") .config(SQLConf.CACHE_VECTORIZED_READER_ENABLED.key, (name != "comet-row").toString) .config(SQLConf.CODEGEN_FACTORY_MODE.key, "CODEGEN_ONLY") - .config(CometConf.COMET_ENABLED.key, "false") + .config(CometConf.COMET_ENABLED.key, "true") + .config(CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key, "true") .config(CometConf.COMET_EXEC_ENABLED.key, "false") + .config(CometConf.COMET_SHUFFLE_ENABLED.key, "false") .withExtensions(new CometSparkSessionExtensions) .getOrCreate() spark.sparkContext.setLogLevel("WARN") From 49b7ec0b52999524d619c7136d69b6bd6bcd0400 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Tue, 15 Sep 2026 17:03:54 +0800 Subject: [PATCH 06/12] fix: align cache fusion with Spark codegen settings --- .../comet/rules/CometCacheColumnarRule.scala | 2 ++ .../execution/arrow/CachedBatchRowIterator.scala | 11 +++++++---- .../comet/exec/CometInMemoryCacheSuite.scala | 14 +++++++++----- .../benchmark/CometCacheRowReaderBenchmark.scala | 9 ++++++--- 4 files changed, 24 insertions(+), 12 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/rules/CometCacheColumnarRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometCacheColumnarRule.scala index fe9383869c7..ad777509646 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometCacheColumnarRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometCacheColumnarRule.scala @@ -26,6 +26,7 @@ import org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer import org.apache.spark.sql.execution.{CodegenSupport, ColumnarToRowExec, ColumnarToRowTransition, SparkPlan, WholeStageCodegenExec} import org.apache.spark.sql.execution.adaptive.QueryStageExec import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec +import org.apache.spark.sql.internal.SQLConf import org.apache.comet.CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED import org.apache.comet.CometSparkSessionExtensions.isCometLoaded @@ -58,6 +59,7 @@ object CometCacheColumnarRule extends Rule[SparkPlan] { override def apply(plan: SparkPlan): SparkPlan = { if (!isCometLoaded(conf) || !COMET_EXEC_IN_MEMORY_CACHE_ENABLED.get(conf)) return plan if (!conf.wholeStageEnabled) return plan + if (conf.getConf(SQLConf.CODEGEN_FACTORY_MODE).toString == "NO_CODEGEN") return plan plan.transformUp { case parent: CodegenSupport diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchRowIterator.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchRowIterator.scala index 907236999aa..d953f109c31 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchRowIterator.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchRowIterator.scala @@ -62,6 +62,7 @@ private[arrow] class CachedBatchRowIterator(attributes: Seq[Attribute]) } } val projection = GenerateUnsafeProjection.createCode(ctx, fields) + val batchesRef = ctx.addReferenceObj("batches", batches, "scala.collection.Iterator") val bindColumns = columns.zipWithIndex .map { case (column, i) => s"$column = batch.column($i);" @@ -69,17 +70,19 @@ private[arrow] class CachedBatchRowIterator(attributes: Seq[Attribute]) .mkString("\n") val code = s""" public Object generate(Object[] references) { - return new SpecificCachedBatchRowIterator((scala.collection.Iterator) references[0]); + return new SpecificCachedBatchRowIterator(references); } class SpecificCachedBatchRowIterator extends scala.collection.AbstractIterator { + private final Object[] references; private final scala.collection.Iterator batches; private int rowId = 0; private int numRows = 0; ${ctx.declareMutableStates()} - public SpecificCachedBatchRowIterator(scala.collection.Iterator batches) { - this.batches = batches; + public SpecificCachedBatchRowIterator(Object[] references) { + this.references = references; + this.batches = $batchesRef; ${ctx.initMutableStates()} } @@ -106,7 +109,7 @@ private[arrow] class CachedBatchRowIterator(attributes: Seq[Attribute]) """ val (compiled, _) = CodeGenerator.compile(new CodeAndComment(code, ctx.getPlaceHolderToComments())) - compiled.generate(Array[Any](batches)).asInstanceOf[Iterator[InternalRow]] + compiled.generate(ctx.references.toArray).asInstanceOf[Iterator[InternalRow]] } override protected def createInterpretedObject( diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala index 2944ee508ca..f75d758f716 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -377,12 +377,14 @@ class CometInMemoryCacheSuite extends CometTestBase { } } - test("Spark generated cache consumers respect runtime Comet disable switches") { + test("Spark generated cache consumers respect runtime enable and codegen settings") { for { adaptive <- Seq(false, true) - enabledKey <- Seq( - CometConf.COMET_ENABLED.key, - CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key) + disabledSetting <- Seq( + CometConf.COMET_ENABLED.key -> "false", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "false", + SQLConf.CODEGEN_FACTORY_MODE.key -> "NO_CODEGEN", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false") } { withSQLConf( CometConf.COMET_ENABLED.key -> "true", @@ -392,6 +394,7 @@ class CometInMemoryCacheSuite extends CometTestBase { SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "true", SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true", + SQLConf.CODEGEN_FACTORY_MODE.key -> "CODEGEN_ONLY", SQLConf.COLUMN_BATCH_SIZE.key -> "7", SQLConf.SHUFFLE_PARTITIONS.key -> "2") { val source = spark @@ -410,7 +413,8 @@ class CometInMemoryCacheSuite extends CometTestBase { .cacheBuilder // Materialize with fusion enabled, then disable and re-enable it on the same cache. Seq(true, false, true).zipWithIndex.foreach { case (enabled, index) => - withSQLConf(enabledKey -> enabled.toString) { + val settings = if (enabled) Seq.empty else Seq(disabledSetting) + withSQLConf(settings: _*) { val cold = index == 0 val df = query val plan = df.queryExecution.executedPlan diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometCacheRowReaderBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometCacheRowReaderBenchmark.scala index 3197f061a89..4473b8cf073 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometCacheRowReaderBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometCacheRowReaderBenchmark.scala @@ -35,8 +35,9 @@ import org.apache.comet.{CometConf, CometSparkSessionExtensions} * Compare Spark consumers of Comet and Spark caches (issue #5485). * * Arguments: [spark|comet|comet-row|all] [rows] [iterations] [all|mixed|numeric]. Run one format - * per JVM in alternating order on main and the patch. comet-row disables vectorized cache reading - * to isolate the row iterator. Cache creation and validation are outside timing. + * per JVM in alternating order. comet requires the fused columnar reader; comet-row disables + * vectorized cache reading to isolate the row iterator. Cache creation and validation are outside + * timing. */ object CometCacheRowReaderBenchmark extends BenchmarkBase { private val warmups = 5 @@ -167,7 +168,9 @@ object CometCacheRowReaderBenchmark extends BenchmarkBase { val scan = scans.head assert(scan.attributes.map(_.name).toSet == selected.toSet, s"Wrong projection:\n$plan") val columnar = plan.exists(_.isInstanceOf[ColumnarToRowExec]) - if (format != "comet") { + if (format == "comet") { + assert(columnar, s"Expected the fused columnar cache reader:\n$plan") + } else { assert(!columnar, s"Expected the cache row reader:\n$plan") } assert(!plan.exists(_.getClass.getName.startsWith("org.apache.spark.sql.comet."))) From ee0e24f3c1cb057d04c61c51d13dce5856a5185f Mon Sep 17 00:00:00 2001 From: peterxcli Date: Tue, 15 Sep 2026 21:55:24 +0800 Subject: [PATCH 07/12] ci: retry after runner and dependency download failures From 3aa603065226390f655b9eb468f94bd714c13536 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Tue, 29 Sep 2026 02:15:16 +0800 Subject: [PATCH 08/12] perf: split the generated cache row reader for wide projections The generated CachedBatchRowIterator passed every column read to GenerateUnsafeProjection through ctx.currentVars, which disables its method splitting, so next() held one write per column. Past about 100 nullable columns it exceeded HotSpot's 8000-byte HugeMethodLimit and was never JIT-compiled, and near 1500 columns it failed to compile at all. Read each column through a codegen-only leaf expression instead, so GenerateUnsafeProjection splits the field writes into bounded methods as it does for any projection, and bind a batch's columns in a loop. The largest generated method is now 371 bytes at 200 nullable bigint columns and 1905 bytes at 1500 (was 16462 bytes and a compile failure). Also check the ByteCodeStats that CodeGenerator.compile returns, as WholeStageCodegenExec does: when a method exceeds min(spark.sql.codegen.hugeMethodLimit, 8000), read the batches with UnsafeProjection over batch.getRow in an indexed loop without copy(). The interpreted path uses the same loop with InterpretedUnsafeProjection. --- .../arrow/CachedBatchRowIterator.scala | 153 ++++++++++++------ .../arrow/CachedBatchRowIteratorSuite.scala | 123 +++++++++++--- 2 files changed, 203 insertions(+), 73 deletions(-) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchRowIterator.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchRowIterator.scala index d953f109c31..82f73d25901 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchRowIterator.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchRowIterator.scala @@ -20,9 +20,11 @@ package org.apache.spark.sql.comet.execution.arrow import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.expressions.{Attribute, BoundReference, CodeGeneratorWithInterpretedFallback, InterpretedUnsafeProjection} +import org.apache.spark.sql.catalyst.expressions.{Attribute, BoundReference, CodeGeneratorWithInterpretedFallback, InterpretedUnsafeProjection, LeafExpression, UnsafeProjection} import org.apache.spark.sql.catalyst.expressions.codegen._ import org.apache.spark.sql.catalyst.expressions.codegen.Block._ +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.DataType import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} /** @@ -30,6 +32,11 @@ import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} * batches and releases them on advancement or task completion. As with Spark's cache reader, * callers must copy rows they retain across next(), but the returned row owns its variable-width * values and remains valid when hasNext() releases the batch that supplied them. + * + * The generated reader hands each column read to GenerateUnsafeProjection as an expression, so it + * splits the field writes of a wide projection into methods of bounded size, as it does for any + * Spark projection. If a generated method still exceeds the huge-method limit, the reader backs + * off the way WholeStageCodegenExec does, here to Spark's projection of each batch row. */ private[arrow] class CachedBatchRowIterator(attributes: Seq[Attribute]) extends CodeGeneratorWithInterpretedFallback[Iterator[ColumnarBatch], Iterator[InternalRow]] { @@ -41,33 +48,21 @@ private[arrow] class CachedBatchRowIterator(attributes: Seq[Attribute]) override protected def createCodeGeneratedObject( batches: Iterator[ColumnarBatch]): Iterator[InternalRow] = { val ctx = new CodegenContext - val columns = attributes.indices.map { i => - ctx.addMutableState(classOf[ColumnVector].getName, s"column$i") - } - ctx.currentVars = attributes.zip(columns).map { case (attr, column) => - val value = JavaCode.variable(ctx.freshName("value"), attr.dataType) - val getter = CodeGenerator.getValueFromVector(column, attr.dataType, "rowId") - val javaType = CodeGenerator.javaType(attr.dataType) - if (attr.nullable) { - val isNull = JavaCode.isNullVariable(ctx.freshName("isNull")) - ExprCode( - code""" - boolean $isNull = $column.isNullAt(rowId); - $javaType $value = $isNull ? ${CodeGenerator.defaultValue(attr.dataType)} : ($getter); - """, - isNull, - value) - } else { - ExprCode(code"$javaType $value = $getter;", FalseLiteral, value) - } + val vectorClass = classOf[ColumnVector].getName + val batchClass = classOf[ColumnarBatch].getName + val columns = ctx.addMutableState( + s"$vectorClass[]", + "columns", + v => s"$v = new $vectorClass[${attributes.length}];", + forceInline = true) + val rowId = ctx.addMutableState(CodeGenerator.JAVA_INT, "rowId", forceInline = true) + val reads = attributes.zipWithIndex.map { case (attr, i) => + VectorValue(s"$columns[$i]", rowId, attr.dataType, attr.nullable) } - val projection = GenerateUnsafeProjection.createCode(ctx, fields) + // With ctx.currentVars unset, GenerateUnsafeProjection splits the field writes into methods + // that take the input row as their argument. The reads above ignore it. + val projection = GenerateUnsafeProjection.createCode(ctx, reads) val batchesRef = ctx.addReferenceObj("batches", batches, "scala.collection.Iterator") - val bindColumns = columns.zipWithIndex - .map { case (column, i) => - s"$column = batch.column($i);" - } - .mkString("\n") val code = s""" public Object generate(Object[] references) { return new SpecificCachedBatchRowIterator(references); @@ -76,7 +71,6 @@ private[arrow] class CachedBatchRowIterator(attributes: Seq[Attribute]) class SpecificCachedBatchRowIterator extends scala.collection.AbstractIterator { private final Object[] references; private final scala.collection.Iterator batches; - private int rowId = 0; private int numRows = 0; ${ctx.declareMutableStates()} @@ -87,45 +81,106 @@ private[arrow] class CachedBatchRowIterator(attributes: Seq[Attribute]) } public boolean hasNext() { - while (rowId >= numRows && batches.hasNext()) { - ${classOf[ColumnarBatch].getName} batch = - (${classOf[ColumnarBatch].getName}) batches.next(); + while ($rowId >= numRows && batches.hasNext()) { + $batchClass batch = ($batchClass) batches.next(); numRows = batch.numRows(); - rowId = 0; - $bindColumns + $rowId = 0; + for (int ordinal = 0; ordinal < $columns.length; ordinal++) { + $columns[ordinal] = batch.column(ordinal); + } } - return rowId < numRows; + return $rowId < numRows; } public InternalRow next() { if (!hasNext()) throw new java.util.NoSuchElementException(); + InternalRow ${ctx.INPUT_ROW} = null; ${projection.code} - rowId++; + $rowId++; return ${projection.value}; } ${ctx.declareAddedFunctions()} } """ - val (compiled, _) = + val (compiled, stats) = CodeGenerator.compile(new CodeAndComment(code, ctx.getPlaceHolderToComments())) - compiled.generate(ctx.references.toArray).asInstanceOf[Iterator[InternalRow]] + // Honor spark.sql.codegen.hugeMethodLimit as whole-stage codegen does, but never go above + // HotSpot's own limit: the config defaults to the largest method the JVM accepts, while this + // runs once per row and HotSpot never JIT-compiles a method longer than + // DEFAULT_JVM_HUGE_METHOD_LIMIT bytes. + val limit = + math.min(SQLConf.get.hugeMethodLimit, CodeGenerator.DEFAULT_JVM_HUGE_METHOD_LIMIT) + if (stats.maxMethodCodeSize > limit) { + logInfo( + s"Generated cache reader for ${attributes.length} columns has a " + + s"${stats.maxMethodCodeSize}-byte method, above the $limit-byte limit; " + + "projecting cached rows with UnsafeProjection instead") + new ProjectedRows(batches, UnsafeProjection.create(fields)) + } else { + compiled.generate(ctx.references.toArray).asInstanceOf[Iterator[InternalRow]] + } } override protected def createInterpretedObject( - batches: Iterator[ColumnarBatch]): Iterator[InternalRow] = { - val toUnsafe = InterpretedUnsafeProjection.createProjection(fields) - batches.flatMap { batch => - new Iterator[InternalRow] { - private var rowId = 0 - override def hasNext: Boolean = rowId < batch.numRows() - override def next(): InternalRow = { - if (!hasNext) throw new NoSuchElementException - val row = toUnsafe(batch.getRow(rowId)) - rowId += 1 - row - } - } + batches: Iterator[ColumnarBatch]): Iterator[InternalRow] = + new ProjectedRows(batches, InterpretedUnsafeProjection.createProjection(fields)) +} + +/** + * Projects each batch row through `projection`, which reuses one UnsafeRow and owns the values it + * writes, under the same contract as the generated reader. + */ +private[arrow] class ProjectedRows( + batches: Iterator[ColumnarBatch], + private[arrow] val projection: UnsafeProjection) + extends Iterator[InternalRow] { + private var batch: ColumnarBatch = _ + private var rowId = 0 + private var numRows = 0 + + override def hasNext: Boolean = { + while (rowId >= numRows && batches.hasNext) { + batch = batches.next() + numRows = batch.numRows() + rowId = 0 + } + rowId < numRows + } + + override def next(): InternalRow = { + if (!hasNext) throw new NoSuchElementException + val row = projection(batch.getRow(rowId)) + rowId += 1 + row + } +} + +/** + * The current row of one column of the batch a generated reader is reading. It exists only to be + * code generated, as an expression rather than through ctx.currentVars, which would stop + * GenerateUnsafeProjection from splitting the writer and leave every field in next(). + */ +private case class VectorValue( + column: String, + rowId: String, + dataType: DataType, + nullable: Boolean) + extends LeafExpression { + + override def eval(input: InternalRow): Any = + throw new UnsupportedOperationException(s"$nodeName is only code generated") + + override protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = { + val javaType = CodeGenerator.javaType(dataType) + val value = CodeGenerator.getValueFromVector(column, dataType, rowId) + if (nullable) { + ev.copy(code = code""" + boolean ${ev.isNull} = $column.isNullAt($rowId); + $javaType ${ev.value} = ${ev.isNull} ? ${CodeGenerator.defaultValue(dataType)} : ($value); + """) + } else { + ev.copy(code = code"$javaType ${ev.value} = $value;", isNull = FalseLiteral) } } } diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchRowIteratorSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchRowIteratorSuite.scala index 1bea5c5e3e1..6d2a5385e81 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchRowIteratorSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchRowIteratorSuite.scala @@ -25,24 +25,26 @@ import org.scalatest.funsuite.AnyFunSuite import org.apache.arrow.memory.RootAllocator import org.apache.arrow.vector.VarCharVector -import org.apache.spark.sql.catalyst.expressions.{AttributeReference, UnsafeRow} +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{AttributeReference, InterpretedUnsafeProjection, UnsafeRow} import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.types.{IntegerType, StringType} +import org.apache.spark.sql.types.{IntegerType, LongType, StringType} import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} import org.apache.comet.vector.CometPlainVector class CachedBatchRowIteratorSuite extends AnyFunSuite { Seq("CODEGEN_ONLY", "NO_CODEGEN").foreach { mode => - def withMode(f: => Unit): Unit = { + def withMode(settings: (String, String)*)(f: => Unit): Unit = { val conf = new SQLConf conf.setConfString(SQLConf.CODEGEN_FACTORY_MODE.key, mode) + settings.foreach { case (key, value) => conf.setConfString(key, value) } SQLConf.withExistingConf(conf)(f) } test(s"$mode: rows own Arrow values across batch release and reuse the output buffer") { - withMode { + withMode() { val allocator = new RootAllocator(Long.MaxValue) val vectors = Seq(Seq("first", null), Seq("字" * 1000, "last")).map { values => val vector = new VarCharVector("s", allocator) @@ -92,7 +94,7 @@ class CachedBatchRowIteratorSuite extends AnyFunSuite { } test(s"$mode: empty input, empty batches, and zero-column rows") { - withMode { + withMode() { val factory = new CachedBatchRowIterator(Seq.empty) val empty = factory.createObject(Iterator.empty) assert(!empty.hasNext) @@ -109,29 +111,102 @@ class CachedBatchRowIteratorSuite extends AnyFunSuite { } } - test(s"$mode: wide projections preserve nullable and required columns") { - withMode { - val attributes = (0 until 150).map { i => - AttributeReference(s"c$i", IntegerType, nullable = i % 2 == 0)() - } - val columns = attributes.indices.map { i => - val column = new OnHeapColumnVector(2, IntegerType) - column.putInt(0, i) - if (i % 2 == 0) column.putNull(1) else column.putInt(1, -i) - column + // Before the generated reader let GenerateUnsafeProjection split its writer, next() passed + // HotSpot's 8000-byte JIT limit near 100 columns and Janino's 64 KB method limit near 1500. + // The generated reader is only kept when none of its methods is above the former. + val reader = if (mode == "CODEGEN_ONLY") "the generated reader" else "the interpreted reader" + Seq(100, 200, 1500).foreach { width => + test(s"$mode: $width-column projections use $reader") { + withMode() { + val (attributes, batches) = wideInput(width) + try { + val rows = new CachedBatchRowIterator(attributes).createObject(batches.iterator) + rows match { + case projected: ProjectedRows => + assert(mode == "NO_CODEGEN", "The generated reader has a method too large to JIT") + assert(projected.projection.isInstanceOf[InterpretedUnsafeProjection]) + case _ => assert(mode == "CODEGEN_ONLY") + } + checkWideRows(attributes, rows) + } finally batches.foreach(_.close()) } - val batch = new ColumnarBatch(columns.toArray[ColumnVector], 2) + } + } + + test( + s"$mode: a generated method above the huge-method limit falls back to UnsafeProjection") { + // Below any generated method, as if the reader had grown past HotSpot's limit. + withMode(SQLConf.WHOLESTAGE_HUGE_METHOD_LIMIT.key -> "1") { + val (attributes, batches) = wideInput(200) try { - val rows = new CachedBatchRowIterator(attributes).createObject(Iterator.single(batch)) - val first = rows.next().copy() - val second = rows.next() - attributes.indices.foreach { i => - assert(first.getInt(i) == i) - if (i % 2 == 0) assert(second.isNullAt(i)) else assert(second.getInt(i) == -i) + val rows = new CachedBatchRowIterator(attributes).createObject(batches.iterator) + val projection = rows.asInstanceOf[ProjectedRows].projection + assert(projection.isInstanceOf[InterpretedUnsafeProjection] == (mode == "NO_CODEGEN")) + checkWideRows(attributes, rows) + } finally batches.foreach(_.close()) + } + } + } + + private val wideTypes = Seq(IntegerType, LongType, StringType) + + /** + * Nullable and required int, bigint and string columns, two rows in each of two batches, so + * that every column is bound again at the batch boundary. Odd rows are null where a column + * allows it. + */ + private def wideInput(width: Int): (Seq[AttributeReference], Seq[ColumnarBatch]) = { + val attributes = (0 until width).map { i => + AttributeReference(s"c$i", wideTypes(i % 3), nullable = i % 2 == 0)() + } + val batches = Seq(0, 2).map { firstRow => + val columns = attributes.zipWithIndex.map { case (attr, i) => + val column = new OnHeapColumnVector(2, attr.dataType) + Seq(0, 1).foreach { r => + val row = firstRow + r + if (attr.nullable && row % 2 == 1) { + column.putNull(r) + } else { + attr.dataType match { + case IntegerType => column.putInt(r, wideValue(i, row).asInstanceOf[Int]) + case LongType => column.putLong(r, wideValue(i, row).asInstanceOf[Long]) + case _ => column.putByteArray(r, wideValue(i, row).toString.getBytes(UTF_8)) + } + } + } + column + } + new ColumnarBatch(columns.toArray[ColumnVector], 2) + } + (attributes, batches) + } + + private def wideValue(column: Int, row: Int): Any = wideTypes(column % 3) match { + case IntegerType => column * 10 + row + case LongType => column * 10000000000L + row + case _ => s"$column:$row" + } + + private def checkWideRows( + attributes: Seq[AttributeReference], + rows: Iterator[InternalRow]): Unit = { + (0 until 4).foreach { row => + assert(rows.hasNext) + val actual = rows.next() + assert(actual.numFields == attributes.size) + attributes.zipWithIndex.foreach { case (attr, i) => + if (attr.nullable && row % 2 == 1) { + assert(actual.isNullAt(i), s"c$i in row $row") + } else { + val value = attr.dataType match { + case IntegerType => actual.getInt(i) + case LongType => actual.getLong(i) + case _ => actual.getUTF8String(i).toString } - assert(!rows.hasNext) - } finally batch.close() + assert(value == wideValue(i, row), s"c$i in row $row") + } } } + assert(!rows.hasNext) } } From 351503a9bddcac6f5094a95f5f28a30dd0285f3b Mon Sep 17 00:00:00 2001 From: peterxcli Date: Tue, 29 Sep 2026 02:15:16 +0800 Subject: [PATCH 09/12] fix: do not fuse cache reads in plan-only mode Since #5394 the post-columnar rules also run in plan-only mode, where Comet only reports the plan it would execute and Spark executes its own. CometCacheColumnarRule still inserted a ColumnarToRowExec there. Return early when plan-only mode applies, using the predicate CometRule uses, except in the plan-only preview, which should still show the transition. In the runtime-settings test, add a plan-only case and use QueryTest.checkAnswer with checkToRDD = false so the cold iteration really reads a cold cache. --- .../comet/rules/CometCacheColumnarRule.scala | 7 +++- .../org/apache/comet/rules/CometRule.scala | 37 ++++++++++++------- .../comet/exec/CometInMemoryCacheSuite.scala | 33 ++++++++++++----- 3 files changed, 54 insertions(+), 23 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/rules/CometCacheColumnarRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometCacheColumnarRule.scala index ad777509646..df35763820a 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometCacheColumnarRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometCacheColumnarRule.scala @@ -54,10 +54,15 @@ import org.apache.comet.CometSparkSessionExtensions.isCometLoaded * | Arrow vectors | * +------------------------+ * }}} + * + * @param preview + * true in the plan-only preview, which shows the plan Comet would execute. Otherwise the rule + * leaves plans alone in plan-only mode, where Spark executes each query unchanged. */ -object CometCacheColumnarRule extends Rule[SparkPlan] { +case class CometCacheColumnarRule(preview: Boolean = false) extends Rule[SparkPlan] { override def apply(plan: SparkPlan): SparkPlan = { if (!isCometLoaded(conf) || !COMET_EXEC_IN_MEMORY_CACHE_ENABLED.get(conf)) return plan + if (!preview && CometRule.planOnlyApplies(conf, plan)) return plan if (!conf.wholeStageEnabled) return plan if (conf.getConf(SQLConf.CODEGEN_FACTORY_MODE).toString == "NO_CODEGEN") return plan diff --git a/spark/src/main/scala/org/apache/comet/rules/CometRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometRule.scala index 8bc1658ca89..17ef4a6ce1b 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometRule.scala @@ -29,6 +29,7 @@ import org.apache.spark.sql.execution.{ApplyColumnarRulesAndInsertTransitions, B import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, InsertAdaptiveSparkPlan, QueryStageExec} import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, Exchange} import org.apache.spark.sql.execution.reuse.ReuseExchangeAndSubquery +import org.apache.spark.sql.internal.SQLConf import org.apache.comet.{CometConf, ExtendedExplainInfo} import org.apache.comet.CometSparkSessionExtensions.isCometLoaded @@ -36,12 +37,29 @@ import org.apache.comet.shims.ShimCometStreaming object CometRule { - /** Comet's post-columnar rules, shared by `CometColumnar` and the plan-only preview. */ - def postColumnarRules(session: SparkSession, wholePlan: Boolean = false): Seq[Rule[SparkPlan]] = + /** + * Comet's post-columnar rules, shared by `CometColumnar` and the plan-only preview. + * + * @param preview + * true for the plan-only preview, which holds the whole plan and shows the plan Comet would + * execute. + */ + def postColumnarRules(session: SparkSession, preview: Boolean = false): Seq[Rule[SparkPlan]] = Seq( - RevertNativeForTransitionHeavyStages(session, wholePlan), + RevertNativeForTransitionHeavyStages(session, wholePlan = preview), EliminateRedundantTransitions(session), - CometCacheColumnarRule) + CometCacheColumnarRule(preview)) + + /** + * Whether plan-only mode applies to `plan`, so that Comet only reports the plan it would + * execute and Spark executes `plan` unchanged. Mirrors the conversion rules' own guards; + * plan-only is scoped to exec being enabled. + */ + private[comet] def planOnlyApplies(conf: SQLConf, plan: SparkPlan): Boolean = + CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get(conf) && + isCometLoaded(conf) && + !ShimCometStreaming.isStreamingPlan(plan) && + CometConf.COMET_EXEC_ENABLED.get(conf) /** * Canonical hashes of the subquery plans reported for the query this thread is preparing. Spark @@ -141,7 +159,7 @@ case class CometRule(session: SparkSession, queryStagePrep: Boolean = false) private val execRule = CometExecRule(session) override def apply(plan: SparkPlan): SparkPlan = { - if (planOnlyApplies(plan)) { + if (CometRule.planOnlyApplies(conf, plan)) { reportPlanOnlyCoverage(plan) plan } else { @@ -151,13 +169,6 @@ case class CometRule(session: SparkSession, queryStagePrep: Boolean = false) private def convert(plan: SparkPlan): SparkPlan = execRule.apply(scanRule.apply(plan)) - /** Mirrors the conversion rules' own guards; plan-only is scoped to exec being enabled. */ - private def planOnlyApplies(plan: SparkPlan): Boolean = - CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get(conf) && - isCometLoaded(conf) && - !ShimCometStreaming.isStreamingPlan(plan) && - CometConf.COMET_EXEC_ENABLED.get(conf) - /** Logs the Comet plan for `plan` unless already reported. Never fails the query. */ private def reportPlanOnlyCoverage(plan: SparkPlan): Unit = { try { @@ -184,7 +195,7 @@ case class CometRule(session: SparkSession, queryStagePrep: Boolean = false) val withTransitions = ApplyColumnarRulesAndInsertTransitions(Seq.empty, outputsColumnar = false).apply(converted) val preview = CometRule - .postColumnarRules(session, wholePlan = true) + .postColumnarRules(session, preview = true) .foldLeft(withTransitions) { case (p, rule) => rule(p) } if (topLevel) ReuseExchangeAndSubquery.apply(preview) else preview } diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala index 1dca4ba868b..b96f9b37227 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -37,7 +37,7 @@ import org.apache.spark.sql.columnar.{CachedBatch, SimpleMetricsCachedBatch} import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometInMemoryTableScanExec, CometSortExec, CometSortMergeJoinExec} import org.apache.spark.sql.comet.execution.arrow.{ArrowCachedBatchSerializer, CometCachedBatchHelper} import org.apache.spark.sql.comet.util.Utils -import org.apache.spark.sql.execution.{ColumnarToRowExec, RowToColumnarExec, SortExec} +import org.apache.spark.sql.execution.{ColumnarToRowExec, FilterExec, RowToColumnarExec, SortExec} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, QueryStageExec, ShuffleQueryStageExec} import org.apache.spark.sql.execution.columnar.{CometInMemoryRelationHelper, InMemoryRelation, InMemoryTableScanExec} import org.apache.spark.sql.execution.exchange.{Exchange, ReusedExchangeExec, ShuffleExchangeLike} @@ -434,13 +434,18 @@ class CometInMemoryCacheSuite extends CometTestBase { } test("Spark generated cache consumers respect runtime enable and codegen settings") { + val planOnly = Seq( + CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true", + // Plan-only mode applies only while native execution is enabled. + CometConf.COMET_EXEC_ENABLED.key -> "true") for { adaptive <- Seq(false, true) - disabledSetting <- Seq( - CometConf.COMET_ENABLED.key -> "false", - CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "false", - SQLConf.CODEGEN_FACTORY_MODE.key -> "NO_CODEGEN", - SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false") + disabledSettings <- Seq( + Seq(CometConf.COMET_ENABLED.key -> "false"), + Seq(CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "false"), + Seq(SQLConf.CODEGEN_FACTORY_MODE.key -> "NO_CODEGEN"), + Seq(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false"), + planOnly) } { withSQLConf( CometConf.COMET_ENABLED.key -> "true", @@ -469,14 +474,16 @@ class CometInMemoryCacheSuite extends CometTestBase { .cacheBuilder // Materialize with fusion enabled, then disable and re-enable it on the same cache. Seq(true, false, true).zipWithIndex.foreach { case (enabled, index) => - val settings = if (enabled) Seq.empty else Seq(disabledSetting) + val settings = if (enabled) Seq.empty else disabledSettings withSQLConf(settings: _*) { val cold = index == 0 val df = query val plan = df.queryExecution.executedPlan // Planning must not materialize the cache or replace AQE's cache-stage metadata. assert(builder.isCachedColumnBuffersLoaded != cold, plan.toString) - checkAnswer(df, expected) + // checkToRDD = false keeps checkAnswer from loading the cache with a query of its + // own, so the cold run's table-cache stage materializes and AQE re-plans above it. + QueryTest.checkAnswer(df, expected, checkToRDD = false) assert(builder.isCachedColumnBuffersLoaded) val transitions = collect(plan) { case c: ColumnarToRowExec if collect(c.child) { case s: InMemoryTableScanExec => @@ -498,8 +505,16 @@ class CometInMemoryCacheSuite extends CometTestBase { // A cache scan can also be the root of a columnar request or already have a // transition. Applying the rule again must preserve those input/output contracts. Seq(scan, ColumnarToRowExec(scan), RowToColumnarExec(scan)).foreach { boundary => - assert(CometCacheColumnarRule(boundary).fastEquals(boundary)) + assert(CometCacheColumnarRule()(boundary).fastEquals(boundary)) } + // The plan-only preview shows the plan Comet would execute, so it still fuses a + // generated consumer that the executed plan leaves alone in plan-only mode. + val consumer = FilterExec(Literal.TrueLiteral, scan) + val fusedConsumer = FilterExec(Literal.TrueLiteral, ColumnarToRowExec(scan)) + assert(CometCacheColumnarRule()(consumer).fastEquals(fusedConsumer) == enabled) + assert( + CometCacheColumnarRule(preview = true)(consumer).fastEquals(fusedConsumer) == + (enabled || disabledSettings == planOnly)) } } } finally source.unpersist(blocking = true) From d4ad3c42800a0b4aef7f4aa6bc1485886ae5f5e3 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Tue, 29 Sep 2026 02:15:16 +0800 Subject: [PATCH 10/12] test: check fused cache reads with Comet on and native execution off The row-consumer test ran with Comet disabled, so no case took the fused path and the test could not notice. Run it with Comet on and native execution and shuffle off, with AQE on and off, add a filter that reads every column through codegen, and assert per query whether the fused transition is present. --- .../comet/exec/CometInMemoryCacheSuite.scala | 59 +++++++++++-------- 1 file changed, 35 insertions(+), 24 deletions(-) diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala index b96f9b37227..ca77452cb30 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -354,12 +354,18 @@ class CometInMemoryCacheSuite extends CometTestBase { test("Spark row consumers of Comet cache preserve values across batches") { for { + adaptive <- Seq(false, true) mode <- Seq("CODEGEN_ONLY", "NO_CODEGEN") vectorized <- Seq(false, true) } { + // Comet on with native execution off, so Spark operators consume the cache scan and the + // generated ones among them read its vectors through the fused transition. withSQLConf( - CometConf.COMET_ENABLED.key -> "false", - SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "false", + CometConf.COMET_SHUFFLE_ENABLED.key -> "false", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> vectorized.toString, SQLConf.COLUMN_BATCH_SIZE.key -> "7", SQLConf.CODEGEN_FACTORY_MODE.key -> mode, @@ -396,15 +402,20 @@ class CometInMemoryCacheSuite extends CometTestBase { "if(id % 3 = 0, null, map('k', array(cast(id AS STRING), null))) AS m", "null AS n")): _*) - def queries(df: DataFrame): Seq[DataFrame] = Seq( - df.select("*"), - df.selectExpr("s AS renamed", "key", "b", "a", "st", "m"), - df.orderBy($"s".desc, $"key"), - df.join(spark.range(41).toDF("join_key"), $"key" === $"join_key").select(df("*")), - df.selectExpr("count(*)"), - df.limit(1)) - - val expected = queries(source).map(_.collect().toSeq) + // Each query, and whether a generated Spark operator consumes the cache scan directly. + // The other consumers (the query root, exchanges and limits) read the row iterator. + def queries(df: DataFrame): Seq[(DataFrame, Boolean)] = Seq( + // The generated filter reads every column, so this covers the whole type matrix. + df.filter($"key" >= 0) -> true, + df.select("*") -> false, + df.selectExpr("s AS renamed", "key", "b", "a", "st", "m") -> true, + df.orderBy($"s".desc, $"key") -> false, + df.join(spark.range(41).toDF("join_key"), $"key" === $"join_key") + .select(df("*")) -> false, + df.selectExpr("count(*)") -> true, + df.limit(1) -> false) + + val expected = queries(source).map(_._1.collect().toSeq) source.cache() try { assert(source.count() == 41) @@ -413,20 +424,20 @@ class CometInMemoryCacheSuite extends CometTestBase { val buffers = relation.cacheBuilder.cachedColumnBuffers.collect() assert(buffers.length > 2) assert(buffers.forall(_.getClass.getSimpleName == "CometCachedBatch")) - queries(source).zip(expected).foreach { case (df, answer) => - val scans = - df.queryExecution.executedPlan.collect { case scan: InMemoryTableScanExec => - scan - } - assert( - scans.nonEmpty && scans.forall(_.supportsColumnar == vectorized), - df.queryExecution.executedPlan.toString) - if (!vectorized || mode == "NO_CODEGEN") { - assert( - !df.queryExecution.executedPlan.exists(_.isInstanceOf[ColumnarToRowExec]), - df.queryExecution.executedPlan.toString) - } + queries(source).zip(expected).foreach { case ((df, generatedConsumer), answer) => + val plan = df.queryExecution.executedPlan checkAnswer(df, answer) + // Inspected after execution, when an adaptive plan is final. + val scans = collect(plan) { case scan: InMemoryTableScanExec => scan } + assert(scans.nonEmpty && scans.forall(_.supportsColumnar == vectorized), plan) + val transitions = collect(plan) { + case c: ColumnarToRowExec if collect(c.child) { case s: InMemoryTableScanExec => + s + }.nonEmpty => + c + } + val fused = generatedConsumer && vectorized && mode == "CODEGEN_ONLY" + assert(transitions.size == (if (fused) 1 else 0), plan) } } finally source.unpersist(blocking = true) } From f6f245f9c5f35cf7521313785a500af55646129c Mon Sep 17 00:00:00 2001 From: peterxcli Date: Tue, 29 Sep 2026 02:15:16 +0800 Subject: [PATCH 11/12] test: measure the fused cache reader in CometInMemoryCacheBenchmark Fold the fused-consumer case into CometInMemoryCacheBenchmark's Spark-operator section instead of a separate harness. Its Comet-off arm now measures the row reader, so name it that, add an arm with Comet on and native execution off (on-heap enabled, so Comet loads and the read fuses), and check each arm's reader in its plan. Add wide relations of 100, 200 and 1500 columns read through the row readers. --- .../CometCacheRowReaderBenchmark.scala | 207 ------------------ .../CometInMemoryCacheBenchmark.scala | 187 ++++++++++++---- 2 files changed, 145 insertions(+), 249 deletions(-) delete mode 100644 spark/src/test/scala/org/apache/spark/sql/benchmark/CometCacheRowReaderBenchmark.scala diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometCacheRowReaderBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometCacheRowReaderBenchmark.scala deleted file mode 100644 index 4473b8cf073..00000000000 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometCacheRowReaderBenchmark.scala +++ /dev/null @@ -1,207 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you under the Apache License, Version 2.0 (the - * "License"); you may not use this file except in compliance - * with the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, - * software distributed under the License is distributed on an - * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY - * KIND, either express or implied. See the License for the - * specific language governing permissions and limitations - * under the License. - */ - -package org.apache.spark.sql.benchmark - -import java.nio.charset.StandardCharsets - -import org.apache.spark.benchmark.BenchmarkBase -import org.apache.spark.sql.{DataFrame, Row, SparkSession} -import org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer -import org.apache.spark.sql.execution.ColumnarToRowExec -import org.apache.spark.sql.execution.columnar.{CometInMemoryRelationHelper, DefaultCachedBatch, DefaultCachedBatchSerializer, InMemoryRelation, InMemoryTableScanExec} -import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.storage.StorageLevel - -import org.apache.comet.{CometConf, CometSparkSessionExtensions} - -/** - * Compare Spark consumers of Comet and Spark caches (issue #5485). - * - * Arguments: [spark|comet|comet-row|all] [rows] [iterations] [all|mixed|numeric]. Run one format - * per JVM in alternating order. comet requires the fused columnar reader; comet-row disables - * vectorized cache reading to isolate the row iterator. Cache creation and validation are outside - * timing. - */ -object CometCacheRowReaderBenchmark extends BenchmarkBase { - private val warmups = 5 - - override def runBenchmarkSuite(args: Array[String]): Unit = { - require(args.length <= 4, "Expected format, rows, iterations, schema") - val format = args.headOption.getOrElse("all") - val rows = args.lift(1).map(_.toLong).getOrElse(5000000L) - val iterations = args.lift(2).map(_.toInt).getOrElse(15) - val schema = args.lift(3).getOrElse("all") - require(Set("all", "spark", "comet", "comet-row").contains(format)) - require(Set("all", "mixed", "numeric").contains(schema)) - require(rows > 0 && iterations > 0) - - emit("CACHE_SAMPLE,format,schema,query,rows,iteration,elapsed_ns") - val formats = - if (format == "all") Seq("spark", "comet", "comet-row") else Seq(format) - val schemas = if (schema == "all") Seq("mixed", "numeric") else Seq(schema) - formats.foreach { name => - CometInMemoryRelationHelper.clearSerializer() - SparkSession.clearActiveSession() - SparkSession.clearDefaultSession() - val serializer = if (name == "spark") { - classOf[DefaultCachedBatchSerializer].getName - } else { - classOf[ArrowCachedBatchSerializer].getName - } - val spark = SparkSession - .builder() - .master("local[1]") - .appName(getClass.getSimpleName) - .config("spark.ui.enabled", "false") - .config("spark.sql.cache.serializer", serializer) - .config("spark.sql.shuffle.partitions", "1") - .config("spark.sql.inMemoryColumnarStorage.batchSize", "10000") - .config("spark.sql.inMemoryColumnarStorage.compressed", "true") - .config("spark.io.compression.codec", "lz4") - .config(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key, "false") - .config(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key, "true") - .config(SQLConf.CACHE_VECTORIZED_READER_ENABLED.key, (name != "comet-row").toString) - .config(SQLConf.CODEGEN_FACTORY_MODE.key, "CODEGEN_ONLY") - .config(CometConf.COMET_ENABLED.key, "true") - .config(CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key, "true") - .config(CometConf.COMET_EXEC_ENABLED.key, "false") - .config(CometConf.COMET_SHUFFLE_ENABLED.key, "false") - .withExtensions(new CometSparkSessionExtensions) - .getOrCreate() - spark.sparkContext.setLogLevel("WARN") - try { - emit(s"CACHE_ENV,$name,Spark=${spark.version},Java=${System.getProperty("java.version")}") - schemas.foreach(runSchema(spark, name, _, rows, iterations, serializer)) - } finally { - spark.stop() - SparkSession.clearActiveSession() - SparkSession.clearDefaultSession() - CometInMemoryRelationHelper.clearSerializer() - } - } - } - - private def runSchema( - spark: SparkSession, - format: String, - schema: String, - rows: Long, - iterations: Int, - serializer: String): Unit = { - val mixed = schema == "mixed" - val first = Seq("id", "id % 1000 AS k", "id + 1 AS v") - val rest = if (mixed) { - Seq( - "concat('str_a_', cast(id % 100000 as string)) AS s1", - "concat('str_b_', cast(id % 7919 as string)) AS s2", - "concat('str_c_', cast(id as string)) AS s3") - } else { - Seq("id % 100000 AS n1", "id % 7919 AS n2", "id * 3 AS n3") - } - val source = spark.range(0, rows, 1, 16).selectExpr((first ++ rest): _*) - val columns = source.columns.toSeq - val three = if (mixed) Seq("id", "s1", "s2") else columns.take(3) - val projections = Seq("count" -> Seq.empty[String], "long" -> Seq("id")) ++ - (if (mixed) Seq("string" -> Seq("s1")) else Seq.empty) ++ - Seq("three" -> three, "all" -> columns) - def expressions(selected: Seq[String]): Seq[String] = { - if (selected.isEmpty) Seq("count(*)") - else - selected.map { name => - if (name.startsWith("s")) s"sum(length($name))" else s"sum($name)" - } - } - // Obtain the expected values before the relation is cached, using Spark's ordinary row plan. - val expected = projections.map { case (_, selected) => - source.selectExpr(expressions(selected): _*).collect() - } - val cached = source.persist(StorageLevel.MEMORY_ONLY) - try { - assert(cached.count() == rows) - val relation = cached.queryExecution.withCachedData.collectFirst { - case relation: InMemoryRelation => relation - }.get - val builder = relation.cacheBuilder - assert(builder.serializer.getClass.getName == serializer) - val batches = builder.cachedColumnBuffers - val batchSummary = batches - .map { batch => - // Spark's sizeInBytes comes from statistics; measure its encoded column buffers. - val bytes = batch match { - case b: DefaultCachedBatch => b.buffers.map(_.length.toLong).sum - case _ => batch.sizeInBytes - } - (batch.getClass.getSimpleName, batch.numRows.toLong, bytes) - } - .collect() - val expectedClass = if (format == "spark") "DefaultCachedBatch" else "CometCachedBatch" - assert(batchSummary.forall(_._1 == expectedClass), "Wrong cached payload format") - assert(batchSummary.map(_._2).sum == rows) - val storage = spark.sparkContext.getRDDStorageInfo.find(_.id == batches.id).get - assert(storage.numCachedPartitions == batches.getNumPartitions && storage.diskSize == 0) - emit( - s"CACHE_STORAGE,$format,$schema,${batchSummary.length}," + - s"${batchSummary.map(_._3).sum},${storage.memSize}") - - projections.zip(expected).foreach { case ((name, selected), answer) => - val query = cached.selectExpr(expressions(selected): _*) - val plan = query.queryExecution.executedPlan - val scans = plan.collect { case scan: InMemoryTableScanExec => scan } - assert(scans.size == 1, s"Expected one Spark cache scan:\n$plan") - val scan = scans.head - assert(scan.attributes.map(_.name).toSet == selected.toSet, s"Wrong projection:\n$plan") - val columnar = plan.exists(_.isInstanceOf[ColumnarToRowExec]) - if (format == "comet") { - assert(columnar, s"Expected the fused columnar cache reader:\n$plan") - } else { - assert(!columnar, s"Expected the cache row reader:\n$plan") - } - assert(!plan.exists(_.getClass.getName.startsWith("org.apache.spark.sql.comet."))) - val reader = if (columnar) "columnar" else "row" - emit(s"CACHE_PLAN,$format,$schema,$name,columns=${selected.size},reader=$reader\n$plan") - runQuery(query, answer, format, schema, name, rows, iterations) - } - } finally cached.unpersist(blocking = true) - } - - private def runQuery( - query: DataFrame, - expected: Array[Row], - format: String, - schema: String, - name: String, - rows: Long, - iterations: Int): Unit = { - (0 until warmups).foreach { _ => assert(query.collect().sameElements(expected)) } - (0 until iterations).foreach { i => - val start = System.nanoTime() - val actual = query.collect() - val elapsed = System.nanoTime() - start - assert(actual.sameElements(expected), s"Wrong result for $format/$schema/$name") - emit(s"CACHE_SAMPLE,$format,$schema,$name,$rows,$i,$elapsed") - } - emit(s"CACHE_RESULT,$format,$schema,$name,${expected.mkString(";")}") - } - - private def emit(line: String): Unit = { - println(line) - output.foreach(_.write((line + "\n").getBytes(StandardCharsets.UTF_8))) - } -} diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometInMemoryCacheBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometInMemoryCacheBenchmark.scala index f33a53af299..9317b5b6bf6 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometInMemoryCacheBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometInMemoryCacheBenchmark.scala @@ -27,6 +27,7 @@ import org.apache.spark.sql.SparkSession import org.apache.spark.sql.catalyst.expressions.AttributeReference import org.apache.spark.sql.comet.CometInMemoryTableScanExec import org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer +import org.apache.spark.sql.execution.ColumnarToRowExec import org.apache.spark.sql.execution.columnar.{CometInMemoryRelationHelper, DefaultCachedBatchSerializer, InMemoryRelation, InMemoryTableScanExec} import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} @@ -234,6 +235,7 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { runCodecBenchmark(flatRelation) runSparkOperatorBenchmark(flatRelation) } + runWideSparkOperatorBenchmark() } /** @@ -339,16 +341,28 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { /** * Reads that feed Spark operators rather than Comet ones, against Spark's own cache format. * - * Comet is off in every case, so this measures Spark consuming the cached data: the shape where - * Comet's format has something to lose, and the reason the feature is off by default. Both - * formats are cached from the same relation, one copy at a time as in runCodecBenchmark, and - * each case checks which serializer cached the relation it reads. + * Native execution is off in every case, so this measures Spark consuming the cached data: the + * shape where Comet's format has something to lose, and the reason the feature is off by + * default. Comet's format is read two ways. With Comet off, every Spark operator reads rows + * from the cache scan's row reader. With Comet on, a generated Spark operator instead reads the + * cached vectors through a ColumnarToRowExec fused into its generated code, which is how these + * reads run when Comet is enabled without native execution. With native execution on, Spark + * operators above the cache read CometInMemoryTableScan's batches through CometColumnarToRow + * instead, which this does not measure. + * + * Both formats are cached from the same relation, one copy at a time as in runCodecBenchmark, + * and each case checks which serializer cached the relation it reads and which reader it uses. */ private def runSparkOperatorBenchmark(relation: CachedRelation): Unit = { val view = s"${relation.table}_spark_operators" - val formats = Seq( - "Spark's cache format" -> classOf[DefaultCachedBatchSerializer].getName, - "Comet's cache format" -> classOf[ArrowCachedBatchSerializer].getName) + val reads = Seq( + SparkOperatorRead("Spark's cache format", sparkSerializer, sparkOperatorConf), + SparkOperatorRead("Comet's cache format, row reader", cometSerializer, sparkOperatorConf), + SparkOperatorRead( + "Comet's cache format, fused reader", + cometSerializer, + fusedReaderConf, + fused = true)) spark.catalog.clearCache() withTempTable(view) { @@ -356,19 +370,7 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { .sql(s"SELECT ${relation.columns.mkString(", ")} FROM ${relation.source}") .createOrReplaceTempView(view) - var cachedBy: String = null - def cacheBy(serializer: String): Unit = if (cachedBy != serializer) { - spark.catalog.uncacheTable(view) - cachedBy = null - withCacheSerializer(serializer) { - withSQLConf(sparkOperatorConf: _*) { - spark.catalog.cacheTable(view) - spark.table(view).count() - } - } - cachedBy = serializer - } - + val cache = new OneCachedCopy(view) Seq( ("row count only (0 of 6 columns)", s"SELECT count(*) FROM $view", 0), ("narrow projection (1 of 6 columns)", s"SELECT count(k) FROM $view", 1), @@ -381,23 +383,7 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { s"in-memory cache read by Spark operators, $label", relation.rows, output = output) - formats.foreach { case (name, serializer) => - var verified = false - // Re-caching in this case's format is setup, so it is outside the timer, and it only - // happens on the case's first call, which is a warmup iteration. - benchmark.addTimerCase(name) { timer => - cacheBy(serializer) - withSQLConf(sparkOperatorConf: _*) { - if (!verified) { - verifySparkOperatorRead(query, scanned, serializer) - verified = true - } - timer.startTiming() - spark.sql(query).noop() - timer.stopTiming() - } - } - } + reads.foreach(addSparkOperatorCase(benchmark, cache, query, scanned, _)) benchmark.run() } @@ -405,6 +391,97 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { } } + /** + * Spark operators reading every column of relations wider than the six-column one, at widths + * where a row reader that writes every column in one generated method would stop being JIT + * compiled (past about a hundred columns) or fail to compile at all (about 1500). + * + * Each case writes the relation to the noop sink, which takes rows straight from the cache + * scan, so both formats are read by their serializers' row readers. Past + * spark.sql.codegen.maxFields (100 by default) that is how any Spark operator reads a cached + * relation, because the scan stops offering columnar output and nothing can fuse with it. Every + * width holds about the same number of values, so the cases differ in width rather than in data + * volume. + */ + private def runWideSparkOperatorBenchmark(): Unit = { + val reads = Seq( + SparkOperatorRead("Spark's cache format", sparkSerializer, sparkOperatorConf), + SparkOperatorRead("Comet's cache format, row reader", cometSerializer, sparkOperatorConf)) + + Seq(100, 200, 1500).foreach { width => + val rows = 20 * 1000 * 1000 / width + val view = s"comet_cache_bench_wide_$width" + spark.catalog.clearCache() + withTempTable(view) { + spark + .range(0, rows, 1, 4) + .selectExpr((0 until width).map(i => + s"if(id % 8 = ${i % 8}, null, id + $i) AS c$i"): _*) + .createOrReplaceTempView(view) + + val cache = new OneCachedCopy(view) + val benchmark = new Benchmark( + s"in-memory cache read by Spark operators, all $width columns", + rows, + output = output) + reads.foreach(addSparkOperatorCase(benchmark, cache, s"SELECT * FROM $view", width, _)) + benchmark.run() + + spark.catalog.uncacheTable(view) + } + } + } + + /** How Spark operators read a cached relation: its format and the settings of the read. */ + private case class SparkOperatorRead( + name: String, + serializer: String, + conf: Seq[(String, String)], + fused: Boolean = false) + + /** + * Holds one cached copy of `view` at a time. Two copies could not coexist anyway: the cache + * manager keys on the plan rather than the name, so a second one would find the first. + */ + private class OneCachedCopy(view: String) { + private var cachedBy: String = _ + + def cacheBy(serializer: String): Unit = if (cachedBy != serializer) { + spark.catalog.uncacheTable(view) + cachedBy = null + withCacheSerializer(serializer) { + withSQLConf(sparkOperatorConf: _*) { + spark.catalog.cacheTable(view) + spark.table(view).count() + } + } + cachedBy = serializer + } + } + + private def addSparkOperatorCase( + benchmark: Benchmark, + cache: OneCachedCopy, + query: String, + scanned: Int, + read: SparkOperatorRead): Unit = { + var verified = false + // Re-caching in this case's format is setup, so it is outside the timer, and it only happens + // on the case's first call, which is a warmup iteration. + benchmark.addTimerCase(read.name) { timer => + cache.cacheBy(read.serializer) + withSQLConf(read.conf: _*) { + if (!verified) { + verifySparkOperatorRead(query, scanned, read) + verified = true + } + timer.startTiming() + spark.sql(query).noop() + timer.stopTiming() + } + } + } + // spark.sql.cache.serializer is static, and InMemoryRelation memoizes the serializer it names // for the life of the JVM. It looks the name up in the active session's conf when a relation is // cached, though, so setting it there directly and clearing the memoized instance around one @@ -422,10 +499,14 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { } } - // Pins what a Spark-operator case claims: no Comet operator anywhere, and one cache scan that - // reads the columns its label counts from a relation the named serializer cached. The last is - // what catches both formats silently reading one copy. - private def verifySparkOperatorRead(query: String, scanned: Int, serializer: String): Unit = { + // Pins what a Spark-operator case claims: no Comet operator anywhere, one cache scan that + // reads the columns its label counts from a relation the named serializer cached, and the + // reader the case names. The serializer check is what catches both formats silently reading one + // copy, and the reader check a fused case silently measuring the row reader. + private def verifySparkOperatorRead( + query: String, + scanned: Int, + read: SparkOperatorRead): Unit = { val executed = spark.sql(query).queryExecution.executedPlan val plan = executed.toString() assert(executed.find(_.nodeName.startsWith("Comet")).isEmpty, s"Expected no Comet:\n$plan") @@ -435,7 +516,15 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { scans.head.attributes.length == scanned, s"Expected the scan to read $scanned columns:\n$plan") val actual = scans.head.relation.cacheBuilder.serializer.getClass.getName - assert(actual == serializer, s"Expected a relation cached by $serializer, not $actual") + assert( + actual == read.serializer, + s"Expected a relation cached by ${read.serializer}, not $actual") + val transitions = executed.collect { + case c: ColumnarToRowExec if c.exists(_.isInstanceOf[InMemoryTableScanExec]) => c + } + assert( + transitions.length == (if (read.fused) 1 else 0), + s"Expected the ${if (read.fused) "fused" else "row"} reader:\n$plan") } /** What the cached relation behind `view` occupies, summed over its batches as written. */ @@ -602,4 +691,18 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { CometConf.COMET_ENABLED.key -> "false", CometConf.COMET_EXEC_ENABLED.key -> "false", "spark.sql.inMemoryColumnarStorage.batchSize" -> "10000") + + // Comet on but native execution and shuffle off, so the operators are still Spark's and the + // generated ones read Comet's cache through the fused reader. On-heap mode has to be enabled + // explicitly, as in cacheConf, or Comet stays unloaded and the reader is never fused. + private val fusedReaderConf: Seq[(String, String)] = Seq( + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "false", + CometConf.COMET_SHUFFLE_ENABLED.key -> "false", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", + "spark.comet.exec.onHeap.enabled" -> "true", + "spark.sql.inMemoryColumnarStorage.batchSize" -> "10000") + + private val sparkSerializer = classOf[DefaultCachedBatchSerializer].getName + private val cometSerializer = classOf[ArrowCachedBatchSerializer].getName } From 384ea00a814c1d0ffeec72f58a538c0e3d351590 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Tue, 29 Sep 2026 02:23:03 +0800 Subject: [PATCH 12/12] docs: describe how Spark operators read Comet's cache format The Limitations section still said Spark-operator reads were slower for reasons not yet established, with a table measured through the row path this PR replaced. Explain the three ways a Spark operator reads a relation cached in Comet's format (under native execution, through the fused ColumnarToRowExec and when it applies, and through the row reader otherwise) and replace the table with one that measures both readers. --- .../user-guide/latest/in-memory-cache.md | 55 ++++++++++++++----- 1 file changed, 41 insertions(+), 14 deletions(-) diff --git a/docs/source/user-guide/latest/in-memory-cache.md b/docs/source/user-guide/latest/in-memory-cache.md index c37fa2dc32d..57d09d900eb 100644 --- a/docs/source/user-guide/latest/in-memory-cache.md +++ b/docs/source/user-guide/latest/in-memory-cache.md @@ -53,7 +53,8 @@ relation whose format could change mid-session could not be read back reliably. codec is a runtime config, but each batch records the codec it was written with, so data cached under one setting stays readable after the setting changes. Turning `spark.comet.exec.inMemoryCache.enabled` off at runtime only sends cached scans back to Spark's -execution path; the cached data stays readable either way. +execution path, where Spark operators read them through the row reader described under +[Limitations](#limitations); the cached data stays readable either way. ## Storage format @@ -175,19 +176,45 @@ registrator. ## Limitations -Reads that feed **Spark** operators rather than Comet ones are slower than Spark's own cache -format, and the narrower the read, the wider the gap. Measured by the same benchmark over the same -5M-row relation, with Comet off so that Spark operators consume the cached data: - -| Read shape | Spark's cache format | Comet's cache format | Slowdown | -| ----------------------- | -------------------: | -------------------: | -------: | -| Row count only (0 of 6) | 35 ms | 183 ms | 5.2x | -| 1 of 6 columns | 54 ms | 257 ms | 4.8x | -| 3 of 6 columns | 98 ms | 331 ms | 3.4x | -| 6 of 6 columns | 410 ms | 623 ms | 1.5x | - -This is why the feature is off by default. The cause is not yet established; -[#5485](https://github.com/apache/datafusion-comet/issues/5485) tracks it. +Reads that feed **Spark** operators rather than Comet ones can be slower than Spark's own cache +format, because every cached batch is decoded from Arrow before Spark reads it. How a Spark +operator reads a relation cached in Comet's format depends on the scan below it: + +- With native execution enabled (`spark.comet.exec.enabled=true`), the cache is scanned by + `CometInMemoryTableScan`, and a Spark operator above it reads the scan's batches through + `CometColumnarToRow`, as it would above any other Comet operator. +- With Comet enabled but native execution disabled, the cache is scanned by Spark's + `InMemoryTableScanExec`. When the operator directly above the scan takes part in whole-stage code + generation, as filters, projections and aggregates do, Comet puts Spark's `ColumnarToRowExec` + between the two, and the generated code reads the cached Arrow vectors directly, with no + intermediate row. The plan shows this as a `ColumnarToRow` above the `InMemoryTableScan`. It + needs `spark.sql.inMemoryColumnarStorage.enableVectorizedReader` (on by default) and whole-stage + code generation, and applies to relations of at most `spark.sql.codegen.maxFields` fields (100 by + default, counting nested fields), beyond which Spark reads a cached relation only as rows. It is + not applied in plan-only mode (`spark.comet.explain.planOnly.enabled`), where Spark executes its + own plan unchanged. +- Otherwise the scan's row reader decodes each batch and writes its rows into one reused + `UnsafeRow`. That covers Comet or `spark.comet.exec.inMemoryCache.enabled` turned off at runtime, + and operators that do not take part in code generation, such as exchanges and limits, or a query + that returns the cached rows as they are. + +Measured by the same benchmark over the same 5M-row relation, with native execution off so that +Spark operators consume the cached data, Comet disabled for the row reader and enabled for the fused +reader (Apple M4, JDK 17, Spark 4.1; the average of two runs): + +| Read shape | Spark's cache format | Comet's format, row reader | Comet's format, fused reader | +| ----------------------- | -------------------: | -------------------------: | ---------------------------: | +| Row count only (0 of 6) | 63 ms | 57 ms | 35 ms | +| 1 of 6 columns | 63 ms | 77 ms | 54 ms | +| 3 of 6 columns | 113 ms | 176 ms | 133 ms | +| 6 of 6 columns | 306 ms | 500 ms | 334 ms | + +The fused reader is faster than Spark's own format for the narrowest reads and within 20% of it for +the others. The row reader takes up to 1.6 times as long, and it is the only reader for relations +wider than `spark.sql.codegen.maxFields`: reading every column of relations of 100, 200 and 1500 +nullable `bigint` columns took 2.2 to 2.5 times as long as from Spark's format. These gaps are why +the feature is still off by default; +[#5485](https://github.com/apache/datafusion-comet/issues/5485) tracks them. Comet's serializer exists because Spark's own Arrow cache format ([SPARK-57268](https://issues.apache.org/jira/browse/SPARK-57268)) is only available from Spark