comphead commented on code in PR #6726:
URL: https://github.com/apache/datafusion-comet/pull/6726#discussion_r4459279857
##########
spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala:
##########
@@ -4558,43 +4558,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<STRING COLLATE UTF8_LCASE>"),
- DataType.fromDDL("MAP<STRING COLLATE UTF8_LCASE,STRING>"),
- DataType.fromDDL("MAP<STRING,STRING COLLATE UTF8_LCASE>"))
- else Seq.empty)
- unsupported.foreach { dataType =>
+ "ARRAY<STRING COLLATE UTF8_LCASE>",
+ "ARRAY<ARRAY<STRING COLLATE UTF8_LCASE>>",
+ "ARRAY<STRUCT<s: STRING COLLATE UTF8_LCASE>>",
+ "MAP<STRING COLLATE UTF8_LCASE,STRING>",
+ "MAP<STRING,STRING COLLATE UTF8_LCASE>",
+ "MAP<INT,ARRAY<STRING COLLATE UTF8_LCASE>>").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<INT>", "cast(id * 10 + x AS INT)"),
+ array(
+ "ARRAY<DECIMAL(38,10)>",
+ "cast(concat(if(x % 2 = 0, '-', ''), cast(id AS STRING),
'12345678901234567.0123456789') " +
+ "AS DECIMAL(38,10))"),
+ array("ARRAY<TIMESTAMP>", "timestamp_micros(1700000000000000 + id *
1000003 + x)"),
+ array(
+ "ARRAY<BINARY>",
+ "cast(concat('b', cast(id AS STRING), cast(x AS STRING)) AS BINARY)"),
+ array("ARRAY<DOUBLE>", "cast(id AS DOUBLE) / 3 + x"),
+ array("ARRAY<BOOLEAN>", "(id + x) % 3 = 0"),
+ array(
+ "ARRAY<ARRAY<INT>>",
+ "if((x + id) % 4 = 3, null, if((x + id) % 4 = 2, cast(array() AS
ARRAY<INT>), " +
+ "transform(sequence(0, cast(x % 3 AS INT)), y -> cast(id + x + y AS
INT))))"),
+ array("ARRAY<STRUCT<i: INT, s: STRING>>", struct),
+ array("ARRAY<MAP<STRING,INT>>", "map(concat('k', cast(x AS STRING)),
cast(id AS INT))"),
+ map("MAP<INT,INT>", "x", "cast(id * 10 + x AS INT)"),
+ map(
+ "MAP<DATE,DECIMAL(20,2)>",
+ "date_add(DATE'2024-01-01', x)",
+ "cast((id * 1000003 + x) / 100 AS DECIMAL(20,2))"),
+ map(
+ "MAP<DECIMAL(20,2),STRING>",
+ "cast(concat(cast(x AS STRING), '.25') AS DECIMAL(20,2))",
+ "concat('v', cast(id AS STRING), '_', cast(x AS STRING))"),
+ map(
+ "MAP<STRING,ARRAY<INT>>",
+ "concat('k', cast(x AS STRING))",
+ "transform(sequence(0, x), y -> cast(id + y AS INT))"),
+ map("MAP<INT,STRUCT<i: INT, s: STRING>>", "x", struct),
+ map(
+ "MAP<STRUCT<i: INT, s: STRING>,STRING>",
+ "named_struct('i', x, 's', concat('k', cast(x AS STRING)))",
+ "concat('v', cast(id AS STRING))"),
+ map("MAP<ARRAY<INT>,INT>", "array(x, x + 1)", "cast(id + x AS INT)"),
+ map(
+ "MAP<STRING,MAP<INT,STRING>>",
+ "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) =>
Review Comment:
The body of this loop repeats the harness of `SparkToColumnar string maps
cross RDD and Parquet native boundaries` further down (the `rdd`, `parquet-row`
and `parquet-columnar` loop, the same confs, the `CometSparkToColumnarExec`
count, `checkSparkSchema`, the native shuffle check, `limit(1)`). Could that
harness become a helper that takes a schema, rows and a query, so the two
string tests and the 17 shapes share it? The string-array and string-map tests
would then differ only in their extra checks, and `ARRAY<STRING>` and
`MAP<STRING,STRING>` would not be the only shapes with their own copy of the
loop.
##########
spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala:
##########
@@ -4558,43 +4558,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)))
Review Comment:
The three entries here that nest `duplicate` (`ArrayType(duplicate)`,
`MapType(LongType, duplicate)`, `ArrayType(ArrayType(duplicate))`) overlap with
the `CometNativeShuffleSuite` change in this PR, which now asserts the same two
shapes against `CometSparkToColumnarExec` and also checks the `struct with
duplicate field names` reason. This test only checks that some reason was
recorded. Could the duplicate-struct entries live in one of the two places?
##########
spark/src/test/scala/org/apache/comet/exec/CometShuffleInputConversionSuite.scala:
##########
@@ -313,32 +329,37 @@ 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<bigint>
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)
}
}
}
- convertTest("a converted wide-decimal shuffle joins the JVM columnar
shuffle") {
- withShuffledJoins {
- Seq("true", "false").foreach { aqe =>
- withSQLConf(
- SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe,
- SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "false",
- SQLConf.SKEW_JOIN_ENABLED.key -> "false") {
- val left = rowsDf(100).select(col("k").cast(DecimalType(38,
0)).as("lk"), col("l"))
- // The array column keeps the right input on the JVM columnar
shuffle.
- val right = arraysDf().select(col("k").cast(DecimalType(38,
0)).as("rk"), col("xs"))
- val (_, plan) = checkSparkAnswer(left.join(right, col("lk") ===
col("rk")))
- assert(convertedShuffles(plan).length == 1, s"AQE $aqe:\n$plan")
- assert(
- cometShuffles(plan).count(_.shuffleType == CometColumnarShuffle)
== 1,
- s"AQE $aqe:\n$plan")
+ Seq("xs", "m").foreach { column =>
+ convertTest(s"a converted wide-decimal shuffle joins the JVM columnar
shuffle with $column") {
Review Comment:
The `xs` and `m` variants differ only in the payload column of the converted
side. The join key and both shuffles are the same, and the hash is over `lk`,
so I do not expect the payload type to change what is being checked. Would one
test that selects both `xs` and `m` from `collectionsDf()` cover the same
ground with half of the joins (each runs with AQE on and off)?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
