andygrove commented on code in PR #6726:
URL: https://github.com/apache/datafusion-comet/pull/6726#discussion_r4542647965
##########
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:
Agreed. The three duplicate-struct entries are out of this list, and
`CometNativeShuffleSuite` keeps them, since it also checks the `struct with
duplicate field names` reason.
##########
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:
Done. `checkSparkToColumnarSources(schema, rows, ...)(query)(check)` runs
the RDD, `parquet-row` and `parquet-columnar` sources with the answer, the one
conversion, the schema, the native shuffle and the `limit(1)` checks. The 17
shapes and the string map test use it, and the map test passes its own checks
(the `RDDScanExec` child, the filter on the map and the conversions-off check)
as the last argument. The string array test keeps its own loop, because it also
reads JSON and the V2 file sources, which the helper does not, and a comment
there says so.
##########
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:
You are right that the payload does not change what is hashed. It is one
test now, with the array and the map column both on the converted side, so it
runs half the joins.
--
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]
