diff --git a/docs/source/user-guide/latest/datasources.md b/docs/source/user-guide/latest/datasources.md index f351aff5b6..a04da32381 100644 --- a/docs/source/user-guide/latest/datasources.md +++ b/docs/source/user-guide/latest/datasources.md @@ -73,12 +73,17 @@ operator in `spark.comet.sparkToColumnar.supportedOperatorList` by its Spark cla ### Spark-to-Comet conversion types -Spark-to-Comet conversion supports `ARRAY` and `MAP` with binary -string semantics, both as top-level fields and inside supported structs. Arrays and maps may -be null; array elements and map values may also be null. Map keys must be non-null. -This applies to Spark row and columnar inputs when conversion is enabled for the source. -Other array element types, other map key/value types, nested collections, and non-binary -string collations remain unsupported at this conversion boundary. Source defaults are unchanged. +Spark-to-Comet conversion supports booleans, integers, floats, decimals, strings, binary, dates +and timestamps, and arrays, maps and structs of them, nested to any depth. Arrays and maps may be +null; array elements and map values may also be null. Map keys must be non-null. This applies to +Spark row and columnar inputs when conversion is enabled for the source. A column with any other +type is not converted, and neither is one with a string that has a non-binary collation, or a +struct that has duplicate field names or no fields, anywhere inside it. The operators above such a +column stay on Spark. Source defaults are unchanged. + +Calendar intervals are not converted, at the top level or inside an array, map or struct. Arrow +holds the time part of an interval in nanoseconds, so it cannot hold one of more than about 292 +years, which Spark can ([#5279](https://github.com/apache/datafusion-comet/issues/5279)). This includes row-backed `ExistingRDD` inputs when `spark.comet.convert.rdd.enabled=true`. Spark still produces the RDD rows; conversion lets eligible downstream operators execute in Comet. diff --git a/docs/source/user-guide/latest/in-memory-cache.md b/docs/source/user-guide/latest/in-memory-cache.md index f0c2cd6a31..faac3441cd 100644 --- a/docs/source/user-guide/latest/in-memory-cache.md +++ b/docs/source/user-guide/latest/in-memory-cache.md @@ -169,11 +169,8 @@ several times faster, at six times the memory. The two relations are not comparable to each other — different row counts, and a struct column carries several values per row. -Array and map columns are deliberately absent from the benchmark, not from the format — the cache -stores and projects them, and `CometInMemoryCacheSuite` covers them. They cannot be measured _here_ -because the left column would not exist: it needs Spark's cache scan to bridge into Comet operators, -and `CometSparkToColumnarExec` declines `ArrayType` and `MapType`, so a query projecting one falls -back to Spark row execution above the scan and the two columns stop measuring the same boundary. +Array and map columns are not in the benchmark, though the cache stores and projects them, and +`CometInMemoryCacheSuite` covers them. Read what this compares carefully. Comet execution is on in both columns, so the aggregation runs on Comet either way and only the cache-scan boundary moves: on the left, Spark's diff --git a/docs/source/user-guide/latest/tuning/shuffle.md b/docs/source/user-guide/latest/tuning/shuffle.md index b7ad3d2f09..6a8d22d54a 100644 --- a/docs/source/user-guide/latest/tuning/shuffle.md +++ b/docs/source/user-guide/latest/tuning/shuffle.md @@ -70,10 +70,9 @@ Setting `spark.comet.convert.shuffleInput.enabled=true` makes such a shuffle use converts the child's rows to Arrow, as the `spark.comet.convert.*` settings do for Spark scans, and shuffles the batches natively. With hash partitioning that is usually faster than columnar shuffle, while with range partitioning the two come out about even. This applies only where native shuffle supports the partitioning and Comet can convert -all of the child's columns, which rules out calendar intervals, and arrays and maps other than `array` and -`map`. A shuffle that hashes a decimal wider than 18 digits, or a string or a value computed from a -string, stays on columnar shuffle, so that it partitions rows as Spark does. The converted batches are allocated on -the JVM outside Comet's memory pools. The setting is disabled by default. +all of the child's columns, which rules out calendar intervals. A shuffle that hashes a decimal wider than 18 digits, or +a string or a value computed from a string, stays on columnar shuffle, so that it partitions rows as Spark does. The +converted batches are allocated on the JVM outside Comet's memory pools. The setting is disabled by default. ### Automatic Revert to Spark Shuffle diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometSparkToColumnarExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometSparkToColumnarExec.scala index b2d72fb25c..203f98bd23 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometSparkToColumnarExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometSparkToColumnarExec.scala @@ -33,7 +33,7 @@ import org.apache.spark.sql.comet.execution.arrow.{CometArrowStream, CometNative import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.execution.{RowToColumnarTransition, SparkPlan} import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} -import org.apache.spark.sql.types._ +import org.apache.spark.sql.types.{CalendarIntervalType, DataType} import org.apache.spark.sql.vectorized.ColumnarBatch import org.apache.comet.{CometConf, DataTypeSupport} @@ -138,13 +138,20 @@ object CometSparkToColumnarExec extends CometSink[SparkPlan] with DataTypeSuppor CometScanWrapper(nativeOp, CometSparkToColumnarExec(op)) } + /** + * Declines calendar intervals at any depth. Arrow holds the time part of an interval in + * nanoseconds, so `IntervalMonthDayNanoWriter` overflows on one of more than about 292 years, + * which Spark accepts (#5279). + */ override def isTypeSupported( dt: DataType, name: String, fallbackReasons: ListBuffer[String]): Boolean = dt match { - case ArrayType(StringType, _) => true - case MapType(StringType, StringType, _) => true - case _: ArrayType | _: MapType => false + case CalendarIntervalType => + fallbackReasons += s"Unsupported $name of type $dt: Arrow holds the time part of an " + + "interval in nanoseconds, which overflows past about 292 years " + + "(https://github.com/apache/datafusion-comet/issues/5279)" + false case _ => super.isTypeSupported(dt, name, fallbackReasons) } } 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 478c62f4ed..b876967fe8 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala @@ -50,7 +50,7 @@ import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.internal.SQLConf.SESSION_LOCAL_TIMEZONE import org.apache.spark.sql.types._ -import org.apache.spark.unsafe.types.UTF8String +import org.apache.spark.unsafe.types.{CalendarInterval, UTF8String} import org.apache.comet.{CometConf, CometExecIterator, CometNativeException, ExtendedExplainInfo} import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus, isSpark41Plus, isSpark42Plus} @@ -4195,43 +4195,263 @@ class CometExecSuite extends CometTestBase { }) } - test("SparkToColumnar admits only binary string arrays and maps through the collection gate") { - for (nullable <- Seq(false, true); containsNull <- Seq(false, true)) { - for (dataType <- Seq( - ArrayType(StringType, containsNull), - MapType(StringType, StringType, containsNull))) { - val field = StructField("tags", dataType, nullable) - Seq( - StructType(Seq(field)), - StructType(Seq(StructField("nested", StructType(Seq(field)))))) - .foreach { schema => - assert(CometSparkToColumnarExec.isSchemaSupported(schema, ListBuffer.empty)) - } - } - } - val unsupported = Seq( + test("SparkToColumnar admits arrays and maps of every type it converts") { + val struct = StructType(Seq(StructField("i", IntegerType), StructField("s", StringType))) + val admitted = Seq( + ArrayType(StringType), ArrayType(IntegerType), + ArrayType(DecimalType(38, 10)), + ArrayType(TimestampType), ArrayType(BinaryType), ArrayType(ArrayType(StringType)), - ArrayType(StructType(Seq(StructField("s", StringType)))), - MapType(StringType, IntegerType), - MapType(IntegerType, StringType), - MapType(StringType, ArrayType(StringType)), - MapType(StringType, StructType(Seq(StructField("s", StringType))))) ++ - (if (isSpark40Plus) + ArrayType(struct), + ArrayType(MapType(StringType, IntegerType)), + MapType(StringType, StringType), + MapType(IntegerType, IntegerType), + MapType(DateType, DecimalType(20, 2)), + MapType(StringType, ArrayType(IntegerType)), + MapType(IntegerType, struct), + MapType(struct, StringType), + MapType(ArrayType(IntegerType), IntegerType), + MapType(StringType, MapType(IntegerType, StringType))) + for (nullable <- Seq(false, true); containsNull <- Seq(false, true); dataType <- admitted) { + val collection = dataType match { + case a: ArrayType => a.copy(containsNull = containsNull) + case m: MapType => m.copy(valueContainsNull = containsNull) + } + val field = StructField("c", collection, nullable) + Seq(StructType(Seq(field)), StructType(Seq(StructField("nested", StructType(Seq(field)))))) + .foreach { schema => + val reasons = ListBuffer.empty[String] + assert(CometSparkToColumnarExec.isSchemaSupported(schema, reasons), s"$collection") + assert(reasons.isEmpty, reasons) + } + } + + // The shared rule looks inside a collection and declines these for what it finds there, and + // so does the conversion for a calendar interval. A decline of the whole collection would + // record no reason, which is how the two differ. + val duplicate = StructType(Seq(StructField("a", LongType), StructField("a", LongType))) + val declined = Seq( + ArrayType(duplicate), + MapType(LongType, duplicate), + ArrayType(ArrayType(duplicate)), + ArrayType(NullType), + MapType(IntegerType, ArrayType(NullType)), + ArrayType(DayTimeIntervalType()), + MapType(StringType, YearMonthIntervalType()), + ArrayType(CalendarIntervalType), + MapType(StringType, CalendarIntervalType)) ++ + (if (isSpark40Plus) { Seq( - DataType.fromDDL("ARRAY"), - DataType.fromDDL("MAP"), - DataType.fromDDL("MAP")) - else Seq.empty) - unsupported.foreach { dataType => + "ARRAY", + "ARRAY>", + "ARRAY>", + "MAP", + "MAP", + "MAP>").map(DataType.fromDDL) + } else Seq.empty) + declined.foreach { dataType => val schema = StructType(Seq(StructField("value", dataType))) - assert(!CometSparkToColumnarExec.isSchemaSupported(schema, ListBuffer.empty), dataType) - assert( - !CometSparkToColumnarExec.isSchemaSupported( - StructType(Seq(StructField("nested", schema))), - ListBuffer.empty), - dataType) + Seq(schema, StructType(Seq(StructField("nested", schema)))).foreach { declinedSchema => + val reasons = ListBuffer.empty[String] + assert(!CometSparkToColumnarExec.isSchemaSupported(declinedSchema, reasons), dataType) + assert(reasons.nonEmpty, s"$dataType was declined without a reason") + } + } + + // The shared rule admits a calendar interval, but Arrow holds the time part of one in + // nanoseconds, so the conversion declines it at the top level too (#5279). + val reasons = ListBuffer.empty[String] + val interval = StructType(Seq(StructField("i", CalendarIntervalType))) + assert(!CometSparkToColumnarExec.isSchemaSupported(interval, reasons)) + assert(reasons.exists(_.contains("issues/5279")), reasons) + } + + // A column of each array and map shape the conversion admits, as its type, an expression over + // `id` that builds it, and one that reads an element of it. The column has, in some rows, a + // null collection, an empty one, and elements that are null. + private val sparkToColumnarShapes: Seq[(String, String, String)] = { + val positions = "sequence(0, cast(id % 4 AS INT))" + def collection(ddl: String, empty: String, body: String): String = + "cast(if(id % 7 = 0, null, " + + s"if(id % 7 = 1, cast($empty AS $ddl), cast($body AS $ddl))) AS $ddl)" + def array(ddl: String, element: String): (String, String, String) = + ( + ddl, + collection( + ddl, + "array()", + s"transform($positions, x -> if((x + id) % 5 = 4, null, $element))"), + "if(size(c) > 0, c[0], null)") + def map(ddl: String, key: String, value: String): (String, String, String) = + ( + ddl, + collection( + ddl, + "map()", + s"map_from_arrays(transform($positions, x -> $key), " + + s"transform($positions, x -> if((x + id) % 5 = 4, null, $value)))"), + "if(size(c) > 0, element_at(map_values(c), 1), null)") + val struct = "named_struct('i', cast(id + x AS INT), 's', concat('s', cast(x AS STRING)))" + Seq( + array("ARRAY", "cast(id * 10 + x AS INT)"), + array( + "ARRAY", + "cast(concat(if(x % 2 = 0, '-', ''), cast(id AS STRING), '12345678901234567.0123456789') " + + "AS DECIMAL(38,10))"), + array("ARRAY", "timestamp_micros(1700000000000000 + id * 1000003 + x)"), + array( + "ARRAY", + "cast(concat('b', cast(id AS STRING), cast(x AS STRING)) AS BINARY)"), + array("ARRAY", "cast(id AS DOUBLE) / 3 + x"), + array("ARRAY", "(id + x) % 3 = 0"), + array( + "ARRAY>", + "if((x + id) % 4 = 3, null, if((x + id) % 4 = 2, cast(array() AS ARRAY), " + + "transform(sequence(0, cast(x % 3 AS INT)), y -> cast(id + x + y AS INT))))"), + array("ARRAY>", struct), + array("ARRAY>", "map(concat('k', cast(x AS STRING)), cast(id AS INT))"), + map("MAP", "x", "cast(id * 10 + x AS INT)"), + map( + "MAP", + "date_add(DATE'2024-01-01', x)", + "cast((id * 1000003 + x) / 100 AS DECIMAL(20,2))"), + map( + "MAP", + "cast(concat(cast(x AS STRING), '.25') AS DECIMAL(20,2))", + "concat('v', cast(id AS STRING), '_', cast(x AS STRING))"), + map( + "MAP>", + "concat('k', cast(x AS STRING))", + "transform(sequence(0, x), y -> cast(id + y AS INT))"), + map("MAP>", "x", struct), + map( + "MAP,STRING>", + "named_struct('i', x, 's', concat('k', cast(x AS STRING)))", + "concat('v', cast(id AS STRING))"), + map("MAP,INT>", "array(x, x + 1)", "cast(id + x AS INT)"), + map( + "MAP>", + "concat('k', cast(x AS STRING))", + "map_from_arrays(array(x, x + 1), array('a', concat('b', cast(id AS STRING))))")) + } + + sparkToColumnarShapes.foreach { case (ddl, expression, element) => + test(s"SparkToColumnar converts $ddl columns that native operators read") { + withTempPath { dir => + var schema: StructType = null + var rows: Seq[Row] = Seq.empty + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + val data = spark.range(100).selectExpr("id", s"$expression AS c") + schema = data.schema + rows = data.collect().toSeq + data.repartition(1).write.parquet(dir.toString) + } + for (sourceType <- Seq("rdd", "parquet-row", "parquet-columnar")) { + val vectorized = sourceType == "parquet-columnar" + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.PARQUET_VECTORIZED_READER_ENABLED.key -> vectorized.toString, + SQLConf.PARQUET_VECTORIZED_READER_NESTED_COLUMN_ENABLED.key -> "true", + // Spark batches smaller than Comet's, so an Arrow batch spans several of them. + SQLConf.PARQUET_VECTORIZED_READER_BATCH_SIZE.key -> "30", + CometConf.COMET_NATIVE_SCAN_ENABLED.key -> "false", + CometConf.COMET_BATCH_SIZE.key -> "16", + CometConf.COMET_SHUFFLE_MODE.key -> "native", + CometConf.COMET_CONVERT_FROM_PARQUET_ENABLED.key -> "true", + CometConf.COMET_CONVERT_FROM_RDD_ENABLED.key -> "true") { + def source = + if (sourceType == "rdd") { + spark.createDataFrame(spark.sparkContext.parallelize(rows, 2), schema) + } else { + spark.read.parquet(dir.toString) + } + def query = + source.filter("id > 0").selectExpr("id", "c", "size(c) AS n", s"$element AS e") + val (_, plan) = checkSparkAnswerAndOperator( + query, + includeClasses = Seq( + classOf[CometSparkToColumnarExec], + classOf[CometProjectExec], + classOf[CometFilterExec])) + val conversions = collect(plan) { case c: CometSparkToColumnarExec => c } + assert(conversions.size == 1, plan) + assert(conversions.head.child.supportsColumnar == vectorized, plan) + checkSparkSchema(query) + val (_, shuffled) = checkSparkAnswerAndOperator( + query.repartition(2, col("id")), + includeClasses = Seq(classOf[CometShuffleExchangeExec])) + val exchanges = collect(shuffled) { case s: CometShuffleExchangeExec => s } + assert(exchanges.nonEmpty && exchanges.forall(_.shuffleType == CometNativeShuffle)) + // Stop before draining the input to exercise normal task-completion cleanup. + checkSparkAnswer(query.limit(1)) + } + } + } + } + } + + test("SparkToColumnar converts arrays and maps whose elements and values are not null") { + // A typed Seq[Int] or Map[Int, Int] field declares its elements and values not null, and so + // does an RDD with that schema. CometLocalTableScanExec widens such children to nullable, + // because native kernels failed on them (#4789). This conversion keeps them as they are. + val data = (0 until 60).map { i => + val n = i % 4 + ( + i.toLong, + if (i % 7 == 0) null else Seq.tabulate(n)(j => i * 10 + j), + if (i % 7 == 1) null else Seq.tabulate(n)(j => j -> (i + j)).toMap) + } + val schema = new StructType() + .add("id", LongType, nullable = false) + .add("xs", ArrayType(IntegerType, containsNull = false)) + .add("m", MapType(IntegerType, IntegerType, valueContainsNull = false)) + val rows = data.map(Row.fromTuple) + val sources: Seq[(String, () => DataFrame)] = Seq( + "typed Dataset" -> (() => + spark.sparkContext.parallelize(data, 2).toDS().toDF(schema.fieldNames.toSeq: _*)), + "RDD" -> (() => spark.createDataFrame(spark.sparkContext.parallelize(rows, 2), schema))) + val expressions = Seq("id", "size(xs) AS n", "slice(xs, 2, 2) AS s", "map_entries(m) AS e") ++ + (if (isSpark35Plus) Seq("array_insert(xs, 1, 0) AS i", "array_prepend(xs, -1) AS p") + else Seq.empty) + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_SHUFFLE_MODE.key -> "native", + CometConf.COMET_CONVERT_FROM_RDD_ENABLED.key -> "true", + CometConf.COMET_CONVERT_FROM_TYPED_DATASET_ENABLED.key -> "true") { + sources.foreach { case (name, source) => + assert(source().schema == schema, name) + val query = source().filter("id > 0").selectExpr(expressions: _*) + val (_, plan) = checkSparkAnswerAndOperator( + query, + includeClasses = Seq(classOf[CometSparkToColumnarExec], classOf[CometProjectExec])) + assert(collect(plan) { case c: CometSparkToColumnarExec => c }.size == 1, s"$name: $plan") + val (_, shuffled) = checkSparkAnswerAndOperator( + query.repartition(2, col("id")), + includeClasses = Seq(classOf[CometShuffleExchangeExec])) + val exchanges = collect(shuffled) { case s: CometShuffleExchangeExec => s } + assert(exchanges.nonEmpty && exchanges.forall(_.shuffleType == CometNativeShuffle), name) + } + } + } + + test("SparkToColumnar leaves calendar intervals to Spark") { + // Arrow holds the time part of an interval in nanoseconds, so converting this one, of about + // 29,000 years, would fail the task where Spark returns it (#5279). + val interval = new CalendarInterval(1, 2, Long.MaxValue / 10) + Seq( + CalendarIntervalType -> interval, + ArrayType(CalendarIntervalType) -> Seq(interval, null), + MapType(StringType, CalendarIntervalType) -> Map("k" -> interval)).foreach { + case (dataType, value) => + val schema = new StructType().add("id", IntegerType).add("c", dataType) + val rows = Seq(Row(1, value), Row(2, null)) + withSQLConf(CometConf.COMET_CONVERT_FROM_RDD_ENABLED.key -> "true") { + val df = spark.createDataFrame(spark.sparkContext.parallelize(rows, 1), schema) + val (_, plan) = checkSparkAnswer(df.filter("id > 0")) + assert(collect(plan) { case c: CometSparkToColumnarExec => c }.isEmpty, plan) + } } } diff --git a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala index 0c56e11b9a..54689c33fb 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala @@ -747,9 +747,10 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper assert(reasons.exists(_.contains("struct with duplicate field names")), reasons.toString) } // The gate recurses, so a duplicate struct nested in an array or map is declined too. - for (nested <- Seq(ArrayType(duplicate), MapType(LongType, duplicate))) { + for (sink <- Seq(CometLocalTableScanExec, CometSparkToColumnarExec); + nested <- Seq(ArrayType(duplicate), MapType(LongType, duplicate))) { val reasons = ListBuffer.empty[String] - assert(!CometLocalTableScanExec.isSchemaSupported(schemaWith(nested), reasons), s"$nested") + assert(!sink.isSchemaSupported(schemaWith(nested), reasons), s"$sink accepted $nested") assert(reasons.exists(_.contains("struct with duplicate field names")), reasons.toString) } // Names that differ only by case are distinct to Java Arrow and stay supported. diff --git a/spark/src/test/scala/org/apache/comet/exec/CometShuffleInputConversionSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometShuffleInputConversionSuite.scala index 0318113600..e422687244 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometShuffleInputConversionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometShuffleInputConversionSuite.scala @@ -27,9 +27,9 @@ import org.apache.spark.sql.comet.{CometHashAggregateExec, CometPlan, CometSpark import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometNativeShuffle, CometShuffleExchangeExec} import org.apache.spark.sql.execution.{ColumnarToRowExec, ColumnarToRowTransition, SparkPlan} import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, SortAggregateExec} -import org.apache.spark.sql.functions.{array, avg, col, count, hash, length, max, min, size, spark_partition_id, sum} +import org.apache.spark.sql.functions.{avg, col, count, hash, length, max, min, spark_partition_id, sum} import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.types.{ArrayType, BinaryType, BooleanType, ByteType, CalendarIntervalType, DataType, DataTypes, DateType, DecimalType, DoubleType, FloatType, IntegerType, LongType, MapType, ShortType, StringType, StructType, TimestampNTZType, TimestampType} +import org.apache.spark.sql.types.{ArrayType, BinaryType, BooleanType, ByteType, CalendarIntervalType, DataType, DataTypes, DateType, DecimalType, DoubleType, FloatType, IntegerType, LongType, MapType, NullType, ShortType, StringType, StructType, TimestampNTZType, TimestampType} import org.apache.spark.unsafe.types.CalendarInterval import org.apache.comet.{CometConf, ExtendedExplainInfo} @@ -81,10 +81,13 @@ class CometShuffleInputConversionSuite extends CometTestBase { spark.createDataFrame(spark.sparkContext.parallelize(data, 4), rowSchema) } - /** Rows with an `array` column, which `CometSparkToColumnarExec` does not convert. */ - private def arraysDf(rows: Int = 100): DataFrame = { - val schema = new StructType().add("k", IntegerType).add("xs", ArrayType(IntegerType)) - val data = (0 until rows).map(i => Row(i % 23, Seq(i, i + 1))) + /** + * Rows with a `NullType` column, which `CometSparkToColumnarExec` does not convert and the JVM + * columnar shuffle does take. + */ + private def nullsDf(rows: Int = 100): DataFrame = { + val schema = new StructType().add("k", IntegerType).add("n", NullType) + val data = (0 until rows).map(i => Row(i % 23, null)) spark.createDataFrame(spark.sparkContext.parallelize(data, 4), schema) } @@ -176,11 +179,14 @@ class CometShuffleInputConversionSuite extends CometTestBase { LocalDateTime.of(9999, 12, 31, 23, 59, 59, 999999000)), BinaryType -> Seq(Array.emptyByteArray, Array[Byte](0), Array[Byte](-1, -128, 127))) - /** Four rows of each value and of a NULL, as column `k`, over four input partitions. */ + /** + * Four rows of each value and of a NULL, as column `k`, over four input partitions, next to a + * `NullType` column `n`, which `CometSparkToColumnarExec` does not convert. + */ private def keysDf(keyType: DataType, values: Seq[Any]): DataFrame = { val keys = values :+ null - val data = (0 until keys.length * 4).map(i => Row(keys(i % keys.length), i.toLong)) - val schema = new StructType().add("k", keyType).add("v", LongType) + val data = (0 until keys.length * 4).map(i => Row(keys(i % keys.length), i.toLong, null)) + val schema = new StructType().add("k", keyType).add("v", LongType).add("n", NullType) spark.createDataFrame(spark.sparkContext.parallelize(data, 4), schema) } @@ -292,10 +298,10 @@ class CometShuffleInputConversionSuite extends CometTestBase { val (_, placed) = checkSparkAnswer( keys.repartition(10, col("k")).select(col("k"), col("v"), spark_partition_id())) assert(convertedShuffles(placed).length == 1, placed) - // The left input converts and the right one, with its array column, does not, so + // The left input converts and the right one, with its NullType column, does not, so // matching keys meet only if both shuffles put them in the same partition. val left = keys.select(col("k").as("lk"), col("v")) - val right = keys.select(col("k").as("rk"), array(col("v")).as("xs")) + val right = keys.select(col("k").as("rk"), col("n")) val (_, joined) = checkSparkAnswer(left.join(right, col("lk") === col("rk"))) assert(convertedShuffles(joined).length == 1, joined) assert(cometShuffles(joined).count(_.shuffleType == CometColumnarShuffle) == 1, joined) @@ -305,11 +311,11 @@ class CometShuffleInputConversionSuite extends CometTestBase { convertTest("a shuffle that hashes a wide decimal stays on the JVM columnar shuffle") { // Native shuffle hashes a decimal wider than 18 digits differently from Spark's partitioner - // (#5994). The right input, with its array column, stays on the JVM columnar shuffle, + // (#5994). The right input, with its NullType column, stays on the JVM columnar shuffle, // so the left one has to as well, or matching keys land in different partitions. withShuffledJoins { val left = rowsDf(100).select(col("k").cast(DecimalType(38, 0)).as("lk"), col("l")) - val right = arraysDf().select(col("k").cast(DecimalType(38, 0)).as("rk"), col("xs")) + val right = nullsDf().select(col("k").cast(DecimalType(38, 0)).as("rk"), col("n")) val df = left.join(right, col("lk") === col("rk")) val (_, plan) = checkSparkAnswer(df) assert(conversions(plan).isEmpty, plan) @@ -320,22 +326,20 @@ class CometShuffleInputConversionSuite extends CometTestBase { convertTest("a shuffle that hashes a string stays on the JVM columnar shuffle") { // Spark's partitioner hashes a string's bytes as they are, but native shuffle hashes them // after the import into native has replaced invalid UTF-8. The right input, with its - // array column, stays on the JVM columnar shuffle, so the left one has to as well, or + // NullType column, stays on the JVM columnar shuffle, so the left one has to as well, or // matching keys with invalid UTF-8 land in different partitions. A key computed from a // string, such as its hash, is no different: native shuffle would compute it from the // replaced string. Both shuffles write the rows with invalid UTF-8 replaced, and the join // compares what they wrote, so the two keys stay apart only because each key expression puts // them in different partitions. withShuffledJoins { - val schema = new StructType().add("b", BinaryType).add("v", LongType) + val schema = new StructType().add("b", BinaryType).add("v", LongType).add("n", NullType) val data = (0 until 40).map { i => - Row(Array((if (i % 2 == 0) 0x80 else 0x83).toByte), i.toLong) + Row(Array((if (i % 2 == 0) 0x80 else 0x83).toByte), i.toLong, null) } val keys = spark.createDataFrame(spark.sparkContext.parallelize(data, 4), schema) val left = keys.select(col("b").cast(StringType).as("lk"), col("v")) - val right = keys.select( - col("b").cast(StringType).as("rk"), - array(col("v").cast(IntegerType)).as("xs")) + val right = keys.select(col("b").cast(StringType).as("rk"), col("n")) Seq(col("lk") === col("rk"), hash(col("lk")) === hash(col("rk"))).foreach { condition => val df = left.join(right, condition) val (_, plan) = checkSparkAnswer(df) @@ -388,10 +392,13 @@ class CometShuffleInputConversionSuite extends CometTestBase { } convertTest("columns the conversion does not support keep the JVM columnar shuffle") { - val df = arraysDf().repartition(5, col("k")).groupBy("k").agg(sum(size(col("xs")))) + val df = nullsDf().repartition(5, col("k")).groupBy("k").agg(count(col("n"))) val (_, plan) = checkSparkAnswer(df) assert(conversions(plan).isEmpty, plan) - assert(cometShuffles(plan).forall(_.shuffleType == CometColumnarShuffle), plan) + assert( + cometShuffles(plan).nonEmpty && + cometShuffles(plan).forall(_.shuffleType == CometColumnarShuffle), + plan) } convertTest("range partitioning") { diff --git a/spark/src/test/scala/org/apache/comet/exec/CometTypedDatasetSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometTypedDatasetSuite.scala index c9e1a349e3..c4df8c2abf 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometTypedDatasetSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometTypedDatasetSuite.scala @@ -19,13 +19,14 @@ package org.apache.comet.exec +import java.time.Duration import java.util.concurrent.atomic.AtomicLong import org.apache.spark.sql.{CometTestBase, DataFrame, Dataset} import org.apache.spark.sql.catalyst.expressions.aggregate.Partial import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometHashAggregateExec, CometSparkToColumnarExec} import org.apache.spark.sql.execution.{ColumnarToRowExec, ColumnarToRowTransition, SerializeFromObjectExec, SparkPlan, WholeStageCodegenExec} -import org.apache.spark.sql.functions.{broadcast, col, size, sum} +import org.apache.spark.sql.functions.{broadcast, col, max, sum} import org.apache.spark.sql.internal.SQLConf import org.apache.comet.{CometConf, ExtendedExplainInfo} @@ -37,11 +38,11 @@ case class TypedDsWide(i: Int, s: String, d: java.math.BigDecimal, opt: Option[L case class TypedDsNested(id: Int, inner: TypedDsRec, tags: Seq[String]) -case class TypedDsInts(id: Int, xs: Seq[Int]) +case class TypedDsDurations(id: Int, d: Duration) case class TypedDsDecimal(k: java.math.BigDecimal, v: Long) -case class TypedDsDecimalInts(k: java.math.BigDecimal, xs: Seq[Int]) +case class TypedDsDecimalDurations(k: java.math.BigDecimal, d: Duration) /** Counts calls to a user function. Comet tests run in local mode, so tasks see this object. */ object TypedDsCounter { @@ -198,22 +199,23 @@ class CometTypedDatasetSuite extends CometTestBase { } convertTest("columns the conversion does not support keep the operators above on Spark") { - withParquetTable((0 until 20).map(i => (i, Seq(i, i + 1))), "ints") { - val ds = spark.sql("SELECT _1 AS id, _2 AS xs FROM ints").as[TypedDsInts] - val (_, plan) = checkSparkAnswerAndFallbackReason( - // The aggregate reads `xs`, or Spark would prune it from the serializer. - ds.map(r => TypedDsInts(r.id % 3, r.xs.reverse)).groupBy("id").agg(sum(size(col("xs")))), - "Comet cannot convert the output of a typed Dataset operation to Arrow because it " + - "does not support the type of these columns: xs: array") - assert(conversions(plan).isEmpty, plan) - } + val ds = spark + .range(20) + .map(i => TypedDsDurations((i.longValue % 3).toInt, Duration.ofDays(i.longValue))) + val (_, plan) = checkSparkAnswerAndFallbackReason( + // The aggregate reads `d`, or Spark would prune it from the serializer. + ds.groupBy("id").agg(max(col("d"))), + "Comet cannot convert the output of a typed Dataset operation to Arrow because it " + + "does not support the type of these columns: d: interval day to second") + assert(conversions(plan).isEmpty, plan) } convertTest("a join on wide decimal keys with an input that is not converted") { - // The left input converts and the right one, with its array column, does not. Native - // shuffle hashes a decimal wider than 18 digits differently from Spark's partitioner (#5994), - // so the shuffle above the conversion has to stay on Comet's columnar shuffle like the right - // input's, or matching keys land in different partitions and the join loses rows. + // The left input converts and the right one, with its interval column, does not, so its + // shuffle stays on Spark's. Native shuffle hashes a decimal wider than 18 digits differently + // from Spark's partitioner (#5994), so the shuffle above the conversion has to stay on + // Comet's columnar shuffle, which hashes as Spark does, or matching keys land in different + // partitions and the join loses rows. withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", @@ -224,13 +226,16 @@ class CometTypedDatasetSuite extends CometTestBase { .alias("l") val right = spark .range(0, 100, 1, 2) - .map(i => TypedDsDecimalInts(new java.math.BigDecimal(i.longValue), Seq(i.intValue))) + .map(i => + TypedDsDecimalDurations( + new java.math.BigDecimal(i.longValue), + Duration.ofDays(i.longValue))) .alias("r") - val df = left.join(right, col("l.k") === col("r.k")).select(col("l.v"), col("r.xs")) + val df = left.join(right, col("l.k") === col("r.k")).select(col("l.v"), col("r.d")) val (_, plan) = checkSparkAnswer(df) assert(conversions(plan).nonEmpty, plan) - // One columnar shuffle for each input of the join. - checkCometExchange(df, 2, native = false) + // The left input's shuffle is the only one of Comet's, and it is not native. + checkCometExchange(df, 1, native = false) } } diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometInMemoryCacheBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometInMemoryCacheBenchmark.scala index a7b8412eb7..60741a838c 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometInMemoryCacheBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometInMemoryCacheBenchmark.scala @@ -138,13 +138,10 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { "if(id % 8 = 5, null, concat('str_c_', cast(id as string))) AS s3") .createOrReplaceTempView(sourceTable) - // Struct columns, not arrays or maps. The baseline arm needs Spark's cache scan to bridge into - // Comet operators, and CometSparkToColumnarExec declines ArrayType and MapType outright, so - // for a relation projecting one of those the arm simply does not exist -- the partial - // aggregate stays on Spark and the two cases stop being a scan-boundary comparison. Structs - // are what can be measured here, and they are the shape that matters for the format anyway: a - // struct is where one cached column owns several field nodes and a validity buffer per level. - // Array and map coverage lives in CometInMemoryCacheSuite instead. + // Struct columns, the shape that matters for the format: one cached column owns several + // field nodes and a validity buffer per level. Array and map columns are not measured here. + // The baseline arm would bridge them into Comet operators like any other column, and + // CometInMemoryCacheSuite covers them. // // The structs themselves are non-nullable and carry nullable fields, rather than the other way // round. Comet cannot evaluate `if(c, null, named_struct(...))` at all: the Spark type keeps