From ec92ee5130d77b62626be2f48bf1994b7c3954ac Mon Sep 17 00:00:00 2001 From: peterxcli Date: Fri, 11 Sep 2026 20:27:42 +0800 Subject: [PATCH 1/6] 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 2/6] 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 3/6] 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 4/6] 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 28bb732598bf093539404d4c822c4f09f61a1a2d Mon Sep 17 00:00:00 2001 From: peterxcli Date: Sat, 12 Sep 2026 02:30:23 +0800 Subject: [PATCH 5/6] perf: delta-encode compressible cached longs --- .../arrow/ArrowCachedBatchSerializer.scala | 38 ++++++-- .../apache/spark/sql/comet/util/Utils.scala | 90 +++++++++++++++---- .../comet/exec/CometInMemoryCacheSuite.scala | 58 ++++++++++++ .../CometCacheRowReaderBenchmark.scala | 17 ++-- .../arrow/CometCachedBatchHelper.scala | 3 + 5 files changed, 176 insertions(+), 30 deletions(-) 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 c9f05289c30..59013ff9ead 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 @@ -23,7 +23,8 @@ import java.lang.{Boolean => JBoolean, Byte => JByte, Double => JDouble, Float = import scala.util.control.NonFatal -import org.apache.spark.TaskContext +import org.apache.spark.{SparkEnv, TaskContext} +import org.apache.spark.io.CompressionCodec 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} @@ -39,6 +40,7 @@ import org.apache.spark.unsafe.types.UTF8String import org.apache.spark.util.io.ChunkedByteBuffer import org.apache.comet.CometArrowAllocator +import org.apache.comet.vector.CometVector /** * Cached batch format used when Comet writes Spark in-memory cache data. @@ -47,13 +49,16 @@ import org.apache.comet.CometArrowAllocator * by `Utils.serializeBatchColumns`. Storing columns separately is what lets a scan decode only * the ones it projected; a single stream covering the whole batch would have to be inflated in * full before any projection could be applied. The cache manager still owns storage and eviction; - * this class only changes the cached payload. + * this class only changes the cached payload. `deltaEncoded` marks numeric streams whose values + * need a prefix sum after decoding; the validity bits and logical schema remain the same as the + * source column. */ private case class CometCachedBatch( override val numRows: Int, override val sizeInBytes: Long, override val stats: InternalRow, - columns: Array[ChunkedByteBuffer]) + columns: Array[ChunkedByteBuffer], + deltaEncoded: Array[Boolean]) extends SimpleMetricsCachedBatch /** @@ -351,7 +356,7 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { val (lower, upper, nulls) = gatherColumnStats(batch, attrs) val numRows = batch.numRows() - val columns = if (Utils.isArrowBacked(batch)) { + val (columns, deltaEncoded) = if (Utils.isArrowBacked(batch)) { Utils.serializeBatchColumns(batch) } else { val arrowBatch = @@ -365,7 +370,8 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { numRows = numRows, sizeInBytes = columnSizes.sum, stats = statsRow(lower, upper, nulls, numRows, columnSizes), - columns = columns) + columns = columns, + deltaEncoded = deltaEncoded) } } @@ -463,6 +469,8 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { val indices = selectedIndices(cacheAttributes, selectedAttributes) input.mapPartitions { it => + // Codec factories are reusable; each selected column still gets its own input stream. + lazy val codec = CompressionCodec.createCodec(SparkEnv.get.conf) // A ColumnReaders closes its readers (releasing the vectors they are holding) only when the // batch it produced has been consumed. A consumer that stops early -- LIMIT, take(), or a // cancelled task -- leaves the readers for the batch in flight open, so close them on task @@ -489,7 +497,11 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { // Nothing to decode: the row count is the whole answer, and it is already here. Iterator.single(new ColumnarBatch(Array.empty[ColumnVector], cb.numRows)) } else { - val readers = new ColumnReaders(indices.map(i => cb.columns(i)), cb.numRows) + val readers = new ColumnReaders( + indices.map(i => cb.columns(i)), + indices.map(i => cb.deltaEncoded(i)), + cb.numRows, + codec) current = readers readers.batches } @@ -507,7 +519,11 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { // decoded vectors stay owned by their readers: closing them releases the batch, which is why // this yields a single-element iterator that closes on exhaustion, matching what // ArrowReaderIterator did when the payload was one stream. - private class ColumnReaders(buffers: Array[ChunkedByteBuffer], numRows: Int) { + private class ColumnReaders( + buffers: Array[ChunkedByteBuffer], + deltaEncoded: Array[Boolean], + numRows: Int, + codec: CompressionCodec) { // decodeBatches opens a reader and eagerly decodes its first batch, so it allocates. If a // later column throws, the readers already opened here are unreachable: the task-completion // listener cannot release them because `current` is only assigned once this constructor @@ -517,7 +533,7 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { var i = 0 try { while (i < buffers.length) { - opened(i) = Utils.decodeBatches(buffers(i), "CometCache") + opened(i) = Utils.decodeBatches(buffers(i), "CometCache", codec) i += 1 } } catch { @@ -569,7 +585,11 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { throw new IllegalStateException( s"Cached column stream $i decoded ${decoded.numRows()} rows, expected $numRows") } - columns(i) = decoded.column(0) + val column = decoded.column(0) + if (deltaEncoded(i)) { + Utils.decodeDeltaLongs(column.asInstanceOf[CometVector].getValueVector) + } + columns(i) = column i += 1 } new ColumnarBatch(columns, numRows) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala index d70fdab35e6..7b328aa5806 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala @@ -42,6 +42,7 @@ import org.apache.spark.sql.comet.execution.arrow.{ArrowReaderIterator, Constant import org.apache.spark.sql.execution.vectorized.ConstantColumnVector import org.apache.spark.sql.types._ import org.apache.spark.sql.vectorized.ColumnarBatch +import org.apache.spark.unsafe.Platform import org.apache.spark.util.io.{ChunkedByteBuffer, ChunkedByteBufferOutputStream} import org.apache.comet.Constants.COMET_CONF_DIR_ENV @@ -284,30 +285,82 @@ object Utils extends CometTypeShim with Logging { * [[serializeBatches]] writes one stream covering every column, so a reader has to inflate all * of them before it can project. Comet's in-memory cache stores columns separately instead, so * a scan decodes only the ones it selected. Each stream is self-contained, including its schema - * and any dictionaries the column needs. + * and any dictionaries the column needs. Returns the streams and flags identifying long columns + * stored as deltas. Their readers restore the original values after Arrow decoding, before + * exposing the vectors to consumers. * * The row count is not recoverable from the result when `batch` has no columns, so callers keep * it alongside. As with [[serializeBatches]], the batch's vectors are cleared once written. */ - def serializeBatchColumns(batch: ColumnarBatch): Array[ChunkedByteBuffer] = { + def serializeBatchColumns(batch: ColumnarBatch): (Array[ChunkedByteBuffer], Array[Boolean]) = { val codec = CompressionCodec.createCodec(SparkEnv.get.conf) // Each column is written with the provider it was decoded with, not the batch's first one: // columns decoded from separate streams have independent dictionary ID namespaces. - getBatchFieldVectorsWithProviders(batch).map { case (fieldVector, providerOpt) => - val provider = providerOpt.getOrElse(new CDataDictionaryProvider) - val cbbos = new ChunkedByteBufferOutputStream(1024 * 1024, ByteBuffer.allocate) - val out = new DataOutputStream(codec.compressedOutputStream(cbbos)) + getBatchFieldVectorsWithProviders(batch) + .map { case (fieldVector, providerOpt) => + val provider = providerOpt.getOrElse(new CDataDictionaryProvider) + def writeColumn(vector: FieldVector): ChunkedByteBuffer = { + val cbbos = new ChunkedByteBufferOutputStream(1024 * 1024, ByteBuffer.allocate) + val out = new DataOutputStream(codec.compressedOutputStream(cbbos)) + val root = new VectorSchemaRoot(Seq(vector).asJava) + val writer = new ArrowStreamWriter(root, provider, Channels.newChannel(out)) + writer.start() + writer.writeBatch() + root.clear() + writer.close() + cbbos.toChunkedByteBuffer + } - val root = new VectorSchemaRoot(Seq(fieldVector).asJava) - val writer = new ArrowStreamWriter(root, provider, Channels.newChannel(out)) - writer.start() - writer.writeBatch() - root.clear() - writer.close() + fieldVector match { + case longs: BigIntVector + if longs.getField.getDictionary == null && longs.getValueCount > 0 => + // Copy before writing: a column may be borrowed from a cache scan, and writeColumn + // clears its input. The validity bits and logical schema stay unchanged. + val deltas = new BigIntVector(longs.getField, longs.getAllocator) + try { + val count = longs.getValueCount + deltas.allocateNew(count) + deltas.setValueCount(count) + deltas.getValidityBuffer.setBytes( + 0, + longs.getValidityBuffer, + 0, + BitVectorHelper.getValidityBufferSize(count)) + val source = longs.getDataBuffer.memoryAddress() + val target = deltas.getDataBuffer.memoryAddress() + var previous = 0L + var i = 0 + while (i < count) { + val value = Platform.getLong(null, source + i * 8L) + Platform.putLong(null, target + i * 8L, value - previous) + previous = value + i += 1 + } + val plain = writeColumn(longs) + val encoded = writeColumn(deltas) + // Irregular values can get larger and slower after delta encoding. Require a + // substantial size reduction to pay for reconstructing the values on each read. + if (encoded.size < plain.size * 3 / 4) (encoded, true) else (plain, false) + } finally deltas.close() + case _ => (writeColumn(fieldVector), false) + } + } + .toArray + .unzip + } - cbbos.toChunkedByteBuffer - }.toArray + private[sql] def decodeDeltaLongs(vector: ValueVector): Unit = { + require(vector.isInstanceOf[BigIntVector], "Delta-encoded cache column must contain longs") + val address = vector.getDataBuffer.memoryAddress() + var previous = 0L + var i = 0 + while (i < vector.getValueCount) { + val value = Platform.getLong(null, address + i * 8L) + previous + Platform.putLong(null, address + i * 8L, value) + previous = value + i += 1 + } } /** @@ -337,12 +390,19 @@ object Utils extends CometTypeShim with Logging { * an iterator of ColumnarBatch */ def decodeBatches(bytes: ChunkedByteBuffer, source: String): Iterator[ColumnarBatch] = { + if (bytes.size == 0) Iterator.empty + else decodeBatches(bytes, source, CompressionCodec.createCodec(SparkEnv.get.conf)) + } + + def decodeBatches( + bytes: ChunkedByteBuffer, + source: String, + codec: CompressionCodec): Iterator[ColumnarBatch] = { if (bytes.size == 0) { return Iterator.empty } // use Spark's compression codec (LZ4 by default) and not Comet's compression - val codec = CompressionCodec.createCodec(SparkEnv.get.conf) val cbbis = bytes.toInputStream() val ins = new DataInputStream(codec.compressedInputStream(cbbis)) // batches are in Arrow IPC format 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..52be0409527 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -1675,6 +1675,64 @@ class CometInMemoryCacheSuite extends CometTestBase { } } + test("Comet in-memory cache preserves delta-encoded longs across every reader") { + val random = new java.util.Random(5485) + val rows = (0 until 4096).map { i => + // Cross the signed-long boundary; delta reconstruction must preserve wrapping arithmetic. + val value = Long.MaxValue - 2048 + i + Row(value, if (i % 7 == 0) null else value, random.nextLong(), s"value_${i % 11}") + } + val schema = new StructType() + .add("seq", LongType, nullable = false) + .add("nullable", LongType, nullable = true) + .add("random", LongType, nullable = false) + .add("text", StringType, nullable = false) + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", + "spark.comet.sparkToColumnar.enabled" -> "true") { + spark.catalog.clearCache() + val cached = spark.createDataFrame(spark.sparkContext.parallelize(rows, 2), schema).cache() + try { + assert(cached.count() == rows.length) + val relation = + spark.sharedState.cacheManager.lookupCachedData(cached).get.cachedRepresentation + val batches = relation.cacheBuilder.cachedColumnBuffers.collect() + assert(batches.forall(b => CometCachedBatchHelper.columnsAreDeltaEncoded(b)(0))) + assert(batches.forall(b => !CometCachedBatchHelper.columnsAreDeltaEncoded(b)(2))) + + val expected = + rows.map(r => Row(r.getLong(2), r.getLong(0), r.get(1), r.getLong(0), r.get(3))) + for ((native, vectorized) <- Seq((true, true), (false, true), (false, false))) { + withSQLConf( + CometConf.COMET_EXEC_ENABLED.key -> native.toString, + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> native.toString, + SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> vectorized.toString) { + val projected = + cached.selectExpr("random", "seq", "nullable", "seq AS repeated", "text") + checkAnswer(projected, expected) + val plan = projected.queryExecution.executedPlan + if (native) assert(plan.exists(_.isInstanceOf[CometInMemoryTableScanExec])) + else { + assert(plan.exists(_.isInstanceOf[InMemoryTableScanExec])) + assert(plan.exists(_.isInstanceOf[ColumnarToRowExec]) == vectorized) + } + } + } + + // Re-caching decoded vectors must neither encode them twice nor change the original cache. + withSQLConf( + CometConf.COMET_EXEC_ENABLED.key -> "false", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "false") { + val recached = cached.union(cached).cache() + try checkAnswer(recached, rows ++ rows) + finally recached.unpersist() + checkAnswer(cached, rows) + } + } finally cached.unpersist() + } + } + test("Comet in-memory cache releases opened readers when a later column fails to decode") { // A cached batch is several independent Arrow streams and decodeBatches opens each eagerly. // If a later column throws, the readers already opened are unreachable: the task-completion 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..db0ce00853e 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,19 +34,20 @@ 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. + * Arguments: [spark|comet|comet-row|all] [rows] [iterations] [all|mixed|numeric] [codec]. 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 override def runBenchmarkSuite(args: Array[String]): Unit = { - require(args.length <= 4, "Expected format, rows, iterations, schema") + require(args.length <= 5, "Expected format, rows, iterations, schema, compression codec") 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") + val codec = args.lift(4).getOrElse("lz4") require(Set("all", "spark", "comet", "comet-row").contains(format)) require(Set("all", "mixed", "numeric").contains(schema)) require(rows > 0 && iterations > 0) @@ -73,7 +74,7 @@ object CometCacheRowReaderBenchmark extends BenchmarkBase { .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("spark.io.compression.codec", codec) .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) @@ -84,7 +85,9 @@ object CometCacheRowReaderBenchmark extends BenchmarkBase { .getOrCreate() spark.sparkContext.setLogLevel("WARN") try { - emit(s"CACHE_ENV,$name,Spark=${spark.version},Java=${System.getProperty("java.version")}") + emit( + s"CACHE_ENV,$name,Spark=${spark.version}," + + s"Java=${System.getProperty("java.version")},codec=$codec") schemas.foreach(runSchema(spark, name, _, rows, iterations, serializer)) } finally { spark.stop() @@ -131,7 +134,9 @@ object CometCacheRowReaderBenchmark extends BenchmarkBase { } val cached = source.persist(StorageLevel.MEMORY_ONLY) try { + val buildStart = System.nanoTime() assert(cached.count() == rows) + emit(s"CACHE_BUILD,$format,$schema,$rows,${System.nanoTime() - buildStart}") val relation = cached.queryExecution.withCachedData.collectFirst { case relation: InMemoryRelation => relation }.get diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometCachedBatchHelper.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometCachedBatchHelper.scala index 5548f6dadbd..24dfedf92c7 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometCachedBatchHelper.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometCachedBatchHelper.scala @@ -48,6 +48,9 @@ object CometCachedBatchHelper { def columnStreamSizes(batch: CachedBatch): Seq[Long] = batch.asInstanceOf[CometCachedBatch].columns.map(_.size).toSeq + def columnsAreDeltaEncoded(batch: CachedBatch): Seq[Boolean] = + batch.asInstanceOf[CometCachedBatch].deltaEncoded.toSeq + /** * Replace one column's stream with bytes that cannot be decoded, in place. * From 2e65a02cba9f9a5979fb6dcaf951829c96b5a039 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Sat, 12 Sep 2026 02:53:05 +0800 Subject: [PATCH 6/6] perf: skip recompressing irregular cached longs --- .../org/apache/spark/sql/comet/util/Utils.scala | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala index 7b328aa5806..bd8a1523ccb 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala @@ -330,18 +330,25 @@ object Utils extends CometTypeShim with Logging { val source = longs.getDataBuffer.memoryAddress() val target = deltas.getDataBuffer.memoryAddress() var previous = 0L + var smallDeltas = 0 var i = 0 while (i < count) { val value = Platform.getLong(null, source + i * 8L) - Platform.putLong(null, target + i * 8L, value - previous) + val delta = value - previous + Platform.putLong(null, target + i * 8L, delta) + if (delta == delta.toInt.toLong) smallDeltas += 1 previous = value i += 1 } val plain = writeColumn(longs) - val encoded = writeColumn(deltas) - // Irregular values can get larger and slower after delta encoding. Require a - // substantial size reduction to pay for reconstructing the values on each read. - if (encoded.size < plain.size * 3 / 4) (encoded, true) else (plain, false) + // ponytail: this cheap filter skips full-width random longs; the size comparison + // below still rejects poorly compressing deltas from narrower distributions. + if (smallDeltas.toLong * 2 < count) (plain, false) + else { + val encoded = writeColumn(deltas) + // Require a substantial size reduction to pay for reconstructing each read. + if (encoded.size < plain.size * 3 / 4) (encoded, true) else (plain, false) + } } finally deltas.close() case _ => (writeColumn(fieldVector), false) }