diff --git a/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/HiveTable.scala b/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/HiveTable.scala index 4c2d0168589..8d43f8d6d21 100644 --- a/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/HiveTable.scala +++ b/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/HiveTable.scala @@ -33,14 +33,12 @@ import org.apache.spark.sql.connector.catalog.TableCapability.{BATCH_READ, BATCH import org.apache.spark.sql.connector.expressions.Transform import org.apache.spark.sql.connector.read.ScanBuilder import org.apache.spark.sql.connector.write.{LogicalWriteInfo, WriteBuilder} -import org.apache.spark.sql.execution.datasources.v2.orc.OrcScanBuilder -import org.apache.spark.sql.execution.datasources.v2.parquet.ParquetScanBuilder import org.apache.spark.sql.hive.kyuubi.connector.HiveBridgeHelper.{BucketSpecHelper, LogicalExpressions} import org.apache.spark.sql.types.StructType import org.apache.spark.sql.util.CaseInsensitiveStringMap import org.apache.kyuubi.spark.connector.hive.KyuubiHiveConnectorConf.{READ_CONVERT_METASTORE_ORC, READ_CONVERT_METASTORE_PARQUET} -import org.apache.kyuubi.spark.connector.hive.read.{HiveCatalogFileIndex, HiveScanBuilder} +import org.apache.kyuubi.spark.connector.hive.read.{HiveCatalogFileIndex, HiveScanBuilder, KyuubiOrcScanBuilder, KyuubiParquetScanBuilder} import org.apache.kyuubi.spark.connector.hive.write.HiveWriteBuilder case class HiveTable( @@ -109,10 +107,24 @@ case class HiveTable( override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = { convertedProvider match { case Some("ORC") if sparkSession.sessionState.conf.getConf(READ_CONVERT_METASTORE_ORC) => - OrcScanBuilder(sparkSession, fileIndex, schema, dataSchema, options) + new KyuubiOrcScanBuilder( + sparkSession, + fileIndex, + schema, + dataSchema, + options, + catalogTable, + hiveTableCatalog) case Some("PARQUET") if sparkSession.sessionState.conf.getConf(READ_CONVERT_METASTORE_PARQUET) => - ParquetScanBuilder(sparkSession, fileIndex, schema, dataSchema, options) + new KyuubiParquetScanBuilder( + sparkSession, + fileIndex, + schema, + dataSchema, + options, + catalogTable, + hiveTableCatalog) case _ => HiveScanBuilder(sparkSession, fileIndex, dataSchema, catalogTable) } } diff --git a/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/HiveTableCatalog.scala b/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/HiveTableCatalog.scala index 3d5c27ebb88..be91dec816a 100644 --- a/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/HiveTableCatalog.scala +++ b/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/HiveTableCatalog.scala @@ -77,7 +77,18 @@ class HiveTableCatalog(sparkSession: SparkSession) SupportsNamespaces.PROP_LOCATION, SupportsNamespaces.PROP_OWNER) - private lazy val hadoopConf: Configuration = { + /** + * Cached Hadoop [[Configuration]] snapshot taken at first catalog use. + */ + private lazy val hadoopConf: Configuration = buildHadoopConf() + + /** + * Non-cached Hadoop [[Configuration]] for scan builders. Re-evaluated on + * every call so mid-session confs reach readers. + */ + def newScanHadoopConf(): Configuration = buildHadoopConf() + + private def buildHadoopConf(): Configuration = { val conf = sparkSession.sessionState.newHadoopConf() catalogOptions.asScala.foreach { case (k, v) => conf.set(k, v) } if (catalogOptions.containsKey("hive.metastore.uris")) { diff --git a/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/read/HiveScan.scala b/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/read/HiveScan.scala index 088fa6167e2..7792ee7b449 100644 --- a/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/read/HiveScan.scala +++ b/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/read/HiveScan.scala @@ -28,7 +28,7 @@ import org.apache.spark.sql.catalyst.catalog.{CatalogStorageFormat, CatalogTable import org.apache.spark.sql.catalyst.expressions.{AttributeReference, Expression} import org.apache.spark.sql.catalyst.expressions.codegen.GenerateUnsafeProjection import org.apache.spark.sql.connector.expressions.NamedReference -import org.apache.spark.sql.connector.read.{PartitionReaderFactory, SupportsRuntimeFiltering} +import org.apache.spark.sql.connector.read.{PartitionReaderFactory, Scan, SupportsRuntimeFiltering} import org.apache.spark.sql.execution.datasources.{FilePartition, PartitionedFile} import org.apache.spark.sql.execution.datasources.v2.FileScan import org.apache.spark.sql.hive.kyuubi.connector.HiveBridgeHelper.HiveClientImpl @@ -181,6 +181,27 @@ case class HiveScan( // SupportsRuntimeFiltering implementation // ------------------------------------------------------------------------------- + /** + * The default [[Scan.ColumnarSupportMode.PARTITION_DEFINED]] (SPARK-44505) + * would drive `DataSourceV2ScanExecBase.supportsColumnar` to materialise + * `inputPartitions` during planning (via `HiveScan.partitions` -> + * `HiveCatalogFileIndex.listHiveFiles`), triggering a full-table HDFS + * listing before runtime filters arrive via + * [[SupportsRuntimeFiltering.filter]] and cancelling DPP's end-to-end win. + * + * [[HivePartitionReaderFactory]] only implements the row-based + * `createReader` path (no `supportColumnarReads` / `createColumnarReader`), + * so `HiveScan` is always row-based. Returning `UNSUPPORTED` + * is semantically equivalent to the default behaviour but short-circuits + * `supportsColumnar` without touching `inputPartitions`. + * + * NOTE: If [[HivePartitionReaderFactory]] ever gains columnar support, + * remove this override so `supportsColumnar` reflects reality, otherwise + * columnar-capable partitions would be silently reported as row-based. + */ + override def columnarSupportMode(): Scan.ColumnarSupportMode = + Scan.ColumnarSupportMode.UNSUPPORTED + override def filterAttributes(): Array[NamedReference] = { HiveRuntimeFilterSupport.filterAttributes(readPartitionSchema.fieldNames.toSeq) } diff --git a/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/read/KyuubiOrcScan.scala b/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/read/KyuubiOrcScan.scala new file mode 100644 index 00000000000..8c141df9527 --- /dev/null +++ b/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/read/KyuubiOrcScan.scala @@ -0,0 +1,171 @@ +/* + * 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.kyuubi.spark.connector.hive.read + +import org.apache.hadoop.conf.Configuration +import org.apache.hadoop.fs.Path +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.catalyst.catalog.CatalogTable +import org.apache.spark.sql.catalyst.expressions.Expression +import org.apache.spark.sql.connector.expressions.NamedReference +import org.apache.spark.sql.connector.expressions.aggregate.Aggregation +import org.apache.spark.sql.connector.read.{InputPartition, PartitionReaderFactory, Scan, SupportsRuntimeFiltering} +import org.apache.spark.sql.execution.WholeStageCodegenExec +import org.apache.spark.sql.execution.datasources.PartitioningAwareFileIndex +import org.apache.spark.sql.execution.datasources.orc.OrcUtils +import org.apache.spark.sql.execution.datasources.v2.FileScan +import org.apache.spark.sql.execution.datasources.v2.orc.OrcScan +import org.apache.spark.sql.sources.Filter +import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.util.CaseInsensitiveStringMap + +/** + * A DPP-aware wrapper around Spark's built-in [[OrcScan]] that adds + * [[SupportsRuntimeFiltering]] so Dynamic Partition Pruning can push runtime + * IN predicates down to the Hive partitioned scan. + * + * Implementation notes: + * 1. Only DPP-specific methods ([[filter]] / [[filterAttributes]] / + * [[planInputPartitions]]) contain custom logic, all other methods + * delegate to the wrapped [[OrcScan]]. + * 2. [[equals]] uses a `KyuubiOrcScan` type pattern before delegating to + * `inner.equals`, so a plain [[OrcScan]] never compares equal and is + * never reused in its place during exchange/subquery reuse. [[hashCode]] + * is a constant, matching [[FileScan]]'s default. + * 3. Native engines (Gluten / Comet) identify scans by class name, so this + * wrapper is not recognized and falls back to JVM reads. + */ +class KyuubiOrcScan( + val sparkSession: SparkSession, + val hadoopConf: Configuration, + val fileIndex: PartitioningAwareFileIndex, + val dataSchema: StructType, + val readDataSchema: StructType, + val readPartitionSchema: StructType, + val options: CaseInsensitiveStringMap, + val pushedAggregate: Option[Aggregation], + val pushedFilters: Array[Filter], + val partitionFilters: Seq[Expression], + val dataFilters: Seq[Expression], + val catalogTable: CatalogTable) + extends FileScan + with SupportsRuntimeFiltering { + + private[hive] val inner: OrcScan = OrcScan( + sparkSession, + hadoopConf, + fileIndex, + dataSchema, + readDataSchema, + readPartitionSchema, + options, + pushedAggregate, + pushedFilters, + partitionFilters, + dataFilters) + + private var runtimeFilters: Seq[Expression] = Seq.empty + + private val isCaseSensitive = sparkSession.sessionState.conf.caseSensitiveAnalysis + + /** + * The default [[Scan.ColumnarSupportMode.PARTITION_DEFINED]] (SPARK-44505) + * would drive `DataSourceV2ScanExecBase.supportsColumnar` to materialise + * `inputPartitions` during planning (via `FileScan.partitions` -> + * `HiveCatalogFileIndex.listFiles`), triggering a full-table HDFS listing + * before runtime filters arrive via [[SupportsRuntimeFiltering.filter]] and + * cancelling DPP's end-to-end win. + * + * We instead decide from sqlConf + schema, matching + * `OrcPartitionReaderFactory.supportColumnarReads` in all non-empty cases. + * When DPP prunes every partition, Spark's default would return + * `UNSUPPORTED` on the empty list, we still return `SUPPORTED`, adding a + * harmless `ColumnarToRow` on an empty RDD. + */ + override def columnarSupportMode(): Scan.ColumnarSupportMode = { + val sqlConf = sparkSession.sessionState.conf + val schema = StructType(readDataSchema.fields ++ readPartitionSchema.fields) + val supportsColumnar = sqlConf.orcVectorizedReaderEnabled && + sqlConf.wholeStageEnabled && + !WholeStageCodegenExec.isTooManyFields(sqlConf, schema) && + schema.forall(s => + OrcUtils.supportColumnarReads( + s.dataType, + sqlConf.orcVectorizedReaderNestedColumnEnabled)) + if (supportsColumnar) Scan.ColumnarSupportMode.SUPPORTED + else Scan.ColumnarSupportMode.UNSUPPORTED + } + + override def filterAttributes(): Array[NamedReference] = { + // Under aggregate pushdown, the scan outputs aggregate columns only, so + // partition columns may be absent from its output, runtime filtering on + // them is also meaningless once results are aggregated. + if (pushedAggregate.nonEmpty) Array.empty[NamedReference] + else HiveRuntimeFilterSupport.filterAttributes(readPartitionSchema.fieldNames.toSeq) + } + + override def filter(filters: Array[Filter]): Unit = { + runtimeFilters = HiveRuntimeFilterSupport.toCatalystPartitionFilters( + filters, + fileIndex.partitionSchema, + isCaseSensitive) + if (runtimeFilters.nonEmpty) { + logInfo(s"Received ${runtimeFilters.length} runtime partition filter(s) for " + + s"${catalogTable.identifier}") + logDebug(s"Runtime partition filter(s) for ${catalogTable.identifier}: " + + s"${runtimeFilters.mkString(", ")}") + } + } + + override def planInputPartitions(): Array[InputPartition] = { + if (runtimeFilters.isEmpty) { + inner.planInputPartitions() + } else { + // Delegate planning to a sibling OrcScan carrying the merged + // partitionFilters ++ runtimeFilters so DPP predicates take effect. + val sibling = OrcScan( + sparkSession, + hadoopConf, + fileIndex, + dataSchema, + readDataSchema, + readPartitionSchema, + options, + pushedAggregate, + pushedFilters, + partitionFilters ++ runtimeFilters, + dataFilters) + sibling.planInputPartitions() + } + } + + override def isSplitable(path: Path): Boolean = inner.isSplitable(path) + + override def readSchema(): StructType = inner.readSchema() + + override def getMetaData(): Map[String, String] = inner.getMetaData() + + override def createReaderFactory(): PartitionReaderFactory = inner.createReaderFactory() + + override def equals(obj: Any): Boolean = obj match { + case that: KyuubiOrcScan => this.inner.equals(that.inner) + case _ => false + } + + override def hashCode(): Int = getClass.hashCode() +} diff --git a/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/read/KyuubiOrcScanBuilder.scala b/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/read/KyuubiOrcScanBuilder.scala new file mode 100644 index 00000000000..e4a2e457567 --- /dev/null +++ b/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/read/KyuubiOrcScanBuilder.scala @@ -0,0 +1,126 @@ +/* + * 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.kyuubi.spark.connector.hive.read + +import scala.collection.JavaConverters._ + +import org.apache.hadoop.conf.Configuration +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.catalyst.catalog.CatalogTable +import org.apache.spark.sql.connector.expressions.aggregate.Aggregation +import org.apache.spark.sql.connector.read.SupportsPushDownAggregates +import org.apache.spark.sql.execution.datasources.{AggregatePushDownUtils, PartitioningAwareFileIndex} +import org.apache.spark.sql.execution.datasources.v2.FileScanBuilder +import org.apache.spark.sql.hive.kyuubi.connector.HiveBridgeHelper +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.sources.Filter +import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.util.CaseInsensitiveStringMap + +import org.apache.kyuubi.spark.connector.hive.HiveTableCatalog + +/** + * A ScanBuilder that mirrors Spark's built-in [[OrcScanBuilder]] but builds + * [[KyuubiOrcScan]] instances, which additionally implement + * `SupportsRuntimeFiltering` so that Dynamic Partition Pruning works when + * a Hive ORC table goes through Spark's vectorized ORC reader path. + * + * Filter, aggregate and column pushdown behaviour matches [[OrcScanBuilder]]. + */ +class KyuubiOrcScanBuilder( + sparkSession: SparkSession, + fileIndex: PartitioningAwareFileIndex, + schema: StructType, + dataSchema: StructType, + options: CaseInsensitiveStringMap, + catalogTable: CatalogTable, + hiveTableCatalog: HiveTableCatalog) + extends FileScanBuilder(sparkSession, fileIndex, dataSchema) + with SupportsPushDownAggregates { + + /** + * Cloned from a freshly-built per-catalog Hadoop [[Configuration]] so + * per-catalog settings and mid-session confs are both honored, matching + * Spark's built-in `OrcScanBuilder.hadoopConf`. Cloned so per-scan + * `options` do not pollute the source instance. + */ + lazy val hadoopConf: Configuration = { + val conf = new Configuration(hiveTableCatalog.newScanHadoopConf()) + // Hadoop Configurations are case sensitive. + options.asCaseSensitiveMap.asScala.foreach { case (k, v) => conf.set(k, v) } + conf + } + + private var finalSchema = new StructType() + + private var pushedAggregations = Option.empty[Aggregation] + + override protected val supportsNestedSchemaPruning: Boolean = true + + override def build(): KyuubiOrcScan = { + // the `finalSchema` is either pruned in pushAggregation (if aggregates are + // pushed down), or pruned in readDataSchema() (in regular column pruning). These + // two are mutual exclusive. + if (pushedAggregations.isEmpty) { + finalSchema = readDataSchema() + } + new KyuubiOrcScan( + sparkSession, + hadoopConf, + fileIndex, + dataSchema, + finalSchema, + readPartitionSchema(), + options, + pushedAggregations, + pushedDataFilters, + partitionFilters, + dataFilters, + catalogTable) + } + + override def pushDataFilters(dataFilters: Array[Filter]): Array[Filter] = { + if (sparkSession.sessionState.conf.orcFilterPushDown) { + HiveBridgeHelper.orcConvertibleFilters( + readDataSchema(), + SQLConf.get.caseSensitiveAnalysis, + dataFilters.toSeq).toArray + } else { + Array.empty[Filter] + } + } + + override def pushAggregation(aggregation: Aggregation): Boolean = { + if (!sparkSession.sessionState.conf.orcAggregatePushDown) { + return false + } + + AggregatePushDownUtils.getSchemaForPushedAggregation( + aggregation, + schema, + partitionNameSet, + dataFilters) match { + + case Some(schema) => + finalSchema = schema + this.pushedAggregations = Some(aggregation) + true + case _ => false + } + } +} diff --git a/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/read/KyuubiParquetScan.scala b/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/read/KyuubiParquetScan.scala new file mode 100644 index 00000000000..c517f54bc19 --- /dev/null +++ b/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/read/KyuubiParquetScan.scala @@ -0,0 +1,249 @@ +/* + * 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.kyuubi.spark.connector.hive.read + +import org.apache.hadoop.conf.Configuration +import org.apache.hadoop.fs.Path +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.catalyst.catalog.CatalogTable +import org.apache.spark.sql.catalyst.expressions.Expression +import org.apache.spark.sql.connector.expressions.NamedReference +import org.apache.spark.sql.connector.expressions.aggregate.Aggregation +import org.apache.spark.sql.connector.read.{InputPartition, PartitionReaderFactory, Scan, SupportsRuntimeFiltering} +import org.apache.spark.sql.execution.WholeStageCodegenExec +import org.apache.spark.sql.execution.datasources.PartitioningAwareFileIndex +import org.apache.spark.sql.execution.datasources.parquet.ParquetUtils +import org.apache.spark.sql.execution.datasources.v2.FileScan +import org.apache.spark.sql.execution.datasources.v2.parquet.ParquetScan +import org.apache.spark.sql.sources.Filter +import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.util.CaseInsensitiveStringMap + +import org.apache.kyuubi.util.reflect.{DynClasses, DynConstructors} + +/** + * A DPP-aware wrapper around Spark's built-in [[ParquetScan]] that adds + * [[SupportsRuntimeFiltering]] so Dynamic Partition Pruning can push runtime + * IN predicates down to the Hive partitioned scan. + * + * Implementation notes: + * 1. Only DPP-specific methods ([[filter]] / [[filterAttributes]] / + * [[planInputPartitions]]) contain custom logic, all other methods + * delegate to the wrapped [[ParquetScan]]. + * 2. [[equals]] uses a `KyuubiParquetScan` type pattern before delegating to + * `inner.equals`, so a plain [[ParquetScan]] never compares equal and is + * never reused in its place during exchange/subquery reuse. [[hashCode]] + * is a constant, matching [[FileScan]]'s default. + * 3. Native engines (Gluten / Comet) identify scans by class name, so this + * wrapper is not recognized and falls back to JVM reads. + */ +class KyuubiParquetScan( + val sparkSession: SparkSession, + val hadoopConf: Configuration, + val fileIndex: PartitioningAwareFileIndex, + val dataSchema: StructType, + val readDataSchema: StructType, + val readPartitionSchema: StructType, + val pushedFilters: Array[Filter], + val options: CaseInsensitiveStringMap, + val pushedAggregate: Option[Aggregation], + val partitionFilters: Seq[Expression], + val dataFilters: Seq[Expression], + val catalogTable: CatalogTable) + extends FileScan + with SupportsRuntimeFiltering { + + private[hive] val inner: ParquetScan = KyuubiParquetScan.newParquetScan( + sparkSession, + hadoopConf, + fileIndex, + dataSchema, + readDataSchema, + readPartitionSchema, + pushedFilters, + options, + pushedAggregate, + partitionFilters, + dataFilters) + + private var runtimeFilters: Seq[Expression] = Seq.empty + + private val isCaseSensitive = sparkSession.sessionState.conf.caseSensitiveAnalysis + + /** + * The default [[Scan.ColumnarSupportMode.PARTITION_DEFINED]] (SPARK-44505) + * would drive `DataSourceV2ScanExecBase.supportsColumnar` to materialise + * `inputPartitions` during planning (via `FileScan.partitions` -> + * `HiveCatalogFileIndex.listFiles`), triggering a full-table HDFS listing + * before runtime filters arrive via [[SupportsRuntimeFiltering.filter]] and + * cancelling DPP's end-to-end win. + * + * We instead decide from sqlConf + schema, matching + * `ParquetPartitionReaderFactory.supportColumnarReads` in all non-empty + * cases. When DPP prunes every partition, Spark's default would return + * `UNSUPPORTED` on the empty list, we still return `SUPPORTED`, adding a + * harmless `ColumnarToRow` on an empty RDD. + */ + override def columnarSupportMode(): Scan.ColumnarSupportMode = { + val sqlConf = sparkSession.sessionState.conf + val schema = StructType(readDataSchema.fields ++ readPartitionSchema.fields) + val supportsColumnar = ParquetUtils.isBatchReadSupportedForSchema(sqlConf, schema) && + sqlConf.wholeStageEnabled && + !WholeStageCodegenExec.isTooManyFields(sqlConf, schema) + if (supportsColumnar) Scan.ColumnarSupportMode.SUPPORTED + else Scan.ColumnarSupportMode.UNSUPPORTED + } + + override def filterAttributes(): Array[NamedReference] = { + // Under aggregate pushdown, the scan outputs aggregate columns only, so + // partition columns may be absent from its output, runtime filtering on + // them is also meaningless once results are aggregated. + if (pushedAggregate.nonEmpty) Array.empty[NamedReference] + else HiveRuntimeFilterSupport.filterAttributes(readPartitionSchema.fieldNames.toSeq) + } + + override def filter(filters: Array[Filter]): Unit = { + runtimeFilters = HiveRuntimeFilterSupport.toCatalystPartitionFilters( + filters, + fileIndex.partitionSchema, + isCaseSensitive) + if (runtimeFilters.nonEmpty) { + logInfo(s"Received ${runtimeFilters.length} runtime partition filter(s) for " + + s"${catalogTable.identifier}") + logDebug(s"Runtime partition filter(s) for ${catalogTable.identifier}: " + + s"${runtimeFilters.mkString(", ")}") + } + } + + override def planInputPartitions(): Array[InputPartition] = { + if (runtimeFilters.isEmpty) { + inner.planInputPartitions() + } else { + // Delegate planning to a sibling ParquetScan carrying the merged + // partitionFilters ++ runtimeFilters so DPP predicates take effect. + val sibling = KyuubiParquetScan.newParquetScan( + sparkSession, + hadoopConf, + fileIndex, + dataSchema, + readDataSchema, + readPartitionSchema, + pushedFilters, + options, + pushedAggregate, + partitionFilters ++ runtimeFilters, + dataFilters) + sibling.planInputPartitions() + } + } + + override def isSplitable(path: Path): Boolean = inner.isSplitable(path) + + override def readSchema(): StructType = inner.readSchema() + + override def getMetaData(): Map[String, String] = inner.getMetaData() + + override def createReaderFactory(): PartitionReaderFactory = inner.createReaderFactory() + + override def equals(obj: Any): Boolean = obj match { + case that: KyuubiParquetScan => this.inner.equals(that.inner) + case _ => false + } + + override def hashCode(): Int = getClass.hashCode() +} + +object KyuubiParquetScan { + + // Element type of `Array[VariantExtraction]` in Spark 4.1+, null when absent. + private lazy val variantExtractionCls: Class[_] = DynClasses.builder() + .impl("org.apache.spark.sql.connector.read.VariantExtraction") + .orNull() + .build() + + private lazy val emptyVariantExtractions: AnyRef = + if (variantExtractionCls == null) null + else java.lang.reflect.Array.newInstance(variantExtractionCls, 0).asInstanceOf[AnyRef] + + private lazy val parquetScanCtor: DynConstructors.Ctor[ParquetScan] = + if (variantExtractionCls != null) { + DynConstructors.builder() + .impl( // SPARK-53880 / SPARK-54656 (4.1.0): adds trailing Array[VariantExtraction] + classOf[ParquetScan], + classOf[SparkSession], + classOf[Configuration], + classOf[PartitioningAwareFileIndex], + classOf[StructType], + classOf[StructType], + classOf[StructType], + classOf[Array[Filter]], + classOf[CaseInsensitiveStringMap], + classOf[Option[Aggregation]], + classOf[Seq[Expression]], + classOf[Seq[Expression]], + emptyVariantExtractions.getClass) + .buildChecked[ParquetScan]() + } else { + DynConstructors.builder() + .impl( // Spark 4.0 and previous + classOf[ParquetScan], + classOf[SparkSession], + classOf[Configuration], + classOf[PartitioningAwareFileIndex], + classOf[StructType], + classOf[StructType], + classOf[StructType], + classOf[Array[Filter]], + classOf[CaseInsensitiveStringMap], + classOf[Option[Aggregation]], + classOf[Seq[Expression]], + classOf[Seq[Expression]]) + .buildChecked[ParquetScan]() + } + + // scalastyle:off parameter.number + private[hive] def newParquetScan( + sparkSession: SparkSession, + hadoopConf: Configuration, + fileIndex: PartitioningAwareFileIndex, + dataSchema: StructType, + readDataSchema: StructType, + readPartitionSchema: StructType, + pushedFilters: Array[Filter], + options: CaseInsensitiveStringMap, + pushedAggregate: Option[Aggregation], + partitionFilters: Seq[Expression], + dataFilters: Seq[Expression]): ParquetScan = { + // `DynConstructors` truncates trailing args to match the resolved ctor arity, + // so always passing `emptyVariantExtractions` is safe on pre-4.1 Spark too. + parquetScanCtor.newInstance( + sparkSession, + hadoopConf, + fileIndex, + dataSchema, + readDataSchema, + readPartitionSchema, + pushedFilters, + options, + pushedAggregate, + partitionFilters, + dataFilters, + emptyVariantExtractions) + } + // scalastyle:on parameter.number +} diff --git a/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/read/KyuubiParquetScanBuilder.scala b/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/read/KyuubiParquetScanBuilder.scala new file mode 100644 index 00000000000..5deac91f413 --- /dev/null +++ b/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/read/KyuubiParquetScanBuilder.scala @@ -0,0 +1,127 @@ +/* + * 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.kyuubi.spark.connector.hive.read + +import scala.collection.JavaConverters._ + +import org.apache.hadoop.conf.Configuration +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.catalyst.catalog.CatalogTable +import org.apache.spark.sql.connector.expressions.aggregate.Aggregation +import org.apache.spark.sql.connector.read.SupportsPushDownAggregates +import org.apache.spark.sql.execution.datasources.{AggregatePushDownUtils, PartitioningAwareFileIndex} +import org.apache.spark.sql.execution.datasources.v2.FileScanBuilder +import org.apache.spark.sql.hive.kyuubi.connector.HiveBridgeHelper +import org.apache.spark.sql.sources.Filter +import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.util.CaseInsensitiveStringMap + +import org.apache.kyuubi.spark.connector.hive.HiveTableCatalog + +/** + * A ScanBuilder that mirrors Spark's built-in [[ParquetScanBuilder]] but builds + * [[KyuubiParquetScan]] instances, which additionally implement + * `SupportsRuntimeFiltering` so that Dynamic Partition Pruning works when + * a Hive Parquet table goes through Spark's vectorized Parquet reader path. + * + * Filter, aggregate and column pushdown behaviour matches [[ParquetScanBuilder]]. + * Gap: on Spark 4.1+ [[ParquetScanBuilder]] also mixes in + * `SupportsPushDownVariantExtractions`, which this builder does not support, + * but this is unreachable in practice. Hive-serde Parquet tables have no way + * to declare `variant` data columns (Hive 4.2 adds `variant` only for Iceberg + * tables, see HIVE-29184). + */ +class KyuubiParquetScanBuilder( + sparkSession: SparkSession, + fileIndex: PartitioningAwareFileIndex, + schema: StructType, + dataSchema: StructType, + options: CaseInsensitiveStringMap, + catalogTable: CatalogTable, + hiveTableCatalog: HiveTableCatalog) + extends FileScanBuilder(sparkSession, fileIndex, dataSchema) + with SupportsPushDownAggregates { + + /** + * Cloned from a freshly-built per-catalog Hadoop [[Configuration]] so + * per-catalog settings and mid-session confs are both honored, matching + * Spark's built-in `ParquetScanBuilder.hadoopConf`. Cloned so per-scan + * `options` do not pollute the source instance. + */ + lazy val hadoopConf: Configuration = { + val conf = new Configuration(hiveTableCatalog.newScanHadoopConf()) + // Hadoop Configurations are case sensitive. + options.asCaseSensitiveMap.asScala.foreach { case (k, v) => conf.set(k, v) } + conf + } + + private var finalSchema = new StructType() + + private var pushedAggregations = Option.empty[Aggregation] + + override protected val supportsNestedSchemaPruning: Boolean = true + + override def build(): KyuubiParquetScan = { + // the `finalSchema` is either pruned in pushAggregation (if aggregates are + // pushed down), or pruned in readDataSchema() (in regular column pruning). These + // two are mutual exclusive. + if (pushedAggregations.isEmpty) { + finalSchema = readDataSchema() + } + new KyuubiParquetScan( + sparkSession, + hadoopConf, + fileIndex, + dataSchema, + finalSchema, + readPartitionSchema(), + pushedDataFilters, + options, + pushedAggregations, + partitionFilters, + dataFilters, + catalogTable) + } + + override def pushDataFilters(dataFilters: Array[Filter]): Array[Filter] = { + if (sparkSession.sessionState.conf.parquetFilterPushDown) { + HiveBridgeHelper.parquetConvertibleFilters(readDataSchema(), dataFilters.toSeq).toArray + } else { + Array.empty[Filter] + } + } + + override def pushAggregation(aggregation: Aggregation): Boolean = { + if (!sparkSession.sessionState.conf.parquetAggregatePushDown) { + return false + } + + AggregatePushDownUtils.getSchemaForPushedAggregation( + aggregation, + schema, + partitionNameSet, + dataFilters) match { + + case Some(schema) => + finalSchema = schema + this.pushedAggregations = Some(aggregation) + true + case _ => false + } + } +} diff --git a/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/spark/sql/hive/kyuubi/connector/HiveBridgeHelper.scala b/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/spark/sql/hive/kyuubi/connector/HiveBridgeHelper.scala index 2993a0f12db..afe9a7a6d17 100644 --- a/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/spark/sql/hive/kyuubi/connector/HiveBridgeHelper.scala +++ b/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/spark/sql/hive/kyuubi/connector/HiveBridgeHelper.scala @@ -22,9 +22,14 @@ import scala.collection.mutable import org.apache.spark.SparkContext import org.apache.spark.sql.catalyst.catalog.{BucketSpec, ExternalCatalogEvent} import org.apache.spark.sql.catalyst.expressions.{AttributeReference, Literal} +import org.apache.spark.sql.catalyst.util.RebaseDateTime.RebaseSpec import org.apache.spark.sql.catalyst.util.quoteIfNeeded import org.apache.spark.sql.connector.expressions.{BucketTransform, FieldReference, IdentityTransform, Transform} import org.apache.spark.sql.connector.expressions.LogicalExpressions.{bucket, reference} +import org.apache.spark.sql.execution.datasources.orc.OrcFilters +import org.apache.spark.sql.execution.datasources.parquet.{ParquetFilters, SparkToParquetSchemaConverter} +import org.apache.spark.sql.internal.{LegacyBehaviorPolicy, SQLConf} +import org.apache.spark.sql.sources.Filter import org.apache.spark.sql.types.{DataType, DoubleType, FloatType, StructType} object HiveBridgeHelper { @@ -119,4 +124,38 @@ object HiveBridgeHelper { implicit class NamespaceHelper(namespace: Array[String]) { def quoted: String = namespace.map(quoteIfNeeded).mkString(".") } + + def orcConvertibleFilters( + schema: StructType, + caseSensitive: Boolean, + dataFilters: Seq[Filter]): Seq[Filter] = { + val dataTypeMap = OrcFilters.getSearchableTypeMap(schema, caseSensitive) + OrcFilters.convertibleFilters(dataTypeMap, dataFilters) + } + + def parquetConvertibleFilters( + readDataSchema: StructType, + dataFilters: Seq[Filter]): Seq[Filter] = { + val sqlConf = SQLConf.get + val pushDownDate = sqlConf.parquetFilterPushDownDate + val pushDownTimestamp = sqlConf.parquetFilterPushDownTimestamp + val pushDownDecimal = sqlConf.parquetFilterPushDownDecimal + val pushDownStringPredicate = sqlConf.parquetFilterPushDownStringPredicate + val pushDownInFilterThreshold = sqlConf.parquetFilterPushDownInFilterThreshold + val isCaseSensitive = sqlConf.caseSensitiveAnalysis + val parquetSchema = new SparkToParquetSchemaConverter(sqlConf).convert(readDataSchema) + val rebaseSpec = RebaseSpec(LegacyBehaviorPolicy.CORRECTED) + val parquetFilters = new ParquetFilters( + parquetSchema, + pushDownDate, + pushDownTimestamp, + pushDownDecimal, + pushDownStringPredicate, + pushDownInFilterThreshold, + isCaseSensitive, + // The rebase mode doesn't matter here because the filters are used to determine + // whether they is convertible. + rebaseSpec) + parquetFilters.convertibleFilters(dataFilters) + } } diff --git a/extensions/spark/kyuubi-spark-connector-hive/src/test/scala/org/apache/kyuubi/spark/connector/hive/DynamicPartitionPruningSuite.scala b/extensions/spark/kyuubi-spark-connector-hive/src/test/scala/org/apache/kyuubi/spark/connector/hive/DynamicPartitionPruningSuite.scala index 602c8cc1fc5..b6bcad70284 100644 --- a/extensions/spark/kyuubi-spark-connector-hive/src/test/scala/org/apache/kyuubi/spark/connector/hive/DynamicPartitionPruningSuite.scala +++ b/extensions/spark/kyuubi-spark-connector-hive/src/test/scala/org/apache/kyuubi/spark/connector/hive/DynamicPartitionPruningSuite.scala @@ -19,46 +19,51 @@ package org.apache.kyuubi.spark.connector.hive import scala.annotation.tailrec -import org.apache.spark.sql.{Row, SparkSession} +import org.apache.spark.sql.Row import org.apache.spark.sql.catalyst.expressions.DynamicPruningExpression +import org.apache.spark.sql.connector.read.Scan import org.apache.spark.sql.execution.SparkPlan import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec import org.apache.spark.sql.execution.datasources.v2.BatchScanExec +import org.apache.spark.sql.internal.SQLConf -import org.apache.kyuubi.spark.connector.hive.read.HiveScan +import org.apache.kyuubi.spark.connector.hive.read.{HiveScan, KyuubiOrcScan, KyuubiParquetScan} class DynamicPartitionPruningSuite extends KyuubiHiveTest { - private def findBatchScanExec( - spark: SparkSession, - sql: String, - tableNameHint: String): BatchScanExec = { - // Match on `HiveScan.catalogTable` rather than the node's `toString` because - // `BatchScanExec.toString` shape differs across Spark versions. - def matchesHint(b: BatchScanExec): Boolean = b.scan match { - case h: HiveScan => h.catalogTable.identifier.table == tableNameHint - case _ => false + private def findBatchScanExec(plan: SparkPlan, tableNameHint: String): BatchScanExec = { + // Match on the underlying Hive `catalogTable` rather than the node's `toString` + // because `BatchScanExec.toString` shape differs across Spark versions. + def hiveTableName(b: BatchScanExec): Option[String] = b.scan match { + case h: HiveScan => Some(h.catalogTable.identifier.table) + case o: KyuubiOrcScan => Some(o.catalogTable.identifier.table) + case p: KyuubiParquetScan => Some(p.catalogTable.identifier.table) + case _ => None } @tailrec - def findBatchScan(plan: SparkPlan): Option[BatchScanExec] = plan match { + def findBatchScan(p: SparkPlan): Option[BatchScanExec] = p match { case aqe: AdaptiveSparkPlanExec => findBatchScan(aqe.inputPlan) - case _ => plan.collectFirst { - case b: BatchScanExec if matchesHint(b) => b + case _ => p.collectFirst { + case b: BatchScanExec if hiveTableName(b).contains(tableNameHint) => b } } - val exec = findBatchScan(spark.sql(sql).queryExecution.executedPlan) + val exec = findBatchScan(plan) assert(exec.isDefined) exec.get } - test("HiveScan supports DPP runtime filtering on partition columns") { + private def runDppCase(storedAs: String): Unit = { + // Collect the number of input partitions actually planned under DPP on / off + // and later assert a strict reduction. + val plannedPartitions = scala.collection.mutable.Map.empty[Boolean, Int] + Seq(true, false).foreach { enabled => withSparkSession(Map( "hive.exec.dynamic.partition.mode" -> "nonstrict", - "spark.sql.optimizer.dynamicPartitionPruning.enabled" -> enabled.toString)) { spark => - val suffix = if (enabled) "on" else "off" + SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> enabled.toString)) { spark => + val suffix = s"${storedAs.toLowerCase}_${if (enabled) "on" else "off"}" val fact = s"hive.default.dpp_fact_$suffix" val dim = s"hive.default.dpp_dim_$suffix" @@ -66,8 +71,8 @@ class DynamicPartitionPruningSuite extends KyuubiHiveTest { spark.sql( s""" | CREATE TABLE $fact (id INT, v STRING) PARTITIONED BY (dt STRING) - | STORED AS TEXTFILE - |""".stripMargin).collect() + | STORED AS $storedAs + |""".stripMargin) spark.sql(s"INSERT INTO $fact PARTITION (dt='2026-01-01') VALUES (1, 'a'), (2, 'b')") spark.sql(s"INSERT INTO $fact PARTITION (dt='2026-05-01') VALUES (3, 'c'), (4, 'd')") spark.sql(s"INSERT INTO $fact PARTITION (dt='2026-09-01') VALUES (5, 'e'), (6, 'f')") @@ -75,30 +80,160 @@ class DynamicPartitionPruningSuite extends KyuubiHiveTest { spark.sql( s""" | CREATE TABLE $dim (dt STRING, tag STRING) - | STORED AS TEXTFILE - |""".stripMargin).collect() + | STORED AS $storedAs + |""".stripMargin) spark.sql(s"INSERT INTO $dim VALUES ('2026-05-01', 'target')") - val sql = + val df = spark.sql( s""" | SELECT f.id, f.v, f.dt | FROM $fact f JOIN $dim d ON f.dt = d.dt | WHERE d.tag = 'target' - |""".stripMargin - + |""".stripMargin) checkAnswer( - spark.sql(sql), + df, Seq( Row(3, "c", "2026-05-01"), Row(4, "d", "2026-05-01"))) // DPP being actually applied is observable as a `DynamicPruningExpression` // injected into `BatchScanExec.runtimeFilters`. - val exec = findBatchScanExec(spark, sql, fact.split('.').last) + val exec = findBatchScanExec(df.queryExecution.executedPlan, fact.split('.').last) val hasDpp = exec.runtimeFilters.exists(_.isInstanceOf[DynamicPruningExpression]) assert(hasDpp == enabled) + + val planned = exec.scan.toBatch.planInputPartitions().length + plannedPartitions(enabled) = planned + + exec.scan match { + case _: KyuubiOrcScan | _: KyuubiParquetScan => + assert(exec.scan.columnarSupportMode() == Scan.ColumnarSupportMode.SUPPORTED) + case _: HiveScan => + assert(exec.scan.columnarSupportMode() == Scan.ColumnarSupportMode.UNSUPPORTED) + case other => + fail(s"unexpected scan type: ${other.getClass.getName}") + } + } + } + } + + val planOn = plannedPartitions(true) + val planOff = plannedPartitions(false) + assert( + planOn < planOff, + s"DPP ($storedAs) should plan fewer partitions when enabled") + } + + test("HiveScan supports DPP runtime filtering on partition columns") { + runDppCase(storedAs = "TEXTFILE") + } + + test("KyuubiOrcScan supports DPP runtime filtering on partition columns") { + runDppCase(storedAs = "ORC") + } + + test("KyuubiParquetScan supports DPP runtime filtering on partition columns") { + runDppCase(storedAs = "PARQUET") + } + + /** + * Build a fact table and assert `columnarSupportMode()` matches + * `expectedColumnarMode` under the given `extraConf`. + */ + private def runColumnarModeCase( + storedAs: String, + extraConf: Map[String, String], + expectedColumnarMode: Scan.ColumnarSupportMode): Unit = { + withSparkSession(extraConf ++ Map( + "hive.exec.dynamic.partition.mode" -> "nonstrict")) { spark => + val fact = s"hive.default.mode_fact_${storedAs.toLowerCase}" + dropTableAfter(fact) { + spark.sql( + s""" + | CREATE TABLE $fact (id INT, v STRING) PARTITIONED BY (dt STRING) + | STORED AS $storedAs + |""".stripMargin) + spark.sql(s"INSERT INTO $fact PARTITION (dt='2026-05-01') VALUES (1, 'a')") + + val df = spark.sql(s"SELECT id, v, dt FROM $fact") + val exec = findBatchScanExec(df.queryExecution.executedPlan, fact.split('.').last) + exec.scan match { + case _: KyuubiOrcScan | _: KyuubiParquetScan => + assert(exec.scan.columnarSupportMode() == expectedColumnarMode) + case other => + fail(s"unexpected scan type: ${other.getClass.getName}") + } + } + } + } + + test("KyuubiOrcScan returns UNSUPPORTED when orc vectorized reader is disabled") { + runColumnarModeCase( + storedAs = "ORC", + extraConf = Map(SQLConf.ORC_VECTORIZED_READER_ENABLED.key -> "false"), + expectedColumnarMode = Scan.ColumnarSupportMode.UNSUPPORTED) + } + + test("KyuubiParquetScan returns UNSUPPORTED when parquet vectorized reader is disabled") { + runColumnarModeCase( + storedAs = "PARQUET", + extraConf = Map(SQLConf.PARQUET_VECTORIZED_READER_ENABLED.key -> "false"), + expectedColumnarMode = Scan.ColumnarSupportMode.UNSUPPORTED) + } + + private def runAllPrunedCase(storedAs: String): Unit = { + withSparkSession(Map( + "hive.exec.dynamic.partition.mode" -> "nonstrict", + SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "true")) { spark => + val suffix = storedAs.toLowerCase + val fact = s"hive.default.pruned_fact_$suffix" + val dim = s"hive.default.pruned_dim_$suffix" + + dropTableAfter(fact, dim) { + spark.sql( + s""" + | CREATE TABLE $fact (id INT, v STRING) PARTITIONED BY (dt STRING) + | STORED AS $storedAs + |""".stripMargin) + spark.sql(s"INSERT INTO $fact PARTITION (dt='2026-01-01') VALUES (1, 'a')") + spark.sql(s"INSERT INTO $fact PARTITION (dt='2026-05-01') VALUES (2, 'b')") + + spark.sql( + s""" + | CREATE TABLE $dim (dt STRING, tag STRING) + | STORED AS $storedAs + |""".stripMargin) + // Dim key matches no fact partition, so DPP prunes every fact partition. + spark.sql(s"INSERT INTO $dim VALUES ('1999-12-31', 'target')") + + val df = spark.sql( + s""" + | SELECT f.id, f.v, f.dt + | FROM $fact f JOIN $dim d ON f.dt = d.dt + | WHERE d.tag = 'target' + |""".stripMargin) + // Trigger full execution so `BatchScanExec.filteredPartitions` pushes + // runtime filters into the wrapped scan (via `SupportsRuntimeFiltering`) + // and `planInputPartitions()` below reflects DPP-pruned partitions. + assert(df.collect().isEmpty) + + val exec = findBatchScanExec(df.queryExecution.executedPlan, fact.split('.').last) + exec.scan match { + case _: KyuubiOrcScan | _: KyuubiParquetScan => + assert(exec.scan.toBatch.planInputPartitions().isEmpty) + assert(exec.scan.columnarSupportMode() == Scan.ColumnarSupportMode.SUPPORTED) + case other => + fail(s"unexpected scan type: ${other.getClass.getName}") } } } } + + test("KyuubiOrcScan returns SUPPORTED when DPP prunes every partition") { + runAllPrunedCase(storedAs = "ORC") + } + + test("KyuubiParquetScan returns SUPPORTED when DPP prunes every partition") { + runAllPrunedCase(storedAs = "PARQUET") + } } diff --git a/extensions/spark/kyuubi-spark-connector-hive/src/test/scala/org/apache/kyuubi/spark/connector/hive/HiveCatalogSuite.scala b/extensions/spark/kyuubi-spark-connector-hive/src/test/scala/org/apache/kyuubi/spark/connector/hive/HiveCatalogSuite.scala index 0e276bdef71..bd39844c620 100644 --- a/extensions/spark/kyuubi-spark-connector-hive/src/test/scala/org/apache/kyuubi/spark/connector/hive/HiveCatalogSuite.scala +++ b/extensions/spark/kyuubi-spark-connector-hive/src/test/scala/org/apache/kyuubi/spark/connector/hive/HiveCatalogSuite.scala @@ -32,15 +32,14 @@ import org.apache.spark.sql.catalyst.catalog.CatalogTableType import org.apache.spark.sql.catalyst.parser.CatalystSqlParser import org.apache.spark.sql.connector.catalog.{Identifier, SupportsNamespaces, TableCatalog} import org.apache.spark.sql.connector.expressions.Transform -import org.apache.spark.sql.execution.datasources.v2.orc.OrcScan -import org.apache.spark.sql.execution.datasources.v2.parquet.ParquetScan import org.apache.spark.sql.hive.kyuubi.connector.HiveBridgeHelper._ +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{IntegerType, StringType, StructType} import org.apache.spark.sql.util.CaseInsensitiveStringMap import org.apache.kyuubi.spark.connector.hive.HiveTableCatalog.IdentifierHelper import org.apache.kyuubi.spark.connector.hive.KyuubiHiveConnectorConf.{READ_CONVERT_METASTORE_ORC, READ_CONVERT_METASTORE_PARQUET} -import org.apache.kyuubi.spark.connector.hive.read.HiveScan +import org.apache.kyuubi.spark.connector.hive.read.{HiveScan, KyuubiOrcScan, KyuubiParquetScan} class HiveCatalogSuite extends KyuubiHiveTest { @@ -509,9 +508,9 @@ class HiveCatalogSuite extends KyuubiHiveTest { val parScan = value match { case "true" => assert( - scan.isInstanceOf[ParquetScan], - s"Expected ParquetScan, got ${scan.getClass.getSimpleName}") - scan.asInstanceOf[ParquetScan] + scan.isInstanceOf[KyuubiParquetScan], + s"Expected KyuubiParquetScan, got ${scan.getClass.getSimpleName}") + scan.asInstanceOf[KyuubiParquetScan] case "false" => assert( scan.isInstanceOf[HiveScan], @@ -538,9 +537,9 @@ class HiveCatalogSuite extends KyuubiHiveTest { val orcScan = value match { case "true" => assert( - scan.isInstanceOf[OrcScan], - s"Expected OrcScan, got ${scan.getClass.getSimpleName}") - scan.asInstanceOf[OrcScan] + scan.isInstanceOf[KyuubiOrcScan], + s"Expected KyuubiOrcScan, got ${scan.getClass.getSimpleName}") + scan.asInstanceOf[KyuubiOrcScan] case "false" => assert( scan.isInstanceOf[HiveScan], @@ -554,4 +553,40 @@ class HiveCatalogSuite extends KyuubiHiveTest { } } } + + test("KyuubiParquetScan and KyuubiOrcScan pick up mid-session confs and catalog overlay") { + // Confs set in the session must take effect for the reader. + val fieldIdRead = SQLConf.PARQUET_FIELD_ID_READ_ENABLED.key + val fieldIdIgnoreMissing = SQLConf.IGNORE_MISSING_PARQUET_FIELD_ID.key + // Injected in `newCatalog` as camelCase, overlay writes it as lowercase. + val overlayKey = "javax.jdo.option.connectionurl" + + Seq("orc", "parquet").foreach { provider => + withSparkSession() { spark => + val props = new util.HashMap[String, String]() + props.put(TableCatalog.PROP_PROVIDER, provider) + val ident = Identifier.of(testNs, s"scan_conf_$provider") + + try { + val table = catalog.createTable(ident, schema, Array.empty[Transform], props) + spark.sessionState.conf.setConfString(fieldIdRead, "true") + spark.sessionState.conf.setConfString(fieldIdIgnoreMissing, "true") + + val hadoopConf = table.asInstanceOf[HiveTable] + .newScanBuilder(CaseInsensitiveStringMap.empty()).build() match { + case s: KyuubiParquetScan => s.hadoopConf + case s: KyuubiOrcScan => s.hadoopConf + case other => fail(s"unexpected scan type: ${other.getClass.getName}") + } + // Mid-session confs reach the reader (no snapshot freeze). + assert(hadoopConf.get(fieldIdRead) == "true") + assert(hadoopConf.get(fieldIdIgnoreMissing) == "true") + // Per-catalog overlay applied (lowercase key). + assert(hadoopConf.get(overlayKey) != null) + } finally { + catalog.dropTable(ident) + } + } + } + } }