comphead commented on code in PR #6932: URL: https://github.com/apache/datafusion-comet/pull/6932#discussion_r4265451611
########## contrib/delta-spark/src/test/scala/org/apache/spark/sql/comet/CometDeltaScanPlanDataKeySuite.scala: ########## @@ -0,0 +1,107 @@ +/* + * 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 scala.jdk.CollectionConverters._ + +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.execution.UnionExec +import org.apache.spark.sql.internal.SQLConf + +import org.apache.comet.contrib.delta.{CometDeltaTestBase, DeltaSparkScanEnvelope} +import org.apache.comet.serde.OperatorOuterClass + +/** + * Planning data keys of Delta scans built through the real conversion path. Lives in this package + * because [[PlanDataInjector.findAllPlanData]] is `private[comet]`. + */ +class CometDeltaScanPlanDataKeySuite extends CometDeltaTestBase { Review Comment: Question on placement. This adds a new suite class, so a new Delta Spark session, for one test, and the closest existing test is `self-join of the same delta table keeps scans distinct` in `CometDeltaNativeScanSuite`, which uses different data filters and so never reaches the equal-hash case. The only reason for the new package is that `findAllPlanData` is `private[comet]`, and `CometScanWithPlanDataSuite` already pins what it does with equal and different data under one key. Could the Delta-specific part (two partition-filtered scans get different keys ending in their `plan_id`, each key's `perPartitionData` holds only its own files) go into `CometDeltaNativeScanSuite` using the public `sourceKey` and `perPartitionData`? If you want to keep the end to end `findAllPlanData` call, a short note on why would help the next reader. ########## contrib/delta-spark/src/test/scala/org/apache/spark/sql/comet/CometDeltaScanPlanDataKeySuite.scala: ########## @@ -0,0 +1,107 @@ +/* + * 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 scala.jdk.CollectionConverters._ + +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.execution.UnionExec +import org.apache.spark.sql.internal.SQLConf + +import org.apache.comet.contrib.delta.{CometDeltaTestBase, DeltaSparkScanEnvelope} +import org.apache.comet.serde.OperatorOuterClass + +/** + * Planning data keys of Delta scans built through the real conversion path. Lives in this package + * because [[PlanDataInjector.findAllPlanData]] is `private[comet]`. + */ +class CometDeltaScanPlanDataKeySuite extends CometDeltaTestBase { + + private def deltaScan(df: DataFrame): CometDeltaNativeScanExec = { + val scans = deltaNativeScans(df) + assert(scans.size == 1, s"Expected one native Delta scan in:\n${df.queryExecution}") + scans.head.asInstanceOf[CometDeltaNativeScanExec] + } + + private def filePaths(perPartition: Array[Array[Byte]]): Seq[String] = + perPartition.toSeq.flatMap { bytes => + OperatorOuterClass.DeltaSparkScan + .parseFrom(bytes) + .getFilePartition + .getPartitionedFileList + .asScala + .map(_.getFile.getFilePath) + } + + test("scans of one table with different partition filters keep their own planning data") { + // Without AQE the executed plan is a plain tree, so the scans are taken straight from it. + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + withTempPath { dir => + val path = dir.getAbsolutePath + spark + .range(0, 20) + .selectExpr("id", "id % 2 as p") + .write + .format("delta") + .partitionBy("p") + .save(path) + val table = spark.read.format("delta").load(path) + + val p0 = deltaScan(table.where("p = 0")) + val p1 = deltaScan(table.where("p = 1")) + // Same table, schema and projection with no data filters: the common halves match, so + // only the plan id can tell the two scans apart. + assert( + DeltaSparkScanEnvelope.unpack(p0.nativeOp).getCommon == + DeltaSparkScanEnvelope.unpack(p1.nativeOp).getCommon) + assert(p0.sourceKey != p1.sourceKey) + Seq(p0, p1).foreach { scan => + assert(scan.sourceKey.endsWith(s"#${scan.nativeOp.getPlanId}"), scan.sourceKey) + } + + // A parent holding both scans, such as a join fused into one native block. + val (commonByKey, perPartitionByKey) = + PlanDataInjector.findAllPlanData(UnionExec(Seq(p0, p1))) + assert(commonByKey.keySet == Set(p0.sourceKey, p1.sourceKey)) + assert(perPartitionByKey.keySet == Set(p0.sourceKey, p1.sourceKey)) + val p0Files = filePaths(perPartitionByKey(p0.sourceKey)) + val p1Files = filePaths(perPartitionByKey(p1.sourceKey)) + assert(p0Files.nonEmpty && p0Files.forall(_.contains("p=0")), p0Files) + assert(p1Files.nonEmpty && p1Files.forall(_.contains("p=1")), p1Files) + + // The executor reads back the key the driver stored, and rewrapping the same op keeps it. + val injector = new DeltaPlanDataInjector + Seq(p0, p1).foreach { scan => + assert(injector.getKey(scan.nativeOp).contains(scan.sourceKey)) + assert( + CometDeltaNativeScanExec(scan.nativeOp, scan.originalPlan).sourceKey == + scan.sourceKey) + } Review Comment: This block and the `sameResult` block at the end of the test cover code this PR does not change. `injector.getKey(scan.nativeOp).contains(scan.sourceKey)` re-checks that `getKey` returns the stored `source_key`, which `DeltaPlanDataInjectorSuite` already pins (`getKey(op).contains("delta_k")`). The rewrap assert only checks that `CometDeltaNativeScanExec.apply` reads the same proto field that `sourceKey` was set from. The `sameResult` asserts check canonicalization, but `equals` and `doCanonicalize` do not use `sourceKey` and this PR leaves them alone. What fails without the fix is the key inequality, the `#<plan_id>` suffix and the `findAllPlanData` result, so I would drop these two blocks. -- 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]
