unikdahal commented on code in PR #5318:
URL: https://github.com/apache/datafusion-comet/pull/5318#discussion_r4126246593
##########
spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala:
##########
@@ -933,6 +927,230 @@ class CometIcebergWriteActionSuite
|WHEN NOT MATCHED THEN INSERT (id, region, amount) VALUES (s.id,
s.region, s.amount)
|""".stripMargin)
}
+
+ withSQLConf(CometConf.COMET_EXEC_MERGE_ROWS_ENABLED.key -> "true") {
+ if (isSpark41Plus) {
+ assertNativeWriteDoesNotEngage("native_cow_merge", Seq(1, 2,
3))(runMerge())
+ } else {
+ val snapshot = withNativeEnabled {
+ captureWrite("native_cow_merge")(runMerge())
+ }
+ assert(
+ snapshot.snapshotDelta == 1L,
+ s"expected exactly one Iceberg snapshot from native MERGE, got
${snapshot.snapshotDelta}")
+ val mergeExecs = snapshot.plans.flatMap { plan =>
+ collectWithSubqueries(plan) { case e: CometMergeRowsExec => e }
+ }
+ assert(
+ mergeExecs.nonEmpty,
+ "expected Iceberg MERGE to execute through CometMergeRowsExec.
Plans:\n" +
+ snapshot.plans.mkString("\n--\n"))
+ val nativeWrites = snapshot.plans.flatMap { plan =>
+ collectWithSubqueries(plan) { case e: CometIcebergWriteExec => e }
+ }
+ assert(
+ nativeWrites.nonEmpty,
+ "expected Iceberg MERGE to feed CometIcebergWriteExec. Plans:\n" +
+ snapshot.plans.mkString("\n--\n"))
+ assertRows("native_cow_merge", expectedIds = Seq(1, 2, 3))
+ }
+ }
+
+ val updated = spark
+ .sql(s"SELECT amount FROM $catalog.$ns.native_cow_merge WHERE id = 2")
+ .collect()
+ assert(updated.length == 1 && updated.head.getDouble(0) == 200.0)
+ }
+ }
+
+ test("native MergeRows matches Spark on partitioned Iceberg copy-on-write
and merge-on-read") {
+ assumeNativeAcceleration()
+ assume(
+ isSpark35Plus && !isSpark41Plus,
+ "native MergeRows is registered only on Spark 3.5 and 4.0")
+ withIcebergCatalog { warehouseDir =>
+ spark
+ .range(0, 20000, 1, 8)
+ .selectExpr(
+ "CAST(id AS INT) AS id",
+ "concat('r', CAST(id % 8 AS STRING)) AS region",
+ "CAST(id AS DOUBLE) AS amount")
+ .createOrReplaceTempView("merge_parity_seed")
+ spark
+ .range(10000, 30000, 1, 8)
+ .selectExpr(
+ "CAST(id AS INT) AS id",
+ "concat('s', CAST(id % 8 AS STRING)) AS region",
+ "CASE WHEN id % 11 = 0 THEN CAST(NULL AS DOUBLE) ELSE CAST(id * 2 AS
DOUBLE) END AS amount")
+ .createOrReplaceTempView("merge_parity_source")
+
+ def merge(table: String): Unit = {
+ spark.sql(s"""
+ |MERGE INTO $catalog.$ns.$table t
+ |USING merge_parity_source s
+ |ON t.id = s.id
+ |WHEN MATCHED AND CAST(NULL AS BOOLEAN) THEN UPDATE SET t.amount =
-999.0
+ |WHEN MATCHED AND s.amount IS NULL THEN DELETE
+ |WHEN MATCHED AND s.id % 5 = 0 THEN DELETE
+ |WHEN MATCHED THEN UPDATE SET t.region = s.region, t.amount =
s.amount + 0.5
+ |WHEN NOT MATCHED AND s.id IS NOT NULL THEN
+ | INSERT (id, region, amount) VALUES (s.id, s.region, s.amount)
+ |WHEN NOT MATCHED BY SOURCE AND t.id % 7 = 0 THEN DELETE
+ |WHEN NOT MATCHED BY SOURCE THEN UPDATE SET t.amount = t.amount + 1.0
+ |""".stripMargin)
+ }
+
+ def rows(table: String): Seq[Row] =
+ spark
+ .sql(s"SELECT id, region, amount FROM $catalog.$ns.$table ORDER BY
id")
+ .collect()
+ .toSeq
+
+ Seq("copy-on-write" -> "cow", "merge-on-read" -> "mor").foreach { case
(mode, suffix) =>
+ val nativeTable = s"merge_parity_${suffix}_native"
+ val sparkTable = s"merge_parity_${suffix}_spark"
+ Seq(nativeTable, sparkTable).foreach { table =>
+ createTable(
+ warehouseDir,
+ table,
+ partitionSpec = "PARTITIONED BY (region)",
+ properties = Some(s"'format-version'='2',
'write.merge.mode'='$mode'"))
+ }
+
+ withSQLConf(
+ CometConf.COMET_ENABLED.key -> "false",
+ "spark.sql.adaptive.coalescePartitions.enabled" -> "false",
+ "spark.sql.shuffle.partitions" -> "8") {
+ spark.sql(
+ s"INSERT INTO $catalog.$ns.$nativeTable SELECT id, region, amount
FROM merge_parity_seed")
+ spark.sql(
+ s"INSERT INTO $catalog.$ns.$sparkTable SELECT id, region, amount
FROM merge_parity_seed")
+ }
+
+ val inputFiles = spark
+ .sql(s"SELECT count(*) FROM $catalog.$ns.$nativeTable.data_files")
+ .collect()
+ .head
+ .getLong(0)
+ assert(inputFiles > 1L, s"expected a multi-file target, got
$inputFiles file(s)")
+
+ var snapshot: Option[WriteSnapshot] = None
+ withSQLConf(
+ CometConf.COMET_EXEC_MERGE_ROWS_ENABLED.key -> "true",
Review Comment:
Added. The native parity path now runs with AQE disabled too.
--
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]
