From b2f2ada90faf267a1a77dcd0b57d1167342c62e5 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 4 Nov 2026 17:36:48 -0700 Subject: [PATCH 1/2] chore: remove the native columnar-to-row implementation Remove CometNativeColumnarToRowExec, NativeColumnarToRowConverter, the Rust converter in columnar_to_row.rs and its JNI entry points. Comet now always converts Arrow batches to Spark rows with CometColumnarToRowExec, which was already the default. spark.comet.exec.columnarToRow.native.enabled is kept as a deprecated no-op, as the versioning policy requires, and the driver logs a warning when it is set. The Spark SQL test diffs are regenerated without the removed operator. --- .github/labeler.yml | 3 - .github/workflows/pr_build_linux.yml | 1 - .github/workflows/pr_build_macos.yml | 1 - dev/diffs/3.4.3.diff | 38 +- dev/diffs/3.5.9.diff | 38 +- dev/diffs/4.0.4.diff | 57 +- dev/diffs/4.1.3.diff | 71 +- dev/diffs/4.2.0.diff | 71 +- .../contributor-guide/plugin_overview.md | 4 +- .../user-guide/latest/migration-guide.md | 13 + .../user-guide/latest/tuning/transitions.md | 3 +- .../latest/understanding-comet-plans.md | 17 +- native/core/src/execution/columnar_to_row.rs | 2726 ----------------- native/core/src/execution/jni_api.rs | 252 +- native/core/src/execution/mod.rs | 1 - .../apache/comet/NativeColumnarToRowInfo.java | 75 - .../scala/org/apache/comet/CometConf.scala | 10 +- .../apache/comet/ExtendedExplainInfo.scala | 4 +- .../main/scala/org/apache/comet/Native.scala | 46 - .../comet/NativeColumnarToRowConverter.scala | 145 - .../apache/comet/rules/CometExecRule.scala | 4 +- ...metPlanAdaptiveDynamicPruningFilters.scala | 1 - .../rules/EliminateRedundantTransitions.scala | 31 +- .../comet/rules/MakeDateEvaluation.scala | 2 +- ...RevertNativeForTransitionHeavyStages.scala | 12 +- .../main/scala/org/apache/spark/Plugins.scala | 12 + .../comet/CometNativeColumnarToRowExec.scala | 260 -- .../apache/comet/CometCodegenFuzzSuite.scala | 4 +- .../org/apache/comet/CometFuzzTestBase.scala | 16 +- .../comet/CometVariantProjectionSuite.scala | 2 - .../apache/comet/CometVariantTypeSuite.scala | 7 - .../exec/CometDirectColumnarToRowSuite.scala | 1 - .../apache/comet/exec/CometExecSuite.scala | 8 +- .../comet/exec/CometInMemoryCacheSuite.scala | 6 +- .../exec/CometNativeColumnarToRowSuite.scala | 550 ---- .../parquet/CometParquetWriterSuite.scala | 6 +- .../EliminateRedundantTransitionsSuite.scala | 3 +- ...tNativeForTransitionHeavyStagesSuite.scala | 24 +- .../org/apache/spark/CometPluginsSuite.scala | 24 + .../sql/benchmark/CometC2RIsolatedBench.scala | 48 +- .../CometColumnarToRowBenchmark.scala | 21 +- .../CometStringFfiImportBenchmark.scala | 52 +- .../spark/sql/comet/CometPlanChecker.scala | 2 +- 43 files changed, 210 insertions(+), 4462 deletions(-) delete mode 100644 native/core/src/execution/columnar_to_row.rs delete mode 100644 spark/src/main/java/org/apache/comet/NativeColumnarToRowInfo.java delete mode 100644 spark/src/main/scala/org/apache/comet/NativeColumnarToRowConverter.scala delete mode 100644 spark/src/main/scala/org/apache/spark/sql/comet/CometNativeColumnarToRowExec.scala delete mode 100644 spark/src/test/scala/org/apache/comet/exec/CometNativeColumnarToRowSuite.scala diff --git a/.github/labeler.yml b/.github/labeler.yml index 83db8b4c6ac..01c9c747dbe 100644 --- a/.github/labeler.yml +++ b/.github/labeler.yml @@ -95,12 +95,9 @@ area:ffi: - changed-files: - any-glob-to-any-file: - 'native/jni-bridge/**' - - 'native/core/src/execution/columnar_to_row.rs' - 'spark/src/main/**/vector/**' - 'common/src/main/**/vector/**' - - 'spark/src/main/scala/org/apache/comet/NativeColumnarToRowConverter.scala' - 'spark/src/main/scala/org/apache/spark/sql/comet/CometColumnarToRowExec.scala' - - 'spark/src/main/scala/org/apache/spark/sql/comet/CometNativeColumnarToRowExec.scala' - 'spark/src/main/scala/org/apache/spark/sql/comet/CometSparkToColumnarExec.scala' area:memory: diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 7f70f81122c..06d05c71799 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -556,7 +556,6 @@ jobs: value: | org.apache.comet.exec.CometShuffleSuite org.apache.comet.exec.CometShuffle4_0Suite - org.apache.comet.exec.CometNativeColumnarToRowSuite org.apache.comet.exec.CometDirectColumnarToRowSuite org.apache.comet.exec.CometNativeShuffleSuite org.apache.comet.exec.CometShuffleInputConversionSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index 44b70c9e2ca..3ca3643191f 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -191,7 +191,6 @@ jobs: value: | org.apache.comet.exec.CometShuffleSuite org.apache.comet.exec.CometShuffle4_0Suite - org.apache.comet.exec.CometNativeColumnarToRowSuite org.apache.comet.exec.CometDirectColumnarToRowSuite org.apache.comet.exec.CometNativeShuffleSuite org.apache.comet.exec.CometShuffleInputConversionSuite diff --git a/dev/diffs/3.4.3.diff b/dev/diffs/3.4.3.diff index fb3d7371a77..4f5600c8659 100644 --- a/dev/diffs/3.4.3.diff +++ b/dev/diffs/3.4.3.diff @@ -1229,14 +1229,14 @@ index 18123a4d6ec..0fe185baa33 100644 test("non-matching optional group") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala -index 75eabcb96f2..c68f4827bf1 100644 +index 75eabcb96f2..abe87ee7903 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala @@ -21,10 +21,11 @@ import scala.collection.mutable.ArrayBuffer import org.apache.spark.sql.catalyst.expressions.SubqueryExpression import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Join, LogicalPlan, Project, Sort, Union} -+import org.apache.spark.sql.comet.{CometColumnarToRowExec, CometNativeColumnarToRowExec, CometNativeScanExec, CometScanExec} ++import org.apache.spark.sql.comet.{CometColumnarToRowExec, CometNativeScanExec, CometScanExec} import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecution} import org.apache.spark.sql.execution.datasources.FileScanRDD @@ -1245,7 +1245,7 @@ index 75eabcb96f2..c68f4827bf1 100644 import org.apache.spark.sql.execution.joins.{BaseJoinExec, BroadcastHashJoinExec, BroadcastNestedLoopJoinExec} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession -@@ -1543,6 +1544,29 @@ class SubquerySuite extends QueryTest +@@ -1543,6 +1544,24 @@ class SubquerySuite extends QueryTest fs.inputRDDs().forall( _.asInstanceOf[FileScanRDD].filePartitions.forall( _.files.forall(_.urlEncodedPath.contains("p=0")))) @@ -1266,16 +1266,11 @@ index 75eabcb96f2..c68f4827bf1 100644 + fs.partitionFilters.exists(ExecSubqueryExpression.hasSubquery) && + fs.inputRDDs().forall( + _.asInstanceOf[FileScanRDD].filePartitions.forall( -+ _.files.forall(_.urlEncodedPath.contains("p=0")))) -+ case CometNativeColumnarToRowExec(fs: CometNativeScanExec) => -+ fs.partitionFilters.exists(ExecSubqueryExpression.hasSubquery) && -+ fs.inputRDDs().forall( -+ _.asInstanceOf[FileScanRDD].filePartitions.forall( + _.files.forall(_.urlEncodedPath.contains("p=0")))) case _ => false }) } -@@ -2108,7 +2132,7 @@ class SubquerySuite extends QueryTest +@@ -2108,7 +2127,7 @@ class SubquerySuite extends QueryTest df.collect() val exchanges = collect(df.queryExecution.executedPlan) { @@ -1704,18 +1699,18 @@ index eec396b2e39..bf3f1c769d6 100644 nums.createOrReplaceTempView("nums") diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/SparkPlanSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/SparkPlanSuite.scala -index b14f4a405f6..90bed10eca9 100644 +index b14f4a405f6..ab7baf434a5 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/SparkPlanSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/SparkPlanSuite.scala @@ -23,6 +23,7 @@ import org.apache.spark.sql.QueryTest import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference} import org.apache.spark.sql.catalyst.plans.logical.Deduplicate -+import org.apache.spark.sql.comet.{CometColumnarToRowExec, CometNativeColumnarToRowExec} ++import org.apache.spark.sql.comet.CometColumnarToRowExec import org.apache.spark.sql.execution.datasources.v2.BatchScanExec import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession -@@ -131,7 +132,11 @@ class SparkPlanSuite extends QueryTest with SharedSparkSession { +@@ -131,7 +132,10 @@ class SparkPlanSuite extends QueryTest with SharedSparkSession { spark.range(1).write.parquet(path.getAbsolutePath) val df = spark.read.parquet(path.getAbsolutePath) val columnarToRowExec = @@ -1723,7 +1718,6 @@ index b14f4a405f6..90bed10eca9 100644 + df.queryExecution.executedPlan.collectFirst { + case p: ColumnarToRowExec => p + case p: CometColumnarToRowExec => p -+ case p: CometNativeColumnarToRowExec => p + }.get try { spark.range(1).foreach { _ => @@ -2825,7 +2819,7 @@ index fdc633f3556..e6990389aa0 100644 } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/sources/BucketedReadSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/sources/BucketedReadSuite.scala -index 266bb343526..50408a27a0a 100644 +index 266bb343526..642208f4b85 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/sources/BucketedReadSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/sources/BucketedReadSuite.scala @@ -19,15 +19,17 @@ package org.apache.spark.sql.sources @@ -2883,7 +2877,7 @@ index 266bb343526..50408a27a0a 100644 val bucketColumnType = bucketedDataFrame.schema.apply(bucketColumnIndex).dataType val rowsWithInvalidBuckets = fileScan.execute().filter(row => { -@@ -451,28 +464,54 @@ abstract class BucketedReadSuite extends QueryTest with SQLTestUtils with Adapti +@@ -451,28 +464,52 @@ abstract class BucketedReadSuite extends QueryTest with SQLTestUtils with Adapti val joinOperator = if (joined.sqlContext.conf.adaptiveExecutionEnabled) { val executedPlan = joined.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec].executedPlan @@ -2914,9 +2908,7 @@ index 266bb343526..50408a27a0a 100644 + case s: SortMergeJoinExec => s + case o => fail(s"expected SortMergeJoinExec, but found\n$o") + } -+ case CometNativeColumnarToRowExec(child) => -+ child.asInstanceOf[CometSortMergeJoinExec].originalPlan match { -+ case s: SortMergeJoinExec => s ++ case s: SortMergeJoinExec => s + case o => fail(s"expected SortMergeJoinExec, but found\n$o") + } + case o => fail(s"expected SortMergeJoinExec, but found\n$o") @@ -2946,7 +2938,7 @@ index 266bb343526..50408a27a0a 100644 s"expected sort in the right child to be $sortRight but found\n${joinOperator.right}") // check the output partitioning -@@ -835,11 +874,11 @@ abstract class BucketedReadSuite extends QueryTest with SQLTestUtils with Adapti +@@ -835,11 +872,11 @@ abstract class BucketedReadSuite extends QueryTest with SQLTestUtils with Adapti df1.write.format("parquet").bucketBy(8, "i").saveAsTable("bucketed_table") val scanDF = spark.table("bucketed_table").select("j") @@ -2960,7 +2952,7 @@ index 266bb343526..50408a27a0a 100644 checkAnswer(aggDF, df1.groupBy("j").agg(max("k"))) } } -@@ -894,7 +933,10 @@ abstract class BucketedReadSuite extends QueryTest with SQLTestUtils with Adapti +@@ -894,7 +931,10 @@ abstract class BucketedReadSuite extends QueryTest with SQLTestUtils with Adapti } test("SPARK-29655 Read bucketed tables obeys spark.sql.shuffle.partitions") { @@ -2971,7 +2963,7 @@ index 266bb343526..50408a27a0a 100644 SQLConf.SHUFFLE_PARTITIONS.key -> "5", SQLConf.COALESCE_PARTITIONS_INITIAL_PARTITION_NUM.key -> "7") { val bucketSpec = Some(BucketSpec(6, Seq("i", "j"), Nil)) -@@ -913,7 +955,10 @@ abstract class BucketedReadSuite extends QueryTest with SQLTestUtils with Adapti +@@ -913,7 +953,10 @@ abstract class BucketedReadSuite extends QueryTest with SQLTestUtils with Adapti } test("SPARK-32767 Bucket join should work if SHUFFLE_PARTITIONS larger than bucket number") { @@ -2982,7 +2974,7 @@ index 266bb343526..50408a27a0a 100644 SQLConf.SHUFFLE_PARTITIONS.key -> "9", SQLConf.COALESCE_PARTITIONS_INITIAL_PARTITION_NUM.key -> "10") { -@@ -943,7 +988,10 @@ abstract class BucketedReadSuite extends QueryTest with SQLTestUtils with Adapti +@@ -943,7 +986,10 @@ abstract class BucketedReadSuite extends QueryTest with SQLTestUtils with Adapti } test("bucket coalescing eliminates shuffle") { @@ -2993,7 +2985,7 @@ index 266bb343526..50408a27a0a 100644 SQLConf.COALESCE_BUCKETS_IN_JOIN_ENABLED.key -> "true", SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { // The side with bucketedTableTestSpec1 will be coalesced to have 4 output partitions. -@@ -1026,15 +1074,26 @@ abstract class BucketedReadSuite extends QueryTest with SQLTestUtils with Adapti +@@ -1026,15 +1072,26 @@ abstract class BucketedReadSuite extends QueryTest with SQLTestUtils with Adapti expectedNumShuffles: Int, expectedCoalescedNumBuckets: Option[Int]): Unit = { val plan = sql(query).queryExecution.executedPlan diff --git a/dev/diffs/3.5.9.diff b/dev/diffs/3.5.9.diff index e9baa8b4650..6e40bc64dc8 100644 --- a/dev/diffs/3.5.9.diff +++ b/dev/diffs/3.5.9.diff @@ -1387,14 +1387,14 @@ index e827396009d..efbee3975e3 100644 } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala -index 04702201f82..4d38d8d6e51 100644 +index 04702201f82..65a29a3bd01 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala @@ -22,10 +22,11 @@ import scala.collection.mutable.ArrayBuffer import org.apache.spark.sql.catalyst.expressions.SubqueryExpression import org.apache.spark.sql.catalyst.plans.{LeftAnti, LeftSemi} import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Join, LogicalPlan, Project, Sort, Union} -+import org.apache.spark.sql.comet.{CometColumnarToRowExec, CometNativeColumnarToRowExec, CometNativeScanExec, CometScanExec} ++import org.apache.spark.sql.comet.{CometColumnarToRowExec, CometNativeScanExec, CometScanExec} import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecution} import org.apache.spark.sql.execution.datasources.FileScanRDD @@ -1403,7 +1403,7 @@ index 04702201f82..4d38d8d6e51 100644 import org.apache.spark.sql.execution.joins.{BaseJoinExec, BroadcastHashJoinExec, BroadcastNestedLoopJoinExec} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession -@@ -1599,6 +1600,29 @@ class SubquerySuite extends QueryTest +@@ -1599,6 +1600,24 @@ class SubquerySuite extends QueryTest fs.inputRDDs().forall( _.asInstanceOf[FileScanRDD].filePartitions.forall( _.files.forall(_.urlEncodedPath.contains("p=0")))) @@ -1424,16 +1424,11 @@ index 04702201f82..4d38d8d6e51 100644 + fs.partitionFilters.exists(ExecSubqueryExpression.hasSubquery) && + fs.inputRDDs().forall( + _.asInstanceOf[FileScanRDD].filePartitions.forall( -+ _.files.forall(_.urlEncodedPath.contains("p=0")))) -+ case CometNativeColumnarToRowExec(fs: CometNativeScanExec) => -+ fs.partitionFilters.exists(ExecSubqueryExpression.hasSubquery) && -+ fs.inputRDDs().forall( -+ _.asInstanceOf[FileScanRDD].filePartitions.forall( + _.files.forall(_.urlEncodedPath.contains("p=0")))) case _ => false }) } -@@ -2164,7 +2188,7 @@ class SubquerySuite extends QueryTest +@@ -2164,7 +2183,7 @@ class SubquerySuite extends QueryTest df.collect() val exchanges = collect(df.queryExecution.executedPlan) { @@ -1902,18 +1897,18 @@ index eec396b2e39..bf3f1c769d6 100644 nums.createOrReplaceTempView("nums") diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/SparkPlanSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/SparkPlanSuite.scala -index b14f4a405f6..90bed10eca9 100644 +index b14f4a405f6..ab7baf434a5 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/SparkPlanSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/SparkPlanSuite.scala @@ -23,6 +23,7 @@ import org.apache.spark.sql.QueryTest import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference} import org.apache.spark.sql.catalyst.plans.logical.Deduplicate -+import org.apache.spark.sql.comet.{CometColumnarToRowExec, CometNativeColumnarToRowExec} ++import org.apache.spark.sql.comet.CometColumnarToRowExec import org.apache.spark.sql.execution.datasources.v2.BatchScanExec import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession -@@ -131,7 +132,11 @@ class SparkPlanSuite extends QueryTest with SharedSparkSession { +@@ -131,7 +132,10 @@ class SparkPlanSuite extends QueryTest with SharedSparkSession { spark.range(1).write.parquet(path.getAbsolutePath) val df = spark.read.parquet(path.getAbsolutePath) val columnarToRowExec = @@ -1921,7 +1916,6 @@ index b14f4a405f6..90bed10eca9 100644 + df.queryExecution.executedPlan.collectFirst { + case p: ColumnarToRowExec => p + case p: CometColumnarToRowExec => p -+ case p: CometNativeColumnarToRowExec => p + }.get try { spark.range(1).foreach { _ => @@ -3201,7 +3195,7 @@ index 67206e9c655..7ce027efad1 100644 } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/sources/BucketedReadSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/sources/BucketedReadSuite.scala -index 746f289c393..f2e0d1c93e5 100644 +index 746f289c393..e8ce3633bf8 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/sources/BucketedReadSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/sources/BucketedReadSuite.scala @@ -19,16 +19,18 @@ package org.apache.spark.sql.sources @@ -3260,7 +3254,7 @@ index 746f289c393..f2e0d1c93e5 100644 val bucketColumnType = bucketedDataFrame.schema.apply(bucketColumnIndex).dataType val rowsWithInvalidBuckets = fileScan.execute().filter(row => { -@@ -452,28 +465,54 @@ abstract class BucketedReadSuite extends QueryTest with SQLTestUtils with Adapti +@@ -452,28 +465,52 @@ abstract class BucketedReadSuite extends QueryTest with SQLTestUtils with Adapti val joinOperator = if (joined.sqlContext.conf.adaptiveExecutionEnabled) { val executedPlan = joined.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec].executedPlan @@ -3291,9 +3285,7 @@ index 746f289c393..f2e0d1c93e5 100644 + case s: SortMergeJoinExec => s + case o => fail(s"expected SortMergeJoinExec, but found\n$o") + } -+ case CometNativeColumnarToRowExec(child) => -+ child.asInstanceOf[CometSortMergeJoinExec].originalPlan match { -+ case s: SortMergeJoinExec => s ++ case s: SortMergeJoinExec => s + case o => fail(s"expected SortMergeJoinExec, but found\n$o") + } + case o => fail(s"expected SortMergeJoinExec, but found\n$o") @@ -3323,7 +3315,7 @@ index 746f289c393..f2e0d1c93e5 100644 s"expected sort in the right child to be $sortRight but found\n${joinOperator.right}") // check the output partitioning -@@ -836,11 +875,11 @@ abstract class BucketedReadSuite extends QueryTest with SQLTestUtils with Adapti +@@ -836,11 +873,11 @@ abstract class BucketedReadSuite extends QueryTest with SQLTestUtils with Adapti df1.write.format("parquet").bucketBy(8, "i").saveAsTable("bucketed_table") val scanDF = spark.table("bucketed_table").select("j") @@ -3337,7 +3329,7 @@ index 746f289c393..f2e0d1c93e5 100644 checkAnswer(aggDF, df1.groupBy("j").agg(max("k"))) } } -@@ -895,7 +934,10 @@ abstract class BucketedReadSuite extends QueryTest with SQLTestUtils with Adapti +@@ -895,7 +932,10 @@ abstract class BucketedReadSuite extends QueryTest with SQLTestUtils with Adapti } test("SPARK-29655 Read bucketed tables obeys spark.sql.shuffle.partitions") { @@ -3348,7 +3340,7 @@ index 746f289c393..f2e0d1c93e5 100644 SQLConf.SHUFFLE_PARTITIONS.key -> "5", SQLConf.COALESCE_PARTITIONS_INITIAL_PARTITION_NUM.key -> "7") { val bucketSpec = Some(BucketSpec(6, Seq("i", "j"), Nil)) -@@ -914,7 +956,10 @@ abstract class BucketedReadSuite extends QueryTest with SQLTestUtils with Adapti +@@ -914,7 +954,10 @@ abstract class BucketedReadSuite extends QueryTest with SQLTestUtils with Adapti } test("SPARK-32767 Bucket join should work if SHUFFLE_PARTITIONS larger than bucket number") { @@ -3359,7 +3351,7 @@ index 746f289c393..f2e0d1c93e5 100644 SQLConf.SHUFFLE_PARTITIONS.key -> "9", SQLConf.COALESCE_PARTITIONS_INITIAL_PARTITION_NUM.key -> "10") { -@@ -944,7 +989,10 @@ abstract class BucketedReadSuite extends QueryTest with SQLTestUtils with Adapti +@@ -944,7 +987,10 @@ abstract class BucketedReadSuite extends QueryTest with SQLTestUtils with Adapti } test("bucket coalescing eliminates shuffle") { @@ -3370,7 +3362,7 @@ index 746f289c393..f2e0d1c93e5 100644 SQLConf.COALESCE_BUCKETS_IN_JOIN_ENABLED.key -> "true", SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { // The side with bucketedTableTestSpec1 will be coalesced to have 4 output partitions. -@@ -1029,15 +1077,24 @@ abstract class BucketedReadSuite extends QueryTest with SQLTestUtils with Adapti +@@ -1029,15 +1075,24 @@ abstract class BucketedReadSuite extends QueryTest with SQLTestUtils with Adapti Seq(true, false).foreach { aqeEnabled => withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqeEnabled.toString) { val plan = sql(query).queryExecution.executedPlan diff --git a/dev/diffs/4.0.4.diff b/dev/diffs/4.0.4.diff index ede5e39dd36..118fa5c75ea 100644 --- a/dev/diffs/4.0.4.diff +++ b/dev/diffs/4.0.4.diff @@ -1653,14 +1653,14 @@ index 0df7f806272..9cdfe8b8f46 100644 df.collect() } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala -index 2e33f6505ab..f0e64fd58ca 100644 +index 2e33f6505ab..1e2d21aafd0 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala @@ -23,10 +23,11 @@ import org.apache.spark.SparkRuntimeException import org.apache.spark.sql.catalyst.expressions.SubqueryExpression import org.apache.spark.sql.catalyst.plans.{LeftAnti, LeftSemi} import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, Join, LogicalPlan, Project, Sort, Union} -+import org.apache.spark.sql.comet.{CometColumnarToRowExec, CometNativeColumnarToRowExec, CometNativeScanExec, CometScanExec} ++import org.apache.spark.sql.comet.{CometColumnarToRowExec, CometNativeScanExec, CometScanExec} import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecution} import org.apache.spark.sql.execution.datasources.FileScanRDD @@ -1669,7 +1669,7 @@ index 2e33f6505ab..f0e64fd58ca 100644 import org.apache.spark.sql.execution.joins.{BaseJoinExec, BroadcastHashJoinExec, BroadcastNestedLoopJoinExec} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession -@@ -1529,6 +1530,30 @@ class SubquerySuite extends QueryTest +@@ -1529,6 +1530,24 @@ class SubquerySuite extends QueryTest fs.inputRDDs().forall( _.asInstanceOf[FileScanRDD].filePartitions.forall( _.files.forall(_.urlEncodedPath.contains("p=0")))) @@ -1690,17 +1690,11 @@ index 2e33f6505ab..f0e64fd58ca 100644 + fs.partitionFilters.exists(ExecSubqueryExpression.hasSubquery) && + fs.inputRDDs().forall( + _.asInstanceOf[FileScanRDD].filePartitions.forall( -+ _.files.forall(_.urlEncodedPath.contains("p=0")))) -+ case CometNativeColumnarToRowExec( -+ fs: CometNativeScanExec) => -+ fs.partitionFilters.exists(ExecSubqueryExpression.hasSubquery) && -+ fs.inputRDDs().forall( -+ _.asInstanceOf[FileScanRDD].filePartitions.forall( + _.files.forall(_.urlEncodedPath.contains("p=0")))) case _ => false }) } -@@ -2094,7 +2119,7 @@ class SubquerySuite extends QueryTest +@@ -2094,7 +2113,7 @@ class SubquerySuite extends QueryTest df.collect() val exchanges = collect(df.queryExecution.executedPlan) { @@ -1709,7 +1703,7 @@ index 2e33f6505ab..f0e64fd58ca 100644 } assert(exchanges.size === 1) } -@@ -2678,18 +2703,29 @@ class SubquerySuite extends QueryTest +@@ -2678,18 +2697,29 @@ class SubquerySuite extends QueryTest def checkFileSourceScan(query: String, answer: Seq[Row]): Unit = { val df = sql(query) checkAnswer(df, answer) @@ -1747,28 +1741,6 @@ index 2e33f6505ab..f0e64fd58ca 100644 } withTable("t1", "t2") { -diff --git a/sql/core/src/test/scala/org/apache/spark/sql/VariantEndToEndSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/VariantEndToEndSuite.scala -index a40e34d94d0..abc1f035d15 100644 ---- a/sql/core/src/test/scala/org/apache/spark/sql/VariantEndToEndSuite.scala -+++ b/sql/core/src/test/scala/org/apache/spark/sql/VariantEndToEndSuite.scala -@@ -21,6 +21,7 @@ import org.apache.spark.sql.QueryTest.sameRows - import org.apache.spark.sql.catalyst.InternalRow - import org.apache.spark.sql.catalyst.expressions.{Cast, Literal} - import org.apache.spark.sql.catalyst.expressions.variant.{ToVariantObject, VariantExpressionEvalUtils} -+import org.apache.spark.sql.comet.CometNativeColumnarToRowExec - import org.apache.spark.sql.execution.WholeStageCodegenExec - import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector - import org.apache.spark.sql.functions._ -@@ -358,7 +359,8 @@ class VariantEndToEndSuite extends QueryTest with SharedSparkSession { - s"cast(to_variant_object(s) as ${schema(2).dataType.sql})") - checkAnswer(df, input) - val plan = df.queryExecution.executedPlan -- assert(plan.isInstanceOf[WholeStageCodegenExec] == (codegenMode == "CODEGEN_ONLY")) -+ assert(plan.isInstanceOf[WholeStageCodegenExec] == (codegenMode == "CODEGEN_ONLY") -+ || plan.isInstanceOf[CometNativeColumnarToRowExec]) - } - } - } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationSuite.scala index 11e9547dfc5..327e0b1bc5b 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationSuite.scala @@ -2369,18 +2341,18 @@ index aa619c5cde8..4b5a42e683f 100644 nums.createOrReplaceTempView("nums") diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/SparkPlanSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/SparkPlanSuite.scala -index aed11badb71..1a365b5aacf 100644 +index aed11badb71..ab7e9456e26 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/SparkPlanSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/SparkPlanSuite.scala @@ -23,6 +23,7 @@ import org.apache.spark.sql.QueryTest import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference} import org.apache.spark.sql.catalyst.plans.logical.Deduplicate -+import org.apache.spark.sql.comet.{CometColumnarToRowExec, CometNativeColumnarToRowExec} ++import org.apache.spark.sql.comet.CometColumnarToRowExec import org.apache.spark.sql.execution.datasources.v2.BatchScanExec import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession -@@ -134,7 +135,11 @@ class SparkPlanSuite extends QueryTest with SharedSparkSession { +@@ -134,7 +135,10 @@ class SparkPlanSuite extends QueryTest with SharedSparkSession { spark.range(1).write.parquet(path.getAbsolutePath) val df = spark.read.parquet(path.getAbsolutePath) val columnarToRowExec = @@ -2388,7 +2360,6 @@ index aed11badb71..1a365b5aacf 100644 + df.queryExecution.executedPlan.collectFirst { + case p: ColumnarToRowExec => p + case p: CometColumnarToRowExec => p -+ case p: CometNativeColumnarToRowExec => p + }.get try { spark.range(1).foreach { _ => @@ -4040,7 +4011,7 @@ index 800a58f0c1d..2506d27b94d 100644 } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/sources/BucketedReadSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/sources/BucketedReadSuite.scala -index c4b09c4b289..75c3437788e 100644 +index c4b09c4b289..7aafa8a74d9 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/sources/BucketedReadSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/sources/BucketedReadSuite.scala @@ -26,10 +26,11 @@ import org.apache.spark.sql.catalyst.expressions @@ -4092,7 +4063,7 @@ index c4b09c4b289..75c3437788e 100644 val bucketColumnType = bucketedDataFrame.schema.apply(bucketColumnIndex).dataType val rowsWithInvalidBuckets = fileScan.execute().filter(row => { -@@ -454,28 +466,54 @@ abstract class BucketedReadSuite extends QueryTest with SQLTestUtils with Adapti +@@ -454,28 +466,52 @@ abstract class BucketedReadSuite extends QueryTest with SQLTestUtils with Adapti val joinOperator = if (joined.sparkSession.sessionState.conf.adaptiveExecutionEnabled) { val executedPlan = joined.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec].executedPlan @@ -4123,9 +4094,7 @@ index c4b09c4b289..75c3437788e 100644 + case s: SortMergeJoinExec => s + case o => fail(s"expected SortMergeJoinExec, but found\n$o") + } -+ case CometNativeColumnarToRowExec(child) => -+ child.asInstanceOf[CometSortMergeJoinExec].originalPlan match { -+ case s: SortMergeJoinExec => s ++ case s: SortMergeJoinExec => s + case o => fail(s"expected SortMergeJoinExec, but found\n$o") + } + case o => fail(s"expected SortMergeJoinExec, but found\n$o") @@ -4155,7 +4124,7 @@ index c4b09c4b289..75c3437788e 100644 s"expected sort in the right child to be $sortRight but found\n${joinOperator.right}") // check the output partitioning -@@ -838,11 +876,11 @@ abstract class BucketedReadSuite extends QueryTest with SQLTestUtils with Adapti +@@ -838,11 +874,11 @@ abstract class BucketedReadSuite extends QueryTest with SQLTestUtils with Adapti df1.write.format("parquet").bucketBy(8, "i").saveAsTable("bucketed_table") val scanDF = spark.table("bucketed_table").select("j") @@ -4169,7 +4138,7 @@ index c4b09c4b289..75c3437788e 100644 checkAnswer(aggDF, df1.groupBy("j").agg(max("k"))) } } -@@ -1031,15 +1069,24 @@ abstract class BucketedReadSuite extends QueryTest with SQLTestUtils with Adapti +@@ -1031,15 +1067,24 @@ abstract class BucketedReadSuite extends QueryTest with SQLTestUtils with Adapti Seq(true, false).foreach { aqeEnabled => withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqeEnabled.toString) { val plan = sql(query).queryExecution.executedPlan diff --git a/dev/diffs/4.1.3.diff b/dev/diffs/4.1.3.diff index a76ac1fd0a8..194a8962d71 100644 --- a/dev/diffs/4.1.3.diff +++ b/dev/diffs/4.1.3.diff @@ -1824,14 +1824,14 @@ index 7bfc8cf4fa6..4bd387801db 100644 df.collect() } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala -index 3ba48da0e32..a7f9513b9c8 100644 +index 3ba48da0e32..53e87857365 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala @@ -23,10 +23,11 @@ import org.apache.spark.SparkRuntimeException import org.apache.spark.sql.catalyst.expressions.{EqualTo, NamedExpression, OuterReference, SubqueryExpression} import org.apache.spark.sql.catalyst.plans.{LeftAnti, LeftSemi} import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, Join, LogicalPlan, Project, Sort, Union} -+import org.apache.spark.sql.comet.{CometColumnarToRowExec, CometNativeColumnarToRowExec, CometNativeScanExec, CometScanExec} ++import org.apache.spark.sql.comet.{CometColumnarToRowExec, CometNativeScanExec, CometScanExec} import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecution} import org.apache.spark.sql.execution.datasources.FileScanRDD @@ -1840,7 +1840,7 @@ index 3ba48da0e32..a7f9513b9c8 100644 import org.apache.spark.sql.execution.joins.{BaseJoinExec, BroadcastHashJoinExec, BroadcastNestedLoopJoinExec} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession -@@ -1529,6 +1530,30 @@ class SubquerySuite extends QueryTest +@@ -1529,6 +1530,24 @@ class SubquerySuite extends QueryTest fs.inputRDDs().forall( _.asInstanceOf[FileScanRDD].filePartitions.forall( _.files.forall(_.urlEncodedPath.contains("p=0")))) @@ -1861,17 +1861,11 @@ index 3ba48da0e32..a7f9513b9c8 100644 + fs.partitionFilters.exists(ExecSubqueryExpression.hasSubquery) && + fs.inputRDDs().forall( + _.asInstanceOf[FileScanRDD].filePartitions.forall( -+ _.files.forall(_.urlEncodedPath.contains("p=0")))) -+ case CometNativeColumnarToRowExec( -+ fs: CometNativeScanExec) => -+ fs.partitionFilters.exists(ExecSubqueryExpression.hasSubquery) && -+ fs.inputRDDs().forall( -+ _.asInstanceOf[FileScanRDD].filePartitions.forall( + _.files.forall(_.urlEncodedPath.contains("p=0")))) case _ => false }) } -@@ -2094,7 +2119,7 @@ class SubquerySuite extends QueryTest +@@ -2094,7 +2113,7 @@ class SubquerySuite extends QueryTest df.collect() val exchanges = collect(df.queryExecution.executedPlan) { @@ -1880,7 +1874,7 @@ index 3ba48da0e32..a7f9513b9c8 100644 } assert(exchanges.size === 1) } -@@ -2678,18 +2703,29 @@ class SubquerySuite extends QueryTest +@@ -2678,18 +2697,29 @@ class SubquerySuite extends QueryTest def checkFileSourceScan(query: String, answer: Seq[Row]): Unit = { val df = sql(query) checkAnswer(df, answer) @@ -1918,28 +1912,6 @@ index 3ba48da0e32..a7f9513b9c8 100644 } withTable("t1", "t2") { -diff --git a/sql/core/src/test/scala/org/apache/spark/sql/VariantEndToEndSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/VariantEndToEndSuite.scala -index 8a0e2c29653..d276a51cbc6 100644 ---- a/sql/core/src/test/scala/org/apache/spark/sql/VariantEndToEndSuite.scala -+++ b/sql/core/src/test/scala/org/apache/spark/sql/VariantEndToEndSuite.scala -@@ -21,6 +21,7 @@ import org.apache.spark.sql.QueryTest.sameRows - import org.apache.spark.sql.catalyst.InternalRow - import org.apache.spark.sql.catalyst.expressions.{Cast, Literal} - import org.apache.spark.sql.catalyst.expressions.variant.{ToVariantObject, VariantExpressionEvalUtils} -+import org.apache.spark.sql.comet.CometNativeColumnarToRowExec - import org.apache.spark.sql.execution.WholeStageCodegenExec - import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector - import org.apache.spark.sql.functions._ -@@ -358,7 +359,8 @@ class VariantEndToEndSuite extends QueryTest with SharedSparkSession { - s"cast(to_variant_object(s) as ${schema(2).dataType.sql})") - checkAnswer(df, input) - val plan = df.queryExecution.executedPlan -- assert(plan.isInstanceOf[WholeStageCodegenExec] == (codegenMode == "CODEGEN_ONLY")) -+ assert(plan.isInstanceOf[WholeStageCodegenExec] == (codegenMode == "CODEGEN_ONLY") -+ || plan.isInstanceOf[CometNativeColumnarToRowExec]) - } - } - } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationSuite.scala index 8f7a68bcbe6..c09c5d74309 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationSuite.scala @@ -2540,18 +2512,18 @@ index aa619c5cde8..4b5a42e683f 100644 nums.createOrReplaceTempView("nums") diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/SparkPlanSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/SparkPlanSuite.scala -index aed11badb71..1a365b5aacf 100644 +index aed11badb71..ab7e9456e26 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/SparkPlanSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/SparkPlanSuite.scala @@ -23,6 +23,7 @@ import org.apache.spark.sql.QueryTest import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference} import org.apache.spark.sql.catalyst.plans.logical.Deduplicate -+import org.apache.spark.sql.comet.{CometColumnarToRowExec, CometNativeColumnarToRowExec} ++import org.apache.spark.sql.comet.CometColumnarToRowExec import org.apache.spark.sql.execution.datasources.v2.BatchScanExec import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession -@@ -134,7 +135,11 @@ class SparkPlanSuite extends QueryTest with SharedSparkSession { +@@ -134,7 +135,10 @@ class SparkPlanSuite extends QueryTest with SharedSparkSession { spark.range(1).write.parquet(path.getAbsolutePath) val df = spark.read.parquet(path.getAbsolutePath) val columnarToRowExec = @@ -2559,7 +2531,6 @@ index aed11badb71..1a365b5aacf 100644 + df.queryExecution.executedPlan.collectFirst { + case p: ColumnarToRowExec => p + case p: CometColumnarToRowExec => p -+ case p: CometNativeColumnarToRowExec => p + }.get try { spark.range(1).foreach { _ => @@ -3316,7 +3287,7 @@ index 188a28ff1c0..5be29a7377d 100644 assert(o1.semanticEquals(o2), "Different output column order after AQE optimization") } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/CachedBatchSerializerSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/CachedBatchSerializerSuite.scala -index 47b935a2880..65ee66c1975 100644 +index 47b935a2880..576218d5285 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/CachedBatchSerializerSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/CachedBatchSerializerSuite.scala @@ -212,6 +212,18 @@ class CachedBatchSerializerNoUnwrapSuite extends QueryTest @@ -3338,7 +3309,7 @@ index 47b935a2880..65ee66c1975 100644 test("Do not unwrap ColumnarToRowExec") { withTempPath { workDir => val workDirPath = workDir.getAbsolutePath -@@ -230,9 +242,21 @@ class CachedBatchSerializerNoUnwrapSuite extends QueryTest +@@ -230,9 +242,19 @@ class CachedBatchSerializerNoUnwrapSuite extends QueryTest assert(cachedPlans.length == 2) cachedPlans.foreach { cachedPlan => @@ -3346,12 +3317,10 @@ index 47b935a2880..65ee66c1975 100644 - assert(cachedPlan.asInstanceOf[WholeStageCodegenExec] - .child.isInstanceOf[ColumnarToRowExec]) + // Comet replaces ColumnarToRowExec with its own columnar-to-row operator -+ // (CometNativeColumnarToRow / CometColumnarToRow). Accept either the -+ // Spark shape or the equivalent Comet shape, since the important property -+ // under this serializer is that the row conversion is not unwrapped. -+ // CometColumnarToRowExec supports codegen, so it appears as the child of a -+ // WholeStageCodegenExec, while CometNativeColumnarToRowExec does not and is -+ // the cached plan itself. ++ // (CometColumnarToRow). Accept either the Spark shape or the equivalent ++ // Comet shape, since the important property under this serializer is that ++ // the row conversion is not unwrapped. CometColumnarToRowExec supports ++ // codegen, so it appears as the child of a WholeStageCodegenExec. + val transition = cachedPlan match { + case w: WholeStageCodegenExec => w.child + case other => other @@ -4349,7 +4318,7 @@ index 800a58f0c1d..2506d27b94d 100644 } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/sources/BucketedReadSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/sources/BucketedReadSuite.scala -index 83ebd24384c..32511091bb2 100644 +index 83ebd24384c..dae97acf8f8 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/sources/BucketedReadSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/sources/BucketedReadSuite.scala @@ -26,10 +26,11 @@ import org.apache.spark.sql.catalyst.expressions @@ -4401,7 +4370,7 @@ index 83ebd24384c..32511091bb2 100644 val bucketColumnType = bucketedDataFrame.schema.apply(bucketColumnIndex).dataType val rowsWithInvalidBuckets = fileScan.execute().filter(row => { -@@ -454,28 +466,54 @@ abstract class BucketedReadSuite extends QueryTest with SQLTestUtils with Adapti +@@ -454,28 +466,52 @@ abstract class BucketedReadSuite extends QueryTest with SQLTestUtils with Adapti val joinOperator = if (joined.sparkSession.sessionState.conf.adaptiveExecutionEnabled) { val executedPlan = joined.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec].executedPlan @@ -4432,9 +4401,7 @@ index 83ebd24384c..32511091bb2 100644 + case s: SortMergeJoinExec => s + case o => fail(s"expected SortMergeJoinExec, but found\n$o") + } -+ case CometNativeColumnarToRowExec(child) => -+ child.asInstanceOf[CometSortMergeJoinExec].originalPlan match { -+ case s: SortMergeJoinExec => s ++ case s: SortMergeJoinExec => s + case o => fail(s"expected SortMergeJoinExec, but found\n$o") + } + case o => fail(s"expected SortMergeJoinExec, but found\n$o") @@ -4464,7 +4431,7 @@ index 83ebd24384c..32511091bb2 100644 s"expected sort in the right child to be $sortRight but found\n${joinOperator.right}") // check the output partitioning -@@ -838,11 +876,11 @@ abstract class BucketedReadSuite extends QueryTest with SQLTestUtils with Adapti +@@ -838,11 +874,11 @@ abstract class BucketedReadSuite extends QueryTest with SQLTestUtils with Adapti df1.write.format("parquet").bucketBy(8, "i").saveAsTable("bucketed_table") val scanDF = spark.table("bucketed_table").select("j") @@ -4478,7 +4445,7 @@ index 83ebd24384c..32511091bb2 100644 checkAnswer(aggDF, df1.groupBy("j").agg(max("k"))) } } -@@ -1031,15 +1069,24 @@ abstract class BucketedReadSuite extends QueryTest with SQLTestUtils with Adapti +@@ -1031,15 +1067,24 @@ abstract class BucketedReadSuite extends QueryTest with SQLTestUtils with Adapti Seq(true, false).foreach { aqeEnabled => withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqeEnabled.toString) { val plan = sql(query).queryExecution.executedPlan diff --git a/dev/diffs/4.2.0.diff b/dev/diffs/4.2.0.diff index 37fc70b53db..802ac905819 100644 --- a/dev/diffs/4.2.0.diff +++ b/dev/diffs/4.2.0.diff @@ -1894,14 +1894,14 @@ index 63589472854..f8c07a9b037 100644 df.collect() } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala -index cd3e389d765..354c5b97e8b 100644 +index cd3e389d765..ccdcdddb8de 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala @@ -23,10 +23,11 @@ import org.apache.spark.SparkRuntimeException import org.apache.spark.sql.catalyst.expressions.{EqualTo, NamedExpression, OuterReference, SubqueryExpression} import org.apache.spark.sql.catalyst.plans.{LeftAnti, LeftSemi} import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, Join, LogicalPlan, Project, Sort, Union} -+import org.apache.spark.sql.comet.{CometColumnarToRowExec, CometNativeColumnarToRowExec, CometNativeScanExec, CometScanExec} ++import org.apache.spark.sql.comet.{CometColumnarToRowExec, CometNativeScanExec, CometScanExec} import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecution} import org.apache.spark.sql.execution.datasources.FileScanRDD @@ -1910,7 +1910,7 @@ index cd3e389d765..354c5b97e8b 100644 import org.apache.spark.sql.execution.joins.{BaseJoinExec, BroadcastHashJoinExec, BroadcastNestedLoopJoinExec} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession -@@ -1528,6 +1529,30 @@ class SubquerySuite extends SharedSparkSession +@@ -1528,6 +1529,24 @@ class SubquerySuite extends SharedSparkSession fs.inputRDDs().forall( _.asInstanceOf[FileScanRDD].filePartitions.forall( _.files.forall(_.urlEncodedPath.contains("p=0")))) @@ -1920,12 +1920,6 @@ index cd3e389d765..354c5b97e8b 100644 + fs.inputRDDs().forall( + _.asInstanceOf[FileScanRDD].filePartitions.forall( + _.files.forall(_.urlEncodedPath.contains("p=0")))) -+ case CometNativeColumnarToRowExec( -+ fs: CometNativeScanExec) => -+ fs.partitionFilters.exists(ExecSubqueryExpression.hasSubquery) && -+ fs.inputRDDs().forall( -+ _.asInstanceOf[FileScanRDD].filePartitions.forall( -+ _.files.forall(_.urlEncodedPath.contains("p=0")))) + case WholeStageCodegenExec(CometColumnarToRowExec(InputAdapter( + fs @ CometScanExec(_, _, _, partitionFilters, _, _, _, _, _, _)))) => + partitionFilters.exists(ExecSubqueryExpression.hasSubquery) && @@ -1941,7 +1935,7 @@ index cd3e389d765..354c5b97e8b 100644 case _ => false }) } -@@ -2093,7 +2118,7 @@ class SubquerySuite extends SharedSparkSession +@@ -2093,7 +2112,7 @@ class SubquerySuite extends SharedSparkSession df.collect() val exchanges = collect(df.queryExecution.executedPlan) { @@ -1950,7 +1944,7 @@ index cd3e389d765..354c5b97e8b 100644 } assert(exchanges.size === 1) } -@@ -2464,18 +2489,29 @@ class SubquerySuite extends SharedSparkSession +@@ -2464,18 +2483,29 @@ class SubquerySuite extends SharedSparkSession def checkFileSourceScan(query: String, answer: Seq[Row]): Unit = { val df = sql(query) checkAnswer(df, answer) @@ -1988,28 +1982,6 @@ index cd3e389d765..354c5b97e8b 100644 } withTable("t1", "t2") { -diff --git a/sql/core/src/test/scala/org/apache/spark/sql/VariantEndToEndSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/VariantEndToEndSuite.scala -index 2d26356890d..2c5994f5fbc 100644 ---- a/sql/core/src/test/scala/org/apache/spark/sql/VariantEndToEndSuite.scala -+++ b/sql/core/src/test/scala/org/apache/spark/sql/VariantEndToEndSuite.scala -@@ -21,6 +21,7 @@ import org.apache.spark.sql.QueryTest.sameRows - import org.apache.spark.sql.catalyst.InternalRow - import org.apache.spark.sql.catalyst.expressions.{Cast, Literal} - import org.apache.spark.sql.catalyst.expressions.variant.{ToVariantObject, VariantExpressionEvalUtils} -+import org.apache.spark.sql.comet.CometNativeColumnarToRowExec - import org.apache.spark.sql.execution.WholeStageCodegenExec - import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector - import org.apache.spark.sql.functions._ -@@ -391,7 +392,8 @@ class VariantEndToEndSuite extends SharedSparkSession { - s"cast(to_variant_object(s) as ${schema(2).dataType.sql})") - checkAnswer(df, input) - val plan = df.queryExecution.executedPlan -- assert(plan.isInstanceOf[WholeStageCodegenExec] == (codegenMode == "CODEGEN_ONLY")) -+ assert(plan.isInstanceOf[WholeStageCodegenExec] == (codegenMode == "CODEGEN_ONLY") -+ || plan.isInstanceOf[CometNativeColumnarToRowExec]) - } - } - } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationSuite.scala index 37684c7fce3..a0a2f72b6c9 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationSuite.scala @@ -2596,18 +2568,18 @@ index e8a39eb8c53..a621c7a7923 100644 nums.createOrReplaceTempView("nums") diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/SparkPlanSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/SparkPlanSuite.scala -index b167dd13dcb..b2b6b88ce66 100644 +index b167dd13dcb..52b33a7b472 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/SparkPlanSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/SparkPlanSuite.scala @@ -31,6 +31,7 @@ import org.apache.spark.sql.catalyst.expressions.{ import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, ExprCode} import org.apache.spark.sql.catalyst.plans.logical.Deduplicate import org.apache.spark.sql.catalyst.trees.LeafLike -+import org.apache.spark.sql.comet.{CometColumnarToRowExec, CometNativeColumnarToRowExec} ++import org.apache.spark.sql.comet.CometColumnarToRowExec import org.apache.spark.sql.execution.datasources.v2.BatchScanExec import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession -@@ -166,7 +167,11 @@ class SparkPlanSuite extends SharedSparkSession { +@@ -166,7 +167,10 @@ class SparkPlanSuite extends SharedSparkSession { spark.range(1).write.parquet(path.getAbsolutePath) val df = spark.read.parquet(path.getAbsolutePath) val columnarToRowExec = @@ -2615,7 +2587,6 @@ index b167dd13dcb..b2b6b88ce66 100644 + df.queryExecution.executedPlan.collectFirst { + case p: ColumnarToRowExec => p + case p: CometColumnarToRowExec => p -+ case p: CometNativeColumnarToRowExec => p + }.get try { spark.range(1).foreach { _ => @@ -3397,7 +3368,7 @@ index d6d19d21e65..7b3942d936b 100644 assert(o1.semanticEquals(o2), "Different output column order after AQE optimization") } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/CachedBatchSerializerSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/CachedBatchSerializerSuite.scala -index 88be4adb6a4..23ec5374f8a 100644 +index 88be4adb6a4..cb727a663e5 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/CachedBatchSerializerSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/CachedBatchSerializerSuite.scala @@ -210,6 +210,18 @@ class CachedBatchSerializerNoUnwrapSuite extends SharedSparkSession with Adaptiv @@ -3419,7 +3390,7 @@ index 88be4adb6a4..23ec5374f8a 100644 test("Do not unwrap ColumnarToRowExec") { withTempPath { workDir => val workDirPath = workDir.getAbsolutePath -@@ -228,9 +240,21 @@ class CachedBatchSerializerNoUnwrapSuite extends SharedSparkSession with Adaptiv +@@ -228,9 +240,19 @@ class CachedBatchSerializerNoUnwrapSuite extends SharedSparkSession with Adaptiv assert(cachedPlans.length == 2) cachedPlans.foreach { cachedPlan => @@ -3427,12 +3398,10 @@ index 88be4adb6a4..23ec5374f8a 100644 - assert(cachedPlan.asInstanceOf[WholeStageCodegenExec] - .child.isInstanceOf[ColumnarToRowExec]) + // Comet replaces ColumnarToRowExec with its own columnar-to-row operator -+ // (CometNativeColumnarToRow / CometColumnarToRow). Accept either the -+ // Spark shape or the equivalent Comet shape, since the important property -+ // under this serializer is that the row conversion is not unwrapped. -+ // CometColumnarToRowExec supports codegen, so it appears as the child of a -+ // WholeStageCodegenExec, while CometNativeColumnarToRowExec does not and is -+ // the cached plan itself. ++ // (CometColumnarToRow). Accept either the Spark shape or the equivalent ++ // Comet shape, since the important property under this serializer is that ++ // the row conversion is not unwrapped. CometColumnarToRowExec supports ++ // codegen, so it appears as the child of a WholeStageCodegenExec. + val transition = cachedPlan match { + case w: WholeStageCodegenExec => w.child + case other => other @@ -4553,7 +4522,7 @@ index d3da8e40c85..505f48092ea 100644 import testImplicits._ diff --git a/sql/core/src/test/scala/org/apache/spark/sql/sources/BucketedReadSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/sources/BucketedReadSuite.scala -index 452b7b9dedb..b2a73a8b772 100644 +index 452b7b9dedb..2e09caef9df 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/sources/BucketedReadSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/sources/BucketedReadSuite.scala @@ -26,10 +26,11 @@ import org.apache.spark.sql.catalyst.expressions @@ -4605,7 +4574,7 @@ index 452b7b9dedb..b2a73a8b772 100644 val bucketColumnType = bucketedDataFrame.schema.apply(bucketColumnIndex).dataType val rowsWithInvalidBuckets = fileScan.execute().filter(row => { -@@ -454,28 +466,54 @@ abstract class BucketedReadSuite extends QueryTest with AdaptiveSparkPlanHelper +@@ -454,28 +466,52 @@ abstract class BucketedReadSuite extends QueryTest with AdaptiveSparkPlanHelper val joinOperator = if (joined.sparkSession.sessionState.conf.adaptiveExecutionEnabled) { val executedPlan = joined.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec].executedPlan @@ -4636,9 +4605,7 @@ index 452b7b9dedb..b2a73a8b772 100644 + case s: SortMergeJoinExec => s + case o => fail(s"expected SortMergeJoinExec, but found\n$o") + } -+ case CometNativeColumnarToRowExec(child) => -+ child.asInstanceOf[CometSortMergeJoinExec].originalPlan match { -+ case s: SortMergeJoinExec => s ++ case s: SortMergeJoinExec => s + case o => fail(s"expected SortMergeJoinExec, but found\n$o") + } + case o => fail(s"expected SortMergeJoinExec, but found\n$o") @@ -4668,7 +4635,7 @@ index 452b7b9dedb..b2a73a8b772 100644 s"expected sort in the right child to be $sortRight but found\n${joinOperator.right}") // check the output partitioning -@@ -838,11 +876,11 @@ abstract class BucketedReadSuite extends QueryTest with AdaptiveSparkPlanHelper +@@ -838,11 +874,11 @@ abstract class BucketedReadSuite extends QueryTest with AdaptiveSparkPlanHelper df1.write.format("parquet").bucketBy(8, "i").saveAsTable("bucketed_table") val scanDF = spark.table("bucketed_table").select("j") @@ -4682,7 +4649,7 @@ index 452b7b9dedb..b2a73a8b772 100644 checkAnswer(aggDF, df1.groupBy("j").agg(max("k"))) } } -@@ -1031,15 +1069,24 @@ abstract class BucketedReadSuite extends QueryTest with AdaptiveSparkPlanHelper +@@ -1031,15 +1067,24 @@ abstract class BucketedReadSuite extends QueryTest with AdaptiveSparkPlanHelper Seq(true, false).foreach { aqeEnabled => withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqeEnabled.toString) { val plan = sql(query).queryExecution.executedPlan diff --git a/docs/source/contributor-guide/plugin_overview.md b/docs/source/contributor-guide/plugin_overview.md index 08bca47e627..a522f5fa6b2 100644 --- a/docs/source/contributor-guide/plugin_overview.md +++ b/docs/source/contributor-guide/plugin_overview.md @@ -120,9 +120,7 @@ if it has more columnar-to-row transitions than `spark.comet.exec.transitionReve of the transitions can outweigh the benefits of running part of the stage in Comet. It keeps the stage in Comet when reverting it is not safe, for example when Spark and Comet would have to exchange aggregate state. Then `EliminateRedundantTransitions` removes transitions that are not needed and replaces the rest that read Comet output -with `CometColumnarToRowExec`. It uses the Rust-implemented `CometNativeColumnarToRowExec` instead when -`spark.comet.exec.columnarToRow.native.enabled=true` (off by default) and that converter supports the schema. A schema -that contains `VariantType` keeps Spark's `ColumnarToRowExec`. +with `CometColumnarToRowExec`. A schema that contains `VariantType` keeps Spark's `ColumnarToRowExec`. ## Query Execution diff --git a/docs/source/user-guide/latest/migration-guide.md b/docs/source/user-guide/latest/migration-guide.md index feaa15aab06..10d8598dc96 100644 --- a/docs/source/user-guide/latest/migration-guide.md +++ b/docs/source/user-guide/latest/migration-guide.md @@ -57,6 +57,19 @@ Treat setting one of these keys as a temporary measure. If you find you cannot s legacy behavior, please open an issue describing your use case so it can be considered before the key is removed. +## Upgrading to Comet 1.3.0 + +Comet `1.3.0` makes no behavior changes that need a `spark.comet.legacy.*` key. The changes below +need none either, but check whether any of them applies to your deployment. + +### Deprecated and Removed Settings + +`spark.comet.exec.columnarToRow.native.enabled` is deprecated and will be removed in a future major +release. Comet `1.3.0` removes the Rust columnar-to-row converter that it enabled, so the setting no +longer has any effect. Comet always converts Arrow batches to Spark rows on the JVM, which was +already the default. The driver logs a warning when the setting is present in the Spark +configuration. + ## Upgrading to Comet 1.2.0 Comet `1.2.0` makes no behavior changes that need a `spark.comet.legacy.*` key. The changes below diff --git a/docs/source/user-guide/latest/tuning/transitions.md b/docs/source/user-guide/latest/tuning/transitions.md index 6acd307d8d1..d0f15e9cda0 100644 --- a/docs/source/user-guide/latest/tuning/transitions.md +++ b/docs/source/user-guide/latest/tuning/transitions.md @@ -41,8 +41,7 @@ expensive conversions again and again. ## Experimental: Direct Columnar-to-Row Conversion -When the JVM columnar-to-row operator is in use (`spark.comet.exec.columnarToRow.native.enabled=false`), -setting `spark.comet.exec.columnarToRow.direct.enabled=true` enables an experimental converter that writes +Setting `spark.comet.exec.columnarToRow.direct.enabled=true` enables an experimental converter that writes values straight from Arrow buffers into Spark's row format without allocating an object per value. This is most beneficial for decimal-heavy schemas, where the default conversion allocates a `Decimal` object per value (and considerably more for decimals with precision above 18); microbenchmarks show up to 2x faster diff --git a/docs/source/user-guide/latest/understanding-comet-plans.md b/docs/source/user-guide/latest/understanding-comet-plans.md index d86f8076bde..e406ed57d99 100644 --- a/docs/source/user-guide/latest/understanding-comet-plans.md +++ b/docs/source/user-guide/latest/understanding-comet-plans.md @@ -52,8 +52,7 @@ the same plan is shown in the Spark SQL UI. When reading a plan, look for: - **Node prefix.** `Comet*` nodes are accelerated by Comet. Anything without the prefix is unmodified Spark. -- **Transitions.** `CometColumnarToRow`, `CometNativeColumnarToRow`, and - `CometSparkRowToColumnar` mark boundaries between columnar Comet execution +- **Transitions.** `CometColumnarToRow` and `CometSparkRowToColumnar` mark boundaries between columnar Comet execution and row-based Spark execution. Frequent transitions usually indicate fallback inside the plan. - **Exchange type.** `CometExchange` is the native shuffle path, @@ -183,8 +182,7 @@ the config has no effect there. Not every node in the plan is an eligible operator. The following are excluded from both operator counts: -- Transition nodes (`CometColumnarToRow`, `CometNativeColumnarToRow`, - `CometSparkRowToColumnar`, `CometSparkColumnarToColumnar`, `ColumnarToRow`, +- Transition nodes (`CometColumnarToRow`, `CometSparkRowToColumnar`, `CometSparkColumnarToColumnar`, `ColumnarToRow`, `RowToColumnar`), which are reported separately as the transition count. `CometSparkRowToColumnar` and `CometSparkColumnarToColumnar` are the two names a single operator renders under, depending on whether its child already @@ -392,12 +390,11 @@ Comet inserts these nodes wherever data has to cross the columnar/row boundary. Multiple implementations exist because the optimal strategy depends on what produced the columnar data. -| Node | Direction | Notes | -| ------------------------------ | ------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| `CometColumnarToRow` | columnar → row | JVM-based row conversion. A fork of Spark's `ColumnarToRowExec` that includes the SPARK-50235 fix. This is the default. | -| `CometNativeColumnarToRow` | columnar → row | Rust-based row conversion via `NativeColumnarToRowConverter`. Disabled by default; enable with `spark.comet.exec.columnarToRow.native.enabled=true`. It carries a fixed JNI cost per batch and is slower than the JVM conversion for small batches. | -| `CometSparkColumnarToColumnar` | columnar → columnar | Converts a Spark columnar input (a non-Comet `ColumnarBatch`) into Comet's Arrow batches. | -| `CometSparkRowToColumnar` | row → columnar | Converts a Spark row input into Comet's Arrow batches. | +| Node | Direction | Notes | +| ------------------------------ | ------------------- | -------------------------------------------------------------------------------------------------- | +| `CometColumnarToRow` | columnar → row | JVM-based row conversion. A fork of Spark's `ColumnarToRowExec` that includes the SPARK-50235 fix. | +| `CometSparkColumnarToColumnar` | columnar → columnar | Converts a Spark columnar input (a non-Comet `ColumnarBatch`) into Comet's Arrow batches. | +| `CometSparkRowToColumnar` | row → columnar | Converts a Spark row input into Comet's Arrow batches. | The two `CometSpark*` names come from a single `CometSparkToColumnarExec` operator that picks the node name based on whether its child supports diff --git a/native/core/src/execution/columnar_to_row.rs b/native/core/src/execution/columnar_to_row.rs deleted file mode 100644 index 74defde6729..00000000000 --- a/native/core/src/execution/columnar_to_row.rs +++ /dev/null @@ -1,2726 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -//! Native implementation of columnar to row conversion for Spark UnsafeRow format. -//! -//! This module converts Arrow columnar data to Spark's UnsafeRow format, which is used -//! for row-based operations in Spark. The conversion is done in native code for better -//! performance compared to the JVM implementation. -//! -//! # UnsafeRow Format -//! -//! ```text -//! ┌─────────────────────────────────────────────────────────────┐ -//! │ Null Bitset: ((numFields + 63) / 64) * 8 bytes │ -//! ├─────────────────────────────────────────────────────────────┤ -//! │ Fixed-width portion: 8 bytes per field │ -//! │ - Primitives: value stored directly (in lowest bytes) │ -//! │ - Variable-length: (offset << 32) | length │ -//! ├─────────────────────────────────────────────────────────────┤ -//! │ Variable-length data: 8-byte aligned │ -//! └─────────────────────────────────────────────────────────────┘ -//! ``` - -use crate::errors::{CometError, CometResult}; -use arrow::array::types::{ - ArrowDictionaryKeyType, Decimal128Type, Int16Type, Int32Type, Int64Type, Int8Type, UInt16Type, - UInt32Type, UInt64Type, UInt8Type, -}; -use arrow::array::*; -use arrow::compute::kernels::arity::unary; -use arrow::compute::{cast_with_options, CastOptions}; -use arrow::datatypes::{ArrowNativeType, DataType, TimeUnit}; -use std::sync::Arc; - -/// Maximum digits for decimal that can fit in a long (8 bytes). -const MAX_LONG_DIGITS: u8 = 18; - -/// Helper macro for downcasting arrays with consistent error messages. -macro_rules! downcast_array { - ($array:expr, $array_type:ty) => { - $array - .as_any() - .downcast_ref::<$array_type>() - .ok_or_else(|| { - CometError::Internal(format!( - "Failed to downcast to {}, actual type: {:?}", - stringify!($array_type), - $array.data_type() - )) - }) - }; -} - -/// Macro to implement is_null for typed array enums. -/// Generates a complete match expression for all variants that have an array as first field. -macro_rules! impl_is_null { - ($self:expr, $row_idx:expr, [$($variant:ident),+ $(,)?]) => { - match $self { - $(Self::$variant(arr, ..) => arr.is_null($row_idx),)+ - } - }; - // Version with special handling for Null variant - ($self:expr, $row_idx:expr, null_always_true, [$($variant:ident),+ $(,)?]) => { - match $self { - Self::Null => true, - $(Self::$variant(arr, ..) => arr.is_null($row_idx),)+ - } - }; -} - -/// Macro to generate TypedElements::from_array match arms for primitive types. -macro_rules! typed_elements_from_primitive { - ($array:expr, $element_type:expr, $(($dt:pat, $variant:ident, $arr_type:ty)),+ $(,)?) => { - match $element_type { - $( - $dt => { - if let Some(arr) = $array.as_any().downcast_ref::<$arr_type>() { - return TypedElements::$variant(arr); - } - } - )+ - _ => {} - } - }; -} - -/// Macro for write_column_fixed_width arms - handles downcast + loop pattern. -macro_rules! write_fixed_column_primitive { - ($self:expr, $array:expr, $row_size:expr, $field_offset:expr, $num_rows:expr, - $arr_type:ty, $to_i64:expr) => {{ - let arr = downcast_array!($array, $arr_type)?; - for row_idx in 0..$num_rows { - if !arr.is_null(row_idx) { - let offset = row_idx * $row_size + $field_offset; - let value: i64 = $to_i64(arr.value(row_idx)); - $self.buffer[offset..offset + 8].copy_from_slice(&value.to_le_bytes()); - } - } - Ok(()) - }}; -} - -/// Macro for get_field_value arms - handles downcast + value extraction. -macro_rules! get_field_value_primitive { - ($array:expr, $row_idx:expr, $arr_type:ty, $to_i64:expr) => {{ - let arr = downcast_array!($array, $arr_type)?; - Ok($to_i64(arr.value($row_idx))) - }}; -} - -/// Macro for write_struct_to_buffer fixed-width field extraction. -macro_rules! extract_fixed_value { - ($column:expr, $row_idx:expr, $(($dt:pat, $arr_type:ty, $to_i64:expr)),+ $(,)?) => { - match $column.data_type() { - $( - $dt => { - let arr = downcast_array!($column, $arr_type)?; - Some($to_i64(arr.value($row_idx))) - } - )+ - _ => None, - } - }; -} - -/// Writes bytes to buffer with 8-byte alignment padding. -/// Returns the unpadded length. -#[inline] -fn write_bytes_padded(buffer: &mut Vec, bytes: &[u8]) -> usize { - let len = bytes.len(); - buffer.extend_from_slice(bytes); - let padding = round_up_to_8(len) - len; - buffer.extend(std::iter::repeat_n(0u8, padding)); - len -} - -/// Pre-downcast array reference to avoid type dispatch in inner loops. -/// This enum holds references to concrete array types, allowing direct access -/// without repeated downcast_ref calls. -enum TypedArray<'a> { - Null, - Boolean(&'a BooleanArray), - Int8(&'a Int8Array), - Int16(&'a Int16Array), - Int32(&'a Int32Array), - Int64(&'a Int64Array), - Float32(&'a Float32Array), - Float64(&'a Float64Array), - Date32(&'a Date32Array), - TimestampMicro(&'a TimestampMicrosecondArray), - Time64Nano(&'a Time64NanosecondArray), - Decimal128(&'a Decimal128Array, u8), // array + precision - String(&'a StringArray), - LargeString(&'a LargeStringArray), - Binary(&'a BinaryArray), - LargeBinary(&'a LargeBinaryArray), - FixedSizeBinary(&'a FixedSizeBinaryArray), - Struct( - &'a StructArray, - arrow::datatypes::Fields, - Vec>, - ), - List(&'a ListArray, arrow::datatypes::FieldRef), - LargeList(&'a LargeListArray, arrow::datatypes::FieldRef), - Map(&'a MapArray, arrow::datatypes::FieldRef), - Dictionary(&'a ArrayRef, DataType), // fallback for dictionary types -} - -impl<'a> TypedArray<'a> { - /// Pre-downcast an ArrayRef to a TypedArray. - fn from_array(array: &'a ArrayRef) -> CometResult { - let actual_type = array.data_type(); - match actual_type { - DataType::Null => { - // Verify the array is actually a NullArray, but we don't need to store the reference - // since all values are null by definition - downcast_array!(array, NullArray)?; - Ok(TypedArray::Null) - } - DataType::Boolean => Ok(TypedArray::Boolean(downcast_array!(array, BooleanArray)?)), - DataType::Int8 => Ok(TypedArray::Int8(downcast_array!(array, Int8Array)?)), - DataType::Int16 => Ok(TypedArray::Int16(downcast_array!(array, Int16Array)?)), - DataType::Int32 => Ok(TypedArray::Int32(downcast_array!(array, Int32Array)?)), - DataType::Int64 => Ok(TypedArray::Int64(downcast_array!(array, Int64Array)?)), - DataType::Float32 => Ok(TypedArray::Float32(downcast_array!(array, Float32Array)?)), - DataType::Float64 => Ok(TypedArray::Float64(downcast_array!(array, Float64Array)?)), - DataType::Date32 => Ok(TypedArray::Date32(downcast_array!(array, Date32Array)?)), - DataType::Timestamp(TimeUnit::Microsecond, _) => Ok(TypedArray::TimestampMicro( - downcast_array!(array, TimestampMicrosecondArray)?, - )), - DataType::Time64(TimeUnit::Nanosecond) => Ok(TypedArray::Time64Nano(downcast_array!( - array, - Time64NanosecondArray - )?)), - DataType::Decimal128(p, _) => Ok(TypedArray::Decimal128( - downcast_array!(array, Decimal128Array)?, - *p, - )), - DataType::Utf8 => Ok(TypedArray::String(downcast_array!(array, StringArray)?)), - DataType::LargeUtf8 => Ok(TypedArray::LargeString(downcast_array!( - array, - LargeStringArray - )?)), - DataType::Binary => Ok(TypedArray::Binary(downcast_array!(array, BinaryArray)?)), - DataType::LargeBinary => Ok(TypedArray::LargeBinary(downcast_array!( - array, - LargeBinaryArray - )?)), - DataType::FixedSizeBinary(_) => Ok(TypedArray::FixedSizeBinary(downcast_array!( - array, - FixedSizeBinaryArray - )?)), - DataType::Struct(fields) => { - let struct_arr = downcast_array!(array, StructArray)?; - // Pre-downcast all struct fields once - let typed_fields: Vec = fields - .iter() - .enumerate() - .map(|(idx, field)| { - TypedElements::from_array(struct_arr.column(idx), field.data_type()) - }) - .collect(); - Ok(TypedArray::Struct(struct_arr, fields.clone(), typed_fields)) - } - DataType::List(field) => Ok(TypedArray::List( - downcast_array!(array, ListArray)?, - Arc::clone(field), - )), - DataType::LargeList(field) => Ok(TypedArray::LargeList( - downcast_array!(array, LargeListArray)?, - Arc::clone(field), - )), - DataType::Map(field, _) => Ok(TypedArray::Map( - downcast_array!(array, MapArray)?, - Arc::clone(field), - )), - DataType::Dictionary(_, _) => Ok(TypedArray::Dictionary(array, actual_type.clone())), - _ => Err(CometError::Internal(format!( - "Unsupported data type for pre-downcast: {:?}", - actual_type - ))), - } - } - - /// Check if the value at the given index is null. - #[inline] - fn is_null(&self, row_idx: usize) -> bool { - impl_is_null!( - self, - row_idx, - null_always_true, - [ - Boolean, - Int8, - Int16, - Int32, - Int64, - Float32, - Float64, - Date32, - TimestampMicro, - Time64Nano, - Decimal128, - String, - LargeString, - Binary, - LargeBinary, - FixedSizeBinary, - Struct, - List, - LargeList, - Map, - Dictionary - ] - ) - } - - /// Get the fixed-width value as i64 (for types that fit in 8 bytes). - #[inline] - fn get_fixed_value(&self, row_idx: usize) -> i64 { - match self { - TypedArray::Boolean(arr) => arr.value(row_idx) as i64, - TypedArray::Int8(arr) => arr.value(row_idx) as i64, - TypedArray::Int16(arr) => arr.value(row_idx) as i64, - TypedArray::Int32(arr) => arr.value(row_idx) as i64, - TypedArray::Int64(arr) => arr.value(row_idx), - TypedArray::Float32(arr) => arr.value(row_idx).to_bits() as i64, - TypedArray::Float64(arr) => arr.value(row_idx).to_bits() as i64, - TypedArray::Date32(arr) => arr.value(row_idx) as i64, - TypedArray::TimestampMicro(arr) => arr.value(row_idx), - TypedArray::Time64Nano(arr) => arr.value(row_idx), - TypedArray::Decimal128(arr, precision) if *precision <= MAX_LONG_DIGITS => { - arr.value(row_idx) as i64 - } - TypedArray::Decimal128(_, _) => 0, // Variable-length decimal, handled elsewhere - // Variable-length types return 0, actual value written separately - _ => 0, - } - } - - /// Check if this is a variable-length type. - #[inline] - fn is_variable_length(&self) -> bool { - match self { - TypedArray::Null - | TypedArray::Boolean(_) - | TypedArray::Int8(_) - | TypedArray::Int16(_) - | TypedArray::Int32(_) - | TypedArray::Int64(_) - | TypedArray::Float32(_) - | TypedArray::Float64(_) - | TypedArray::Date32(_) - | TypedArray::TimestampMicro(_) - | TypedArray::Time64Nano(_) => false, - TypedArray::Decimal128(_, precision) => *precision > MAX_LONG_DIGITS, - _ => true, - } - } - - /// Write variable-length data to buffer. Returns actual length (0 if not variable-length). - fn write_variable_to_buffer(&self, buffer: &mut Vec, row_idx: usize) -> CometResult { - match self { - TypedArray::String(arr) => { - Ok(write_bytes_padded(buffer, arr.value(row_idx).as_bytes())) - } - TypedArray::LargeString(arr) => { - Ok(write_bytes_padded(buffer, arr.value(row_idx).as_bytes())) - } - TypedArray::Binary(arr) => Ok(write_bytes_padded(buffer, arr.value(row_idx))), - TypedArray::LargeBinary(arr) => Ok(write_bytes_padded(buffer, arr.value(row_idx))), - TypedArray::FixedSizeBinary(arr) => Ok(write_bytes_padded(buffer, arr.value(row_idx))), - TypedArray::Decimal128(arr, precision) if *precision > MAX_LONG_DIGITS => { - let bytes = i128_to_spark_decimal_bytes(arr.value(row_idx)); - Ok(write_bytes_padded(buffer, &bytes)) - } - TypedArray::Struct(arr, fields, typed_fields) => { - write_struct_to_buffer_typed(buffer, arr, row_idx, fields, typed_fields) - } - TypedArray::List(arr, field) => write_list_to_buffer(buffer, arr, row_idx, field), - TypedArray::LargeList(arr, field) => { - write_large_list_to_buffer(buffer, arr, row_idx, field) - } - TypedArray::Map(arr, field) => write_map_to_buffer(buffer, arr, row_idx, field), - TypedArray::Dictionary(arr, schema_type) => { - if let DataType::Dictionary(key_type, value_type) = schema_type { - write_dictionary_to_buffer( - buffer, - arr, - row_idx, - key_type.as_ref(), - value_type.as_ref(), - ) - } else { - Err(CometError::Internal(format!( - "Expected Dictionary type but got {:?}", - schema_type - ))) - } - } - _ => Ok(0), // Fixed-width types - } - } -} - -/// Pre-downcast element array for list/array types. -/// This allows direct access to element values without per-row allocation. -enum TypedElements<'a> { - Boolean(&'a BooleanArray), - Int8(&'a Int8Array), - Int16(&'a Int16Array), - Int32(&'a Int32Array), - Int64(&'a Int64Array), - Float32(&'a Float32Array), - Float64(&'a Float64Array), - Date32(&'a Date32Array), - TimestampMicro(&'a TimestampMicrosecondArray), - Time64Nano(&'a Time64NanosecondArray), - Decimal128(&'a Decimal128Array, u8), - String(&'a StringArray), - LargeString(&'a LargeStringArray), - Binary(&'a BinaryArray), - LargeBinary(&'a LargeBinaryArray), - FixedSizeBinary(&'a FixedSizeBinaryArray), - // For nested types, fall back to ArrayRef - Other(&'a ArrayRef, DataType), -} - -impl<'a> TypedElements<'a> { - /// Create from an ArrayRef and element type. - fn from_array(array: &'a ArrayRef, element_type: &DataType) -> Self { - // Try primitive types first using macro - typed_elements_from_primitive!( - array, - element_type, - (DataType::Boolean, Boolean, BooleanArray), - (DataType::Int8, Int8, Int8Array), - (DataType::Int16, Int16, Int16Array), - (DataType::Int32, Int32, Int32Array), - (DataType::Int64, Int64, Int64Array), - (DataType::Float32, Float32, Float32Array), - (DataType::Float64, Float64, Float64Array), - (DataType::Date32, Date32, Date32Array), - (DataType::Utf8, String, StringArray), - (DataType::LargeUtf8, LargeString, LargeStringArray), - (DataType::Binary, Binary, BinaryArray), - (DataType::LargeBinary, LargeBinary, LargeBinaryArray), - ); - - // Handle special cases that need extra processing - match element_type { - DataType::Timestamp(TimeUnit::Microsecond, _) => { - if let Some(arr) = array.as_any().downcast_ref::() { - return TypedElements::TimestampMicro(arr); - } - } - DataType::Time64(TimeUnit::Nanosecond) => { - if let Some(arr) = array.as_any().downcast_ref::() { - return TypedElements::Time64Nano(arr); - } - } - DataType::Decimal128(p, _) => { - if let Some(arr) = array.as_any().downcast_ref::() { - return TypedElements::Decimal128(arr, *p); - } - } - DataType::FixedSizeBinary(_) => { - if let Some(arr) = array.as_any().downcast_ref::() { - return TypedElements::FixedSizeBinary(arr); - } - } - _ => {} - } - TypedElements::Other(array, element_type.clone()) - } - - /// Get element size for UnsafeArrayData format. - fn element_size(&self) -> usize { - match self { - TypedElements::Boolean(_) => 1, - TypedElements::Int8(_) => 1, - TypedElements::Int16(_) => 2, - TypedElements::Int32(_) | TypedElements::Date32(_) | TypedElements::Float32(_) => 4, - TypedElements::Int64(_) - | TypedElements::TimestampMicro(_) - | TypedElements::Time64Nano(_) - | TypedElements::Float64(_) => 8, - TypedElements::Decimal128(_, p) if *p <= MAX_LONG_DIGITS => 8, - _ => 8, // Variable-length uses 8 bytes for offset+length - } - } - - /// Check if this is a fixed-width primitive type that supports bulk copy. - fn supports_bulk_copy(&self) -> bool { - matches!( - self, - TypedElements::Int8(_) - | TypedElements::Int16(_) - | TypedElements::Int32(_) - | TypedElements::Int64(_) - | TypedElements::Float32(_) - | TypedElements::Float64(_) - | TypedElements::Date32(_) - | TypedElements::TimestampMicro(_) - | TypedElements::Time64Nano(_) - ) - } - - /// Check if value at given index is null. - #[inline] - fn is_null_at(&self, idx: usize) -> bool { - impl_is_null!( - self, - idx, - [ - Boolean, - Int8, - Int16, - Int32, - Int64, - Float32, - Float64, - Date32, - TimestampMicro, - Time64Nano, - Decimal128, - String, - LargeString, - Binary, - LargeBinary, - FixedSizeBinary, - Other - ] - ) - } - - /// Check if this is a fixed-width type (value fits in 8-byte slot). - #[inline] - fn is_fixed_width(&self) -> bool { - match self { - TypedElements::Boolean(_) - | TypedElements::Int8(_) - | TypedElements::Int16(_) - | TypedElements::Int32(_) - | TypedElements::Int64(_) - | TypedElements::Float32(_) - | TypedElements::Float64(_) - | TypedElements::Date32(_) - | TypedElements::TimestampMicro(_) - | TypedElements::Time64Nano(_) => true, - TypedElements::Decimal128(_, p) => *p <= MAX_LONG_DIGITS, - _ => false, - } - } - - /// Get fixed-width value as i64 for the 8-byte field slot. - #[inline] - fn get_fixed_value(&self, idx: usize) -> i64 { - match self { - TypedElements::Boolean(arr) => arr.value(idx) as i64, - TypedElements::Int8(arr) => arr.value(idx) as i64, - TypedElements::Int16(arr) => arr.value(idx) as i64, - TypedElements::Int32(arr) => arr.value(idx) as i64, - TypedElements::Int64(arr) => arr.value(idx), - TypedElements::Float32(arr) => (arr.value(idx).to_bits() as i32) as i64, - TypedElements::Float64(arr) => arr.value(idx).to_bits() as i64, - TypedElements::Date32(arr) => arr.value(idx) as i64, - TypedElements::TimestampMicro(arr) => arr.value(idx), - TypedElements::Time64Nano(arr) => arr.value(idx), - TypedElements::Decimal128(arr, _) => arr.value(idx) as i64, - _ => 0, // Should not be called for variable-length types - } - } - - /// Write variable-length data to buffer. Returns length written (0 for fixed-width). - fn write_variable_value(&self, buffer: &mut Vec, idx: usize) -> CometResult { - match self { - TypedElements::String(arr) => Ok(write_bytes_padded(buffer, arr.value(idx).as_bytes())), - TypedElements::LargeString(arr) => { - Ok(write_bytes_padded(buffer, arr.value(idx).as_bytes())) - } - TypedElements::Binary(arr) => Ok(write_bytes_padded(buffer, arr.value(idx))), - TypedElements::LargeBinary(arr) => Ok(write_bytes_padded(buffer, arr.value(idx))), - TypedElements::FixedSizeBinary(arr) => Ok(write_bytes_padded(buffer, arr.value(idx))), - TypedElements::Decimal128(arr, precision) if *precision > MAX_LONG_DIGITS => { - let bytes = i128_to_spark_decimal_bytes(arr.value(idx)); - Ok(write_bytes_padded(buffer, &bytes)) - } - TypedElements::Other(arr, element_type) => { - write_nested_variable_to_buffer(buffer, element_type, arr, idx) - } - _ => Ok(0), // Fixed-width types - } - } - - /// Write a range of elements to buffer in UnsafeArrayData format. - /// Returns the total bytes written (including header). - fn write_range_to_buffer( - &self, - buffer: &mut Vec, - start_idx: usize, - num_elements: usize, - ) -> CometResult { - let element_size = self.element_size(); - let array_start = buffer.len(); - let element_bitset_width = ColumnarToRowContext::calculate_bitset_width(num_elements); - - // Write number of elements - buffer.extend_from_slice(&(num_elements as i64).to_le_bytes()); - - // Reserve space for null bitset - let null_bitset_start = buffer.len(); - buffer.resize(null_bitset_start + element_bitset_width, 0); - - // Reserve space for element values - let elements_start = buffer.len(); - let elements_total_size = round_up_to_8(num_elements * element_size); - buffer.resize(elements_start + elements_total_size, 0); - - // Try bulk copy for primitive types - if self.supports_bulk_copy() { - self.bulk_copy_range( - buffer, - null_bitset_start, - elements_start, - start_idx, - num_elements, - ); - return Ok(buffer.len() - array_start); - } - - // Handle other types element by element - self.write_elements_slow( - buffer, - array_start, - null_bitset_start, - elements_start, - element_size, - start_idx, - num_elements, - ) - } - - /// Bulk copy primitive values from a range. - fn bulk_copy_range( - &self, - buffer: &mut [u8], - null_bitset_start: usize, - elements_start: usize, - start_idx: usize, - num_elements: usize, - ) { - macro_rules! bulk_copy_range { - ($arr:expr, $elem_size:expr) => {{ - let values_slice = $arr.values(); - let byte_len = num_elements * $elem_size; - let src_start = start_idx * $elem_size; - debug_assert!( - src_start + byte_len <= values_slice.len() * $elem_size, - "bulk_copy_range: source slice out of bounds: src_start={}, byte_len={}, values_len={}, elem_size={}", - src_start, byte_len, values_slice.len() * $elem_size, $elem_size - ); - debug_assert!( - elements_start + byte_len <= buffer.len(), - "bulk_copy_range: destination slice out of bounds: elements_start={}, byte_len={}, buffer_len={}", - elements_start, byte_len, buffer.len() - ); - let src_bytes = unsafe { - std::slice::from_raw_parts( - (values_slice.as_ptr() as *const u8).add(src_start), - byte_len, - ) - }; - buffer[elements_start..elements_start + byte_len].copy_from_slice(src_bytes); - - // Set null bits - if $arr.null_count() > 0 { - for i in 0..num_elements { - if $arr.is_null(start_idx + i) { - let word_idx = i / 64; - let bit_idx = i % 64; - let word_offset = null_bitset_start + word_idx * 8; - let mut word = i64::from_le_bytes( - buffer[word_offset..word_offset + 8].try_into().unwrap(), - ); - word |= 1i64 << bit_idx; - buffer[word_offset..word_offset + 8] - .copy_from_slice(&word.to_le_bytes()); - } - } - } - }}; - } - - match self { - TypedElements::Int8(arr) => bulk_copy_range!(arr, 1), - TypedElements::Int16(arr) => bulk_copy_range!(arr, 2), - TypedElements::Int32(arr) => bulk_copy_range!(arr, 4), - TypedElements::Int64(arr) => bulk_copy_range!(arr, 8), - TypedElements::Float32(arr) => bulk_copy_range!(arr, 4), - TypedElements::Float64(arr) => bulk_copy_range!(arr, 8), - TypedElements::Date32(arr) => bulk_copy_range!(arr, 4), - TypedElements::TimestampMicro(arr) => bulk_copy_range!(arr, 8), - TypedElements::Time64Nano(arr) => bulk_copy_range!(arr, 8), - _ => {} // Should not reach here due to supports_bulk_copy check - } - } - - /// Slow path for non-bulk-copyable types. - #[allow(clippy::too_many_arguments)] - fn write_elements_slow( - &self, - buffer: &mut Vec, - array_start: usize, - null_bitset_start: usize, - elements_start: usize, - element_size: usize, - start_idx: usize, - num_elements: usize, - ) -> CometResult { - match self { - TypedElements::Boolean(arr) => { - for i in 0..num_elements { - let src_idx = start_idx + i; - if arr.is_null(src_idx) { - set_null_bit(buffer, null_bitset_start, i); - } else { - buffer[elements_start + i] = if arr.value(src_idx) { 1 } else { 0 }; - } - } - } - TypedElements::Decimal128(arr, precision) if *precision <= MAX_LONG_DIGITS => { - for i in 0..num_elements { - let src_idx = start_idx + i; - if arr.is_null(src_idx) { - set_null_bit(buffer, null_bitset_start, i); - } else { - let slot_offset = elements_start + i * 8; - let value = arr.value(src_idx) as i64; - buffer[slot_offset..slot_offset + 8].copy_from_slice(&value.to_le_bytes()); - } - } - } - TypedElements::Decimal128(arr, _) => { - // Large decimal - variable length - for i in 0..num_elements { - let src_idx = start_idx + i; - if arr.is_null(src_idx) { - set_null_bit(buffer, null_bitset_start, i); - } else { - let bytes = i128_to_spark_decimal_bytes(arr.value(src_idx)); - let len = write_bytes_padded(buffer, &bytes); - let data_offset = buffer.len() - round_up_to_8(len) - array_start; - let offset_and_len = ((data_offset as i64) << 32) | (len as i64); - let slot_offset = elements_start + i * 8; - buffer[slot_offset..slot_offset + 8] - .copy_from_slice(&offset_and_len.to_le_bytes()); - } - } - } - TypedElements::String(arr) => { - for i in 0..num_elements { - let src_idx = start_idx + i; - if arr.is_null(src_idx) { - set_null_bit(buffer, null_bitset_start, i); - } else { - let len = write_bytes_padded(buffer, arr.value(src_idx).as_bytes()); - let data_offset = buffer.len() - round_up_to_8(len) - array_start; - let offset_and_len = ((data_offset as i64) << 32) | (len as i64); - let slot_offset = elements_start + i * 8; - buffer[slot_offset..slot_offset + 8] - .copy_from_slice(&offset_and_len.to_le_bytes()); - } - } - } - TypedElements::LargeString(arr) => { - for i in 0..num_elements { - let src_idx = start_idx + i; - if arr.is_null(src_idx) { - set_null_bit(buffer, null_bitset_start, i); - } else { - let len = write_bytes_padded(buffer, arr.value(src_idx).as_bytes()); - let data_offset = buffer.len() - round_up_to_8(len) - array_start; - let offset_and_len = ((data_offset as i64) << 32) | (len as i64); - let slot_offset = elements_start + i * 8; - buffer[slot_offset..slot_offset + 8] - .copy_from_slice(&offset_and_len.to_le_bytes()); - } - } - } - TypedElements::Binary(arr) => { - for i in 0..num_elements { - let src_idx = start_idx + i; - if arr.is_null(src_idx) { - set_null_bit(buffer, null_bitset_start, i); - } else { - let len = write_bytes_padded(buffer, arr.value(src_idx)); - let data_offset = buffer.len() - round_up_to_8(len) - array_start; - let offset_and_len = ((data_offset as i64) << 32) | (len as i64); - let slot_offset = elements_start + i * 8; - buffer[slot_offset..slot_offset + 8] - .copy_from_slice(&offset_and_len.to_le_bytes()); - } - } - } - TypedElements::LargeBinary(arr) => { - for i in 0..num_elements { - let src_idx = start_idx + i; - if arr.is_null(src_idx) { - set_null_bit(buffer, null_bitset_start, i); - } else { - let len = write_bytes_padded(buffer, arr.value(src_idx)); - let data_offset = buffer.len() - round_up_to_8(len) - array_start; - let offset_and_len = ((data_offset as i64) << 32) | (len as i64); - let slot_offset = elements_start + i * 8; - buffer[slot_offset..slot_offset + 8] - .copy_from_slice(&offset_and_len.to_le_bytes()); - } - } - } - TypedElements::Other(arr, element_type) => { - // Fall back to old method for nested types - for i in 0..num_elements { - let src_idx = start_idx + i; - if arr.is_null(src_idx) { - set_null_bit(buffer, null_bitset_start, i); - } else { - let slot_offset = elements_start + i * element_size; - let var_len = - write_nested_variable_to_buffer(buffer, element_type, arr, src_idx)?; - - if var_len > 0 { - let padded_len = round_up_to_8(var_len); - let data_offset = buffer.len() - padded_len - array_start; - let offset_and_len = ((data_offset as i64) << 32) | (var_len as i64); - buffer[slot_offset..slot_offset + 8] - .copy_from_slice(&offset_and_len.to_le_bytes()); - } else { - let value = get_field_value(element_type, arr, src_idx)?; - write_array_element(buffer, element_type, value, slot_offset); - } - } - } - } - _ => { - // Should not reach here - all cases covered above - } - } - Ok(buffer.len() - array_start) - } -} - -/// Helper to set a null bit in the buffer. -#[inline] -fn set_null_bit(buffer: &mut [u8], null_bitset_start: usize, idx: usize) { - let word_idx = idx / 64; - let bit_idx = idx % 64; - let word_offset = null_bitset_start + word_idx * 8; - let mut word = i64::from_le_bytes(buffer[word_offset..word_offset + 8].try_into().unwrap()); - word |= 1i64 << bit_idx; - buffer[word_offset..word_offset + 8].copy_from_slice(&word.to_le_bytes()); -} - -/// Check if a data type is fixed-width for UnsafeRow purposes. -/// Fixed-width types are stored directly in the 8-byte field slot. -#[inline] -fn is_fixed_width(data_type: &DataType) -> bool { - match data_type { - DataType::Boolean - | DataType::Int8 - | DataType::Int16 - | DataType::Int32 - | DataType::Int64 - | DataType::Float32 - | DataType::Float64 - | DataType::Date32 - | DataType::Timestamp(TimeUnit::Microsecond, _) - | DataType::Time64(TimeUnit::Nanosecond) => true, - DataType::Decimal128(p, _) => *p <= MAX_LONG_DIGITS, - _ => false, - } -} - -/// Check if all columns in a schema are fixed-width. -#[inline] -fn is_all_fixed_width(schema: &[DataType]) -> bool { - schema.iter().all(is_fixed_width) -} - -/// Context for columnar to row conversion. -/// -/// This struct maintains the output buffer and schema information needed for -/// converting Arrow columnar data to Spark UnsafeRow format. The buffer is -/// reused across multiple `convert` calls to minimize allocations. -pub struct ColumnarToRowContext { - /// The Arrow data types for each column. - schema: Vec, - /// The output buffer containing converted rows. - /// Layout: [Row0][Row1]...[RowN] where each row is an UnsafeRow. - buffer: Vec, - /// Byte offset where each row starts in the buffer. - offsets: Vec, - /// Byte length of each row. - lengths: Vec, - /// Pre-calculated null bitset width in bytes. - null_bitset_width: usize, - /// Pre-calculated fixed-width portion size in bytes (null bitset + 8 bytes per field). - fixed_width_size: usize, - /// Maximum batch size for pre-allocation. - _batch_size: usize, - /// Whether all columns are fixed-width (enables fast path). - all_fixed_width: bool, -} - -impl ColumnarToRowContext { - /// Creates a new ColumnarToRowContext with the given schema. - /// - /// # Arguments - /// - /// * `schema` - The Arrow data types for each column. - /// * `batch_size` - Maximum number of rows expected per batch (for pre-allocation). - pub fn new(schema: Vec, batch_size: usize) -> Self { - let num_fields = schema.len(); - let null_bitset_width = Self::calculate_bitset_width(num_fields); - let fixed_width_size = null_bitset_width + num_fields * 8; - let all_fixed_width = is_all_fixed_width(&schema); - - // Pre-allocate buffer for maximum batch size - // For fixed-width schemas, we know exact size; otherwise estimate - let estimated_row_size = if all_fixed_width { - fixed_width_size - } else { - fixed_width_size + 64 // Conservative estimate for variable-length data - }; - let initial_capacity = batch_size * estimated_row_size; - - Self { - schema, - buffer: Vec::with_capacity(initial_capacity), - offsets: Vec::with_capacity(batch_size), - lengths: Vec::with_capacity(batch_size), - null_bitset_width, - fixed_width_size, - _batch_size: batch_size, - all_fixed_width, - } - } - - /// Calculate the width of the null bitset in bytes. - /// This matches Spark's `UnsafeRow.calculateBitSetWidthInBytes`. - #[inline] - pub const fn calculate_bitset_width(num_fields: usize) -> usize { - num_fields.div_ceil(64) * 8 - } - - /// Round up to the nearest multiple of 8 for alignment. - #[inline] - const fn round_up_to_8(value: usize) -> usize { - value.div_ceil(8) * 8 - } - - /// Converts Arrow arrays to Spark UnsafeRow format. - /// - /// # Arguments - /// - /// * `arrays` - The Arrow arrays to convert, one per column. - /// * `num_rows` - The number of rows to convert. - /// - /// # Returns - /// - /// A tuple containing: - /// - A pointer to the output buffer - /// - A reference to the offsets array - /// - A reference to the lengths array - pub fn convert( - &mut self, - arrays: &[ArrayRef], - num_rows: usize, - ) -> CometResult<(*const u8, &[i32], &[i32])> { - if arrays.len() != self.schema.len() { - return Err(CometError::Internal(format!( - "Column count mismatch: expected {}, got {}", - self.schema.len(), - arrays.len() - ))); - } - - // Unpack any dictionary arrays to their underlying value type - // This is needed because Parquet may return dictionary-encoded arrays - // even when the schema expects a specific type like Decimal128 - let arrays: Vec = arrays - .iter() - .zip(self.schema.iter()) - .map(|(arr, schema_type)| Self::maybe_cast_to_schema_type(arr, schema_type)) - .collect::>>()?; - let arrays = arrays.as_slice(); - - // Clear previous data - self.buffer.clear(); - self.offsets.clear(); - self.lengths.clear(); - - // Reserve space for offsets and lengths - self.offsets.reserve(num_rows); - self.lengths.reserve(num_rows); - - // Use fast path for fixed-width-only schemas - if self.all_fixed_width { - return self.convert_fixed_width(arrays, num_rows); - } - - // Pre-downcast all arrays to avoid type dispatch in inner loop - let typed_arrays: Vec = arrays - .iter() - .map(TypedArray::from_array) - .collect::>>()?; - - // Pre-compute variable-length column indices (once per batch, not per row) - let var_len_indices: Vec = typed_arrays - .iter() - .enumerate() - .filter(|(_, arr)| arr.is_variable_length()) - .map(|(idx, _)| idx) - .collect(); - - // Process each row (general path for variable-length data) - for row_idx in 0..num_rows { - let row_start = self.buffer.len(); - self.offsets.push(row_start as i32); - - // Write fixed-width portion (null bitset + field values) - self.write_row_typed(&typed_arrays, &var_len_indices, row_idx)?; - - let row_end = self.buffer.len(); - self.lengths.push((row_end - row_start) as i32); - } - - Ok((self.buffer.as_ptr(), &self.offsets, &self.lengths)) - } - - /// Casts an array to match the expected schema type if needed. - /// This handles cases where: - /// 1. Parquet returns dictionary-encoded arrays but the schema expects a non-dictionary type - /// 2. Parquet returns NullArray when all values are null, but the schema expects a typed array - /// 3. Parquet returns Int32/Int64 for small-precision decimals but schema expects Decimal128 - fn maybe_cast_to_schema_type( - array: &ArrayRef, - schema_type: &DataType, - ) -> CometResult { - let actual_type = array.data_type(); - - // If types already match, no cast needed - if actual_type == schema_type { - return Ok(Arc::clone(array)); - } - - match (actual_type, schema_type) { - (DataType::Dictionary(_, _), schema) - if !matches!(schema, DataType::Dictionary(_, _)) => - { - // Unpack dictionary if the schema type is not a dictionary - let options = CastOptions::default(); - cast_with_options(array, schema_type, &options).map_err(|e| { - CometError::Internal(format!( - "Failed to unpack dictionary array from {:?} to {:?}: {}", - actual_type, schema_type, e - )) - }) - } - (DataType::Null, _) => { - // Cast NullArray to the expected schema type - // This happens when all values in a column are null - let options = CastOptions::default(); - cast_with_options(array, schema_type, &options).map_err(|e| { - CometError::Internal(format!( - "Failed to cast NullArray to {:?}: {}", - schema_type, e - )) - }) - } - (DataType::Int32, DataType::Decimal128(precision, scale)) => { - // Parquet stores small-precision decimals as Int32 for efficiency, and the - // reader may surface them as the physical Int32 type. The value is already - // scaled (e.g., -1 means -0.01 for scale 2). Reinterpret (not cast) to - // Decimal128 preserving the value. `arity::unary` widens the value and reuses - // the input null buffer zero-copy (an Arrow cast would rescale the value). - let int_array = array.as_primitive::(); - let decimal_array = unary::<_, _, Decimal128Type>(int_array, |x| x as i128) - .with_precision_and_scale(*precision, *scale) - .map_err(|e| { - CometError::Internal(format!("Invalid decimal precision/scale: {}", e)) - })?; - Ok(Arc::new(decimal_array)) - } - (DataType::Int64, DataType::Decimal128(precision, scale)) => { - // Same as Int32 but for medium-precision decimals stored as Int64. - let int_array = array.as_primitive::(); - let decimal_array = unary::<_, _, Decimal128Type>(int_array, |x| x as i128) - .with_precision_and_scale(*precision, *scale) - .map_err(|e| { - CometError::Internal(format!("Invalid decimal precision/scale: {}", e)) - })?; - Ok(Arc::new(decimal_array)) - } - _ => { - // For any other type mismatch, attempt an Arrow cast. - // This handles cases like Int32 → Date32 (which can happen when Spark - // generates default column values using the physical storage type rather - // than the logical type). - let options = CastOptions::default(); - cast_with_options(array, schema_type, &options).map_err(|e| { - CometError::Internal(format!( - "Failed to cast array from {:?} to {:?}: {}", - actual_type, schema_type, e - )) - }) - } - } - } - - /// Fast path for schemas with only fixed-width columns. - /// Pre-allocates entire buffer and processes more efficiently. - fn convert_fixed_width( - &mut self, - arrays: &[ArrayRef], - num_rows: usize, - ) -> CometResult<(*const u8, &[i32], &[i32])> { - let row_size = self.fixed_width_size; - let total_size = row_size * num_rows; - let null_bitset_width = self.null_bitset_width; - - // Pre-allocate entire buffer at once (all zeros) - self.buffer.resize(total_size, 0); - - // Pre-fill offsets and lengths (constant for fixed-width) - let row_size_i32 = row_size as i32; - for row_idx in 0..num_rows { - self.offsets.push((row_idx * row_size) as i32); - self.lengths.push(row_size_i32); - } - - // Process column by column for better cache locality - for (col_idx, array) in arrays.iter().enumerate() { - let field_offset_in_row = null_bitset_width + col_idx * 8; - - // Write values for all rows in this column - self.write_column_fixed_width( - array, - &self.schema[col_idx].clone(), - col_idx, - field_offset_in_row, - row_size, - num_rows, - )?; - } - - Ok((self.buffer.as_ptr(), &self.offsets, &self.lengths)) - } - - /// Write a fixed-width column's values for all rows. - /// Processes column-by-column for better cache locality. - fn write_column_fixed_width( - &mut self, - array: &ArrayRef, - data_type: &DataType, - col_idx: usize, - field_offset_in_row: usize, - row_size: usize, - num_rows: usize, - ) -> CometResult<()> { - // Handle nulls first - set null bits - if array.null_count() > 0 { - let word_idx = col_idx / 64; - let bit_idx = col_idx % 64; - let bit_mask = 1i64 << bit_idx; - - for row_idx in 0..num_rows { - if array.is_null(row_idx) { - let row_start = row_idx * row_size; - let word_offset = row_start + word_idx * 8; - let mut word = i64::from_le_bytes( - self.buffer[word_offset..word_offset + 8] - .try_into() - .unwrap(), - ); - word |= bit_mask; - self.buffer[word_offset..word_offset + 8].copy_from_slice(&word.to_le_bytes()); - } - } - } - - // Write non-null values using type-specific fast paths - match data_type { - DataType::Boolean => { - // Boolean is special: writes single byte, not 8-byte i64 - let arr = downcast_array!(array, BooleanArray)?; - for row_idx in 0..num_rows { - if !arr.is_null(row_idx) { - let offset = row_idx * row_size + field_offset_in_row; - self.buffer[offset] = if arr.value(row_idx) { 1 } else { 0 }; - } - } - Ok(()) - } - DataType::Int8 => write_fixed_column_primitive!( - self, - array, - row_size, - field_offset_in_row, - num_rows, - Int8Array, - |v: i8| v as i64 - ), - DataType::Int16 => write_fixed_column_primitive!( - self, - array, - row_size, - field_offset_in_row, - num_rows, - Int16Array, - |v: i16| v as i64 - ), - DataType::Int32 => write_fixed_column_primitive!( - self, - array, - row_size, - field_offset_in_row, - num_rows, - Int32Array, - |v: i32| v as i64 - ), - DataType::Int64 => write_fixed_column_primitive!( - self, - array, - row_size, - field_offset_in_row, - num_rows, - Int64Array, - |v: i64| v - ), - DataType::Float32 => write_fixed_column_primitive!( - self, - array, - row_size, - field_offset_in_row, - num_rows, - Float32Array, - |v: f32| v.to_bits() as i64 - ), - DataType::Float64 => write_fixed_column_primitive!( - self, - array, - row_size, - field_offset_in_row, - num_rows, - Float64Array, - |v: f64| v.to_bits() as i64 - ), - DataType::Date32 => write_fixed_column_primitive!( - self, - array, - row_size, - field_offset_in_row, - num_rows, - Date32Array, - |v: i32| v as i64 - ), - DataType::Timestamp(TimeUnit::Microsecond, _) => write_fixed_column_primitive!( - self, - array, - row_size, - field_offset_in_row, - num_rows, - TimestampMicrosecondArray, - |v: i64| v - ), - DataType::Time64(TimeUnit::Nanosecond) => write_fixed_column_primitive!( - self, - array, - row_size, - field_offset_in_row, - num_rows, - Time64NanosecondArray, - |v: i64| v - ), - DataType::Decimal128(precision, _) if *precision <= MAX_LONG_DIGITS => { - write_fixed_column_primitive!( - self, - array, - row_size, - field_offset_in_row, - num_rows, - Decimal128Array, - |v: i128| v as i64 - ) - } - _ => Err(CometError::Internal(format!( - "Unexpected non-fixed-width type in fast path: {:?}", - data_type - ))), - } - } - - /// Writes a complete row using pre-downcast TypedArrays. - /// This avoids type dispatch overhead in the inner loop. - fn write_row_typed( - &mut self, - typed_arrays: &[TypedArray], - var_len_indices: &[usize], - row_idx: usize, - ) -> CometResult<()> { - let row_start = self.buffer.len(); - let null_bitset_width = self.null_bitset_width; - let fixed_width_size = self.fixed_width_size; - - // Extend buffer for fixed-width portion - self.buffer.resize(row_start + fixed_width_size, 0); - - // First pass: write null bits and fixed-width values - for (col_idx, typed_arr) in typed_arrays.iter().enumerate() { - let is_null = typed_arr.is_null(row_idx); - - if is_null { - // Set null bit - let word_idx = col_idx / 64; - let bit_idx = col_idx % 64; - let word_offset = row_start + word_idx * 8; - - let mut word = i64::from_le_bytes( - self.buffer[word_offset..word_offset + 8] - .try_into() - .unwrap(), - ); - word |= 1i64 << bit_idx; - self.buffer[word_offset..word_offset + 8].copy_from_slice(&word.to_le_bytes()); - } else if !typed_arr.is_variable_length() { - // Write fixed-width field value (skip variable-length, they're handled in pass 2) - let field_offset = row_start + null_bitset_width + col_idx * 8; - let value = typed_arr.get_fixed_value(row_idx); - self.buffer[field_offset..field_offset + 8].copy_from_slice(&value.to_le_bytes()); - } - } - - // Second pass: write variable-length data (only iterate over var-len columns) - for &col_idx in var_len_indices { - let typed_arr = &typed_arrays[col_idx]; - if typed_arr.is_null(row_idx) { - continue; - } - - // Write variable-length data directly to buffer - let actual_len = typed_arr.write_variable_to_buffer(&mut self.buffer, row_idx)?; - if actual_len > 0 { - // Calculate offset: buffer grew by padded_len, but we need offset to start of data - let padded_len = Self::round_up_to_8(actual_len); - let current_offset = self.buffer.len() - padded_len - row_start; - - // Update the field slot with (offset << 32) | length - let offset_and_len = ((current_offset as i64) << 32) | (actual_len as i64); - let field_offset = row_start + null_bitset_width + col_idx * 8; - self.buffer[field_offset..field_offset + 8] - .copy_from_slice(&offset_and_len.to_le_bytes()); - } - } - - Ok(()) - } - - /// Returns a pointer to the buffer. - pub fn buffer_ptr(&self) -> *const u8 { - self.buffer.as_ptr() - } - - /// Returns the schema. - pub fn schema(&self) -> &[DataType] { - &self.schema - } -} - -/// Gets the fixed-width value for a field as i64. -#[inline] -fn get_field_value(data_type: &DataType, array: &ArrayRef, row_idx: usize) -> CometResult { - // Use the actual array type for dispatching to handle type mismatches - let actual_type = array.data_type(); - - match actual_type { - DataType::Boolean => { - let arr = downcast_array!(array, BooleanArray)?; - Ok(arr.value(row_idx) as i64) - } - DataType::Int8 => get_field_value_primitive!(array, row_idx, Int8Array, |v: i8| v as i64), - DataType::Int16 => { - get_field_value_primitive!(array, row_idx, Int16Array, |v: i16| v as i64) - } - DataType::Int32 => { - get_field_value_primitive!(array, row_idx, Int32Array, |v: i32| v as i64) - } - DataType::Int64 => get_field_value_primitive!(array, row_idx, Int64Array, |v: i64| v), - DataType::Float32 => { - get_field_value_primitive!(array, row_idx, Float32Array, |v: f32| v.to_bits() as i64) - } - DataType::Float64 => { - get_field_value_primitive!(array, row_idx, Float64Array, |v: f64| v.to_bits() as i64) - } - DataType::Date32 => { - get_field_value_primitive!(array, row_idx, Date32Array, |v: i32| v as i64) - } - DataType::Timestamp(TimeUnit::Microsecond, _) => { - get_field_value_primitive!(array, row_idx, TimestampMicrosecondArray, |v: i64| v) - } - DataType::Time64(TimeUnit::Nanosecond) => { - get_field_value_primitive!(array, row_idx, Time64NanosecondArray, |v: i64| v) - } - DataType::Decimal128(precision, _) if *precision <= MAX_LONG_DIGITS => { - get_field_value_primitive!(array, row_idx, Decimal128Array, |v: i128| v as i64) - } - // Variable-length types use placeholder (will be overwritten by get_variable_length_data) - DataType::Utf8 - | DataType::LargeUtf8 - | DataType::Binary - | DataType::LargeBinary - | DataType::Decimal128(_, _) - | DataType::Struct(_) - | DataType::List(_) - | DataType::LargeList(_) - | DataType::Map(_, _) - | DataType::Dictionary(_, _) => Ok(0i64), - _ => { - // Check if the schema type is a known type that we should handle - match data_type { - DataType::Boolean - | DataType::Int8 - | DataType::Int16 - | DataType::Int32 - | DataType::Int64 - | DataType::Float32 - | DataType::Float64 - | DataType::Date32 - | DataType::Timestamp(TimeUnit::Microsecond, _) - | DataType::Decimal128(_, _) => Err(CometError::Internal(format!( - "Type mismatch in get_field_value: schema expects {:?} but actual array type is {:?}", - data_type, actual_type - ))), - // If schema is also a variable-length type, return placeholder - DataType::Utf8 - | DataType::LargeUtf8 - | DataType::Binary - | DataType::LargeBinary - | DataType::Struct(_) - | DataType::List(_) - | DataType::LargeList(_) - | DataType::Map(_, _) - | DataType::Dictionary(_, _) => Ok(0i64), - _ => Err(CometError::Internal(format!( - "Unsupported data type for columnar to row conversion: schema={:?}, actual={:?}", - data_type, actual_type - ))), - } - } - } -} - -/// Writes dictionary-encoded value directly to buffer. -#[inline] -fn write_dictionary_to_buffer( - buffer: &mut Vec, - array: &ArrayRef, - row_idx: usize, - key_type: &DataType, - value_type: &DataType, -) -> CometResult { - match key_type { - DataType::Int8 => { - write_dictionary_to_buffer_with_key::(buffer, array, row_idx, value_type) - } - DataType::Int16 => { - write_dictionary_to_buffer_with_key::(buffer, array, row_idx, value_type) - } - DataType::Int32 => { - write_dictionary_to_buffer_with_key::(buffer, array, row_idx, value_type) - } - DataType::Int64 => { - write_dictionary_to_buffer_with_key::(buffer, array, row_idx, value_type) - } - DataType::UInt8 => { - write_dictionary_to_buffer_with_key::(buffer, array, row_idx, value_type) - } - DataType::UInt16 => { - write_dictionary_to_buffer_with_key::(buffer, array, row_idx, value_type) - } - DataType::UInt32 => { - write_dictionary_to_buffer_with_key::(buffer, array, row_idx, value_type) - } - DataType::UInt64 => { - write_dictionary_to_buffer_with_key::(buffer, array, row_idx, value_type) - } - _ => Err(CometError::Internal(format!( - "Unsupported dictionary key type: {:?}", - key_type - ))), - } -} - -/// Writes dictionary value directly to buffer with specific key type. -#[inline] -fn write_dictionary_to_buffer_with_key( - buffer: &mut Vec, - array: &ArrayRef, - row_idx: usize, - value_type: &DataType, -) -> CometResult { - let dict_array = array - .as_any() - .downcast_ref::>() - .ok_or_else(|| { - CometError::Internal(format!( - "Failed to downcast to DictionaryArray<{:?}>", - std::any::type_name::() - )) - })?; - - let values = dict_array.values(); - let key_idx = dict_array.keys().value(row_idx).to_usize().ok_or_else(|| { - CometError::Internal("Dictionary key index out of usize range".to_string()) - })?; - - match value_type { - DataType::Utf8 => { - let string_values = downcast_array!(values, StringArray)?; - Ok(write_bytes_padded( - buffer, - string_values.value(key_idx).as_bytes(), - )) - } - DataType::LargeUtf8 => { - let string_values = downcast_array!(values, LargeStringArray)?; - Ok(write_bytes_padded( - buffer, - string_values.value(key_idx).as_bytes(), - )) - } - DataType::Binary => { - let binary_values = downcast_array!(values, BinaryArray)?; - Ok(write_bytes_padded(buffer, binary_values.value(key_idx))) - } - DataType::LargeBinary => { - let binary_values = downcast_array!(values, LargeBinaryArray)?; - Ok(write_bytes_padded(buffer, binary_values.value(key_idx))) - } - _ => Err(CometError::Internal(format!( - "Unsupported dictionary value type for direct buffer write: {:?}", - value_type - ))), - } -} - -/// Converts i128 to Spark's big-endian decimal byte format. -fn i128_to_spark_decimal_bytes(value: i128) -> Vec { - // Spark uses big-endian format for large decimals - let bytes = value.to_be_bytes(); - - // Find the minimum number of bytes needed (excluding leading sign-extension bytes) - let is_negative = value < 0; - let sign_byte = if is_negative { 0xFF } else { 0x00 }; - - let mut start = 0; - while start < 15 && bytes[start] == sign_byte { - // Check if the next byte's sign bit matches - let next_byte = bytes[start + 1]; - let has_correct_sign = if is_negative { - (next_byte & 0x80) != 0 - } else { - (next_byte & 0x80) == 0 - }; - if has_correct_sign { - start += 1; - } else { - break; - } - } - - bytes[start..].to_vec() -} - -/// Round up to the nearest multiple of 8 for alignment. -#[inline] -const fn round_up_to_8(value: usize) -> usize { - value.div_ceil(8) * 8 -} - -/// Writes a primitive value with the correct size for UnsafeArrayData. -#[inline] -fn write_array_element(buffer: &mut [u8], data_type: &DataType, value: i64, offset: usize) { - match data_type { - DataType::Boolean => { - buffer[offset] = if value != 0 { 1 } else { 0 }; - } - DataType::Int8 => { - buffer[offset] = value as u8; - } - DataType::Int16 => { - buffer[offset..offset + 2].copy_from_slice(&(value as i16).to_le_bytes()); - } - DataType::Int32 | DataType::Date32 => { - buffer[offset..offset + 4].copy_from_slice(&(value as i32).to_le_bytes()); - } - DataType::Float32 => { - buffer[offset..offset + 4].copy_from_slice(&(value as u32).to_le_bytes()); - } - // All 8-byte types - _ => { - buffer[offset..offset + 8].copy_from_slice(&value.to_le_bytes()); - } - } -} - -// ============================================================================= -// Optimized direct-write functions for complex types -// These write directly to the output buffer to avoid intermediate allocations. -// ============================================================================= - -/// Writes a struct value directly to the buffer using pre-downcast typed fields. -/// Returns the unpadded length written. -/// -/// This version uses pre-downcast TypedElements for each field, eliminating -/// per-row type dispatch overhead. -#[inline] -fn write_struct_to_buffer_typed( - buffer: &mut Vec, - _struct_array: &StructArray, - row_idx: usize, - _fields: &arrow::datatypes::Fields, - typed_fields: &[TypedElements], -) -> CometResult { - let num_fields = typed_fields.len(); - let nested_bitset_width = ColumnarToRowContext::calculate_bitset_width(num_fields); - let nested_fixed_size = nested_bitset_width + num_fields * 8; - - // Remember where this struct starts in the buffer - let struct_start = buffer.len(); - - // Reserve space for fixed-width portion (zeros for null bits and field slots) - buffer.resize(struct_start + nested_fixed_size, 0); - - // Write each field using pre-downcast types - for (field_idx, typed_field) in typed_fields.iter().enumerate() { - if typed_field.is_null_at(row_idx) { - // Set null bit in nested struct - set_null_bit(buffer, struct_start, field_idx); - } else { - let field_offset = struct_start + nested_bitset_width + field_idx * 8; - - if typed_field.is_fixed_width() { - // Fixed-width field - use pre-downcast accessor - let value = typed_field.get_fixed_value(row_idx); - buffer[field_offset..field_offset + 8].copy_from_slice(&value.to_le_bytes()); - } else { - // Variable-length field - use pre-downcast writer - let var_len = typed_field.write_variable_value(buffer, row_idx)?; - if var_len > 0 { - let padded_len = round_up_to_8(var_len); - let data_offset = buffer.len() - padded_len - struct_start; - let offset_and_len = ((data_offset as i64) << 32) | (var_len as i64); - buffer[field_offset..field_offset + 8] - .copy_from_slice(&offset_and_len.to_le_bytes()); - } - } - } - } - - Ok(buffer.len() - struct_start) -} - -/// Writes a struct value directly to the buffer. -/// Returns the unpadded length written. -/// -/// Processes each field using inline type dispatch to avoid allocation overhead. -/// This is used for nested structs where we don't have pre-downcast fields. -fn write_struct_to_buffer( - buffer: &mut Vec, - struct_array: &StructArray, - row_idx: usize, - fields: &arrow::datatypes::Fields, -) -> CometResult { - let num_fields = fields.len(); - let nested_bitset_width = ColumnarToRowContext::calculate_bitset_width(num_fields); - let nested_fixed_size = nested_bitset_width + num_fields * 8; - - // Remember where this struct starts in the buffer - let struct_start = buffer.len(); - - // Reserve space for fixed-width portion (zeros for null bits and field slots) - buffer.resize(struct_start + nested_fixed_size, 0); - - // Write each field with inline type handling (no allocation) - for (field_idx, field) in fields.iter().enumerate() { - let column = struct_array.column(field_idx); - let data_type = field.data_type(); - - if column.is_null(row_idx) { - // Set null bit in nested struct - set_null_bit(buffer, struct_start, field_idx); - } else { - let field_offset = struct_start + nested_bitset_width + field_idx * 8; - - // Inline type dispatch for fixed-width types (most common case) - let value: Option = extract_fixed_value!( - column, - row_idx, - (DataType::Boolean, BooleanArray, |v: bool| if v { - 1i64 - } else { - 0i64 - }), - (DataType::Int8, Int8Array, |v: i8| v as i64), - (DataType::Int16, Int16Array, |v: i16| v as i64), - (DataType::Int32, Int32Array, |v: i32| v as i64), - (DataType::Int64, Int64Array, |v: i64| v), - (DataType::Float32, Float32Array, |v: f32| v.to_bits() as i64), - (DataType::Float64, Float64Array, |v: f64| v.to_bits() as i64), - (DataType::Date32, Date32Array, |v: i32| v as i64), - ( - DataType::Timestamp(TimeUnit::Microsecond, _), - TimestampMicrosecondArray, - |v: i64| v - ), - ); - // Handle Decimal128 with precision guard separately - let value: Option = match (value, data_type) { - (Some(v), _) => Some(v), - (None, DataType::Decimal128(p, _)) if *p <= MAX_LONG_DIGITS => { - let arr = downcast_array!(column, Decimal128Array)?; - Some(arr.value(row_idx) as i64) - } - _ => None, - }; - - if let Some(v) = value { - // Fixed-width field - buffer[field_offset..field_offset + 8].copy_from_slice(&v.to_le_bytes()); - } else { - // Variable-length field - let var_len = write_nested_variable_to_buffer(buffer, data_type, column, row_idx)?; - if var_len > 0 { - let padded_len = round_up_to_8(var_len); - let data_offset = buffer.len() - padded_len - struct_start; - let offset_and_len = ((data_offset as i64) << 32) | (var_len as i64); - buffer[field_offset..field_offset + 8] - .copy_from_slice(&offset_and_len.to_le_bytes()); - } - } - } - } - - Ok(buffer.len() - struct_start) -} - -/// Writes a list value directly to the buffer in UnsafeArrayData format. -/// Returns the unpadded length written. -/// -/// This uses offsets directly to avoid per-row ArrayRef allocation. -#[inline] -fn write_list_to_buffer( - buffer: &mut Vec, - list_array: &ListArray, - row_idx: usize, - element_field: &arrow::datatypes::FieldRef, -) -> CometResult { - // Get offsets directly to avoid creating a sliced ArrayRef - let offsets = list_array.value_offsets(); - let start_offset = offsets[row_idx] as usize; - let end_offset = offsets[row_idx + 1] as usize; - let num_elements = end_offset - start_offset; - - // Pre-downcast the element array once - let element_array = list_array.values(); - let element_type = element_field.data_type(); - let typed_elements = TypedElements::from_array(element_array, element_type); - - // Write the range of elements - typed_elements.write_range_to_buffer(buffer, start_offset, num_elements) -} - -/// Writes a large list value directly to the buffer in UnsafeArrayData format. -/// Returns the unpadded length written. -/// -/// This uses offsets directly to avoid per-row ArrayRef allocation. -#[inline] -fn write_large_list_to_buffer( - buffer: &mut Vec, - list_array: &LargeListArray, - row_idx: usize, - element_field: &arrow::datatypes::FieldRef, -) -> CometResult { - // Get offsets directly to avoid creating a sliced ArrayRef - let offsets = list_array.value_offsets(); - let start_offset = offsets[row_idx] as usize; - let end_offset = offsets[row_idx + 1] as usize; - let num_elements = end_offset - start_offset; - - // Pre-downcast the element array once - let element_array = list_array.values(); - let element_type = element_field.data_type(); - let typed_elements = TypedElements::from_array(element_array, element_type); - - // Write the range of elements - typed_elements.write_range_to_buffer(buffer, start_offset, num_elements) -} - -/// Writes a map value directly to the buffer in UnsafeMapData format. -/// Returns the unpadded length written. -/// -/// This uses offsets directly to avoid per-row ArrayRef allocation. -fn write_map_to_buffer( - buffer: &mut Vec, - map_array: &MapArray, - row_idx: usize, - entries_field: &arrow::datatypes::FieldRef, -) -> CometResult { - // UnsafeMapData format: - // [key array size: 8 bytes][key array data][value array data] - let map_start = buffer.len(); - - // Get offsets directly to avoid creating a sliced ArrayRef - let offsets = map_array.value_offsets(); - let start_offset = offsets[row_idx] as usize; - let end_offset = offsets[row_idx + 1] as usize; - let num_entries = end_offset - start_offset; - - // Get keys and values from the underlying entries struct - let entries_array = map_array.entries(); - let keys = entries_array.column(0); - let values = entries_array.column(1); - - let (key_type, value_type) = if let DataType::Struct(fields) = entries_field.data_type() { - (fields[0].data_type().clone(), fields[1].data_type().clone()) - } else { - return Err(CometError::Internal(format!( - "Map entries field is not a struct: {:?}", - entries_field.data_type() - ))); - }; - - // Pre-downcast keys and values once - let typed_keys = TypedElements::from_array(keys, &key_type); - let typed_values = TypedElements::from_array(values, &value_type); - - // Placeholder for key array size - let key_size_offset = buffer.len(); - buffer.extend_from_slice(&0i64.to_le_bytes()); - - // Write key array using range - let key_array_start = buffer.len(); - typed_keys.write_range_to_buffer(buffer, start_offset, num_entries)?; - let key_array_size = (buffer.len() - key_array_start) as i64; - buffer[key_size_offset..key_size_offset + 8].copy_from_slice(&key_array_size.to_le_bytes()); - - // Write value array using range - typed_values.write_range_to_buffer(buffer, start_offset, num_entries)?; - - Ok(buffer.len() - map_start) -} - -/// Writes variable-length data for a nested field directly to buffer. -/// Used by struct, list, and map writers for their nested elements. -/// Returns the unpadded length written (0 if not variable-length). -fn write_nested_variable_to_buffer( - buffer: &mut Vec, - data_type: &DataType, - array: &ArrayRef, - row_idx: usize, -) -> CometResult { - let actual_type = array.data_type(); - - match actual_type { - DataType::Utf8 => { - let arr = downcast_array!(array, StringArray)?; - Ok(write_bytes_padded(buffer, arr.value(row_idx).as_bytes())) - } - DataType::LargeUtf8 => { - let arr = downcast_array!(array, LargeStringArray)?; - Ok(write_bytes_padded(buffer, arr.value(row_idx).as_bytes())) - } - DataType::Binary => { - let arr = downcast_array!(array, BinaryArray)?; - Ok(write_bytes_padded(buffer, arr.value(row_idx))) - } - DataType::LargeBinary => { - let arr = downcast_array!(array, LargeBinaryArray)?; - Ok(write_bytes_padded(buffer, arr.value(row_idx))) - } - DataType::Decimal128(precision, _) if *precision > MAX_LONG_DIGITS => { - let arr = downcast_array!(array, Decimal128Array)?; - let bytes = i128_to_spark_decimal_bytes(arr.value(row_idx)); - Ok(write_bytes_padded(buffer, &bytes)) - } - DataType::Struct(fields) => { - let struct_array = downcast_array!(array, StructArray)?; - write_struct_to_buffer(buffer, struct_array, row_idx, fields) - } - DataType::List(field) => { - let list_array = downcast_array!(array, ListArray)?; - write_list_to_buffer(buffer, list_array, row_idx, field) - } - DataType::LargeList(field) => { - let list_array = downcast_array!(array, LargeListArray)?; - write_large_list_to_buffer(buffer, list_array, row_idx, field) - } - DataType::Map(field, _) => { - let map_array = downcast_array!(array, MapArray)?; - write_map_to_buffer(buffer, map_array, row_idx, field) - } - DataType::Dictionary(key_type, value_type) => { - write_dictionary_to_buffer(buffer, array, row_idx, key_type, value_type) - } - // Check if schema type expects variable-length but actual type doesn't match - _ => match data_type { - DataType::Utf8 - | DataType::LargeUtf8 - | DataType::Binary - | DataType::LargeBinary - | DataType::Struct(_) - | DataType::List(_) - | DataType::LargeList(_) - | DataType::Map(_, _) => Err(CometError::Internal(format!( - "Type mismatch in nested write: schema expects {:?} but actual array type is {:?}", - data_type, actual_type - ))), - DataType::Decimal128(precision, _) if *precision > MAX_LONG_DIGITS => { - Err(CometError::Internal(format!( - "Type mismatch for large decimal: schema expects {:?} but actual is {:?}", - data_type, actual_type - ))) - } - _ => Ok(0), // Not a variable-length type - }, - } -} - -#[cfg(test)] -mod tests { - use super::*; - use std::sync::Arc; - - #[test] - fn test_bitset_width_calculation() { - assert_eq!(ColumnarToRowContext::calculate_bitset_width(0), 0); - assert_eq!(ColumnarToRowContext::calculate_bitset_width(1), 8); - assert_eq!(ColumnarToRowContext::calculate_bitset_width(64), 8); - assert_eq!(ColumnarToRowContext::calculate_bitset_width(65), 16); - assert_eq!(ColumnarToRowContext::calculate_bitset_width(128), 16); - assert_eq!(ColumnarToRowContext::calculate_bitset_width(129), 24); - } - - #[test] - fn test_round_up_to_8() { - assert_eq!(ColumnarToRowContext::round_up_to_8(0), 0); - assert_eq!(ColumnarToRowContext::round_up_to_8(1), 8); - assert_eq!(ColumnarToRowContext::round_up_to_8(7), 8); - assert_eq!(ColumnarToRowContext::round_up_to_8(8), 8); - assert_eq!(ColumnarToRowContext::round_up_to_8(9), 16); - } - - #[test] - fn test_convert_int_array() { - let schema = vec![DataType::Int32]; - let mut ctx = ColumnarToRowContext::new(schema, 100); - - let array: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), Some(2), None, Some(4)])); - let arrays = vec![array]; - - let (ptr, offsets, lengths) = ctx.convert(&arrays, 4).unwrap(); - - assert!(!ptr.is_null()); - assert_eq!(offsets.len(), 4); - assert_eq!(lengths.len(), 4); - - // Each row should have: 8 bytes null bitset + 8 bytes for one field = 16 bytes - for len in lengths { - assert_eq!(*len, 16); - } - } - - #[test] - fn test_convert_multiple_columns() { - let schema = vec![DataType::Int32, DataType::Int64, DataType::Float64]; - let mut ctx = ColumnarToRowContext::new(schema, 100); - - let array1: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); - let array2: ArrayRef = Arc::new(Int64Array::from(vec![100i64, 200, 300])); - let array3: ArrayRef = Arc::new(Float64Array::from(vec![1.1, 2.2, 3.3])); - let arrays = vec![array1, array2, array3]; - - let (ptr, offsets, lengths) = ctx.convert(&arrays, 3).unwrap(); - - assert!(!ptr.is_null()); - assert_eq!(offsets.len(), 3); - assert_eq!(lengths.len(), 3); - - // Each row should have: 8 bytes null bitset + 24 bytes for three fields = 32 bytes - for len in lengths { - assert_eq!(*len, 32); - } - } - - #[test] - fn test_fixed_width_fast_path() { - // Test that the fixed-width fast path produces correct results - let schema = vec![DataType::Int32, DataType::Int64, DataType::Float64]; - let mut ctx = ColumnarToRowContext::new(schema.clone(), 100); - - // Verify that the context detects this as all fixed-width - assert!( - ctx.all_fixed_width, - "Schema should be detected as all fixed-width" - ); - - let array1: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), None, Some(3)])); - let array2: ArrayRef = Arc::new(Int64Array::from(vec![Some(100i64), Some(200), None])); - let array3: ArrayRef = Arc::new(Float64Array::from(vec![1.5, 2.5, 3.5])); - let arrays = vec![array1, array2, array3]; - - let (ptr, offsets, lengths) = ctx.convert(&arrays, 3).unwrap(); - - assert!(!ptr.is_null()); - assert_eq!(offsets.len(), 3); - assert_eq!(lengths.len(), 3); - - // Each row: 8 bytes null bitset + 24 bytes for three fields = 32 bytes - let row_size = 32; - for (i, len) in lengths.iter().enumerate() { - assert_eq!( - *len, row_size as i32, - "Row {} should be {} bytes", - i, row_size - ); - } - - // Verify the actual data - let buffer = unsafe { std::slice::from_raw_parts(ptr, row_size * 3) }; - - // Row 0: int32=1 (not null), int64=100 (not null), float64=1.5 (not null) - let null_bitset_0 = i64::from_le_bytes(buffer[0..8].try_into().unwrap()); - assert_eq!(null_bitset_0, 0, "Row 0 should have no nulls"); - let val0_0 = i64::from_le_bytes(buffer[8..16].try_into().unwrap()); - assert_eq!(val0_0, 1, "Row 0, col 0 should be 1"); - let val0_1 = i64::from_le_bytes(buffer[16..24].try_into().unwrap()); - assert_eq!(val0_1, 100, "Row 0, col 1 should be 100"); - let val0_2 = f64::from_bits(u64::from_le_bytes(buffer[24..32].try_into().unwrap())); - assert!((val0_2 - 1.5).abs() < 0.001, "Row 0, col 2 should be 1.5"); - - // Row 1: int32=null, int64=200 (not null), float64=2.5 (not null) - let null_bitset_1 = i64::from_le_bytes(buffer[32..40].try_into().unwrap()); - assert_eq!(null_bitset_1 & 1, 1, "Row 1, col 0 should be null"); - let val1_1 = i64::from_le_bytes(buffer[48..56].try_into().unwrap()); - assert_eq!(val1_1, 200, "Row 1, col 1 should be 200"); - - // Row 2: int32=3 (not null), int64=null, float64=3.5 (not null) - let null_bitset_2 = i64::from_le_bytes(buffer[64..72].try_into().unwrap()); - assert_eq!(null_bitset_2 & 2, 2, "Row 2, col 1 should be null"); - let val2_0 = i64::from_le_bytes(buffer[72..80].try_into().unwrap()); - assert_eq!(val2_0, 3, "Row 2, col 0 should be 3"); - } - - #[test] - fn test_mixed_schema_uses_general_path() { - // Test that schemas with variable-length types use the general path - let schema = vec![DataType::Int32, DataType::Utf8]; - let ctx = ColumnarToRowContext::new(schema, 100); - - // Should NOT be detected as all fixed-width - assert!( - !ctx.all_fixed_width, - "Schema with Utf8 should not be all fixed-width" - ); - } - - #[test] - fn test_convert_string_array() { - let schema = vec![DataType::Utf8]; - let mut ctx = ColumnarToRowContext::new(schema, 100); - - let array: ArrayRef = Arc::new(StringArray::from(vec!["hello", "world"])); - let arrays = vec![array]; - - let (ptr, offsets, lengths) = ctx.convert(&arrays, 2).unwrap(); - - assert!(!ptr.is_null()); - assert_eq!(offsets.len(), 2); - assert_eq!(lengths.len(), 2); - - // Row 0: 8 (bitset) + 8 (field slot) + 8 (aligned "hello") = 24 - // Row 1: 8 (bitset) + 8 (field slot) + 8 (aligned "world") = 24 - assert_eq!(lengths[0], 24); - assert_eq!(lengths[1], 24); - } - - #[test] - fn test_i128_to_spark_decimal_bytes() { - // Test positive number - let bytes = i128_to_spark_decimal_bytes(12345); - assert!(bytes.len() <= 16); - - // Test negative number - let bytes = i128_to_spark_decimal_bytes(-12345); - assert!(bytes.len() <= 16); - - // Test zero - let bytes = i128_to_spark_decimal_bytes(0); - assert!(!bytes.is_empty()); - } - - #[test] - fn test_list_data_conversion() { - use arrow::datatypes::Field; - - // Create a list with elements [0, 1, 2, 3, 4] - let values = Int32Array::from(vec![0, 1, 2, 3, 4]); - let offsets = arrow::buffer::OffsetBuffer::new(vec![0, 5].into()); - - let list_field = Arc::new(Field::new("item", DataType::Int32, true)); - let list_array = ListArray::new(Arc::clone(&list_field), offsets, Arc::new(values), None); - - // Convert the list for row 0 using the new direct-write function - let mut buffer = Vec::new(); - write_list_to_buffer(&mut buffer, &list_array, 0, &list_field).expect("conversion failed"); - let result = &buffer; - - // UnsafeArrayData format for Int32: - // [0..8]: numElements = 5 - // [8..16]: null bitset (8 bytes for up to 64 elements) - // [16..20]: element 0 (4 bytes for Int32) - // [20..24]: element 1 (4 bytes for Int32) - // ... (total 20 bytes for 5 elements, rounded up to 24 for 8-byte alignment) - - let num_elements = i64::from_le_bytes(result[0..8].try_into().unwrap()); - assert_eq!(num_elements, 5, "should have 5 elements"); - - let bitset_width = ColumnarToRowContext::calculate_bitset_width(5); - assert_eq!(bitset_width, 8); - - // Read each element value (Int32 uses 4 bytes per element) - let element_size = 4; // Int32 - for i in 0..5 { - let slot_offset = 8 + bitset_width + i * element_size; - let value = - i32::from_le_bytes(result[slot_offset..slot_offset + 4].try_into().unwrap()); - assert_eq!(value, i as i32, "element {} should be {}", i, i); - } - } - - #[test] - fn test_list_data_conversion_multiple_rows() { - use arrow::datatypes::Field; - - // Create multiple lists: - // Row 0: [0] - // Row 1: [0, 1] - // Row 2: [0, 1, 2] - let values = Int32Array::from(vec![ - 0, // row 0 - 0, 1, // row 1 - 0, 1, 2, // row 2 - ]); - let offsets = arrow::buffer::OffsetBuffer::new(vec![0, 1, 3, 6].into()); - - let list_field = Arc::new(Field::new("item", DataType::Int32, true)); - let list_array = ListArray::new(Arc::clone(&list_field), offsets, Arc::new(values), None); - - // Test row 1 which has elements [0, 1] - let mut buffer = Vec::new(); - write_list_to_buffer(&mut buffer, &list_array, 1, &list_field).expect("conversion failed"); - let result = &buffer; - - let num_elements = i64::from_le_bytes(result[0..8].try_into().unwrap()); - assert_eq!(num_elements, 2, "row 1 should have 2 elements"); - - // Int32 uses 4 bytes per element - let element_size = 4; - let bitset_width = ColumnarToRowContext::calculate_bitset_width(2); - let slot0_offset = 8 + bitset_width; - let slot1_offset = slot0_offset + element_size; - - let value0 = i32::from_le_bytes(result[slot0_offset..slot0_offset + 4].try_into().unwrap()); - let value1 = i32::from_le_bytes(result[slot1_offset..slot1_offset + 4].try_into().unwrap()); - - assert_eq!(value0, 0, "row 1, element 0 should be 0"); - assert_eq!(value1, 1, "row 1, element 1 should be 1"); - - // Also verify that list slicing is working correctly - let list_values = list_array.value(1); - assert_eq!(list_values.len(), 2, "row 1 should have 2 elements"); - let int_arr = list_values.as_any().downcast_ref::().unwrap(); - assert_eq!(int_arr.value(0), 0, "row 1 value[0] via Arrow should be 0"); - assert_eq!(int_arr.value(1), 1, "row 1 value[1] via Arrow should be 1"); - } - - #[test] - fn test_map_data_conversion() { - use arrow::datatypes::{Field, Fields}; - - // Create a map with 3 entries: {"key_0": 0, "key_1": 10, "key_2": 20} - let keys = StringArray::from(vec!["key_0", "key_1", "key_2"]); - let values = Int32Array::from(vec![0, 10, 20]); - - let entries_field = Field::new( - "entries", - DataType::Struct(Fields::from(vec![ - Field::new("key", DataType::Utf8, false), - Field::new("value", DataType::Int32, true), - ])), - false, - ); - - let entries = StructArray::from(vec![ - ( - Arc::new(Field::new("key", DataType::Utf8, false)), - Arc::new(keys) as ArrayRef, - ), - ( - Arc::new(Field::new("value", DataType::Int32, true)), - Arc::new(values) as ArrayRef, - ), - ]); - - let map_array = MapArray::new( - Arc::new(entries_field.clone()), - arrow::buffer::OffsetBuffer::new(vec![0, 3].into()), - entries, - None, - false, - ); - - // Convert the map for row 0 - let mut buffer = Vec::new(); - write_map_to_buffer(&mut buffer, &map_array, 0, &Arc::new(entries_field)) - .expect("conversion failed"); - let result = &buffer; - - // Verify the structure: - // - [0..8]: key array size - // - [8..key_end]: key array data (UnsafeArrayData format) - // - [key_end..]: value array data (UnsafeArrayData format) - - let key_array_size = i64::from_le_bytes(result[0..8].try_into().unwrap()); - assert!(key_array_size > 0, "key array size should be positive"); - - let value_array_start = (8 + key_array_size) as usize; - assert!( - value_array_start < result.len(), - "value array should start within buffer" - ); - - // Read value array - let value_num_elements = i64::from_le_bytes( - result[value_array_start..value_array_start + 8] - .try_into() - .unwrap(), - ); - assert_eq!(value_num_elements, 3, "should have 3 values"); - - // Value array layout for Int32 (4 bytes per element): - // [0..8]: numElements = 3 - // [8..16]: null bitset (8 bytes for up to 64 elements) - // [16..20]: element 0 (4 bytes) - // [20..24]: element 1 (4 bytes) - // [24..28]: element 2 (4 bytes) - let value_bitset_width = ColumnarToRowContext::calculate_bitset_width(3); - assert_eq!(value_bitset_width, 8); - - let element_size = 4; // Int32 - let slot0_offset = value_array_start + 8 + value_bitset_width; - let slot1_offset = slot0_offset + element_size; - let slot2_offset = slot1_offset + element_size; - - let value0 = i32::from_le_bytes(result[slot0_offset..slot0_offset + 4].try_into().unwrap()); - let value1 = i32::from_le_bytes(result[slot1_offset..slot1_offset + 4].try_into().unwrap()); - let value2 = i32::from_le_bytes(result[slot2_offset..slot2_offset + 4].try_into().unwrap()); - - assert_eq!(value0, 0, "first value should be 0"); - assert_eq!(value1, 10, "second value should be 10"); - assert_eq!(value2, 20, "third value should be 20"); - } - - #[test] - fn test_map_data_conversion_multiple_rows() { - use arrow::datatypes::{Field, Fields}; - - // Create multiple maps: - // Row 0: {"key_0": 0} - // Row 1: {"key_0": 0, "key_1": 10} - // Row 2: {"key_0": 0, "key_1": 10, "key_2": 20} - // All entries are concatenated in the underlying arrays - let keys = StringArray::from(vec![ - "key_0", // row 0 - "key_0", "key_1", // row 1 - "key_0", "key_1", "key_2", // row 2 - ]); - let values = Int32Array::from(vec![ - 0, // row 0 - 0, 10, // row 1 - 0, 10, 20, // row 2 - ]); - - let entries_field = Field::new( - "entries", - DataType::Struct(Fields::from(vec![ - Field::new("key", DataType::Utf8, false), - Field::new("value", DataType::Int32, true), - ])), - false, - ); - - let entries = StructArray::from(vec![ - ( - Arc::new(Field::new("key", DataType::Utf8, false)), - Arc::new(keys) as ArrayRef, - ), - ( - Arc::new(Field::new("value", DataType::Int32, true)), - Arc::new(values) as ArrayRef, - ), - ]); - - // Offsets: row 0 has 1 entry, row 1 has 2 entries, row 2 has 3 entries - let map_array = MapArray::new( - Arc::new(entries_field.clone()), - arrow::buffer::OffsetBuffer::new(vec![0, 1, 3, 6].into()), - entries, - None, - false, - ); - - // Test row 1 which has 2 entries - let mut buffer = Vec::new(); - write_map_to_buffer(&mut buffer, &map_array, 1, &Arc::new(entries_field.clone())) - .expect("conversion failed"); - let result = &buffer; - - let key_array_size = i64::from_le_bytes(result[0..8].try_into().unwrap()); - let value_array_start = (8 + key_array_size) as usize; - - let value_num_elements = i64::from_le_bytes( - result[value_array_start..value_array_start + 8] - .try_into() - .unwrap(), - ); - assert_eq!(value_num_elements, 2, "row 1 should have 2 values"); - - // Int32 uses 4 bytes per element - let element_size = 4; - let value_bitset_width = ColumnarToRowContext::calculate_bitset_width(2); - let slot0_offset = value_array_start + 8 + value_bitset_width; - let slot1_offset = slot0_offset + element_size; - - let value0 = i32::from_le_bytes(result[slot0_offset..slot0_offset + 4].try_into().unwrap()); - let value1 = i32::from_le_bytes(result[slot1_offset..slot1_offset + 4].try_into().unwrap()); - - assert_eq!(value0, 0, "row 1, first value should be 0"); - assert_eq!(value1, 10, "row 1, second value should be 10"); - - // Also verify that entries slicing is working correctly - let entries_row1 = map_array.value(1); - assert_eq!(entries_row1.len(), 2, "row 1 should have 2 entries"); - - let entries_values = entries_row1.column(1); - assert_eq!( - entries_values.len(), - 2, - "row 1 values should have 2 elements" - ); - - // Check the actual values from the sliced array - let values_arr = entries_values - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!( - values_arr.value(0), - 0, - "row 1 value[0] via Arrow should be 0" - ); - assert_eq!( - values_arr.value(1), - 10, - "row 1 value[1] via Arrow should be 10" - ); - } - - /// Test map conversion with a sliced MapArray to simulate FFI import behavior. - /// When data comes from FFI, the MapArray might be a slice of a larger array, - /// and the entries' child arrays might have offsets that don't start at 0. - #[test] - fn test_map_data_conversion_sliced_maparray() { - use arrow::datatypes::{Field, Fields}; - - // Create multiple maps (same as above) - let keys = StringArray::from(vec![ - "key_0", // row 0 - "key_0", "key_1", // row 1 - "key_0", "key_1", "key_2", // row 2 - ]); - let values = Int32Array::from(vec![ - 0, // row 0 - 0, 10, // row 1 - 0, 10, 20, // row 2 - ]); - - let entries_field = Field::new( - "entries", - DataType::Struct(Fields::from(vec![ - Field::new("key", DataType::Utf8, false), - Field::new("value", DataType::Int32, true), - ])), - false, - ); - - let entries = StructArray::from(vec![ - ( - Arc::new(Field::new("key", DataType::Utf8, false)), - Arc::new(keys) as ArrayRef, - ), - ( - Arc::new(Field::new("value", DataType::Int32, true)), - Arc::new(values) as ArrayRef, - ), - ]); - - let map_array = MapArray::new( - Arc::new(entries_field.clone()), - arrow::buffer::OffsetBuffer::new(vec![0, 1, 3, 6].into()), - entries, - None, - false, - ); - - // Slice the MapArray to skip row 0 - this simulates what might happen with FFI - let sliced_map = map_array.slice(1, 2); - let sliced_map_array = sliced_map.as_any().downcast_ref::().unwrap(); - - // Now test row 0 of the sliced array (which is row 1 of the original) - let mut buffer = Vec::new(); - write_map_to_buffer( - &mut buffer, - sliced_map_array, - 0, - &Arc::new(entries_field.clone()), - ) - .expect("conversion failed"); - let result = &buffer; - - let key_array_size = i64::from_le_bytes(result[0..8].try_into().unwrap()); - let value_array_start = (8 + key_array_size) as usize; - - let value_num_elements = i64::from_le_bytes( - result[value_array_start..value_array_start + 8] - .try_into() - .unwrap(), - ); - assert_eq!(value_num_elements, 2, "sliced row 0 should have 2 values"); - - let value_bitset_width = ColumnarToRowContext::calculate_bitset_width(2); - let slot0_offset = value_array_start + 8 + value_bitset_width; - let slot1_offset = slot0_offset + 4; // Int32 uses 4 bytes - - let value0 = i32::from_le_bytes(result[slot0_offset..slot0_offset + 4].try_into().unwrap()); - let value1 = i32::from_le_bytes(result[slot1_offset..slot1_offset + 4].try_into().unwrap()); - - assert_eq!(value0, 0, "sliced row 0, first value should be 0"); - assert_eq!(value1, 10, "sliced row 0, second value should be 10"); - } - - #[test] - fn test_large_list_data_conversion() { - use arrow::datatypes::Field; - - // Create a large list with elements [0, 1, 2, 3, 4] - // LargeListArray uses i64 offsets instead of i32 - let values = Int32Array::from(vec![0, 1, 2, 3, 4]); - let offsets = arrow::buffer::OffsetBuffer::new(vec![0i64, 5].into()); - - let list_field = Arc::new(Field::new("item", DataType::Int32, true)); - let list_array = - LargeListArray::new(Arc::clone(&list_field), offsets, Arc::new(values), None); - - // Convert the list for row 0 - let mut buffer = Vec::new(); - write_large_list_to_buffer(&mut buffer, &list_array, 0, &list_field) - .expect("conversion failed"); - let result = &buffer; - - // UnsafeArrayData format for Int32: - // [0..8]: numElements = 5 - // [8..16]: null bitset (8 bytes for up to 64 elements) - // [16..20]: element 0 (4 bytes for Int32) - // ... (total 20 bytes for 5 elements, rounded up to 24 for 8-byte alignment) - - let num_elements = i64::from_le_bytes(result[0..8].try_into().unwrap()); - assert_eq!(num_elements, 5, "should have 5 elements"); - - let bitset_width = ColumnarToRowContext::calculate_bitset_width(5); - assert_eq!(bitset_width, 8); - - // Read each element value (Int32 uses 4 bytes per element) - let element_size = 4; // Int32 - for i in 0..5 { - let slot_offset = 8 + bitset_width + i * element_size; - let value = - i32::from_le_bytes(result[slot_offset..slot_offset + 4].try_into().unwrap()); - assert_eq!(value, i as i32, "element {} should be {}", i, i); - } - } - - #[test] - fn test_convert_fixed_size_binary_array() { - // FixedSizeBinary(3) - each value is exactly 3 bytes - let schema = vec![DataType::FixedSizeBinary(3)]; - let mut ctx = ColumnarToRowContext::new(schema, 100); - - let array: ArrayRef = Arc::new( - FixedSizeBinaryArray::try_from(vec![ - Some(&[1u8, 2, 3][..]), - Some(&[4u8, 5, 6][..]), - None, // Test null handling - ]) - .unwrap(), - ); - let arrays = vec![array]; - - let (ptr, offsets, lengths) = ctx.convert(&arrays, 3).unwrap(); - - assert!(!ptr.is_null()); - assert_eq!(offsets.len(), 3); - assert_eq!(lengths.len(), 3); - - // Row 0: 8 (bitset) + 8 (field slot) + 8 (aligned 3-byte data) = 24 - // Row 1: 8 (bitset) + 8 (field slot) + 8 (aligned 3-byte data) = 24 - // Row 2: 8 (bitset) + 8 (field slot) = 16 (null, no variable data) - assert_eq!(lengths[0], 24); - assert_eq!(lengths[1], 24); - assert_eq!(lengths[2], 16); - - // Verify the data is correct for non-null rows - unsafe { - let row0 = - std::slice::from_raw_parts(ptr.add(offsets[0] as usize), lengths[0] as usize); - // Variable data starts at offset 16 (8 bitset + 8 field slot) - assert_eq!(&row0[16..19], &[1u8, 2, 3]); - - let row1 = - std::slice::from_raw_parts(ptr.add(offsets[1] as usize), lengths[1] as usize); - assert_eq!(&row1[16..19], &[4u8, 5, 6]); - } - } - - #[test] - fn test_convert_dictionary_decimal_array() { - // Test that dictionary-encoded decimals are correctly unpacked and converted - // This tests the fix for casting to schema_type instead of value_type - use arrow::datatypes::Int8Type; - - // Create a dictionary array with Decimal128 values - // Values: [-0.01, -0.02, -0.03] represented as [-1, -2, -3] with scale 2 - let values = Decimal128Array::from(vec![-1i128, -2, -3]) - .with_precision_and_scale(5, 2) - .unwrap(); - - // Keys: [0, 1, 2, 0, 1, 2] - each value appears twice - let keys = Int8Array::from(vec![0i8, 1, 2, 0, 1, 2]); - - let dict_array: ArrayRef = - Arc::new(DictionaryArray::::try_new(keys, Arc::new(values)).unwrap()); - - // Schema expects Decimal128(5, 2) - not a dictionary type - let schema = vec![DataType::Decimal128(5, 2)]; - let mut ctx = ColumnarToRowContext::new(schema, 100); - - let arrays = vec![dict_array]; - let (ptr, offsets, lengths) = ctx.convert(&arrays, 6).unwrap(); - - assert!(!ptr.is_null()); - assert_eq!(offsets.len(), 6); - assert_eq!(lengths.len(), 6); - - // Verify the decimal values are correct (not doubled or otherwise corrupted) - // Fixed-width decimal is stored directly in the 8-byte field slot - unsafe { - for (i, expected) in [-1i64, -2, -3, -1, -2, -3].iter().enumerate() { - let row = - std::slice::from_raw_parts(ptr.add(offsets[i] as usize), lengths[i] as usize); - // Field value starts at offset 8 (after null bitset) - let value = i64::from_le_bytes(row[8..16].try_into().unwrap()); - assert_eq!( - value, *expected, - "Row {} should have value {}, got {}", - i, expected, value - ); - } - } - } - - #[test] - fn test_convert_int32_to_decimal128() { - // Int32 → Decimal128 reinterpret. Values [-1, null, -3] at scale 2 mean - // [-0.01, null, -0.03]. The null covers the arity::unary null-buffer path. - let int_array: ArrayRef = Arc::new(Int32Array::from(vec![Some(-1i32), None, Some(-3)])); - - let schema = vec![DataType::Decimal128(5, 2)]; - let mut ctx = ColumnarToRowContext::new(schema, 100); - - let arrays = vec![int_array]; - let (ptr, offsets, lengths) = ctx.convert(&arrays, 3).unwrap(); - - assert!(!ptr.is_null()); - assert_eq!(offsets.len(), 3); - assert_eq!(lengths.len(), 3); - - // Row layout: [8-byte null bitset][8-byte field slot]. Column 0's null bit - // is bit 0 of the bitset. - unsafe { - let row0 = - std::slice::from_raw_parts(ptr.add(offsets[0] as usize), lengths[0] as usize); - let null_bitset_0 = i64::from_le_bytes(row0[0..8].try_into().unwrap()); - assert_eq!(null_bitset_0 & 1, 0, "row 0 should not be null"); - assert_eq!(i64::from_le_bytes(row0[8..16].try_into().unwrap()), -1); - - let row1 = - std::slice::from_raw_parts(ptr.add(offsets[1] as usize), lengths[1] as usize); - let null_bitset_1 = i64::from_le_bytes(row1[0..8].try_into().unwrap()); - assert_eq!(null_bitset_1 & 1, 1, "row 1 should be null"); - - let row2 = - std::slice::from_raw_parts(ptr.add(offsets[2] as usize), lengths[2] as usize); - let null_bitset_2 = i64::from_le_bytes(row2[0..8].try_into().unwrap()); - assert_eq!(null_bitset_2 & 1, 0, "row 2 should not be null"); - assert_eq!(i64::from_le_bytes(row2[8..16].try_into().unwrap()), -3); - } - } - - #[test] - fn test_convert_int32_to_decimal128_sliced() { - // Sliced input must reinterpret correctly. Arrays imported over FFI often - // carry a non-zero offset; both the value widening and the null buffer - // traversal must honor it. - let full = Int32Array::from(vec![ - Some(-99i32), // discarded prefix - Some(-98), - Some(-1), // start of sliced portion - None, - Some(-3), - Some(-97), // discarded suffix - ]); - let sliced: ArrayRef = Arc::new(full.slice(2, 3)); - - let schema = vec![DataType::Decimal128(5, 2)]; - let mut ctx = ColumnarToRowContext::new(schema, 100); - let arrays = vec![sliced]; - let (ptr, offsets, lengths) = ctx.convert(&arrays, 3).unwrap(); - - assert!(!ptr.is_null()); - assert_eq!(offsets.len(), 3); - - unsafe { - let row0 = - std::slice::from_raw_parts(ptr.add(offsets[0] as usize), lengths[0] as usize); - assert_eq!(i64::from_le_bytes(row0[0..8].try_into().unwrap()) & 1, 0); - assert_eq!(i64::from_le_bytes(row0[8..16].try_into().unwrap()), -1); - - let row1 = - std::slice::from_raw_parts(ptr.add(offsets[1] as usize), lengths[1] as usize); - assert_eq!(i64::from_le_bytes(row1[0..8].try_into().unwrap()) & 1, 1); - - let row2 = - std::slice::from_raw_parts(ptr.add(offsets[2] as usize), lengths[2] as usize); - assert_eq!(i64::from_le_bytes(row2[0..8].try_into().unwrap()) & 1, 0); - assert_eq!(i64::from_le_bytes(row2[8..16].try_into().unwrap()), -3); - } - } - - #[test] - fn test_convert_int64_to_decimal128() { - // Int64 → Decimal128 reinterpret. Values [-100, null, -300] with a null in - // the middle to exercise the arity::unary null-buffer path. - let int_array: ArrayRef = Arc::new(Int64Array::from(vec![Some(-100i64), None, Some(-300)])); - - let schema = vec![DataType::Decimal128(10, 2)]; - let mut ctx = ColumnarToRowContext::new(schema, 100); - - let arrays = vec![int_array]; - let (ptr, offsets, lengths) = ctx.convert(&arrays, 3).unwrap(); - - assert!(!ptr.is_null()); - assert_eq!(offsets.len(), 3); - assert_eq!(lengths.len(), 3); - - unsafe { - let row0 = - std::slice::from_raw_parts(ptr.add(offsets[0] as usize), lengths[0] as usize); - let null_bitset_0 = i64::from_le_bytes(row0[0..8].try_into().unwrap()); - assert_eq!(null_bitset_0 & 1, 0, "row 0 should not be null"); - assert_eq!(i64::from_le_bytes(row0[8..16].try_into().unwrap()), -100); - - let row1 = - std::slice::from_raw_parts(ptr.add(offsets[1] as usize), lengths[1] as usize); - let null_bitset_1 = i64::from_le_bytes(row1[0..8].try_into().unwrap()); - assert_eq!(null_bitset_1 & 1, 1, "row 1 should be null"); - - let row2 = - std::slice::from_raw_parts(ptr.add(offsets[2] as usize), lengths[2] as usize); - let null_bitset_2 = i64::from_le_bytes(row2[0..8].try_into().unwrap()); - assert_eq!(null_bitset_2 & 1, 0, "row 2 should not be null"); - assert_eq!(i64::from_le_bytes(row2[8..16].try_into().unwrap()), -300); - } - } - - #[test] - fn test_convert_int64_to_decimal128_sliced() { - let full = Int64Array::from(vec![ - Some(-9999i64), - Some(-9998), - Some(-100), - None, - Some(-300), - Some(-9997), - ]); - let sliced: ArrayRef = Arc::new(full.slice(2, 3)); - - let schema = vec![DataType::Decimal128(10, 2)]; - let mut ctx = ColumnarToRowContext::new(schema, 100); - let arrays = vec![sliced]; - let (ptr, offsets, lengths) = ctx.convert(&arrays, 3).unwrap(); - - assert!(!ptr.is_null()); - assert_eq!(offsets.len(), 3); - - unsafe { - let row0 = - std::slice::from_raw_parts(ptr.add(offsets[0] as usize), lengths[0] as usize); - assert_eq!(i64::from_le_bytes(row0[0..8].try_into().unwrap()) & 1, 0); - assert_eq!(i64::from_le_bytes(row0[8..16].try_into().unwrap()), -100); - - let row1 = - std::slice::from_raw_parts(ptr.add(offsets[1] as usize), lengths[1] as usize); - assert_eq!(i64::from_le_bytes(row1[0..8].try_into().unwrap()) & 1, 1); - - let row2 = - std::slice::from_raw_parts(ptr.add(offsets[2] as usize), lengths[2] as usize); - assert_eq!(i64::from_le_bytes(row2[0..8].try_into().unwrap()) & 1, 0); - assert_eq!(i64::from_le_bytes(row2[8..16].try_into().unwrap()), -300); - } - } -} diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index d915abfed0b..28cd922d19d 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -59,7 +59,6 @@ use datafusion::{ physical_plan::{display::DisplayableExecutionPlan, SendableRecordBatchStream}, prelude::{SessionConfig, SessionContext}, }; -use datafusion_comet_common::decode_string_arrays; use datafusion_comet_proto::spark_expression::agg_expr::ExprStruct as AggExprStruct; use datafusion_comet_proto::spark_operator::{AggregateMode, Operator, ShuffleScan}; use datafusion_comet_shuffle::decoded_batch_coalescer::DecodedBatchCoalescer; @@ -80,6 +79,7 @@ use datafusion_spark::function::hash::sha1::SparkSha1; use datafusion_spark::function::hash::sha2::SparkSha2; use datafusion_spark::function::map::map_from_entries::MapFromEntries; use datafusion_spark::function::map::str_to_map::SparkStrToMap; +use datafusion_spark::function::math::bin::SparkBin; use datafusion_spark::function::math::expm1::SparkExpm1; use datafusion_spark::function::math::factorial::SparkFactorial; use datafusion_spark::function::math::hex::SparkHex; @@ -91,6 +91,7 @@ use datafusion_spark::function::string::char::CharFunc; use datafusion_spark::function::string::concat::SparkConcat; use datafusion_spark::function::string::length::SparkLengthFunc; use datafusion_spark::function::string::luhn_check::SparkLuhnCheck; +use datafusion_spark::function::string::soundex::SparkSoundex; use datafusion_spark::function::string::space::SparkSpace; use datafusion_spark::function::string::substring::SparkSubstring; use datafusion_spark::function::url::try_url_decode::TryUrlDecode as SparkTryUrlDecode; @@ -2532,170 +2533,6 @@ pub extern "system" fn Java_org_apache_comet_Native_getMemoryUsage( }) } -// ============================================================================ -// Native Columnar to Row Conversion -// ============================================================================ - -use crate::execution::columnar_to_row::ColumnarToRowContext; -use arrow::ffi::{from_ffi, FFI_ArrowArray, FFI_ArrowSchema}; -use datafusion_spark::function::math::bin::SparkBin; -use datafusion_spark::function::string::soundex::SparkSoundex; - -/// Initialize a native columnar to row converter. -/// -/// # Safety -/// This function is inherently unsafe since it deals with raw pointers passed from JNI. -#[no_mangle] -pub unsafe extern "system" fn Java_org_apache_comet_Native_columnarToRowInit( - e: EnvUnowned, - _class: JClass, - serialized_schema: JObjectArray, - batch_size: jint, -) -> jlong { - try_unwrap_or_throw(&e, |env| { - // Deserialize the schema - let schema = convert_datatype_arrays(env, serialized_schema)?; - - // Create the context - let ctx = Box::new(ColumnarToRowContext::new(schema, batch_size as usize)); - - Ok(Box::into_raw(ctx) as jlong) - }) -} - -/// Convert Arrow columnar data to Spark UnsafeRow format. -/// -/// # Safety -/// This function is inherently unsafe since it deals with raw pointers passed from JNI. -#[no_mangle] -pub unsafe extern "system" fn Java_org_apache_comet_Native_columnarToRowConvert( - e: EnvUnowned, - _class: JClass, - c2r_handle: jlong, - array_addrs: JLongArray, - schema_addrs: JLongArray, - num_rows: jint, -) -> jni::sys::jobject { - try_unwrap_or_throw(&e, |env| { - // Get the context - debug_assert!(c2r_handle != 0, "columnarToRowConvert: c2r_handle is null"); - let ctx = (c2r_handle as *mut ColumnarToRowContext) - .as_mut() - .ok_or_else(|| CometError::Internal("Null columnar to row context".to_string()))?; - - // Get array and schema addresses - let array_addrs = unsafe { array_addrs.get_elements(env, ReleaseMode::NoCopyBack)? }; - let schema_addrs = unsafe { schema_addrs.get_elements(env, ReleaseMode::NoCopyBack)? }; - - let (buffer_ptr, offsets, lengths) = - unsafe { columnar_to_row_convert(ctx, &array_addrs, &schema_addrs, num_rows)? }; - - // Create Java int arrays for offsets and lengths - let offsets_array = env.new_int_array(offsets.len())?; - offsets_array.set_region(env, 0, offsets)?; - - let lengths_array = env.new_int_array(lengths.len())?; - lengths_array.set_region(env, 0, lengths)?; - - // Create the NativeColumnarToRowInfo object - let info_class = - env.find_class(jni::jni_str!("org/apache/comet/NativeColumnarToRowInfo"))?; - let info_obj = env.new_object( - info_class, - jni::jni_sig!("(J[I[I)V"), - &[ - jni::objects::JValue::Long(buffer_ptr as jlong), - jni::objects::JValue::Object(&offsets_array), - jni::objects::JValue::Object(&lengths_array), - ], - )?; - - Ok(info_obj.into_raw()) - }) -} - -/// Imports the columns at `array_addrs` and `schema_addrs` and converts their first `num_rows` -/// rows to Spark `UnsafeRow`s in `ctx`'s buffer. Returns the buffer address and the offset and -/// length of each row in it, valid until the next conversion. Core of -/// `Native.columnarToRowConvert`. -/// -/// # Safety -/// Each address must point to an exported `FFI_ArrowArray` / `FFI_ArrowSchema`, whose ownership -/// moves to this call. -unsafe fn columnar_to_row_convert<'a>( - ctx: &'a mut ColumnarToRowContext, - array_addrs: &[i64], - schema_addrs: &[i64], - num_rows: i32, -) -> CometResult<(*const u8, &'a [i32], &'a [i32])> { - if schema_addrs.len() != array_addrs.len() { - return Err(CometError::Internal(format!( - "{} schema addresses for {} array addresses", - schema_addrs.len(), - array_addrs.len() - ))); - } - let num_cols = array_addrs.len(); - - // Import Arrow arrays from FFI - let mut arrays = Vec::with_capacity(num_cols); - for i in 0..num_cols { - let array_ptr = array_addrs[i] as *mut FFI_ArrowArray; - let schema_ptr = schema_addrs[i] as *mut FFI_ArrowSchema; - - debug_assert!( - !array_ptr.is_null(), - "columnarToRowConvert: null array pointer at index {}", - i - ); - debug_assert!( - !schema_ptr.is_null(), - "columnarToRowConvert: null schema pointer at index {}", - i - ); - - // Take ownership of the FFI structures - let ffi_array = unsafe { std::ptr::read(array_ptr) }; - let ffi_schema = unsafe { std::ptr::read(schema_ptr) }; - - // Convert to Arrow ArrayData - let array_data = from_ffi(ffi_array, &ffi_schema) - .map_err(|e| CometError::Internal(format!("Failed to import array: {}", e)))?; - - let imported = arrow::array::make_array(array_data); - arrays.push(decode_string_arrays(&imported)?); - } - - // Convert columnar to row - debug_assert!( - num_rows >= 0, - "columnarToRowConvert: num_rows is negative: {}", - num_rows - ); - ctx.convert(&arrays, num_rows as usize) -} - -/// Close and release the native columnar to row converter. -/// -/// # Safety -/// This function is inherently unsafe since it deals with raw pointers passed from JNI. -#[no_mangle] -pub unsafe extern "system" fn Java_org_apache_comet_Native_columnarToRowClose( - e: EnvUnowned, - _class: JClass, - c2r_handle: jlong, -) { - try_unwrap_or_throw(&e, |_env| { - debug_assert!(c2r_handle != 0, "columnarToRowClose: c2r_handle is null"); - if c2r_handle != 0 { - let _ctx: Box = - Box::from_raw(c2r_handle as *mut ColumnarToRowContext); - // ctx is dropped here, freeing the buffer - } - Ok(()) - }) -} - #[cfg(test)] mod tests { use super::*; @@ -2703,6 +2540,7 @@ mod tests { use crate::execution::planner::TEST_EXEC_CONTEXT_ID; use arrow::array::{ArrayRef, Int32Array}; use arrow::datatypes::{DataType, Field, Schema}; + use arrow::ffi::{from_ffi, FFI_ArrowArray, FFI_ArrowSchema}; use datafusion::execution::memory_pool::{ MemoryConsumer, MemoryReservation, UnboundedMemoryPool, }; @@ -4306,90 +4144,6 @@ mod tests { assert!(error.contains("record length size 2"), "{error}"); } - /// Exports `columns` as Arrow C Data structs and converts them with `columnar_to_row_convert`, - /// returning the offsets, the lengths, and the row bytes. - fn convert_exported( - columns: &[ArrayRef], - num_rows: i32, - ) -> CometResult<(Vec, Vec, Vec)> { - let mut ffi_arrays: Vec> = columns - .iter() - .map(|column| Box::new(FFI_ArrowArray::new(&column.to_data()))) - .collect(); - let mut ffi_schemas: Vec> = columns - .iter() - .map(|column| Box::new(FFI_ArrowSchema::try_from(column.data_type()).unwrap())) - .collect(); - let array_addrs: Vec = ffi_arrays - .iter_mut() - .map(|array| array.as_mut() as *mut FFI_ArrowArray as i64) - .collect(); - let schema_addrs: Vec = ffi_schemas - .iter_mut() - .map(|schema| schema.as_mut() as *mut FFI_ArrowSchema as i64) - .collect(); - let types = columns.iter().map(|c| c.data_type().clone()).collect(); - - let mut ctx = ColumnarToRowContext::new(types, 8); - let converted = - unsafe { columnar_to_row_convert(&mut ctx, &array_addrs, &schema_addrs, num_rows) }; - // The call takes ownership of the exported structs, so they must not be released again. - ffi_arrays.into_iter().for_each(std::mem::forget); - ffi_schemas.into_iter().for_each(std::mem::forget); - - let (buffer, offsets, lengths) = converted?; - let total = offsets - .iter() - .zip(lengths) - .map(|(offset, length)| (offset + length) as usize) - .max() - .unwrap_or(0); - let rows = unsafe { std::slice::from_raw_parts(buffer, total) }.to_vec(); - Ok((offsets.to_vec(), lengths.to_vec(), rows)) - } - - #[test] - fn columnar_to_row_convert_imports_and_converts() { - let columns: Vec = vec![ - Arc::new(Int32Array::from(vec![Some(1), None, Some(3)])), - Arc::new(arrow::array::StringArray::from(vec![ - Some("a"), - Some("a longer string"), - None, - ])), - ]; - let (offsets, lengths, rows) = convert_exported(&columns, 3).unwrap(); - - let types = columns.iter().map(|c| c.data_type().clone()).collect(); - let mut expected_ctx = ColumnarToRowContext::new(types, 8); - let (expected_buffer, expected_offsets, expected_lengths) = - expected_ctx.convert(&columns, 3).unwrap(); - assert_eq!(offsets, expected_offsets); - assert_eq!(lengths, expected_lengths); - assert_eq!(rows, unsafe { - std::slice::from_raw_parts(expected_buffer, rows.len()) - }); - // The string column makes the rows variable-length. - assert_ne!(lengths[0], lengths[1]); - } - - #[test] - fn columnar_to_row_convert_rejects_mismatched_addresses() { - let column: ArrayRef = Arc::new(Int32Array::from(vec![1, 2])); - let mut ffi_array = Box::new(FFI_ArrowArray::new(&column.to_data())); - let array_addrs = [ffi_array.as_mut() as *mut FFI_ArrowArray as i64]; - let mut ctx = ColumnarToRowContext::new(vec![DataType::Int32], 8); - // Rejected before any struct is imported, so the export is still released on drop. - let error = unsafe { columnar_to_row_convert(&mut ctx, &array_addrs, &[], 2) } - .expect_err("mismatched addresses are rejected"); - assert!( - error - .to_string() - .contains("0 schema addresses for 1 array addresses"), - "{error}" - ); - } - /// A range of `num_elements` longs from 0, which needs no JVM input and reserves no memory. fn range_operator(num_elements: i64) -> Operator { use datafusion_comet_proto::spark_operator::RangeScan; diff --git a/native/core/src/execution/mod.rs b/native/core/src/execution/mod.rs index 5088bfd3e46..caf95cbfc62 100644 --- a/native/core/src/execution/mod.rs +++ b/native/core/src/execution/mod.rs @@ -17,7 +17,6 @@ //! PoC of vectorization execution through JNI to Rust. pub mod c_udf; -pub mod columnar_to_row; #[cfg(feature = "delta")] pub mod delta_dv; pub mod expressions; diff --git a/spark/src/main/java/org/apache/comet/NativeColumnarToRowInfo.java b/spark/src/main/java/org/apache/comet/NativeColumnarToRowInfo.java deleted file mode 100644 index 319390255b8..00000000000 --- a/spark/src/main/java/org/apache/comet/NativeColumnarToRowInfo.java +++ /dev/null @@ -1,75 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you under the Apache License, Version 2.0 (the - * "License"); you may not use this file except in compliance - * with the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, - * software distributed under the License is distributed on an - * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY - * KIND, either express or implied. See the License for the - * specific language governing permissions and limitations - * under the License. - */ - -package org.apache.comet; - -/** - * Container for the result of native columnar to row conversion. - * - *

This class holds the memory address of the converted row data buffer and metadata about each - * row (offsets and lengths). The native side allocates and owns the memory buffer, and this class - * provides the JVM with the information needed to read the UnsafeRow data. - * - *

Memory Layout of the buffer: - * - * 

- * ┌─────────────────────────────────────────────────────────────┐
- * │ Row 0: [null bitset][fixed-width values][variable-length]  │
- * ├─────────────────────────────────────────────────────────────┤
- * │ Row 1: [null bitset][fixed-width values][variable-length]  │
- * ├─────────────────────────────────────────────────────────────┤
- * │ ...                                                         │
- * └─────────────────────────────────────────────────────────────┘
- *
- * - *

The offsets array provides the byte offset from memoryAddress where each row starts. The - * lengths array provides the total byte length of each row. - */ -public class NativeColumnarToRowInfo { - /** The memory address of the buffer containing all converted rows. */ - public final long memoryAddress; - - /** The byte offset from memoryAddress where each row starts. */ - public final int[] offsets; - - /** The total byte length of each row. */ - public final int[] lengths; - - /** - * Constructs a NativeColumnarToRowInfo with the given memory address and row metadata. - * - * @param memoryAddress The memory address of the buffer containing converted row data. - * @param offsets The byte offset for each row from the base memory address. - * @param lengths The byte length of each row. - */ - public NativeColumnarToRowInfo(long memoryAddress, int[] offsets, int[] lengths) { - this.memoryAddress = memoryAddress; - this.offsets = offsets; - this.lengths = lengths; - } - - /** - * Returns the number of rows in this result. - * - * @return The number of rows. - */ - public int numRows() { - return offsets != null ? offsets.length : 0; - } -} diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 824118f3efe..48ceee6b5dd 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -513,11 +513,8 @@ object CometConf extends ShimCometConf { val COMET_NATIVE_COLUMNAR_TO_ROW_ENABLED: ConfigEntry[Boolean] = conf(s"$COMET_EXEC_CONFIG_PREFIX.columnarToRow.native.enabled") .category(CATEGORY_EXEC) - .doc( - "Whether to enable native columnar to row conversion. When enabled, Comet will use " + - "native Rust code to convert Arrow columnar data to Spark UnsafeRow format instead " + - "of the JVM implementation. The native conversion carries a fixed JNI cost per batch " + - "and is slower than the JVM implementation for small batches.") + .doc("Deprecated: this config will be removed in a future major release. It has no " + + "effect: Comet always converts Arrow batches to Spark rows on the JVM.") .booleanConf .createWithDefault(false) @@ -530,8 +527,7 @@ object CometConf extends ShimCometConf { "builds). The direct converter writes values straight from Arrow buffers into " + "UnsafeRow format without per-value object allocation, which is significantly " + "faster for decimal-heavy schemas. Schemas with unsupported data types fall back " + - "to the default conversion. Only applies when the JVM columnar to row operator is " + - "used (spark.comet.exec.columnarToRow.native.enabled=false).") + "to the default conversion.") .booleanConf .createWithDefault(false) diff --git a/spark/src/main/scala/org/apache/comet/ExtendedExplainInfo.scala b/spark/src/main/scala/org/apache/comet/ExtendedExplainInfo.scala index 8638787f907..2d10c851720 100644 --- a/spark/src/main/scala/org/apache/comet/ExtendedExplainInfo.scala +++ b/spark/src/main/scala/org/apache/comet/ExtendedExplainInfo.scala @@ -26,7 +26,7 @@ import scala.collection.mutable import org.apache.spark.sql.ExtendedExplainGenerator import org.apache.spark.sql.catalyst.expressions.{Attribute, BoundReference, Expression, Literal, ScalaUDF} import org.apache.spark.sql.catalyst.trees.{TreeNode, TreeNodeTag} -import org.apache.spark.sql.comet.{CometColumnarToRowExec, CometEmptyRelationExec, CometInMemoryTableScanExec, CometNativeColumnarToRowExec, CometPlan, CometSparkToColumnarExec} +import org.apache.spark.sql.comet.{CometColumnarToRowExec, CometEmptyRelationExec, CometInMemoryTableScanExec, CometPlan, CometSparkToColumnarExec} import org.apache.spark.sql.execution.{ColumnarToRowExec, InputAdapter, ReusedSubqueryExec, RowToColumnarExec, SparkPlan, WholeStageCodegenExec} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, QueryStageExec} import org.apache.spark.sql.execution.exchange.{Exchange, ReusedExchangeExec} @@ -167,7 +167,7 @@ class ExtendedExplainInfo extends ExtendedExplainGenerator { // is pure reuse bookkeeping: the subquery it points at is counted where that subquery is // shown, so counting the wrapper too would invent an un-accelerated Spark operator. case _: RowToColumnarExec | _: ColumnarToRowExec | _: CometColumnarToRowExec | - _: CometNativeColumnarToRowExec | _: CometSparkToColumnarExec => + _: CometSparkToColumnarExec => stats.transitions += 1 case _: CometPlan => stats.cometOperators += 1 diff --git a/spark/src/main/scala/org/apache/comet/Native.scala b/spark/src/main/scala/org/apache/comet/Native.scala index c24bd762c3e..b848a18539d 100644 --- a/spark/src/main/scala/org/apache/comet/Native.scala +++ b/spark/src/main/scala/org/apache/comet/Native.scala @@ -391,50 +391,4 @@ class Native extends NativeBase { * `plans` is the number of native plans created and not yet released. */ @native def getMemoryUsage(): Array[Long] - - // Native Columnar to Row conversion methods - - /** - * Initialize a native columnar to row converter. - * - * @param schema - * Array of serialized data types (as byte arrays) for each column in the schema. - * @param batchSize - * The maximum number of rows that will be converted in a single batch. Used to pre-allocate - * the output buffer. - * @return - * A handle to the native converter context. This handle must be passed to subsequent convert - * and close calls. - */ - @native def columnarToRowInit(schema: Array[Array[Byte]], batchSize: Int): Long - - /** - * Convert Arrow columnar data to Spark UnsafeRow format. - * - * @param c2rHandle - * The handle returned by columnarToRowInit. - * @param arrayAddrs - * The addresses of Arrow Array structures for each column. - * @param schemaAddrs - * The addresses of Arrow Schema structures for each column. - * @param numRows - * The number of rows to convert. - * @return - * A NativeColumnarToRowInfo containing the memory address of the row buffer and metadata - * (offsets and lengths) for each row. - */ - @native def columnarToRowConvert( - c2rHandle: Long, - arrayAddrs: Array[Long], - schemaAddrs: Array[Long], - numRows: Int): NativeColumnarToRowInfo - - /** - * Close and release the native columnar to row converter. - * - * @param c2rHandle - * The handle returned by columnarToRowInit. - */ - @native def columnarToRowClose(c2rHandle: Long): Unit - } diff --git a/spark/src/main/scala/org/apache/comet/NativeColumnarToRowConverter.scala b/spark/src/main/scala/org/apache/comet/NativeColumnarToRowConverter.scala deleted file mode 100644 index b6aad967504..00000000000 --- a/spark/src/main/scala/org/apache/comet/NativeColumnarToRowConverter.scala +++ /dev/null @@ -1,145 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you under the Apache License, Version 2.0 (the - * "License"); you may not use this file except in compliance - * with the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, - * software distributed under the License is distributed on an - * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY - * KIND, either express or implied. See the License for the - * specific language governing permissions and limitations - * under the License. - */ - -package org.apache.comet - -import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.expressions.UnsafeRow -import org.apache.spark.sql.types.StructType -import org.apache.spark.sql.vectorized.ColumnarBatch - -import org.apache.comet.serde.QueryPlanSerde -import org.apache.comet.vector.NativeUtil - -/** - * Native converter that converts Arrow columnar data to Spark UnsafeRow format. - * - * This converter maintains a native handle that holds the conversion context and output buffer. - * The buffer is reused across conversions to minimize allocations. - * - * Memory Management: - * - The native side owns the output buffer - * - UnsafeRow objects returned by convert() point directly to native memory (zero-copy) - * - The buffer is valid until the next convert() call or close() - * - Always call close() when done to release native resources - * - * @param schema - * The schema of the data to convert - * @param batchSize - * Maximum number of rows per batch (used for buffer pre-allocation) - */ -class NativeColumnarToRowConverter(schema: StructType, batchSize: Int) extends AutoCloseable { - - private val nativeLib = new Native() - private val nativeUtil = new NativeUtil() - - // serializeDataType is broader than native C2R support. Production construction is gated by - // CometNativeColumnarToRowExec.supportsSchema, which recursively rejects VariantType. - private val serializedSchema: Array[Array[Byte]] = schema.fields.map { field => - QueryPlanSerde.serializeDataType(field.dataType) match { - case Some(dataType) => dataType.toByteArray - case None => - throw new UnsupportedOperationException( - s"Data type ${field.dataType} is not supported for native columnar to row conversion") - } - } - - // Initialize native context - handle is 0 if initialization failed - private var c2rHandle: Long = nativeLib.columnarToRowInit(serializedSchema, batchSize) - - // Reusable UnsafeRow for iteration - private val unsafeRow = new UnsafeRow(schema.fields.length) - - /** - * Converts a ColumnarBatch to an iterator of InternalRows. - * - * The returned iterator yields UnsafeRow objects that point directly to native memory. These - * rows are valid only until the next call to convert() or close(). - * - * @param batch - * The columnar batch to convert - * @return - * An iterator of InternalRows - */ - def convert(batch: ColumnarBatch): Iterator[InternalRow] = { - if (c2rHandle == 0) { - throw new IllegalStateException("NativeColumnarToRowConverter has been closed") - } - - val numRows = batch.numRows() - if (numRows == 0) { - return Iterator.empty - } - - // Export the batch to Arrow FFI and get memory addresses - val (arrayAddrs, schemaAddrs, exportedNumRows) = nativeUtil.exportBatchToAddresses(batch) - - // Call native conversion - val info = nativeLib.columnarToRowConvert(c2rHandle, arrayAddrs, schemaAddrs, exportedNumRows) - - // Return an iterator that yields UnsafeRows pointing to native memory - new NativeRowIterator(info, unsafeRow) - } - - /** - * Checks if this converter is still open and usable. - */ - def isOpen: Boolean = c2rHandle != 0 - - /** - * Closes the converter and releases native resources. - */ - override def close(): Unit = { - if (c2rHandle != 0) { - nativeLib.columnarToRowClose(c2rHandle) - c2rHandle = 0 - } - nativeUtil.close() - } -} - -/** - * Iterator that yields UnsafeRows backed by native memory. - * - * The UnsafeRow is reused across iterations - callers must copy the row if they need to retain it - * beyond the current iteration. - */ -private class NativeRowIterator(info: NativeColumnarToRowInfo, unsafeRow: UnsafeRow) - extends Iterator[InternalRow] { - - private var currentIdx = 0 - private val numRows = info.numRows() - - override def hasNext: Boolean = currentIdx < numRows - - override def next(): InternalRow = { - if (!hasNext) { - throw new NoSuchElementException("No more rows") - } - - // Point the UnsafeRow to the native memory - val rowAddress = info.memoryAddress + info.offsets(currentIdx) - val rowSize = info.lengths(currentIdx) - - unsafeRow.pointTo(null, rowAddress, rowSize) - currentIdx += 1 - - unsafeRow.copy() - } -} diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala index 390d2954161..d5ede46fc77 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -954,10 +954,8 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) // a Comet plan. Strip the row transition to get the columnar Comet plan. // CometColumnarToRowExec is CodegenSupport, so by the time this rule sees the // subquery plan it is compiled into WholeStageCodegenExec(CometColumnarToRowExec( - // InputAdapter(cometPlan))); CometNativeColumnarToRowExec is not CodegenSupport and - // sits directly under the exchange. + // InputAdapter(cometPlan))). val cometChild = b.child match { - case c2r: CometNativeColumnarToRowExec => c2r.child case c2r: CometColumnarToRowExec => c2r.child case WholeStageCodegenExec(c2r: CometColumnarToRowExec) => c2r.child match { diff --git a/spark/src/main/scala/org/apache/comet/rules/CometPlanAdaptiveDynamicPruningFilters.scala b/spark/src/main/scala/org/apache/comet/rules/CometPlanAdaptiveDynamicPruningFilters.scala index b92fd418366..66f3f6ceb26 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometPlanAdaptiveDynamicPruningFilters.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometPlanAdaptiveDynamicPruningFilters.scala @@ -369,7 +369,6 @@ case object CometPlanAdaptiveDynamicPruningFilters */ private def isNativeBuildSide(plan: SparkPlan): Boolean = { val stripped = stripAQEPlan(plan) match { - case c2r: CometNativeColumnarToRowExec => c2r.child case c2r: CometColumnarToRowExec => c2r.child case WholeStageCodegenExec(c2r: CometColumnarToRowExec) => c2r.child match { diff --git a/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala b/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala index 923fa1983e0..62ed06fc85a 100644 --- a/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala +++ b/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala @@ -24,7 +24,7 @@ import java.util.IdentityHashMap import org.apache.spark.sql.SparkSession import org.apache.spark.sql.catalyst.rules.Rule import org.apache.spark.sql.catalyst.util.sideBySide -import org.apache.spark.sql.comet.{CometCollectLimitExec, CometColumnarToRowExec, CometIcebergWriteExec, CometMapInBatchExec, CometNativeColumnarToRowExec, CometNativeWriteExec, CometPlan, CometRangeExec, CometSparkToColumnarExec} +import org.apache.spark.sql.comet.{CometCollectLimitExec, CometColumnarToRowExec, CometIcebergWriteExec, CometMapInBatchExec, CometNativeWriteExec, CometPlan, CometRangeExec, CometSparkToColumnarExec} import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometShuffleExchangeExec} import org.apache.spark.sql.comet.shims.{MapInBatchInfo, ShimCometMapInBatch} import org.apache.spark.sql.comet.util.Utils.containsVariantType @@ -122,12 +122,8 @@ case class EliminateRedundantTransitions(session: SparkSession) op case CometColumnarToRowExec(sparkToColumnar: CometSparkToColumnarExec) => sparkToColumnar.child - case CometNativeColumnarToRowExec(sparkToColumnar: CometSparkToColumnarExec) => - sparkToColumnar.child case CometColumnarToRowExec(range: CometRangeExec) => range.originalPlan - case CometNativeColumnarToRowExec(range: CometRangeExec) => - range.originalPlan case CometSparkToColumnarExec(child: CometSparkToColumnarExec, _) => child // Replace MapInBatchExec (PythonMapInArrowExec / MapInArrowExec / MapInPandasExec) that has // a ColumnarToRow child with CometMapInBatchExec, eliminating the input and output @@ -178,8 +174,6 @@ case class EliminateRedundantTransitions(session: SparkSession) child case CometColumnarToRowExec(child: CometCollectLimitExec) => child - case CometNativeColumnarToRowExec(child: CometCollectLimitExec) => - child case other => other } @@ -259,7 +253,6 @@ case class EliminateRedundantTransitions(session: SparkSession) * runtime with `requires a columnar (Comet native) child`. */ private def stripColumnarToRow(plan: SparkPlan): Option[SparkPlan] = plan match { - case CometNativeColumnarToRowExec(child) => Some(child) case CometColumnarToRowExec(child) => Some(child) case ColumnarToRowExec(child) => Some(child) case _ => None @@ -274,7 +267,6 @@ case class EliminateRedundantTransitions(session: SparkSession) */ private def extractColumnarChild(plan: SparkPlan): Option[SparkPlan] = plan match { case CometColumnarToRowExec(child) => Some(child) - case CometNativeColumnarToRowExec(child) => Some(child) // Chained `mapInArrow(udf1).mapInArrow(udf2)`: by the time the outer operator is visited // (transformUp is bottom-up) the inner one has already become a `CometMapInBatchExec`, which // is itself columnar. There is no row transition between them to strip, so consume its @@ -319,29 +311,18 @@ case class EliminateRedundantTransitions(session: SparkSession) } /** - * Creates an appropriate columnar to row transition operator. - * - * If native columnar to row conversion is enabled and the schema is supported, uses - * CometNativeColumnarToRowExec. Variant uses Spark's conversion; other unsupported schemas use - * CometColumnarToRowExec. + * Creates the columnar to row transition operator: `CometColumnarToRowExec`, or Spark's + * `ColumnarToRowExec` for a schema that contains Variant. */ private def createColumnarToRowExec(child: SparkPlan): SparkPlan = { - val schema = child.schema // TODO: Remove this fallback once Comet columnar-to-row conversion supports Variant getters // and Spark's Variant UnsafeRow encoding. // https://github.com/apache/datafusion-comet/issues/5436 - if (containsVariantType(schema)) { + if (containsVariantType(child.schema)) { return withFallbackReason( ColumnarToRowExec(child), - "Native columnar-to-row conversion does not support type VariantType") - } - val useNative = CometConf.COMET_NATIVE_COLUMNAR_TO_ROW_ENABLED.get() && - CometNativeColumnarToRowExec.supportsSchema(schema) - - if (useNative) { - CometNativeColumnarToRowExec(child) - } else { - CometColumnarToRowExec(child) + "Comet columnar-to-row conversion does not support type VariantType") } + CometColumnarToRowExec(child) } } diff --git a/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala b/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala index 3a6f0d6c196..7811b893b2f 100644 --- a/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala +++ b/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala @@ -444,7 +444,7 @@ private[rules] object MakeDateEvaluation { val rebuilt = node match { case _: RowToColumnarExec | _: CometSparkToColumnarExec if childReason.isDefined => children.head - case _: ColumnarToRowExec | _: CometColumnarToRowExec | _: CometNativeColumnarToRowExec + case _: ColumnarToRowExec | _: CometColumnarToRowExec if childReason.isDefined && !children.head.supportsColumnar => children.head case native: CometExec if fallbackReason.isDefined || children != node.children => diff --git a/spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala b/spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala index 8a565dd080e..30cd466fd80 100644 --- a/spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala +++ b/spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala @@ -23,7 +23,7 @@ import org.apache.spark.internal.Logging import org.apache.spark.sql.SparkSession import org.apache.spark.sql.catalyst.expressions.aggregate.{Final, Partial, PartialMerge} import org.apache.spark.sql.catalyst.rules.Rule -import org.apache.spark.sql.comet.{CometBaseAggregateExec, CometCollectMetricsExec, CometColumnarToRowExec, CometExec, CometNativeColumnarToRowExec, CometSparkToColumnarExec} +import org.apache.spark.sql.comet.{CometBaseAggregateExec, CometCollectMetricsExec, CometColumnarToRowExec, CometExec, CometSparkToColumnarExec} import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometNativeShuffle, CometShuffleExchangeExec} import org.apache.spark.sql.execution.{ColumnarToRowExec, ColumnarToRowTransition, RowToColumnarExec, SparkPlan} import org.apache.spark.sql.execution.adaptive.QueryStageExec @@ -136,7 +136,6 @@ case class RevertNativeForTransitionHeavyStages(session: SparkSession, wholePlan */ private def convertsSparkRows(plan: SparkPlan): Boolean = plan match { case conversion: CometSparkToColumnarExec => - !conversion.child.isInstanceOf[CometNativeColumnarToRowExec] && !conversion.child.isInstanceOf[CometColumnarToRowExec] case _ => false } @@ -225,10 +224,10 @@ case class RevertNativeForTransitionHeavyStages(session: SparkSession, wholePlan /** * Like `transformDown`, never descends stage-boundary children. If the rule rewrites the * current node, re-apply it to the result so stacked transitions such as - * `CometSparkToColumnarExec(CometNativeColumnarToRowExec(x))` are fully unwrapped before - * children are visited. Spark's `transformDown` does not do this; leaving the inner C2R in - * place later calls `CometNativeColumnarToRowExec.withNewChildren` with a reverted row-based - * child, which asserts `child.supportsColumnar`. + * `CometSparkToColumnarExec(CometColumnarToRowExec(x))` are fully unwrapped before children are + * visited. Spark's `transformDown` does not do this; leaving the inner C2R in place later calls + * `CometColumnarToRowExec.withNewChildren` with a reverted row-based child, which asserts + * `child.supportsColumnar`. * * A rewrite can itself be the stage boundary. Unwrapping a transition that sits directly on a * shuffle yields that shuffle, and descending into it strips transitions in the next stage. @@ -298,7 +297,6 @@ case class RevertNativeForTransitionHeavyStages(session: SparkSession, wholePlan private[rules] def revertToSpark(plan: SparkPlan): SparkPlan = { validateOriginalPlanAliases(plan) val stripped = transformStageDown(plan) { - case CometNativeColumnarToRowExec(child) => child case CometColumnarToRowExec(child) => child case ColumnarToRowExec(child) => child case sparkToColumnar: CometSparkToColumnarExec => sparkToColumnar.child diff --git a/spark/src/main/scala/org/apache/spark/Plugins.scala b/spark/src/main/scala/org/apache/spark/Plugins.scala index 70f3ce6cedc..c4c40dd7d51 100644 --- a/spark/src/main/scala/org/apache/spark/Plugins.scala +++ b/spark/src/main/scala/org/apache/spark/Plugins.scala @@ -124,6 +124,7 @@ class CometDriverPlugin private[spark] (clock: Clock) extends DriverPlugin with CometDriverPlugin.warnIfExecutorMemoryOverheadUnset(sc.getConf) CometDriverPlugin.warnIfMemoryPoolFractionSet(sc.getConf) + CometDriverPlugin.warnIfNativeColumnarToRowSet(sc.getConf) extraConfs } @@ -336,6 +337,17 @@ object CometDriverPlugin extends Logging { } } + // The Rust columnar-to-row converter that spark.comet.exec.columnarToRow.native.enabled turned + // on has been removed, so the key no longer does anything. + private[apache] def warnIfNativeColumnarToRowSet(conf: SparkConf): Unit = { + val key = CometConf.COMET_NATIVE_COLUMNAR_TO_ROW_ENABLED.key + conf.getOption(key).foreach { value => + logWarning( + s"$key=$value is deprecated and will be removed in a future major release. It has no " + + "effect: Comet always converts Arrow batches to Spark rows on the JVM.") + } + } + // Reads a deprecated alternative too, such as spark.comet.exec.shuffle.enabled, as a session // would. private def getBooleanConf(conf: SparkConf, entry: ConfigEntry[Boolean]): Boolean = diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeColumnarToRowExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeColumnarToRowExec.scala deleted file mode 100644 index 4d8c913f2cb..00000000000 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeColumnarToRowExec.scala +++ /dev/null @@ -1,260 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you under the Apache License, Version 2.0 (the - * "License"); you may not use this file except in compliance - * with the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, - * software distributed under the License is distributed on an - * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY - * KIND, either express or implied. See the License for the - * specific language governing permissions and limitations - * under the License. - */ - -package org.apache.spark.sql.comet - -import java.util.UUID -import java.util.concurrent.{Future, TimeoutException, TimeUnit} - -import scala.concurrent.Promise -import scala.util.control.NonFatal - -import org.apache.spark.{broadcast, SparkException, TaskContext} -import org.apache.spark.rdd.RDD -import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.expressions.{Attribute, SortOrder} -import org.apache.spark.sql.catalyst.plans.physical.Partitioning -import org.apache.spark.sql.comet.util.{Utils => CometUtils} -import org.apache.spark.sql.errors.QueryExecutionErrors -import org.apache.spark.sql.execution.{ColumnarToRowTransition, SparkPlan, SQLExecution} -import org.apache.spark.sql.execution.adaptive.BroadcastQueryStageExec -import org.apache.spark.sql.execution.exchange.ReusedExchangeExec -import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} -import org.apache.spark.sql.types.StructType -import org.apache.spark.util.{SparkFatalException, Utils} - -import org.apache.comet.{CometConf, NativeColumnarToRowConverter} - -/** - * Native implementation of ColumnarToRowExec that converts Arrow columnar data to Spark UnsafeRow - * format using Rust. - * - * This feature is disabled by default and can be enabled by setting - * `spark.comet.exec.columnarToRow.native.enabled=true`. The native conversion carries a fixed JNI - * cost per batch and is slower than the JVM implementation for small batches (see issue #5112 for - * measurements). - * - * @param child - * The child plan that produces columnar batches - */ -case class CometNativeColumnarToRowExec(child: SparkPlan) - extends ColumnarToRowTransition - with CometPlan { - - // supportsColumnar requires to be only called on driver side, see also SPARK-37779. - assert(Utils.isInRunningSparkTask || child.supportsColumnar) - - override def output: Seq[Attribute] = child.output - - override def outputPartitioning: Partitioning = child.outputPartitioning - - override def outputOrdering: Seq[SortOrder] = child.outputOrdering - - override lazy val metrics: Map[String, SQLMetric] = Map( - "numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output rows"), - "numInputBatches" -> SQLMetrics.createMetric(sparkContext, "number of input batches"), - "convertTime" -> SQLMetrics.createNanoTimingMetric(sparkContext, "time in conversion")) - - @transient - private lazy val promise = Promise[broadcast.Broadcast[Any]]() - - @transient - private val timeout: Long = conf.broadcastTimeout - - private val runId: UUID = UUID.randomUUID - - private lazy val cometBroadcastExchange = findCometBroadcastExchange(child) - - @transient - lazy val relationFuture: Future[broadcast.Broadcast[Any]] = { - SQLExecution.withThreadLocalCaptured[broadcast.Broadcast[Any]]( - session, - CometBroadcastExchangeExec.executionContext) { - try { - // Setup a job group here so later it may get cancelled by groupId if necessary. - sparkContext.setJobGroup( - runId.toString, - s"CometNativeColumnarToRow broadcast exchange (runId $runId)", - interruptOnCancel = true) - - val numOutputRows = longMetric("numOutputRows") - val numInputBatches = longMetric("numInputBatches") - val localSchema = this.schema - val batchSize = CometConf.COMET_BATCH_SIZE.get() - val broadcastColumnar = child.executeBroadcast[Any]() - val serializedBatches = - broadcastColumnar.value.asInstanceOf[CometBroadcastValue].chunks - - // Use native converter to convert columnar data to rows - val converter = new NativeColumnarToRowConverter(localSchema, batchSize) - try { - val rows = serializedBatches.iterator - .flatMap(CometUtils.decodeBatches(_, this.getClass.getSimpleName)) - .flatMap { batch => - numInputBatches += 1 - numOutputRows += batch.numRows().toLong - val result = converter.convert(batch) - // Wrap iterator to close batch after consumption - new Iterator[InternalRow] { - override def hasNext: Boolean = { - val hasMore = result.hasNext - if (!hasMore) { - batch.close() - } - hasMore - } - override def next(): InternalRow = result.next() - } - } - - val mode = cometBroadcastExchange.get.mode - val relation = mode.transform(rows, Some(numOutputRows.value)) - val broadcasted = sparkContext.broadcastInternal(relation, serializedOnly = true) - val executionId = sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY) - SQLMetrics.postDriverMetricUpdates(sparkContext, executionId, metrics.values.toSeq) - promise.trySuccess(broadcasted) - broadcasted - } finally { - converter.close() - } - } catch { - // SPARK-24294: To bypass scala bug: https://github.com/scala/bug/issues/9554, we throw - // SparkFatalException, which is a subclass of Exception. ThreadUtils.awaitResult - // will catch this exception and re-throw the wrapped fatal throwable. - case oe: OutOfMemoryError => - val ex = new SparkFatalException(oe) - promise.tryFailure(ex) - throw ex - case e if !NonFatal(e) => - val ex = new SparkFatalException(e) - promise.tryFailure(ex) - throw ex - case e: Throwable => - promise.tryFailure(e) - throw e - } - } - } - - override def doExecuteBroadcast[T](): broadcast.Broadcast[T] = { - if (cometBroadcastExchange.isEmpty) { - throw new SparkException( - "CometNativeColumnarToRowExec only supports doExecuteBroadcast when child contains a " + - "CometBroadcastExchange, but got " + child) - } - - try { - relationFuture.get(timeout, TimeUnit.SECONDS).asInstanceOf[broadcast.Broadcast[T]] - } catch { - case ex: TimeoutException => - logError(s"Could not execute broadcast in $timeout secs.", ex) - if (!relationFuture.isDone) { - sparkContext.cancelJobGroup(runId.toString) - relationFuture.cancel(true) - } - throw QueryExecutionErrors.executeBroadcastTimeoutError(timeout, Some(ex)) - } - } - - private def findCometBroadcastExchange(op: SparkPlan): Option[CometBroadcastExchangeExec] = { - op match { - case b: CometBroadcastExchangeExec => Some(b) - case b: BroadcastQueryStageExec => findCometBroadcastExchange(b.plan) - case b: ReusedExchangeExec => findCometBroadcastExchange(b.child) - case _ => op.children.collectFirst(Function.unlift(findCometBroadcastExchange)) - } - } - - override def doExecute(): RDD[InternalRow] = { - val numOutputRows = longMetric("numOutputRows") - val numInputBatches = longMetric("numInputBatches") - val convertTime = longMetric("convertTime") - - // Get the schema and batch size for native conversion - val localSchema = child.schema - val batchSize = CometConf.COMET_BATCH_SIZE.get() - - child.executeColumnar().mapPartitionsInternal { batches => - // Create native converter for this partition - val converter = new NativeColumnarToRowConverter(localSchema, batchSize) - - // Register cleanup on task completion - TaskContext.get().addTaskCompletionListener[Unit] { _ => - converter.close() - } - - batches.flatMap { batch => - numInputBatches += 1 - val numRows = batch.numRows() - numOutputRows += numRows.toLong - - val startTime = System.nanoTime() - val result = converter.convert(batch) - convertTime += System.nanoTime() - startTime - - // Wrap iterator to close batch after consumption - new Iterator[InternalRow] { - override def hasNext: Boolean = { - val hasMore = result.hasNext - if (!hasMore) { - batch.close() - } - hasMore - } - override def next(): InternalRow = result.next() - } - } - } - } - - override protected def withNewChildInternal(newChild: SparkPlan): CometNativeColumnarToRowExec = - copy(child = newChild) -} - -object CometNativeColumnarToRowExec { - - /** - * Checks if the given schema is supported by native columnar to row conversion. - * - * Currently supported types: - * - Primitive types: Boolean, Byte, Short, Int, Long, Float, Double - * - Date and Timestamp (microseconds) - * - Decimal (both inline and variable-length) - * - String and Binary - * - Struct, Array, Map (nested types) - */ - def supportsSchema(schema: StructType): Boolean = { - import org.apache.spark.sql.types._ - - def isSupported(dataType: DataType): Boolean = dataType match { - case BooleanType | ByteType | ShortType | IntegerType | LongType => true - case FloatType | DoubleType => true - case DateType => true - case TimestampType => true - case _: DecimalType => true - case StringType | BinaryType => true - case StructType(fields) => fields.forall(f => isSupported(f.dataType)) - case ArrayType(elementType, _) => isSupported(elementType) - case MapType(keyType, valueType, _) => isSupported(keyType) && isSupported(valueType) - case _ => false - } - - schema.fields.forall(f => isSupported(f.dataType)) - } -} diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenFuzzSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenFuzzSuite.scala index 9c0940ad099..2f1c605cb01 100644 --- a/spark/src/test/scala/org/apache/comet/CometCodegenFuzzSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometCodegenFuzzSuite.scala @@ -44,8 +44,8 @@ import org.apache.comet.vector.CometVector * Randomized end-to-end tests for the Arrow-direct codegen dispatcher: schema-driven coverage of * every input vector class against random parquet, plus a decimal precision-scale sweep across * the `Decimal.MAX_LONG_DIGITS=18` boundary at varying null densities. Extends [[CometTestBase]] - * (not [[CometFuzzTestBase]]) because the base's `shuffle` x `nativeC2R` cross-product is - * irrelevant for projection-only queries. + * (not [[CometFuzzTestBase]]) because the base's per-shuffle-mode runs are irrelevant for + * projection-only queries. */ class CometCodegenFuzzSuite extends CometTestBase diff --git a/spark/src/test/scala/org/apache/comet/CometFuzzTestBase.scala b/spark/src/test/scala/org/apache/comet/CometFuzzTestBase.scala index b112a11d548..88c09d0c624 100644 --- a/spark/src/test/scala/org/apache/comet/CometFuzzTestBase.scala +++ b/spark/src/test/scala/org/apache/comet/CometFuzzTestBase.scala @@ -112,16 +112,14 @@ class CometFuzzTestBase extends CometTestBase with AdaptiveSparkPlanHelper { override protected def test(testName: String, testTags: Tag*)(testFun: => Any)(implicit pos: Position): Unit = { - Seq(("native", "false"), ("jvm", "true"), ("jvm", "false")).foreach { - case (shuffleMode, nativeC2R) => - super.test(testName + s" ($shuffleMode shuffle, nativeC2R=$nativeC2R)", testTags: _*) { - withSQLConf( - CometConf.COMET_NATIVE_COLUMNAR_TO_ROW_ENABLED.key -> nativeC2R, - CometConf.COMET_PARQUET_UNSIGNED_SMALL_INT_CHECK.key -> "false", - CometConf.COMET_SHUFFLE_MODE.key -> shuffleMode) { - testFun - } + Seq("native", "jvm").foreach { shuffleMode => + super.test(testName + s" ($shuffleMode shuffle)", testTags: _*) { + withSQLConf( + CometConf.COMET_PARQUET_UNSIGNED_SMALL_INT_CHECK.key -> "false", + CometConf.COMET_SHUFFLE_MODE.key -> shuffleMode) { + testFun } + } } } diff --git a/spark/src/test/scala/org/apache/comet/CometVariantProjectionSuite.scala b/spark/src/test/scala/org/apache/comet/CometVariantProjectionSuite.scala index 8bb9f62a72c..e61510ed966 100644 --- a/spark/src/test/scala/org/apache/comet/CometVariantProjectionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometVariantProjectionSuite.scala @@ -25,7 +25,6 @@ import org.apache.parquet.io.api.Binary import org.apache.parquet.schema.MessageTypeParser import org.apache.spark.SparkConf import org.apache.spark.sql.{CometTestBase, DataFrame, Row} -import org.apache.spark.sql.comet.CometNativeColumnarToRowExec import org.apache.spark.sql.comet.CometNativeScanExec import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.execution.{ColumnarToRowExec, CommandResultExec, ProjectExec, SparkPlan} @@ -71,7 +70,6 @@ class CometVariantProjectionSuite extends CometTestBase { val plan = checkVariantAnswer(df, expected.getOrElse(sparkRows(df))) checkCometOperators(plan, classOf[ColumnarToRowExec]) assert(collect(plan) { case scan: CometNativeScanExec => scan }.nonEmpty, plan.toString) - assert(collect(plan) { case c: CometNativeColumnarToRowExec => c }.isEmpty, plan.toString) } private def checkScanFallbackPlan(df: DataFrame, reason: String): Unit = { diff --git a/spark/src/test/scala/org/apache/comet/CometVariantTypeSuite.scala b/spark/src/test/scala/org/apache/comet/CometVariantTypeSuite.scala index 92c2ec76f13..222dba677ef 100644 --- a/spark/src/test/scala/org/apache/comet/CometVariantTypeSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometVariantTypeSuite.scala @@ -28,7 +28,6 @@ import org.scalatest.funsuite.AnyFunSuite import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType} import org.apache.spark.sql.catalyst.expressions.AttributeReference -import org.apache.spark.sql.comet.CometNativeColumnarToRowExec import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.types.{ArrayType, BinaryType, StructField, StructType} @@ -70,12 +69,6 @@ class CometVariantTypeSuite extends AnyFunSuite { assert(QueryPlanSerde.serializeDataType(ArrayType(variantType)).isDefined) assert(!QueryPlanSerde.supportedDataType(variantType)) assert(!QueryPlanSerde.supportedDataType(ArrayType(variantType), allowComplex = true)) - assert( - !CometNativeColumnarToRowExec.supportsSchema( - StructType(Seq(StructField("v", variantType))))) - assert( - !CometNativeColumnarToRowExec.supportsSchema( - StructType(Seq(StructField("nested", ArrayType(variantType)))))) assert(!CometScanTypeChecker().isTypeSupported(variantType, "v", ListBuffer.empty)) assert( !CometScanTypeChecker() diff --git a/spark/src/test/scala/org/apache/comet/exec/CometDirectColumnarToRowSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometDirectColumnarToRowSuite.scala index 1c5ff110f94..3a4ee1f00a5 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometDirectColumnarToRowSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometDirectColumnarToRowSuite.scala @@ -34,7 +34,6 @@ class CometDirectColumnarToRowSuite extends CometTestBase { private def withDirectConverter(minBatchSize: Int = 0)(f: => Unit): Unit = { withSQLConf( - CometConf.COMET_NATIVE_COLUMNAR_TO_ROW_ENABLED.key -> "false", CometConf.COMET_DIRECT_COLUMNAR_TO_ROW_ENABLED.key -> "true", CometConf.COMET_DIRECT_COLUMNAR_TO_ROW_MIN_BATCH_SIZE.key -> minBatchSize.toString, SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false")(f) diff --git a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala index 1ba17ea2dae..7b30248b833 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala @@ -3143,7 +3143,6 @@ class CometExecSuite extends CometTestBase { var columnarToRowExec: Seq[SparkPlan] = stripAQEPlan(df.queryExecution.executedPlan).collect { case s: CometColumnarToRowExec => s - case s: CometNativeColumnarToRowExec => s } assert(columnarToRowExec.isEmpty) @@ -3158,9 +3157,8 @@ class CometExecSuite extends CometTestBase { // After AQE: CometBroadcastExchange has to be converted to rows to conform to Spark // BroadcastHashJoin. val plan = stripAQEPlan(df.queryExecution.executedPlan) - columnarToRowExec = plan.collect { - case s: CometColumnarToRowExec => s - case s: CometNativeColumnarToRowExec => s + columnarToRowExec = plan.collect { case s: CometColumnarToRowExec => + s } assert(columnarToRowExec.length == 1) @@ -3170,7 +3168,7 @@ class CometExecSuite extends CometTestBase { assert(broadcastJoins.nonEmpty, s"Expected BroadcastHashJoinExec in plan:\n$plan") val hasC2RDescendant = broadcastJoins.exists { join => join.find { - case _: CometColumnarToRowExec | _: CometNativeColumnarToRowExec => true + case _: CometColumnarToRowExec => true case _ => false }.isDefined } 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 1f9b4c1a513..496b80ef51a 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -1023,11 +1023,9 @@ class CometInMemoryCacheSuite extends CometTestBase { val plan = df.queryExecution.executedPlan.toString() assert(plan.contains("CometInMemoryTableScan")) - // Rows come out of a Comet converter rather than Spark's ColumnarToRow. Either variant - // satisfies that; which one is used depends on the default of - // spark.comet.exec.columnarToRow.native.enabled. + // Rows come out of Comet's converter rather than Spark's ColumnarToRow. assert( - plan.contains("CometColumnarToRow") || plan.contains("CometNativeColumnarToRow"), + plan.contains("CometColumnarToRow"), s"expected a Comet columnar-to-row above the cache scan, got:\n$plan") spark.catalog.clearCache() diff --git a/spark/src/test/scala/org/apache/comet/exec/CometNativeColumnarToRowSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometNativeColumnarToRowSuite.scala deleted file mode 100644 index 2dc6fcfd325..00000000000 --- a/spark/src/test/scala/org/apache/comet/exec/CometNativeColumnarToRowSuite.scala +++ /dev/null @@ -1,550 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you under the Apache License, Version 2.0 (the - * "License"); you may not use this file except in compliance - * with the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, - * software distributed under the License is distributed on an - * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY - * KIND, either express or implied. See the License for the - * specific language governing permissions and limitations - * under the License. - */ - -package org.apache.comet.exec - -import java.sql.Date -import java.sql.Timestamp - -import scala.util.Random - -import org.scalactic.source.Position -import org.scalatest.Tag - -import org.apache.spark.sql.{CometTestBase, Row} -import org.apache.spark.sql.comet.CometNativeColumnarToRowExec -import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper -import org.apache.spark.sql.types._ - -import org.apache.comet.{CometConf, NativeColumnarToRowConverter} -import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator, SchemaGenOptions} - -/** - * Test suite for native columnar to row conversion. - * - * These tests verify that CometNativeColumnarToRowExec produces correct results for all supported - * data types. - */ -class CometNativeColumnarToRowSuite extends CometTestBase with AdaptiveSparkPlanHelper { - - override protected def test(testName: String, testTags: Tag*)(testFun: => Any)(implicit - pos: Position): Unit = { - super.test(testName, testTags: _*) { - withSQLConf(CometConf.COMET_NATIVE_COLUMNAR_TO_ROW_ENABLED.key -> "true") { - testFun - } - } - } - - /** - * Helper to verify that CometNativeColumnarToRowExec is present in the plan. - */ - private def assertNativeC2RPresent(df: org.apache.spark.sql.DataFrame): Unit = { - val plan = stripAQEPlan(df.queryExecution.executedPlan) - val nativeC2R = plan.collect { case c: CometNativeColumnarToRowExec => c } - assert( - nativeC2R.nonEmpty, - s"Expected CometNativeColumnarToRowExec in plan but found none.\nPlan: $plan") - } - - test("primitive types: boolean, byte, short, int, long") { - val data = (0 until 100).map { i => - (i % 2 == 0, i.toByte, i.toShort, i, i.toLong) - } - withParquetTable(data, "primitives") { - // Force row output by using a UDF or collect - val df = sql("SELECT * FROM primitives") - assertNativeC2RPresent(df) - checkSparkAnswer(df) - } - } - - test("primitive types: float, double") { - val data = (0 until 100).map { i => - (i.toFloat / 10.0f, i.toDouble / 10.0) - } - withParquetTable(data, "floats") { - val df = sql("SELECT * FROM floats") - assertNativeC2RPresent(df) - checkSparkAnswer(df) - } - } - - test("primitive types with nulls") { - val data = (0 until 100).map { i => - val isNull = i % 5 == 0 - ( - if (isNull) null else (i % 2 == 0), - if (isNull) null else i.toByte, - if (isNull) null else i.toShort, - if (isNull) null else i, - if (isNull) null else i.toLong, - if (isNull) null else i.toFloat, - if (isNull) null else i.toDouble) - } - val df = spark - .createDataFrame( - spark.sparkContext.parallelize(data.map(Row.fromTuple(_))), - StructType( - Seq( - StructField("bool", BooleanType, nullable = true), - StructField("byte", ByteType, nullable = true), - StructField("short", ShortType, nullable = true), - StructField("int", IntegerType, nullable = true), - StructField("long", LongType, nullable = true), - StructField("float", FloatType, nullable = true), - StructField("double", DoubleType, nullable = true)))) - withParquetDataFrame(df) { parquetDf => - assertNativeC2RPresent(parquetDf) - checkSparkAnswer(parquetDf) - } - } - - test("string type") { - val data = (0 until 100).map { i => - (i, s"string_value_$i", if (i % 10 == 0) null else s"nullable_$i") - } - val df = spark - .createDataFrame( - spark.sparkContext.parallelize(data.map(Row.fromTuple(_))), - StructType( - Seq( - StructField("id", IntegerType), - StructField("str", StringType), - StructField("nullable_str", StringType, nullable = true)))) - withParquetDataFrame(df) { parquetDf => - assertNativeC2RPresent(parquetDf) - checkSparkAnswer(parquetDf) - } - } - - test("string type with various lengths") { - val random = new Random(42) - val data = (0 until 100).map { i => - val len = random.nextInt(1000) - (i, random.alphanumeric.take(len).mkString) - } - withParquetTable(data, "strings") { - val df = sql("SELECT * FROM strings") - assertNativeC2RPresent(df) - checkSparkAnswer(df) - } - } - - test("binary type") { - val data = (0 until 100).map { i => - (i, s"binary_$i".getBytes) - } - val df = spark - .createDataFrame( - spark.sparkContext.parallelize(data.map { case (id, bytes) => Row(id, bytes) }), - StructType( - Seq(StructField("id", IntegerType), StructField("bin", BinaryType, nullable = true)))) - withParquetDataFrame(df) { parquetDf => - assertNativeC2RPresent(parquetDf) - checkSparkAnswer(parquetDf) - } - } - - test("date type") { - val baseDate = Date.valueOf("2024-01-01") - val data = (0 until 100).map { i => - (i, new Date(baseDate.getTime + i * 24 * 60 * 60 * 1000L)) - } - withParquetTable(data, "dates") { - val df = sql("SELECT * FROM dates") - assertNativeC2RPresent(df) - checkSparkAnswer(df) - } - } - - test("timestamp type") { - val baseTs = Timestamp.valueOf("2024-01-01 00:00:00") - val data = (0 until 100).map { i => - (i, new Timestamp(baseTs.getTime + i * 1000L)) - } - withParquetTable(data, "timestamps") { - val df = sql("SELECT * FROM timestamps") - assertNativeC2RPresent(df) - checkSparkAnswer(df) - } - } - - test("decimal type - inline precision (precision <= 18)") { - val data = (0 until 100).map { i => - (i, BigDecimal(i * 100 + i) / 100) - } - val df = spark - .createDataFrame( - spark.sparkContext.parallelize(data.map { case (id, dec) => - Row(id, dec.bigDecimal) - }), - StructType( - Seq( - StructField("id", IntegerType), - StructField("dec", DecimalType(10, 2), nullable = true)))) - withParquetDataFrame(df) { parquetDf => - assertNativeC2RPresent(parquetDf) - checkSparkAnswer(parquetDf) - } - } - - test("decimal type - variable length precision (precision > 18)") { - val data = (0 until 100).map { i => - (i, BigDecimal(s"12345678901234567890.$i")) - } - val df = spark - .createDataFrame( - spark.sparkContext.parallelize(data.map { case (id, dec) => - Row(id, dec.bigDecimal) - }), - StructType( - Seq( - StructField("id", IntegerType), - StructField("dec", DecimalType(30, 5), nullable = true)))) - withParquetDataFrame(df) { parquetDf => - assertNativeC2RPresent(parquetDf) - checkSparkAnswer(parquetDf) - } - } - - test("struct type") { - val data = (0 until 100).map { i => - (i, (i * 2, s"nested_$i")) - } - withParquetTable(data, "structs") { - val df = sql("SELECT * FROM structs") - assertNativeC2RPresent(df) - checkSparkAnswer(df) - } - } - - test("array type") { - val data = (0 until 100).map { i => - (i, (0 to i % 10).toArray) - } - withParquetTable(data, "arrays") { - val df = sql("SELECT * FROM arrays") - assertNativeC2RPresent(df) - checkSparkAnswer(df) - } - } - - test("array of strings") { - val data = (0 until 100).map { i => - (i, (0 to i % 5).map(j => s"elem_${i}_$j").toArray) - } - withParquetTable(data, "string_arrays") { - val df = sql("SELECT * FROM string_arrays") - assertNativeC2RPresent(df) - checkSparkAnswer(df) - } - } - - test("map type") { - val data = (0 until 100).map { i => - (i, (0 to i % 5).map(j => (s"key_$j", j * 10)).toMap) - } - withParquetTable(data, "maps") { - val df = sql("SELECT * FROM maps") - assertNativeC2RPresent(df) - checkSparkAnswer(df) - } - } - - test("nested struct in array") { - val data = (0 until 100).map { i => - (i, (0 to i % 5).map(j => (j, s"nested_$j")).toArray) - } - withParquetTable(data, "nested_structs") { - val df = sql("SELECT * FROM nested_structs") - assertNativeC2RPresent(df) - checkSparkAnswer(df) - } - } - - test("deeply nested: array of arrays") { - val data = (0 until 100).map { i => - (i, (0 to i % 3).map(j => (0 to j).toArray).toArray) - } - withParquetTable(data, "nested_arrays") { - val df = sql("SELECT * FROM nested_arrays") - assertNativeC2RPresent(df) - checkSparkAnswer(df) - } - } - - test("deeply nested: map with array values") { - val data = (0 until 100).map { i => - (i, (0 to i % 3).map(j => (s"key_$j", (0 to j).toArray)).toMap) - } - withParquetTable(data, "map_array_values") { - val df = sql("SELECT * FROM map_array_values") - assertNativeC2RPresent(df) - checkSparkAnswer(df) - } - } - - test("deeply nested: struct containing array of maps") { - val data = (0 until 100).map { i => - ( - i, - ((0 to i % 3).map(j => (0 to j % 2).map(k => (s"k$k", k * 10)).toMap).toArray, s"str_$i")) - } - withParquetTable(data, "struct_array_maps") { - val df = sql("SELECT * FROM struct_array_maps") - assertNativeC2RPresent(df) - checkSparkAnswer(df) - } - } - - test("all null values") { - val df = spark - .createDataFrame( - spark.sparkContext.parallelize((0 until 100).map(_ => Row(null, null, null))), - StructType( - Seq( - StructField("int_col", IntegerType, nullable = true), - StructField("str_col", StringType, nullable = true), - StructField("double_col", DoubleType, nullable = true)))) - withParquetDataFrame(df) { parquetDf => - assertNativeC2RPresent(parquetDf) - checkSparkAnswer(parquetDf) - } - } - - test("empty batch") { - val df = spark - .createDataFrame( - spark.sparkContext.parallelize(Seq.empty[Row]), - StructType(Seq(StructField("int_col", IntegerType), StructField("str_col", StringType)))) - withParquetDataFrame(df) { parquetDf => - assertNativeC2RPresent(parquetDf) - checkSparkAnswer(parquetDf) - } - } - - test("mixed types - comprehensive") { - val baseDate = Date.valueOf("2024-01-01") - val baseTs = Timestamp.valueOf("2024-01-01 00:00:00") - - val data = (0 until 100).map { i => - val isNull = i % 7 == 0 - Row( - i, - if (isNull) null else (i % 2 == 0), - if (isNull) null else i.toByte, - if (isNull) null else i.toShort, - if (isNull) null else i.toLong, - if (isNull) null else i.toFloat, - if (isNull) null else i.toDouble, - if (isNull) null else s"string_$i", - if (isNull) null else new Date(baseDate.getTime + i * 24 * 60 * 60 * 1000L), - if (isNull) null else new Timestamp(baseTs.getTime + i * 1000L), - if (isNull) null else BigDecimal((i * 100 + i).toLong, 2).bigDecimal) - } - - val schema = StructType( - Seq( - StructField("id", IntegerType), - StructField("bool_col", BooleanType, nullable = true), - StructField("byte_col", ByteType, nullable = true), - StructField("short_col", ShortType, nullable = true), - StructField("long_col", LongType, nullable = true), - StructField("float_col", FloatType, nullable = true), - StructField("double_col", DoubleType, nullable = true), - StructField("string_col", StringType, nullable = true), - StructField("date_col", DateType, nullable = true), - StructField("timestamp_col", TimestampType, nullable = true), - StructField("decimal_col", DecimalType(10, 2), nullable = true))) - - val df = spark.createDataFrame(spark.sparkContext.parallelize(data), schema) - withParquetDataFrame(df) { parquetDf => - assertNativeC2RPresent(parquetDf) - checkSparkAnswer(parquetDf) - } - } - - test("large batch") { - val data = (0 until 10000).map { i => - (i, s"string_value_$i", i.toDouble) - } - withParquetTable(data, "large_batch") { - val df = sql("SELECT * FROM large_batch") - assertNativeC2RPresent(df) - checkSparkAnswer(df) - } - } - - test("disabled by default") { - withSQLConf(CometConf.COMET_NATIVE_COLUMNAR_TO_ROW_ENABLED.key -> "false") { - val data = (0 until 10).map(i => (i, s"value_$i")) - withParquetTable(data, "test_disabled") { - val df = sql("SELECT * FROM test_disabled") - val plan = stripAQEPlan(df.queryExecution.executedPlan) - val nativeC2R = plan.collect { case c: CometNativeColumnarToRowExec => c } - assert( - nativeC2R.isEmpty, - s"Expected no CometNativeColumnarToRowExec when disabled.\nPlan: $plan") - checkSparkAnswer(df) - } - } - } - - test("fuzz test with nested types") { - val random = new Random(42) - val schemaGenOptions = - SchemaGenOptions(generateArray = true, generateStruct = true, generateMap = true) - val dataGenOptions = - DataGenOptions(generateNegativeZero = false, generateNaN = false, generateInfinity = false) - - // Use generateSchema which creates various nested types including arrays, structs, and maps. - // Not all generated types may be supported by native C2R, so we just verify correctness. - val schema = FuzzDataGenerator.generateSchema(schemaGenOptions) - val df = FuzzDataGenerator.generateDataFrame(random, spark, schema, 100, dataGenOptions) - - withParquetDataFrame(df) { parquetDf => - checkSparkAnswer(parquetDf) - } - } - - test("fuzz test with generateNestedSchema") { - val random = new Random(42) - - // Use only primitive types supported by native C2R (excludes TimestampNTZType) - val supportedPrimitiveTypes: Seq[DataType] = Seq( - DataTypes.BooleanType, - DataTypes.ByteType, - DataTypes.ShortType, - DataTypes.IntegerType, - DataTypes.LongType, - DataTypes.FloatType, - DataTypes.DoubleType, - DataTypes.createDecimalType(10, 2), - DataTypes.DateType, - DataTypes.TimestampType, - DataTypes.StringType, - DataTypes.BinaryType) - - val schemaGenOptions = SchemaGenOptions( - generateArray = true, - generateStruct = true, - generateMap = true, - primitiveTypes = supportedPrimitiveTypes) - val dataGenOptions = - DataGenOptions(generateNegativeZero = false, generateNaN = false, generateInfinity = false) - - // Test with multiple random deeply nested schemas - for (_ <- 1 to 3) { - val schema = FuzzDataGenerator.generateNestedSchema( - random, - numCols = 5, - minDepth = 1, - maxDepth = 3, - options = schemaGenOptions) - - val df = FuzzDataGenerator.generateDataFrame(random, spark, schema, 100, dataGenOptions) - - withParquetDataFrame(df) { parquetDf => - assertNativeC2RPresent(parquetDf) - checkSparkAnswer(parquetDf) - } - } - } - - // Regression test for https://github.com/apache/datafusion-comet/issues/3308 - // Native columnar-to-row returns UnsafeRow pointing into a Rust-owned buffer that is - // cleared/reused on each convert() call. This test directly exercises the converter: - // it converts multiple batches and holds row references from earlier batches, then - // verifies they still contain correct data. Without a fix (e.g., copying rows), - // rows from earlier batches will contain corrupted data from buffer reuse. - test("rows from earlier batches are not corrupted by subsequent convert() calls") { - import org.apache.spark.sql.catalyst.InternalRow - import org.apache.spark.sql.comet.execution.arrow.CometArrowConverters - import org.apache.spark.unsafe.types.UTF8String - - import scala.collection.mutable.ArrayBuffer - - val schema = new StructType().add("id", IntegerType).add("str", StringType) - - // Create multiple small batches using CometArrowConverters - val numBatches = 10 - val rowsPerBatch = 5 - val totalRows = numBatches * rowsPerBatch - - val rows = (0 until totalRows).map { i => - InternalRow(i, UTF8String.fromString(s"value_$i")) - } - - // Each emitted batch needs independent Arrow buffers so the test can hold rows from - // earlier batches while later batches are consumed. - val allocator = - org.apache.comet.CometArrowAllocator.newChildAllocator("c2r-test", 0, Long.MaxValue) - val batchIter = CometArrowConverters - .rowToArrowBatchIter(rows.iterator, schema, rowsPerBatch, "UTC", allocator) - - val converter = new NativeColumnarToRowConverter(schema, rowsPerBatch) - try { - // Collect all rows from all batches into a single array - // The converter returns rows that should be independent copies - val allRows = new ArrayBuffer[InternalRow]() - var batchCount = 0 - - while (batchIter.hasNext) { - val batch = batchIter.next() - batchCount += 1 - // Convert this batch and collect all rows - val rowIter = converter.convert(batch) - while (rowIter.hasNext) { - allRows += rowIter.next() - } - batch.close() - } - - assert(batchCount == numBatches, s"Expected $numBatches batches, got $batchCount") - assert(allRows.length == totalRows, s"Expected $totalRows rows, got ${allRows.length}") - - // Verify that reading through held references produces all expected - // distinct values. If rows weren't copied, all entries would point - // to the same reused UnsafeRow object with stale data. - val distinctIds = allRows.map(_.getInt(0)).toSet - assert( - distinctIds.size == totalRows, - s"UnsafeRow reuse bug: expected $totalRows distinct row IDs but got " + - s"${distinctIds.size} (values: ${distinctIds.toSeq.sorted.mkString(", ")}). " + - "This means rows were not copied and all references point to the same " + - "reused UnsafeRow object.") - } finally { - converter.close() - allocator.close() - } - } - - /** - * Helper to create a parquet table from a DataFrame and run a function with it. - */ - private def withParquetDataFrame(df: org.apache.spark.sql.DataFrame)( - f: org.apache.spark.sql.DataFrame => Unit): Unit = { - withTempPath { path => - df.write.parquet(path.getAbsolutePath) - spark.read.parquet(path.getAbsolutePath).createOrReplaceTempView("test_table") - f(sql("SELECT * FROM test_table")) - } - } -} diff --git a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala index 29a19307768..62e23e1a72e 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala @@ -42,7 +42,7 @@ import org.apache.spark.internal.io.cloud.{BindingParquetOutputCommitter, PathOu import org.apache.spark.memory.MemoryMode import org.apache.spark.sql.{AnalysisException, DataFrame, Row, SaveMode} import org.apache.spark.sql.catalyst.{CatalystTypeConverters, InternalRow, TableIdentifier} -import org.apache.spark.sql.comet.{CometBatchScanExec, CometNativeColumnarToRowExec, CometNativeScanExec, CometNativeWriteExec, CometScanExec, CometSparkToColumnarExec, CometWriteFilesExec} +import org.apache.spark.sql.comet.{CometBatchScanExec, CometColumnarToRowExec, CometNativeScanExec, CometNativeWriteExec, CometScanExec, CometSparkToColumnarExec, CometWriteFilesExec} import org.apache.spark.sql.execution.{ColumnarToRowTransition, FileSourceScanExec, SparkPlan, SQLExecution} import org.apache.spark.sql.execution.command.DataWritingCommandExec import org.apache.spark.sql.execution.datasources.{BasicWriteTaskStats, BasicWriteTaskStatsTracker, FileFormatWriter, SQLHadoopMapReduceCommitProtocol, WriteFilesExec, WriteJobStatsTracker, WriteTaskStats, WriteTaskStatsTracker} @@ -167,8 +167,8 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { case write: CometWriteFilesExec => write } .getOrElse(fail(s"expected a native parquet writer:\n$nativePlan")) - val stageWithTransition = nativeWrite.withNewChildren(Seq( - CometSparkToColumnarExec(CometNativeColumnarToRowExec(nativeWrite.children.head)))) + val stageWithTransition = nativeWrite.withNewChildren( + Seq(CometSparkToColumnarExec(CometColumnarToRowExec(nativeWrite.children.head)))) assert( stageWithTransition.collect { case _: ColumnarToRowTransition => true }.nonEmpty, s"test requires a C2R transition:\n$stageWithTransition") diff --git a/spark/src/test/scala/org/apache/comet/rules/EliminateRedundantTransitionsSuite.scala b/spark/src/test/scala/org/apache/comet/rules/EliminateRedundantTransitionsSuite.scala index 33293c3d342..76f16945578 100644 --- a/spark/src/test/scala/org/apache/comet/rules/EliminateRedundantTransitionsSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/EliminateRedundantTransitionsSuite.scala @@ -20,7 +20,7 @@ package org.apache.comet.rules import org.apache.spark.sql.CometTestBase -import org.apache.spark.sql.comet.{CometColumnarToRowExec, CometNativeColumnarToRowExec, CometPlan} +import org.apache.spark.sql.comet.{CometColumnarToRowExec, CometPlan} import org.apache.spark.sql.execution.{ColumnarToRowExec, SparkPlan, UnionExec} import org.apache.comet.CometConf @@ -37,7 +37,6 @@ class EliminateRedundantTransitionsSuite extends CometTestBase { private def isCometColumnarToRow(plan: SparkPlan): Boolean = plan match { case _: CometColumnarToRowExec => true - case _: CometNativeColumnarToRowExec => true case _ => false } diff --git a/spark/src/test/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStagesSuite.scala b/spark/src/test/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStagesSuite.scala index 6eb7c65a43c..cb46516e5cc 100644 --- a/spark/src/test/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStagesSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStagesSuite.scala @@ -355,7 +355,7 @@ class RevertNativeForTransitionHeavyStagesSuite extends CometTestBase { val leaf = sparkPlan.collectFirst { case node: LeafExecNode => node }.getOrElse { fail(s"expected a leaf node in test plan:\n$sparkPlan") } - val stacked = CometSparkToColumnarExec(CometNativeColumnarToRowExec(cometFilter(leaf))) + val stacked = CometSparkToColumnarExec(CometColumnarToRowExec(cometFilter(leaf))) val write = cometIcebergWrite(stacked) val reverted = RevertNativeForTransitionHeavyStages(spark).revertToSpark(write) @@ -365,7 +365,7 @@ class RevertNativeForTransitionHeavyStagesSuite extends CometTestBase { reverted.collect { case _: FilterExec => true }.size == 1, s"expected exactly one Spark FilterExec:\n$reverted") assert( - reverted.collect { case _: CometNativeColumnarToRowExec | _: CometSparkToColumnarExec => + reverted.collect { case _: CometColumnarToRowExec | _: CometSparkToColumnarExec => true }.isEmpty, s"stacked transitions should be fully unwrapped:\n$reverted") @@ -440,7 +440,7 @@ class RevertNativeForTransitionHeavyStagesSuite extends CometTestBase { fail(s"expected a leaf node in test plan:\n$sparkPlan") } val aliasing = AliasingFallbackCometExec(leaf, leaf) - val stagePlan = CometNativeColumnarToRowExec(aliasing) + val stagePlan = CometColumnarToRowExec(aliasing) val rule = RevertNativeForTransitionHeavyStages(spark) assert(rule.countTransitions(stagePlan) == 1) @@ -462,7 +462,7 @@ class RevertNativeForTransitionHeavyStagesSuite extends CometTestBase { } val cometChild = cometFilter(leaf) val aliasing = AliasingFallbackCometExec(cometChild, cometChild) - val stagePlan = CometNativeColumnarToRowExec(aliasing) + val stagePlan = CometColumnarToRowExec(aliasing) val rule = RevertNativeForTransitionHeavyStages(spark) assert(rule.countTransitions(stagePlan) == 1) @@ -756,12 +756,12 @@ class RevertNativeForTransitionHeavyStagesSuite extends CometTestBase { .getOrElse(fail(s"test requires a native shuffle:\n$cometPlan")) // The transition that must survive lives in the map stage, under the exchange. val preserved = - if (shuffle.child.supportsColumnar) CometNativeColumnarToRowExec(shuffle.child) + if (shuffle.child.supportsColumnar) CometColumnarToRowExec(shuffle.child) else CometSparkToColumnarExec(shuffle.child) val shuffleWithTransition = shuffle.withNewChildren(Seq(preserved)) // Stacked transitions sit directly on the shuffle, the shape #6152 strips through. val onExchange = - CometSparkToColumnarExec(CometNativeColumnarToRowExec(shuffleWithTransition)) + CometSparkToColumnarExec(CometColumnarToRowExec(shuffleWithTransition)) val write = cometIcebergWrite(onExchange) val reverted = RevertNativeForTransitionHeavyStages(spark).revertToSpark(write) @@ -902,8 +902,8 @@ class RevertNativeForTransitionHeavyStagesSuite extends CometTestBase { assert(scan.supportsColumnar, s"test requires a vectorized Spark scan:\n$scan") val rule = RevertNativeForTransitionHeavyStages(spark, wholePlan) - val overFilter = BothFormatsExec( - CometSparkToColumnarExec(CometNativeColumnarToRowExec(cometFilter(scan)))) + val overFilter = + BothFormatsExec(CometSparkToColumnarExec(CometColumnarToRowExec(cometFilter(scan)))) assert(overFilter.supportsColumnar && overFilter.supportsRowBased) assert(rule.countTransitions(overFilter) == 1) val filterSparkPlan = rule.revertToSpark(overFilter) @@ -956,7 +956,7 @@ class RevertNativeForTransitionHeavyStagesSuite extends CometTestBase { val rule = RevertNativeForTransitionHeavyStages(spark) val columnarRoot = - CometSparkToColumnarExec(CometNativeColumnarToRowExec(cometFilter(leaf))) + CometSparkToColumnarExec(CometColumnarToRowExec(cometFilter(leaf))) assert(columnarRoot.supportsColumnar && !columnarRoot.supportsRowBased) assert(rule.countTransitions(columnarRoot) == 1) val columnar = rule(columnarRoot) @@ -995,8 +995,8 @@ class RevertNativeForTransitionHeavyStagesSuite extends CometTestBase { .getOrElse(fail("test requires a native shuffle")) assert(exchange.child.isInstanceOf[CometNativeScanExec]) val input = if (columnarRoot) exchange.child else cometFilter(exchange.child) - val stage = exchange.withNewChildren( - Seq(CometSparkToColumnarExec(CometNativeColumnarToRowExec(input)))) + val stage = + exchange.withNewChildren(Seq(CometSparkToColumnarExec(CometColumnarToRowExec(input)))) val rule = RevertNativeForTransitionHeavyStages(spark) assert(rule.countTransitions(stage.children.head) == 1) @@ -1191,7 +1191,7 @@ class RevertNativeForTransitionHeavyStagesSuite extends CometTestBase { } .getOrElse(fail(s"expected a native partial aggregate:\n$nativePlan")) val producerWithTransition = partial.withNewChildren( - Seq(CometSparkToColumnarExec(CometNativeColumnarToRowExec(partial.child)))) + Seq(CometSparkToColumnarExec(CometColumnarToRowExec(partial.child)))) val reverter = RevertNativeForTransitionHeavyStages(spark) assert(reverter.countTransitions(producerWithTransition) == 1) diff --git a/spark/src/test/scala/org/apache/spark/CometPluginsSuite.scala b/spark/src/test/scala/org/apache/spark/CometPluginsSuite.scala index f85045a57ec..ae091d1b508 100644 --- a/spark/src/test/scala/org/apache/spark/CometPluginsSuite.scala +++ b/spark/src/test/scala/org/apache/spark/CometPluginsSuite.scala @@ -300,6 +300,30 @@ class CometPluginsMemoryPoolFractionWarningSuite extends CometTestBase { } } +class CometPluginsNativeColumnarToRowWarningSuite extends CometTestBase { + + private val warning = "spark.comet.exec.columnarToRow.native.enabled=true is deprecated" + + private def warningsFor(conf: SparkConf): Seq[String] = { + // Logging derives the logger name by stripping the object's trailing '$' + val logger = CometDriverPlugin.getClass.getName.stripSuffix("$") + val appender = new LogAppender("native columnar to row warning") + withLogAppender(appender, Seq(logger), Some(Level.WARN)) { + CometDriverPlugin.warnIfNativeColumnarToRowSet(conf) + } + appender.loggingEvents.map(_.getMessage.getFormattedMessage).toSeq + } + + test("warns when native columnar to row conversion is set") { + val conf = new SparkConf().set("spark.comet.exec.columnarToRow.native.enabled", "true") + assert(warningsFor(conf).exists(_.contains(warning))) + } + + test("does not warn when native columnar to row conversion is unset") { + assert(!warningsFor(new SparkConf()).exists(_.contains("columnarToRow.native"))) + } +} + class CometPluginsUnifiedModeSuite extends CometTestBase { override protected def sparkConf: SparkConf = { val conf = new SparkConf() diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometC2RIsolatedBench.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometC2RIsolatedBench.scala index 07d9f475754..8ac441c9db4 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometC2RIsolatedBench.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometC2RIsolatedBench.scala @@ -27,7 +27,7 @@ import org.apache.spark.sql.types._ import org.apache.spark.sql.vectorized.ColumnarBatch import org.apache.spark.unsafe.types.UTF8String -import org.apache.comet.{DirectColumnarToRowConverter, NativeColumnarToRowConverter} +import org.apache.comet.DirectColumnarToRowConverter /** * Isolated columnar-to-row microbenchmark that excludes the parquet scan entirely. Batches are @@ -35,10 +35,9 @@ import org.apache.comet.{DirectColumnarToRowConverter, NativeColumnarToRowConver * conversion itself is what gets measured. * * The schema scenarios exist because the implementations have very different per-value costs by - * data type. The JVM path allocates a Decimal object per non-null decimal value (and a - * BigInteger/BigDecimal chain for precision > 18), while the native path writes the unscaled - * value directly. The native path also has a column-wise fast path that only engages when every - * column in the schema is fixed-width. + * data type. The `rowIterator` path allocates a Decimal object per non-null decimal value (and a + * BigInteger/BigDecimal chain for precision > 18), while the direct converter writes the unscaled + * value directly. */ object CometC2RIsolatedBench { @@ -192,29 +191,10 @@ object CometC2RIsolatedBench { if (sink == Long.MinValue) println(sink) } - benchmark.addCase("Native converter") { _ => - val converter = new NativeColumnarToRowConverter(scenario.schema, batchSize) - var sink = 0L - try { - var b = 0 - while (b < batches.length) { - val it = converter.convert(batches(b)) - while (it.hasNext) { - sink += scenario.consume(it.next()) - } - b += 1 - } - } finally { - converter.close() - } - if (sink == Long.MinValue) println(sink) - } - benchmark.run() // GC pressure comparison: the JVM path allocates per object-typed value (Decimal, - // UTF8String) but reuses its output row buffer, while the native path allocates a - // byte[] + UnsafeRow per row for the defensive copy in NativeRowIterator. + // UTF8String) but reuses its output row buffer. val jvmAlloc = measureAllocatedBytes { val proj = UnsafeProjection.create(scenario.schema.fields.map(_.dataType)) var sink = 0L @@ -240,25 +220,9 @@ object CometC2RIsolatedBench { } if (sink == Long.MinValue) println(sink) } - val nativeAlloc = measureAllocatedBytes { - val converter = new NativeColumnarToRowConverter(scenario.schema, batchSize) - var sink = 0L - try { - for (batch <- batches) { - val it = converter.convert(batch) - while (it.hasNext) { - sink += scenario.consume(it.next()) - } - } - } finally { - converter.close() - } - if (sink == Long.MinValue) println(sink) - } println( f"JVM heap allocation per row: JVM path ${jvmAlloc.toDouble / totalRows}%.1f bytes, " + - f"direct path ${directAlloc.toDouble / totalRows}%.1f bytes, " + - f"native path ${nativeAlloc.toDouble / totalRows}%.1f bytes%n") + f"direct path ${directAlloc.toDouble / totalRows}%.1f bytes%n") batches.foreach(_.close()) } diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometColumnarToRowBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometColumnarToRowBenchmark.scala index 1a05bdbcd73..a861cbc83a2 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometColumnarToRowBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometColumnarToRowBenchmark.scala @@ -24,7 +24,7 @@ import scala.collection.mutable.ArrayBuffer import org.apache.spark.SparkConf import org.apache.spark.benchmark.Benchmark import org.apache.spark.sql.SparkSession -import org.apache.spark.sql.comet.{CometColumnarToRowExec, CometNativeColumnarToRowExec} +import org.apache.spark.sql.comet.CometColumnarToRowExec import org.apache.spark.sql.execution.{ColumnarToRowExec, QueryExecution, SparkPlan} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.util.QueryExecutionListener @@ -35,7 +35,6 @@ import org.apache.comet.{CometConf, CometSparkSessionExtensions} * Benchmark to compare Columnar to Row conversion performance: * - Spark's default ColumnarToRowExec * - Comet's JVM-based CometColumnarToRowExec - * - Comet's Native CometNativeColumnarToRowExec * * To run this benchmark: * {{{ @@ -109,7 +108,6 @@ object CometColumnarToRowBenchmark extends CometBenchmarkBase { collect(plan) { case c: ColumnarToRowExec => c case c: CometColumnarToRowExec => c - case c: CometNativeColumnarToRowExec => c } } require( @@ -134,15 +132,7 @@ object CometColumnarToRowBenchmark extends CometBenchmarkBase { "Comet JVM (CometColumnarToRowExec)", classOf[CometColumnarToRowExec], CometConf.COMET_ENABLED.key -> "true", - CometConf.COMET_EXEC_ENABLED.key -> "true", - CometConf.COMET_NATIVE_COLUMNAR_TO_ROW_ENABLED.key -> "false") - - addCase( - "Comet Native (CometNativeColumnarToRowExec)", - classOf[CometNativeColumnarToRowExec], - CometConf.COMET_ENABLED.key -> "true", - CometConf.COMET_EXEC_ENABLED.key -> "true", - CometConf.COMET_NATIVE_COLUMNAR_TO_ROW_ENABLED.key -> "true") + CometConf.COMET_EXEC_ENABLED.key -> "true") } /** @@ -177,8 +167,7 @@ object CometColumnarToRowBenchmark extends CometBenchmarkBase { } /** - * Benchmark columnar to row conversion for fixed-width types ONLY (no strings). This tests the - * fast path in native C2R that pre-allocates buffers. + * Benchmark columnar to row conversion for fixed-width types ONLY (no strings). */ def fixedWidthOnlyBenchmark(values: Int): Unit = { val benchmark = @@ -447,8 +436,8 @@ object CometColumnarToRowBenchmark extends CometBenchmarkBase { /** * Benchmark where a JVM operator (Scala UDF projection + aggregate) consumes the rows. Unlike - * the noop() sink, this defeats escape analysis and exercises the WholeStageCodegen fusion - * asymmetry between the CodegenSupport JVM C2R and the iterator-based native C2R. + * the noop() sink, this defeats escape analysis and exercises the WholeStageCodegen fusion of + * the CodegenSupport JVM C2R with its consumer. */ def jvmConsumerBenchmark(values: Int): Unit = { val benchmark = diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometStringFfiImportBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometStringFfiImportBenchmark.scala index 49c3cfde804..6d008072905 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometStringFfiImportBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometStringFfiImportBenchmark.scala @@ -24,7 +24,7 @@ import scala.collection.mutable.ArrayBuffer import org.apache.spark.SparkConf import org.apache.spark.benchmark.Benchmark import org.apache.spark.sql.SparkSession -import org.apache.spark.sql.comet.{CometColumnarToRowExec, CometNativeColumnarToRowExec, CometNativeScanExec, CometSparkToColumnarExec} +import org.apache.spark.sql.comet.{CometNativeScanExec, CometSparkToColumnarExec} import org.apache.spark.sql.execution.{QueryExecution, SparkPlan} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.util.QueryExecutionListener @@ -37,15 +37,13 @@ import org.apache.comet.{CometConf, CometSparkSessionExtensions} * * arrow-rs builds FFI-imported string arrays without checking UTF-8, so Comet validates every * imported `Utf8` column and decodes it if it is invalid. The data here is always valid, so the - * cases measure the validation pass that every query crossing an import site pays. Two import - * sites carry most string data: - * - `ScanExec`, which receives batches from JVM operators. Here it is fed by Spark's Parquet - * reader converted to Arrow. - * - native columnar-to-row conversion, which imports native batches back from the JVM. + * cases measure the validation pass that every query crossing an import site pays. The import + * site here is `ScanExec`, which receives batches from JVM operators. It is fed by Spark's + * Parquet reader converted to Arrow. * - * Each case that crosses an import site is paired with a control that does the same work without - * one: the native Parquet scan for `ScanExec`, and the JVM conversion for columnar-to-row. When - * comparing a build with and without validation, the controls should not move. + * The imported case is paired with a control that does the same work without an import: the + * native Parquet scan. When comparing a build with and without validation, the control should not + * move. * * To run this benchmark: * {{{ @@ -80,7 +78,7 @@ object CometStringFfiImportBenchmark extends CometBenchmarkBase { sparkSession.conf.set(CometConf.COMET_ENABLED.key, "false") sparkSession.conf.set(CometConf.COMET_EXEC_ENABLED.key, "false") // Every value is distinct, so dictionary pages would fall back anyway. Disabling them keeps - // both import sites on plain `Utf8` arrays rather than depending on where the fallback lands. + // the import site on plain `Utf8` arrays rather than depending on where the fallback lands. sparkSession.conf.set("parquet.enable.dictionary", "false") sparkSession @@ -165,36 +163,6 @@ object CometStringFfiImportBenchmark extends CometBenchmarkBase { benchmark.run() } - /** Wide string columns returned to Spark as rows. */ - def columnarToRowImportBenchmark(values: Int, kind: String): Unit = { - val benchmark = - new Benchmark(s"Columnar to row of wide $kind strings", values.toLong, output = output) - val query = "SELECT s FROM parquetV1Table" - val nativeScan = cometConf :+ (CometConf.COMET_NATIVE_SCAN_ENABLED.key -> "true") - - benchmark.addCase("Spark") { _ => - spark.sql(query).noop() - } - - addVerifiedCase( - benchmark, - "Comet, JVM columnar to row (no import)", - query, - nativeScan :+ (CometConf.COMET_NATIVE_COLUMNAR_TO_ROW_ENABLED.key -> "false"), - required = classOf[CometColumnarToRowExec], - excluded = classOf[CometNativeColumnarToRowExec]) - - addVerifiedCase( - benchmark, - "Comet, native columnar to row (imported)", - query, - nativeScan :+ (CometConf.COMET_NATIVE_COLUMNAR_TO_ROW_ENABLED.key -> "true"), - required = classOf[CometNativeColumnarToRowExec], - excluded = classOf[CometColumnarToRowExec]) - - benchmark.run() - } - override def runCometBenchmark(mainArgs: Array[String]): Unit = { val values = 1024 * 1024 // 1M rows @@ -215,10 +183,6 @@ object CometStringFfiImportBenchmark extends CometBenchmarkBase { runBenchmark(s"FFI String Import - Scan ($kind)") { scanImportBenchmark(values, kind) } - - runBenchmark(s"FFI String Import - Columnar to Row ($kind)") { - columnarToRowImportBenchmark(values, kind) - } } } } diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/CometPlanChecker.scala b/spark/src/test/scala/org/apache/spark/sql/comet/CometPlanChecker.scala index c8c4baff4ab..7caac71351d 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/CometPlanChecker.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/CometPlanChecker.scala @@ -46,7 +46,7 @@ trait CometPlanChecker { case _: CometNativeScanExec | _: CometScanExec | _: CometBatchScanExec | _: CometIcebergNativeScanExec => case _: CometSinkPlaceHolder | _: CometScanWrapper => - case _: CometColumnarToRowExec | _: CometNativeColumnarToRowExec => + case _: CometColumnarToRowExec => case _: CometSparkToColumnarExec => case _: CometExec | _: CometShuffleExchangeExec => case _: CometBroadcastExchangeExec => From ca89966b1be279363b69b6481a71126fdbcc3927 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 5 Nov 2026 06:30:53 -0700 Subject: [PATCH 2/2] chore: share the deprecated config warning and drop leftover native C2R references --- .github/workflows/pr_build_linux.yml | 2 +- .github/workflows/pr_build_macos.yml | 2 +- .../latest/understanding-comet-plans.md | 13 +++--- .../comet/DirectColumnarToRowConverter.java | 3 +- .../scala/org/apache/comet/CometConf.scala | 2 +- .../rules/EliminateRedundantTransitions.scala | 37 +++++++-------- .../main/scala/org/apache/spark/Plugins.scala | 45 +++++++++---------- .../EliminateRedundantTransitionsSuite.scala | 7 +-- .../org/apache/spark/CometPluginsSuite.scala | 44 +++++------------- .../sql/benchmark/CometC2RIsolatedBench.scala | 9 ++-- .../CometColumnarToRowBenchmark.scala | 4 +- 11 files changed, 66 insertions(+), 102 deletions(-) diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 06d05c71799..dba17288544 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -617,7 +617,7 @@ jobs: org.apache.spark.CometExecIteratorLifecycleSuite org.apache.spark.CometPluginsDefaultSuite org.apache.spark.CometPluginsMemoryOverheadWarningSuite - org.apache.spark.CometPluginsMemoryPoolFractionWarningSuite + org.apache.spark.CometPluginsDeprecatedConfigWarningSuite org.apache.spark.CometPluginsUnifiedModeSuite org.apache.spark.CometPluginsExtensionOnlySuite org.apache.spark.CometPluginsSparkShuffleManagerSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index 3ca3643191f..8bd8a0aca99 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -252,7 +252,7 @@ jobs: org.apache.spark.CometExecIteratorLifecycleSuite org.apache.spark.CometPluginsDefaultSuite org.apache.spark.CometPluginsMemoryOverheadWarningSuite - org.apache.spark.CometPluginsMemoryPoolFractionWarningSuite + org.apache.spark.CometPluginsDeprecatedConfigWarningSuite org.apache.spark.CometPluginsUnifiedModeSuite org.apache.spark.CometPluginsExtensionOnlySuite org.apache.spark.CometPluginsSparkShuffleManagerSuite diff --git a/docs/source/user-guide/latest/understanding-comet-plans.md b/docs/source/user-guide/latest/understanding-comet-plans.md index e406ed57d99..06a7b1df77b 100644 --- a/docs/source/user-guide/latest/understanding-comet-plans.md +++ b/docs/source/user-guide/latest/understanding-comet-plans.md @@ -52,9 +52,9 @@ the same plan is shown in the Spark SQL UI. When reading a plan, look for: - **Node prefix.** `Comet*` nodes are accelerated by Comet. Anything without the prefix is unmodified Spark. -- **Transitions.** `CometColumnarToRow` and `CometSparkRowToColumnar` mark boundaries between columnar Comet execution - and row-based Spark execution. Frequent transitions usually indicate - fallback inside the plan. +- **Transitions.** `CometColumnarToRow` and `CometSparkRowToColumnar` mark + boundaries between columnar Comet execution and row-based Spark execution. + Frequent transitions usually indicate fallback inside the plan. - **Exchange type.** `CometExchange` is the native shuffle path, `CometColumnarExchange` is the JVM columnar shuffle path, and a plain `Exchange` means Spark shuffle. See [Shuffle Operators](#shuffle-operators) @@ -182,8 +182,9 @@ the config has no effect there. Not every node in the plan is an eligible operator. The following are excluded from both operator counts: -- Transition nodes (`CometColumnarToRow`, `CometSparkRowToColumnar`, `CometSparkColumnarToColumnar`, `ColumnarToRow`, - `RowToColumnar`), which are reported separately as the transition count. +- Transition nodes (`CometColumnarToRow`, `CometSparkRowToColumnar`, + `CometSparkColumnarToColumnar`, `ColumnarToRow`, `RowToColumnar`), which are + reported separately as the transition count. `CometSparkRowToColumnar` and `CometSparkColumnarToColumnar` are the two names a single operator renders under, depending on whether its child already produces columnar data, and both are excluded. @@ -387,8 +388,6 @@ shuffle and choose between the implementations. ### Columnar/Row Transitions Comet inserts these nodes wherever data has to cross the columnar/row boundary. -Multiple implementations exist because the optimal strategy depends on what -produced the columnar data. | Node | Direction | Notes | | ------------------------------ | ------------------- | -------------------------------------------------------------------------------------------------- | diff --git a/spark/src/main/java/org/apache/comet/DirectColumnarToRowConverter.java b/spark/src/main/java/org/apache/comet/DirectColumnarToRowConverter.java index 9a797f0dbec..aae7c3328bd 100644 --- a/spark/src/main/java/org/apache/comet/DirectColumnarToRowConverter.java +++ b/spark/src/main/java/org/apache/comet/DirectColumnarToRowConverter.java @@ -44,8 +44,7 @@ *

When every column is fixed-width (no strings, no decimals above precision 18), the whole batch * is converted column-at-a-time in {@link #setBatch(ColumnarBatch)} into a single buffer with a * constant row stride: one monomorphic loop per column with the null check hoisted out for - * null-free columns, mirroring the native converter's fixed-width fast path. In that mode {@link - * #convertRow(int)} only repoints the reused row. + * null-free columns. In that mode {@link #convertRow(int)} only repoints the reused row. * *

The produced rows are byte-identical to {@code UnsafeProjection} output, which matters because * {@link UnsafeRow} equality and hashing are byte-wise. The returned row and its backing buffer are diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 48ceee6b5dd..6e70026175b 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -522,7 +522,7 @@ object CometConf extends ShimCometConf { conf(s"$COMET_EXEC_CONFIG_PREFIX.columnarToRow.direct.enabled") .category(CATEGORY_EXEC) .doc( - "Experimental: Whether to use the direct columnar to row converter in the JVM " + + "Experimental: Whether to use the direct columnar to row converter in Comet's " + "columnar to row operator's non-codegen paths (including broadcast relation " + "builds). The direct converter writes values straight from Arrow buffers into " + "UnsafeRow format without per-value object allocation, which is significantly " + diff --git a/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala b/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala index 62ed06fc85a..f9b8c2cf2f2 100644 --- a/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala +++ b/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala @@ -112,14 +112,23 @@ case class EliminateRedundantTransitions(session: SparkSession) case ColumnarToRowExec(range: CometRangeExec) => range.originalPlan case c @ ColumnarToRowExec(child) if hasCometNativeChild(child, containsCometPlanMemo) => - val op = createColumnarToRowExec(child) - if (c.logicalLink.isEmpty) { - op.unsetTagValue(SparkPlan.LOGICAL_PLAN_TAG) - op.unsetTagValue(SparkPlan.LOGICAL_PLAN_INHERITED_TAG) + if (containsVariantType(child.schema)) { + // TODO: Remove this fallback once Comet columnar-to-row conversion supports Variant + // getters and Spark's Variant UnsafeRow encoding. + // https://github.com/apache/datafusion-comet/issues/5436 + withFallbackReason( + c, + "Comet columnar-to-row conversion does not support type VariantType") } else { - c.logicalLink.foreach(op.setLogicalLink) + val op = CometColumnarToRowExec(child) + if (c.logicalLink.isEmpty) { + op.unsetTagValue(SparkPlan.LOGICAL_PLAN_TAG) + op.unsetTagValue(SparkPlan.LOGICAL_PLAN_INHERITED_TAG) + } else { + c.logicalLink.foreach(op.setLogicalLink) + } + op } - op case CometColumnarToRowExec(sparkToColumnar: CometSparkToColumnarExec) => sparkToColumnar.child case CometColumnarToRowExec(range: CometRangeExec) => @@ -309,20 +318,4 @@ case class EliminateRedundantTransitions(session: SparkSession) } } } - - /** - * Creates the columnar to row transition operator: `CometColumnarToRowExec`, or Spark's - * `ColumnarToRowExec` for a schema that contains Variant. - */ - private def createColumnarToRowExec(child: SparkPlan): SparkPlan = { - // TODO: Remove this fallback once Comet columnar-to-row conversion supports Variant getters - // and Spark's Variant UnsafeRow encoding. - // https://github.com/apache/datafusion-comet/issues/5436 - if (containsVariantType(child.schema)) { - return withFallbackReason( - ColumnarToRowExec(child), - "Comet columnar-to-row conversion does not support type VariantType") - } - CometColumnarToRowExec(child) - } } diff --git a/spark/src/main/scala/org/apache/spark/Plugins.scala b/spark/src/main/scala/org/apache/spark/Plugins.scala index c4c40dd7d51..1a93c7cc6d7 100644 --- a/spark/src/main/scala/org/apache/spark/Plugins.scala +++ b/spark/src/main/scala/org/apache/spark/Plugins.scala @@ -123,8 +123,7 @@ class CometDriverPlugin private[spark] (clock: Clock) extends DriverPlugin with CometDriverPlugin.registerIcebergWriteReport(sc.conf) CometDriverPlugin.warnIfExecutorMemoryOverheadUnset(sc.getConf) - CometDriverPlugin.warnIfMemoryPoolFractionSet(sc.getConf) - CometDriverPlugin.warnIfNativeColumnarToRowSet(sc.getConf) + CometDriverPlugin.warnIfDeprecatedConfigsSet(sc.getConf) extraConfs } @@ -321,32 +320,32 @@ object CometDriverPlugin extends Logging { .exists(factor => !Try(factor.toDouble).toOption.contains(submitDefault)) } + // Deprecated configs that the driver warns about when they are set, with the reason for each. + // // spark.comet.exec.memoryPool.fraction was documented as holding back part of the off-heap pool // for the native memory that Comet does not reserve. It cannot: Spark hands out the whole pool // to the tasks that ask for it, and the fraction only caps each task's total under // fair_unified. Users who set it for that purpose need to size the memory overhead instead. - private[apache] def warnIfMemoryPoolFractionSet(conf: SparkConf): Unit = { - val key = CometConf.COMET_OFFHEAP_MEMORY_POOL_FRACTION.key - conf.getOption(key).foreach { value => - logWarning( - s"$key=$value is deprecated and will be removed in a future major release. It does " + - "not leave room in spark.memory.offHeap.size for native memory that Comet's memory " + - "pools do not track, because Spark hands out the whole off-heap pool whatever it is " + - s"set to. Size ${EXECUTOR_MEMORY_OVERHEAD.key} for that memory instead. " + - s"${CometConf.TUNING_GUIDE}.") - } - } - - // The Rust columnar-to-row converter that spark.comet.exec.columnarToRow.native.enabled turned - // on has been removed, so the key no longer does anything. - private[apache] def warnIfNativeColumnarToRowSet(conf: SparkConf): Unit = { - val key = CometConf.COMET_NATIVE_COLUMNAR_TO_ROW_ENABLED.key - conf.getOption(key).foreach { value => - logWarning( - s"$key=$value is deprecated and will be removed in a future major release. It has no " + - "effect: Comet always converts Arrow batches to Spark rows on the JVM.") + // + // spark.comet.exec.columnarToRow.native.enabled turned on a Rust columnar-to-row converter that + // has been removed, so the key no longer does anything. + private def deprecatedConfigs: Seq[(ConfigEntry[_], String)] = Seq( + CometConf.COMET_OFFHEAP_MEMORY_POOL_FRACTION -> + ("It does not leave room in spark.memory.offHeap.size for native memory that Comet's " + + "memory pools do not track, because Spark hands out the whole off-heap pool whatever " + + s"it is set to. Size ${EXECUTOR_MEMORY_OVERHEAD.key} for that memory instead. " + + s"${CometConf.TUNING_GUIDE}."), + CometConf.COMET_NATIVE_COLUMNAR_TO_ROW_ENABLED -> + "It has no effect: Comet always converts Arrow batches to Spark rows on the JVM.") + + private[apache] def warnIfDeprecatedConfigsSet(conf: SparkConf): Unit = + deprecatedConfigs.foreach { case (entry, reason) => + conf.getOption(entry.key).foreach { value => + logWarning( + s"${entry.key}=$value is deprecated and will be removed in a future major release. " + + reason) + } } - } // Reads a deprecated alternative too, such as spark.comet.exec.shuffle.enabled, as a session // would. diff --git a/spark/src/test/scala/org/apache/comet/rules/EliminateRedundantTransitionsSuite.scala b/spark/src/test/scala/org/apache/comet/rules/EliminateRedundantTransitionsSuite.scala index 76f16945578..b6625d4fd02 100644 --- a/spark/src/test/scala/org/apache/comet/rules/EliminateRedundantTransitionsSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/EliminateRedundantTransitionsSuite.scala @@ -35,11 +35,6 @@ class EliminateRedundantTransitionsSuite extends CometTestBase { stripAQEPlan(plan) } - private def isCometColumnarToRow(plan: SparkPlan): Boolean = plan match { - case _: CometColumnarToRowExec => true - case _ => false - } - // The rule decides, per `ColumnarToRowExec`, whether the subtree below it is Comet columnar. // That check is memoized across the whole plan, so a Comet branch must not make a sibling // Spark-only branch look Comet, or the other way around. @@ -62,7 +57,7 @@ class EliminateRedundantTransitionsSuite extends CometTestBase { val Seq(cometBranch, sparkBranch) = result.children assert( - isCometColumnarToRow(cometBranch), + cometBranch.isInstanceOf[CometColumnarToRowExec], s"Comet branch should use a Comet columnar-to-row transition:\n${result.treeString}") assert( sparkBranch == sparkTransition, diff --git a/spark/src/test/scala/org/apache/spark/CometPluginsSuite.scala b/spark/src/test/scala/org/apache/spark/CometPluginsSuite.scala index ae091d1b508..e08f9f7a0c4 100644 --- a/spark/src/test/scala/org/apache/spark/CometPluginsSuite.scala +++ b/spark/src/test/scala/org/apache/spark/CometPluginsSuite.scala @@ -276,51 +276,29 @@ class CometPluginsMemoryOverheadWarningSuite extends CometTestBase { } } -class CometPluginsMemoryPoolFractionWarningSuite extends CometTestBase { - - private val warning = "spark.comet.exec.memoryPool.fraction=0.8 is deprecated" +class CometPluginsDeprecatedConfigWarningSuite extends CometTestBase { private def warningsFor(conf: SparkConf): Seq[String] = { // Logging derives the logger name by stripping the object's trailing '$' val logger = CometDriverPlugin.getClass.getName.stripSuffix("$") - val appender = new LogAppender("memory pool fraction warning") + val appender = new LogAppender("deprecated config warning") withLogAppender(appender, Seq(logger), Some(Level.WARN)) { - CometDriverPlugin.warnIfMemoryPoolFractionSet(conf) + CometDriverPlugin.warnIfDeprecatedConfigsSet(conf) } appender.loggingEvents.map(_.getMessage.getFormattedMessage).toSeq } - test("warns when the memory pool fraction is set") { - val conf = new SparkConf().set("spark.comet.exec.memoryPool.fraction", "0.8") - assert(warningsFor(conf).exists(_.contains(warning))) - } - - test("does not warn when the memory pool fraction is unset") { - assert(!warningsFor(new SparkConf()).exists(_.contains("memoryPool.fraction"))) - } -} - -class CometPluginsNativeColumnarToRowWarningSuite extends CometTestBase { - - private val warning = "spark.comet.exec.columnarToRow.native.enabled=true is deprecated" - - private def warningsFor(conf: SparkConf): Seq[String] = { - // Logging derives the logger name by stripping the object's trailing '$' - val logger = CometDriverPlugin.getClass.getName.stripSuffix("$") - val appender = new LogAppender("native columnar to row warning") - withLogAppender(appender, Seq(logger), Some(Level.WARN)) { - CometDriverPlugin.warnIfNativeColumnarToRowSet(conf) + Seq( + "spark.comet.exec.memoryPool.fraction" -> "0.8", + "spark.comet.exec.columnarToRow.native.enabled" -> "true").foreach { case (key, value) => + test(s"warns when $key is set") { + val warnings = warningsFor(new SparkConf().set(key, value)) + assert(warnings.exists(_.contains(s"$key=$value is deprecated")), warnings) } - appender.loggingEvents.map(_.getMessage.getFormattedMessage).toSeq - } - - test("warns when native columnar to row conversion is set") { - val conf = new SparkConf().set("spark.comet.exec.columnarToRow.native.enabled", "true") - assert(warningsFor(conf).exists(_.contains(warning))) } - test("does not warn when native columnar to row conversion is unset") { - assert(!warningsFor(new SparkConf()).exists(_.contains("columnarToRow.native"))) + test("does not warn when no deprecated config is set") { + assert(!warningsFor(new SparkConf()).exists(_.contains("is deprecated"))) } } diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometC2RIsolatedBench.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometC2RIsolatedBench.scala index 8ac441c9db4..064f000f9bc 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometC2RIsolatedBench.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometC2RIsolatedBench.scala @@ -57,12 +57,12 @@ object CometC2RIsolatedBench { .add("c", DoubleType) .add("d", StringType), i => InternalRow(i.toLong, i, i.toDouble, UTF8String.fromString(s"value_$i")), - // All three converters return UnsafeRow, so reading a raw fixed slot avoids allocating a + // Both converters return UnsafeRow, so reading a raw fixed slot avoids allocating a // UTF8String in the benchmark sink. row => row.getLong(0) + row.getLong(3)) // The columns a TPC-H lineitem aggregation reaches C2R with: four compact decimals, a date, - // and two short flag strings. The strings force the native general (row-at-a-time) path. + // and two short flag strings. The strings force the direct converter's row-at-a-time path. private val tpchDecimalsWithStrings = Scenario( "4 x decimal(12,2), date, 2 x string (tpch-like)", new StructType() @@ -85,7 +85,8 @@ object CometC2RIsolatedBench { // A compact decimal is stored directly as an unscaled long in UnsafeRow. row => row.getLong(0) + row.getInt(4)) - // Same decimals without any var-length column, so the native column-wise fast path engages. + // Same decimals without any var-length column, so the direct converter's column-wise fast + // path engages. private val decimalsAllFixedWidth = Scenario( "4 x decimal(12,2), date, long (all fixed-width)", new StructType() @@ -106,7 +107,7 @@ object CometC2RIsolatedBench { row => row.getLong(0) + row.getLong(5)) // Precision > 18: the JVM accessor allocates byte[] -> BigInteger -> java BigDecimal -> - // scala BigDecimal -> Decimal per value; the native side writes 16 bytes. + // scala BigDecimal -> Decimal per value; the direct converter writes 16 bytes. private val highPrecisionDecimals = Scenario( "2 x decimal(38,10), long (high precision)", new StructType() diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometColumnarToRowBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometColumnarToRowBenchmark.scala index a861cbc83a2..3e1893e7478 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometColumnarToRowBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometColumnarToRowBenchmark.scala @@ -129,7 +129,7 @@ object CometColumnarToRowBenchmark extends CometBenchmarkBase { CometConf.COMET_ENABLED.key -> "false") addCase( - "Comet JVM (CometColumnarToRowExec)", + "Comet (CometColumnarToRowExec)", classOf[CometColumnarToRowExec], CometConf.COMET_ENABLED.key -> "true", CometConf.COMET_EXEC_ENABLED.key -> "true") @@ -437,7 +437,7 @@ object CometColumnarToRowBenchmark extends CometBenchmarkBase { /** * Benchmark where a JVM operator (Scala UDF projection + aggregate) consumes the rows. Unlike * the noop() sink, this defeats escape analysis and exercises the WholeStageCodegen fusion of - * the CodegenSupport JVM C2R with its consumer. + * the CodegenSupport Comet C2R with its consumer. */ def jvmConsumerBenchmark(values: Int): Unit = { val benchmark =