From 007194d04594210cf9dee47bfada052bdbd9016e Mon Sep 17 00:00:00 2001 From: lifumao Date: Sat, 15 Aug 2026 08:35:43 +0800 Subject: [PATCH 1/4] [KYUUBI #6943][2/2] OrcScan and ParquetScan support DPP --- .../spark/connector/hive/HiveTable.scala | 22 +- .../spark/connector/hive/read/HiveScan.scala | 23 +- .../connector/hive/read/KyuubiOrcScan.scala | 171 ++++++++++++ .../hive/read/KyuubiOrcScanBuilder.scala | 130 +++++++++ .../hive/read/KyuubiParquetScan.scala | 261 ++++++++++++++++++ .../hive/read/KyuubiParquetScanBuilder.scala | 131 +++++++++ .../kyuubi/connector/HiveBridgeHelper.scala | 37 +++ .../hive/DynamicPartitionPruningSuite.scala | 82 ++++-- .../connector/hive/HiveCatalogSuite.scala | 53 +++- 9 files changed, 872 insertions(+), 38 deletions(-) create mode 100644 extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/read/KyuubiOrcScan.scala create mode 100644 extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/read/KyuubiOrcScanBuilder.scala create mode 100644 extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/read/KyuubiParquetScan.scala create mode 100644 extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/read/KyuubiParquetScanBuilder.scala 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/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..604a74cbb73 --- /dev/null +++ b/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/read/KyuubiOrcScanBuilder.scala @@ -0,0 +1,130 @@ +/* + * 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 { + + /** + * Starts from `hiveTableCatalog.hadoopConfiguration()` so per-catalog Hadoop + * settings are honored. Cloned so per-scan `options` do not pollute the + * shared catalog instance. + * + * Note: the catalog's `hadoopConfiguration()` is a `lazy val` snapshot of + * `sessionState.newHadoopConf()` taken at first use. Session confs changed + * after that snapshot may not reach reader code that reads straight from + * the Hadoop `Configuration`. + */ + lazy val hadoopConf: Configuration = { + val conf = new Configuration(hiveTableCatalog.hadoopConfiguration()) + // 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..e7da1cd281b --- /dev/null +++ b/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/read/KyuubiParquetScan.scala @@ -0,0 +1,261 @@ +/* + * 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] + + // 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 = { + if (variantExtractionCls != null) { + // Spark 4.1+ + DynConstructors.builder() + .impl( + 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() + .invokeChecked[ParquetScan]( + null, + sparkSession, + hadoopConf, + fileIndex, + dataSchema, + readDataSchema, + readPartitionSchema, + pushedFilters, + options, + pushedAggregate, + partitionFilters, + dataFilters, + emptyVariantExtractions) + } else { + // Spark 4.0 and previous + DynConstructors.builder() + .impl( + 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() + .invokeChecked[ParquetScan]( + null, + sparkSession, + hadoopConf, + fileIndex, + dataSchema, + readDataSchema, + readPartitionSchema, + pushedFilters, + options, + pushedAggregate, + partitionFilters, + dataFilters) + } + } + // 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..d3590aca72f --- /dev/null +++ b/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/read/KyuubiParquetScanBuilder.scala @@ -0,0 +1,131 @@ +/* + * 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 { + + /** + * Starts from `hiveTableCatalog.hadoopConfiguration()` so per-catalog Hadoop + * settings are honored. Cloned so per-scan `options` do not pollute the + * shared catalog instance. + * + * Note: the catalog's `hadoopConfiguration()` is a `lazy val` snapshot of + * `sessionState.newHadoopConf()` taken at first use. Session confs changed + * after that snapshot may not reach reader code that reads straight from + * the Hadoop `Configuration`. + */ + lazy val hadoopConf: Configuration = { + val conf = new Configuration(hiveTableCatalog.hadoopConfiguration()) + // 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..7045594c3cb 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,36 @@ 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, + 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..a29ac6888e8 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,50 @@ 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.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" + 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 +70,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,8 +79,8 @@ 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 = @@ -86,19 +90,51 @@ class DynamicPartitionPruningSuite extends KyuubiHiveTest { | WHERE d.tag = 'target' |""".stripMargin + val df = spark.sql(sql) 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," + + s" got on=$planOn off=$planOff") + } + + 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") } } 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..648b8a114b5 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 hadoopConf snapshots at catalog first use") { + // Unlike Spark's built-in ScanBuilders rebuilt per query planning from + // the current sqlConf, KSHC clones a `lazy val` snapshot taken at the + // catalog's first use, so confs SET after that snapshot cannot reach + // readers that consult the Hadoop `Configuration` directly. + val fieldIdRead = SQLConf.PARQUET_FIELD_ID_READ_ENABLED.key + val fieldIdIgnoreMissing = SQLConf.IGNORE_MISSING_PARQUET_FIELD_ID.key + + Seq("orc", "parquet").foreach { provider => + withSparkSession() { spark => + // `catalog` is built in `beforeEach`, its `hadoopConf` snapshot is + // already frozen by the time this test runs, so any conf we SET below + // is "mid-session" from its perspective. + val props = new util.HashMap[String, String]() + props.put(TableCatalog.PROP_PROVIDER, provider) + + try { + val table = catalog.createTable(testIdent, 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}") + } + assert(hadoopConf.get(fieldIdRead) != "true") + assert(hadoopConf.get(fieldIdIgnoreMissing) != "true") + } finally { + catalog.dropTable(testIdent) + } + } + } + } } From 7cc6b0508c07df52a9d96f1bb2e3198c14318397 Mon Sep 17 00:00:00 2001 From: lifumao Date: Mon, 17 Aug 2026 11:54:29 +0800 Subject: [PATCH 2/4] Cache the Ctor in a lazy val --- .../hive/read/KyuubiParquetScan.scala | 80 ++++++++----------- 1 file changed, 34 insertions(+), 46 deletions(-) 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 index e7da1cd281b..c517f54bc19 100644 --- 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 @@ -180,23 +180,10 @@ object KyuubiParquetScan { if (variantExtractionCls == null) null else java.lang.reflect.Array.newInstance(variantExtractionCls, 0).asInstanceOf[AnyRef] - // 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 = { + private lazy val parquetScanCtor: DynConstructors.Ctor[ParquetScan] = if (variantExtractionCls != null) { - // Spark 4.1+ DynConstructors.builder() - .impl( + .impl( // SPARK-53880 / SPARK-54656 (4.1.0): adds trailing Array[VariantExtraction] classOf[ParquetScan], classOf[SparkSession], classOf[Configuration], @@ -210,25 +197,10 @@ object KyuubiParquetScan { classOf[Seq[Expression]], classOf[Seq[Expression]], emptyVariantExtractions.getClass) - .buildChecked() - .invokeChecked[ParquetScan]( - null, - sparkSession, - hadoopConf, - fileIndex, - dataSchema, - readDataSchema, - readPartitionSchema, - pushedFilters, - options, - pushedAggregate, - partitionFilters, - dataFilters, - emptyVariantExtractions) + .buildChecked[ParquetScan]() } else { - // Spark 4.0 and previous DynConstructors.builder() - .impl( + .impl( // Spark 4.0 and previous classOf[ParquetScan], classOf[SparkSession], classOf[Configuration], @@ -241,21 +213,37 @@ object KyuubiParquetScan { classOf[Option[Aggregation]], classOf[Seq[Expression]], classOf[Seq[Expression]]) - .buildChecked() - .invokeChecked[ParquetScan]( - null, - sparkSession, - hadoopConf, - fileIndex, - dataSchema, - readDataSchema, - readPartitionSchema, - pushedFilters, - options, - pushedAggregate, - partitionFilters, - dataFilters) + .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 } From 21242e24d8fdee1423fc8735cb55295f58149121 Mon Sep 17 00:00:00 2001 From: lifumao Date: Tue, 18 Aug 2026 23:12:29 +0800 Subject: [PATCH 3/4] add comments --- .../spark/sql/hive/kyuubi/connector/HiveBridgeHelper.scala | 2 ++ 1 file changed, 2 insertions(+) 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 7045594c3cb..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 @@ -153,6 +153,8 @@ object HiveBridgeHelper { 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) } From b0005139fd4b98163ebbd093f20719d43f6bbf84 Mon Sep 17 00:00:00 2001 From: lifumao Date: Wed, 19 Aug 2026 12:41:38 +0800 Subject: [PATCH 4/4] Fix UT --- .../connector/hive/HiveTableCatalog.scala | 13 +- .../hive/read/KyuubiOrcScanBuilder.scala | 14 +-- .../hive/read/KyuubiParquetScanBuilder.scala | 14 +-- .../hive/DynamicPartitionPruningSuite.scala | 113 ++++++++++++++++-- .../connector/hive/HiveCatalogSuite.scala | 24 ++-- 5 files changed, 140 insertions(+), 38 deletions(-) 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/KyuubiOrcScanBuilder.scala b/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/read/KyuubiOrcScanBuilder.scala index 604a74cbb73..e4a2e457567 100644 --- 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 @@ -54,17 +54,13 @@ class KyuubiOrcScanBuilder( with SupportsPushDownAggregates { /** - * Starts from `hiveTableCatalog.hadoopConfiguration()` so per-catalog Hadoop - * settings are honored. Cloned so per-scan `options` do not pollute the - * shared catalog instance. - * - * Note: the catalog's `hadoopConfiguration()` is a `lazy val` snapshot of - * `sessionState.newHadoopConf()` taken at first use. Session confs changed - * after that snapshot may not reach reader code that reads straight from - * the Hadoop `Configuration`. + * 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.hadoopConfiguration()) + val conf = new Configuration(hiveTableCatalog.newScanHadoopConf()) // Hadoop Configurations are case sensitive. options.asCaseSensitiveMap.asScala.foreach { case (k, v) => conf.set(k, v) } conf 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 index d3590aca72f..5deac91f413 100644 --- 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 @@ -58,17 +58,13 @@ class KyuubiParquetScanBuilder( with SupportsPushDownAggregates { /** - * Starts from `hiveTableCatalog.hadoopConfiguration()` so per-catalog Hadoop - * settings are honored. Cloned so per-scan `options` do not pollute the - * shared catalog instance. - * - * Note: the catalog's `hadoopConfiguration()` is a `lazy val` snapshot of - * `sessionState.newHadoopConf()` taken at first use. Session confs changed - * after that snapshot may not reach reader code that reads straight from - * the Hadoop `Configuration`. + * 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.hadoopConfiguration()) + val conf = new Configuration(hiveTableCatalog.newScanHadoopConf()) // Hadoop Configurations are case sensitive. options.asCaseSensitiveMap.asScala.foreach { case (k, v) => conf.set(k, v) } conf 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 a29ac6888e8..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 @@ -25,6 +25,7 @@ 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, KyuubiOrcScan, KyuubiParquetScan} @@ -61,7 +62,7 @@ class DynamicPartitionPruningSuite extends KyuubiHiveTest { Seq(true, false).foreach { enabled => withSparkSession(Map( "hive.exec.dynamic.partition.mode" -> "nonstrict", - "spark.sql.optimizer.dynamicPartitionPruning.enabled" -> enabled.toString)) { spark => + 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" @@ -83,14 +84,12 @@ class DynamicPartitionPruningSuite extends KyuubiHiveTest { |""".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 - - val df = spark.sql(sql) + |""".stripMargin) checkAnswer( df, Seq( @@ -122,8 +121,7 @@ class DynamicPartitionPruningSuite extends KyuubiHiveTest { val planOff = plannedPartitions(false) assert( planOn < planOff, - s"DPP ($storedAs) should plan fewer partitions when enabled," + - s" got on=$planOn off=$planOff") + s"DPP ($storedAs) should plan fewer partitions when enabled") } test("HiveScan supports DPP runtime filtering on partition columns") { @@ -137,4 +135,105 @@ class DynamicPartitionPruningSuite extends KyuubiHiveTest { 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 648b8a114b5..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 @@ -554,24 +554,21 @@ class HiveCatalogSuite extends KyuubiHiveTest { } } - test("KyuubiParquetScan and KyuubiOrcScan hadoopConf snapshots at catalog first use") { - // Unlike Spark's built-in ScanBuilders rebuilt per query planning from - // the current sqlConf, KSHC clones a `lazy val` snapshot taken at the - // catalog's first use, so confs SET after that snapshot cannot reach - // readers that consult the Hadoop `Configuration` directly. + 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 => - // `catalog` is built in `beforeEach`, its `hadoopConf` snapshot is - // already frozen by the time this test runs, so any conf we SET below - // is "mid-session" from its perspective. 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(testIdent, schema, Array.empty[Transform], props) + val table = catalog.createTable(ident, schema, Array.empty[Transform], props) spark.sessionState.conf.setConfString(fieldIdRead, "true") spark.sessionState.conf.setConfString(fieldIdIgnoreMissing, "true") @@ -581,10 +578,13 @@ class HiveCatalogSuite extends KyuubiHiveTest { case s: KyuubiOrcScan => s.hadoopConf case other => fail(s"unexpected scan type: ${other.getClass.getName}") } - assert(hadoopConf.get(fieldIdRead) != "true") - assert(hadoopConf.get(fieldIdIgnoreMissing) != "true") + // 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(testIdent) + catalog.dropTable(ident) } } }