dwsmith1983 commented on code in PR #6753: URL: https://github.com/apache/datafusion-comet/pull/6753#discussion_r4214890913
########## spark/src/test/scala/org/apache/comet/rules/CometMultiStoreScanSuite.scala: ########## @@ -0,0 +1,290 @@ +/* + * 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.comet.rules + +import java.io.File +import java.net.URI +import java.nio.file.Files +import java.util.UUID + +import scala.jdk.CollectionConverters._ + +import org.apache.commons.io.FileUtils +import org.apache.spark.SparkConf +import org.apache.spark.sql.{CometTestBase, DataFrame, SaveMode} +import org.apache.spark.sql.catalyst.expressions.DynamicPruningExpression +import org.apache.spark.sql.comet.{CometCsvNativeScanExec, CometNativeScanExec, CometScanExec} +import org.apache.spark.sql.execution.{FileSourceScanExec, SparkPlan} +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +import org.apache.spark.sql.execution.datasources.FilePartition +import org.apache.spark.sql.execution.datasources.v2.BatchScanExec +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{IntegerType, StructType} + +import org.apache.comet.CometConf +import org.apache.comet.hadoop.fs.FakeHdfsAuthorityFileSystem + +/** + * Native scans over files in more than one object store, without a cloud store: `hdfs://nn1` and + * `hdfs://nn2` are two native stores backed by the local disk. Native execution cannot read them, + * so claimed scans are checked on their plans, and declined scans are run by Spark. + */ +class CometMultiStoreScanSuite extends CometTestBase with AdaptiveSparkPlanHelper { + + private var rootDir: File = _ + + override protected def sparkConf: SparkConf = { + val conf = super.sparkConf + conf.set("spark.hadoop.fs.hdfs.impl", classOf[FakeHdfsAuthorityFileSystem].getName) + conf.set("spark.hadoop.fs.hdfs.impl.disable.cache", "true") + conf + } + + override def beforeAll(): Unit = { + rootDir = Files.createTempDirectory(s"comet_multi_store_${UUID.randomUUID()}").toFile + super.beforeAll() + } + + protected override def afterAll(): Unit = { + if (rootDir != null) FileUtils.deleteDirectory(rootDir) + super.afterAll() + } + + private val nativeScan = Seq( + CometConf.COMET_NATIVE_SCAN_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true") + + private val nativeCsv = + Seq(CometConf.COMET_CSV_V2_NATIVE_ENABLED.key -> "true", SQLConf.USE_V1_SOURCE_LIST.key -> "") + + // Spark packs every file of the scan into one partition. + private val onePartition = Seq( + SQLConf.FILES_MIN_PARTITION_NUM.key -> "1", + SQLConf.FILES_OPEN_COST_IN_BYTES.key -> "1", + SQLConf.FILES_MAX_PARTITION_BYTES.key -> (128L * 1024 * 1024).toString) + + private def hdfs(nameNode: String, name: String): String = + s"hdfs://$nameNode${rootDir.getAbsolutePath}/$name" + + private def local(name: String): String = s"file://${rootDir.getAbsolutePath}/$name" + + private def storeOf(path: String): String = new URI(path).getAuthority + + private def withoutComet(f: => Unit): Unit = + withSQLConf(CometConf.COMET_ENABLED.key -> "false")(f) + + private def writeIds(path: String, from: Int, format: String = "parquet"): Unit = + withoutComet { + spark + .range(from.toLong, from.toLong + 5) + .selectExpr("cast(id as int) as id") + .coalesce(1) + .write + .mode(SaveMode.Overwrite) + .format(format) + .save(path) + } + + /** The files of each partition of Spark's own scans in `df`, which is planned, not run. */ + private def sparkLayout(df: => DataFrame): Seq[Seq[String]] = { + var layout: Seq[Seq[String]] = Nil + withoutComet { + val plan = df.queryExecution.executedPlan + val partitions = collect(plan) { + case scan: FileSourceScanExec => scan.inputRDD.partitions.toSeq + case scan: BatchScanExec => scan.inputPartitions + }.flatten + layout = partitions.collect { case p: FilePartition => + p.files.map(_.filePath.toString).toSeq + } + } + layout + } + + /** The scans CometScanRule claims in the Spark plan of `df`. */ + private def ruleClaims(df: => DataFrame): Seq[CometScanExec] = { + var sparkPlan: SparkPlan = null + withoutComet { + sparkPlan = df.queryExecution.executedPlan + } + CometScanRule(spark).apply(stripAQEPlan(sparkPlan)).collect { case s: CometScanExec => s } + } + + private def nativeParquetScan(df: DataFrame): CometNativeScanExec = { + val plan = df.queryExecution.executedPlan + val scans = collect(plan) { case scan: CometNativeScanExec => scan } + assert(scans.size == 1, s"expected one native Parquet scan:\n$plan") + scans.head + } + + test("parquet scan over two name nodes packs each name node's files on its own") { + val (a, b) = (hdfs("nn1", "two-nn-a"), hdfs("nn2", "two-nn-b")) + writeIds(a, 0) + writeIds(b, 10) + withSQLConf(nativeScan ++ onePartition: _*) { + val sparkFiles = sparkLayout(spark.read.parquet(a, b)) + assert(sparkFiles.exists(_.map(storeOf).distinct.size > 1), s"Spark's: $sparkFiles") + val scan = nativeParquetScan(spark.read.parquet(a, b)) + val cometFiles = scan.perPartitionFilePaths.toSeq + assert(cometFiles.forall(_.map(storeOf).distinct.size == 1), s"Comet's: $cometFiles") + assert(cometFiles.flatten.sorted == sparkFiles.flatten.sorted) + assert(cometFiles.size != sparkFiles.size) + assert(scan.outputPartitioning.numPartitions == scan.perPartitionData.length) + } + } + + test("parquet scan over one name node keeps Spark's layout") { + val (a, b) = (hdfs("nn1", "one-nn-a"), hdfs("nn1", "one-nn-b")) + writeIds(a, 0) + writeIds(b, 10) + withSQLConf(nativeScan ++ onePartition: _*) { + val sparkFiles = sparkLayout(spark.read.parquet(a, b)) + val scan = nativeParquetScan(spark.read.parquet(a, b)) + assert(scan.perPartitionFilePaths.toSeq == sparkFiles) + } + } + + test("csv scan over two name nodes splits partitions that mix name nodes") { Review Comment: > Could we add a case with a global aggregate over a scan that Spark packs into one partition, and check that Comet either falls back or keeps the partition count? Added in e909c077b. The test first checks that Spark's scan has one input partition with files on both name nodes, then expects either the fallback with its reason (Spark 3.4) or the split scan with an exchange above it (3.5 and later). It fails on 3.4 without the guard. `CsvReadFromS3Suite` runs the same aggregate end to end over two MinIO buckets with AQE on and off, on 3.4 and 4.1, and compares the count with Spark. -- 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]
