+ * When cost-based optimization or plan statistics are disabled, Spark primarily needs the scan's + * size in bytes (for example, for broadcast-join thresholding). This method lets connectors serve + * that size estimate cheaply and avoid computing the full statistics. The default implementation + * delegates to {@link #estimateStatistics()} and returns its {@code sizeInBytes()}, so connectors + * that already compute statistics cheaply do not need to override this method. + * + * @since 4.3.0 + */ + default OptionalLong estimateSizeInBytes() { + Statistics statistics = estimateStatistics(); + return statistics != null ? statistics.sizeInBytes() : OptionalLong.empty(); + } + + /** + * Returns whether the statistics reported by this scan already reflect all filters that were + * fully pushed down to the data source. + *
+ * When {@code true} (the default), the reported statistics describe exactly the data the scan + * will produce. When {@code false}, they do not account for the fully pushed filters + * (for example, they describe the whole table), so Spark may use those fully pushed filters to + * adjust stats. Re-applying those fully pushed filters in Spark should be redundant for query + * results because the data source already evaluates them. + * + * @since 4.3.0 + */ + default boolean reflectsFullyPushedDownFilters() { + return true; + } } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Relation.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Relation.scala index 0b7c829939ff7..7dd5907737089 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Relation.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Relation.scala @@ -25,6 +25,7 @@ import org.apache.spark.sql.catalyst.catalog.{CatalogColumnStat, CatalogStatisti import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeMap, AttributeReference, AttributeSet, Expression, SortOrder, V2ExpressionUtils} import org.apache.spark.sql.catalyst.plans.QueryPlan import org.apache.spark.sql.catalyst.plans.logical.{ColumnStat, ExposesMetadataColumns, Histogram, HistogramBin, LeafNode, LogicalPlan, Statistics} +import org.apache.spark.sql.catalyst.plans.logical.statsEstimation.EstimationUtils import org.apache.spark.sql.catalyst.streaming.{StreamingSourceIdentifyingName, Unassigned} import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes import org.apache.spark.sql.catalyst.util.{removeInternalMetadata, truncatedString, CharVarcharUtils} @@ -34,6 +35,7 @@ import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReferenc import org.apache.spark.sql.connector.read.{Scan, Statistics => V2Statistics, SupportsReportStatistics, SupportsRuntimeV2Filtering} import org.apache.spark.sql.connector.read.colstats.{ColumnStatistics, Histogram => V2Histogram, HistogramBin => V2HistogramBin} import org.apache.spark.sql.connector.read.streaming.{Offset, SparkDataStream} +import org.apache.spark.sql.internal.connector.V2StatisticsUtils import org.apache.spark.sql.types.{DataType, StructType} import org.apache.spark.sql.util.CaseInsensitiveStringMap import org.apache.spark.util.ArrayImplicits._ @@ -95,7 +97,7 @@ abstract class DataSourceV2RelationBase( table.asReadable.newScanBuilder(options).build() match { case r: SupportsReportStatistics => val statistics = r.estimateStatistics() - DataSourceV2Relation.transformV2Stats(statistics, None, conf.defaultSizeInBytes, output) + DataSourceV2Relation.transformV2Stats(statistics, conf.defaultSizeInBytes, output) case _ => Statistics(sizeInBytes = conf.defaultSizeInBytes) } @@ -199,15 +201,32 @@ case class DataSourceV2ScanRelation( } override def computeStats(): Statistics = { - scan match { - case r: SupportsReportStatistics => - val statistics = r.estimateStatistics() - DataSourceV2Relation.transformV2Stats(statistics, None, conf.defaultSizeInBytes, output) - case _ => - Statistics(sizeInBytes = conf.defaultSizeInBytes) + if (conf.cboEnabled || conf.planStatsEnabled) { + computeFullStats() + } else { + computeSizeOnlyStats() + } + } + + private def computeFullStats(): Statistics = { + V2StatisticsUtils.computeStats(scan) match { + case Some(v2Stats) => + DataSourceV2Relation.transformV2Stats(v2Stats, conf.defaultSizeInBytes, output) + case _ => defaultSizeOnlyStats + } + } + + private def computeSizeOnlyStats(): Statistics = { + V2StatisticsUtils.computeSizeInBytes(scan, EstimationUtils.getSizePerRow(output)) match { + case Some(sizeInBytes) => Statistics(sizeInBytes = sizeInBytes) + case _ => defaultSizeOnlyStats } } + private def defaultSizeOnlyStats: Statistics = { + Statistics(sizeInBytes = conf.defaultSizeInBytes) + } + override def doCanonicalize(): DataSourceV2ScanRelation = { this.copy( relation = this.relation.copy( @@ -276,7 +295,7 @@ case class StreamingDataSourceV2ScanRelation( override def computeStats(): Statistics = scan match { case r: SupportsReportStatistics => val statistics = r.estimateStatistics() - DataSourceV2Relation.transformV2Stats(statistics, None, conf.defaultSizeInBytes, output) + DataSourceV2Relation.transformV2Stats(statistics, conf.defaultSizeInBytes, output) case _ => Statistics(sizeInBytes = conf.defaultSizeInBytes) } @@ -419,13 +438,12 @@ object DataSourceV2Relation { */ def transformV2Stats( v2Statistics: V2Statistics, - defaultRowCount: Option[BigInt], defaultSizeInBytes: Long, output: Seq[Attribute] = Seq.empty): Statistics = { val numRows: Option[BigInt] = if (v2Statistics.numRows().isPresent) { Some(v2Statistics.numRows().getAsLong) } else { - defaultRowCount + None } var colStats: Seq[(Attribute, ColumnStat)] = Seq.empty[(Attribute, ColumnStat)] @@ -463,9 +481,20 @@ object DataSourceV2Relation { }) }) } + val attributeStats = AttributeMap(colStats) + // Prefer the source-reported size. Otherwise infer a projection-aware size from the row count + // (numRows * outputRowSize via getOutputSize). Fall back to the default size when neither is + // available. + val sizeInBytes = if (v2Statistics.sizeInBytes().isPresent) { + BigInt(v2Statistics.sizeInBytes().getAsLong) + } else if (numRows.isDefined) { + EstimationUtils.getOutputSize(output, numRows.get, attributeStats) + } else { + BigInt(defaultSizeInBytes) + } Statistics( - sizeInBytes = v2Statistics.sizeInBytes().orElse(defaultSizeInBytes), + sizeInBytes = sizeInBytes, rowCount = numRows, - attributeStats = AttributeMap(colStats)) + attributeStats = attributeStats) } } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/V2StatisticsUtils.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/V2StatisticsUtils.scala new file mode 100644 index 0000000000000..514eceb9a8f17 --- /dev/null +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/V2StatisticsUtils.scala @@ -0,0 +1,62 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.internal.connector + +import java.util.OptionalLong + +import org.apache.spark.sql.connector.read.{Scan, Statistics, SupportsReportStatistics} + +object V2StatisticsUtils { + + def isNotEmpty(stats: Statistics): Boolean = { + stats != null && hasAnyValue(stats) + } + + private def hasAnyValue(stats: Statistics): Boolean = { + stats.sizeInBytes().isPresent || + stats.numRows().isPresent || + (stats.columnStats() != null && !stats.columnStats().isEmpty) + } + + def computeStats(scan: Scan): Option[Statistics] = scan match { + case s: SupportsReportStatistics => Some(s.estimateStatistics()).filter(isNotEmpty) + case _ => None + } + + def computeSizeInBytes( + scan: Scan, + avgRowSize: => BigInt): Option[BigInt] = { + extractSizeInBytes(scan).orElse(extractRowCount(scan).map(_ * avgRowSize)) + } + + private def extractSizeInBytes(scan: Scan): Option[BigInt] = scan match { + case s: SupportsReportStatistics => toBigInt(s.estimateSizeInBytes()) + case _ => None + } + + private def extractRowCount(scan: Scan): Option[BigInt] = scan match { + case s: SupportsReportStatistics => + Option(s.estimateStatistics()).flatMap(stats => toBigInt(stats.numRows())) + case _ => + None + } + + private def toBigInt(value: OptionalLong): Option[BigInt] = { + if (value.isPresent) Some(BigInt(value.getAsLong)) else None + } +} diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2RelationSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2RelationSuite.scala index 46ed870cb3221..f82395a7d9f23 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2RelationSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2RelationSuite.scala @@ -18,18 +18,25 @@ package org.apache.spark.sql.execution.datasources.v2 import java.util +import java.util.OptionalLong import org.apache.spark.SparkFunSuite import org.apache.spark.sql.catalyst.catalog.{CatalogColumnStat, CatalogStatistics} +import org.apache.spark.sql.catalyst.expressions.AttributeReference +import org.apache.spark.sql.catalyst.plans.SQLHelper import org.apache.spark.sql.catalyst.plans.logical.{Histogram, HistogramBin} +import org.apache.spark.sql.catalyst.plans.logical.statsEstimation.EstimationUtils import org.apache.spark.sql.catalyst.util.FieldMetadataUtils.FIELD_ID_METADATA_KEY import org.apache.spark.sql.catalyst.util.INTERNAL_METADATA_KEYS import org.apache.spark.sql.connector.catalog.{Column, Table, TableCapability} -import org.apache.spark.sql.connector.expressions.FieldReference +import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference} +import org.apache.spark.sql.connector.read.{Scan, Statistics => V2Statistics, SupportsReportStatistics} +import org.apache.spark.sql.connector.read.colstats.ColumnStatistics +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{IntegerType, MetadataBuilder, StringType, StructField, StructType} import org.apache.spark.sql.util.CaseInsensitiveStringMap -class DataSourceV2RelationSuite extends SparkFunSuite { +class DataSourceV2RelationSuite extends SparkFunSuite with SQLHelper { test("DataSourceV2Relation.v1StatsToV2Stats") { val schema = StructType(Seq( @@ -114,6 +121,293 @@ class DataSourceV2RelationSuite extends SparkFunSuite { assert(!idV2NoHist.histogram().isPresent) } + private def scanRel( + output: Seq[AttributeReference], + scan: Scan, + table: Table = new FakeTableWithSchema()): DataSourceV2ScanRelation = { + DataSourceV2ScanRelation( + DataSourceV2Relation(table, output, None, None, CaseInsensitiveStringMap.empty()), + scan, + output) + } + + test("DataSourceV2ScanRelation.computeStats uses non-empty scan stats with CBO") { + val idAttr = AttributeReference("id", IntegerType)() + val output = Seq(idAttr) + val scan = new Scan with SupportsReportStatistics { + override def readSchema(): StructType = StructType(Seq(StructField("id", IntegerType))) + override def estimateStatistics(): V2Statistics = new V2Statistics { + override def sizeInBytes(): OptionalLong = OptionalLong.empty() + override def numRows(): OptionalLong = OptionalLong.of(42L) + override def columnStats(): java.util.Map[NamedReference, ColumnStatistics] = { + val stats = new java.util.HashMap[NamedReference, ColumnStatistics]() + stats.put(FieldReference.column("id"), new ColumnStatistics { + override def distinctCount(): OptionalLong = OptionalLong.of(40L) + override def avgLen(): OptionalLong = OptionalLong.of(4L) + }) + stats + } + } + } + + withSQLConf(SQLConf.CBO_ENABLED.key -> "true") { + val stats = scanRel(output, scan).computeStats() + + assert(stats.rowCount.contains(BigInt(42))) + assert(stats.attributeStats.size === 1) + assert(stats.attributeStats(idAttr).distinctCount.contains(BigInt(40))) + assert(stats.attributeStats(idAttr).avgLen.contains(4L)) + assert(stats.sizeInBytes === + EstimationUtils.getOutputSize(output, BigInt(42), stats.attributeStats)) + } + } + + test("DataSourceV2ScanRelation.computeStats derives size 1 for a zero-row scan with CBO") { + val idAttr = AttributeReference("id", IntegerType)() + val output = Seq(idAttr) + val scan = new Scan with SupportsReportStatistics { + override def readSchema(): StructType = StructType(Seq(StructField("id", IntegerType))) + override def estimateStatistics(): V2Statistics = new V2Statistics { + override def sizeInBytes(): OptionalLong = OptionalLong.empty() + override def numRows(): OptionalLong = OptionalLong.of(0L) + } + } + + withSQLConf( + SQLConf.CBO_ENABLED.key -> "true", + SQLConf.DEFAULT_SIZE_IN_BYTES.key -> "12345") { + val stats = scanRel(output, scan).computeStats() + + assert(stats.rowCount.contains(BigInt(0))) + assert(stats.sizeInBytes === BigInt(1)) + } + } + + test("DataSourceV2ScanRelation.computeStats uses full stats with plan stats enabled") { + val idAttr = AttributeReference("id", IntegerType)() + val output = Seq(idAttr) + val scan = new Scan with SupportsReportStatistics { + override def readSchema(): StructType = StructType(Seq(StructField("id", IntegerType))) + override def estimateSizeInBytes(): OptionalLong = OptionalLong.of(50L) + override def estimateStatistics(): V2Statistics = new V2Statistics { + override def sizeInBytes(): OptionalLong = OptionalLong.of(1000L) + override def numRows(): OptionalLong = OptionalLong.of(7L) + } + } + + withSQLConf( + SQLConf.CBO_ENABLED.key -> "false", + SQLConf.PLAN_STATS_ENABLED.key -> "true") { + val stats = scanRel(output, scan).computeStats() + + assert(stats.sizeInBytes === BigInt(1000)) + assert(stats.rowCount.contains(BigInt(7))) + } + } + + test("DataSourceV2ScanRelation.computeStats treats column-only scan stats as non-empty") { + val idAttr = AttributeReference("id", IntegerType)() + val output = Seq(idAttr) + val scan = new Scan with SupportsReportStatistics { + override def readSchema(): StructType = StructType(Seq(StructField("id", IntegerType))) + override def estimateStatistics(): V2Statistics = new V2Statistics { + override def sizeInBytes(): OptionalLong = OptionalLong.empty() + override def numRows(): OptionalLong = OptionalLong.empty() + override def columnStats(): java.util.Map[NamedReference, ColumnStatistics] = { + val stats = new java.util.HashMap[NamedReference, ColumnStatistics]() + stats.put(FieldReference.column("id"), new ColumnStatistics { + override def distinctCount(): OptionalLong = OptionalLong.of(7L) + }) + stats + } + } + } + + withSQLConf( + SQLConf.CBO_ENABLED.key -> "true", + SQLConf.DEFAULT_SIZE_IN_BYTES.key -> "12345") { + val stats = scanRel(output, scan).computeStats() + + assert(stats.sizeInBytes === BigInt(12345)) + assert(stats.rowCount.isEmpty) + assert(stats.attributeStats.size === 1) + assert(stats.attributeStats(idAttr).distinctCount.contains(BigInt(7))) + } + } + + test("DataSourceV2ScanRelation.computeStats uses size-only estimates without CBO") { + val idAttr = AttributeReference("id", IntegerType)() + val output = Seq(idAttr) + val scan = new Scan with SupportsReportStatistics { + override def readSchema(): StructType = StructType(Seq(StructField("id", IntegerType))) + override def estimateSizeInBytes(): OptionalLong = OptionalLong.of(50L) + override def estimateStatistics(): V2Statistics = new V2Statistics { + override def sizeInBytes(): OptionalLong = OptionalLong.of(1000L) + override def numRows(): OptionalLong = OptionalLong.of(5L) + } + } + + withSQLConf( + SQLConf.CBO_ENABLED.key -> "false", + SQLConf.PLAN_STATS_ENABLED.key -> "false") { + val stats = scanRel(output, scan).computeStats() + + assert(stats.sizeInBytes === BigInt(50)) + assert(stats.rowCount.isEmpty) + assert(stats.attributeStats.isEmpty) + } + } + + test("DataSourceV2ScanRelation.computeStats uses default estimateSizeInBytes without CBO") { + val idAttr = AttributeReference("id", IntegerType)() + val output = Seq(idAttr) + val scan = new Scan with SupportsReportStatistics { + override def readSchema(): StructType = StructType(Seq(StructField("id", IntegerType))) + override def estimateStatistics(): V2Statistics = new V2Statistics { + override def sizeInBytes(): OptionalLong = OptionalLong.of(64L) + override def numRows(): OptionalLong = OptionalLong.of(5L) + } + } + + withSQLConf( + SQLConf.CBO_ENABLED.key -> "false", + SQLConf.PLAN_STATS_ENABLED.key -> "false") { + val stats = scanRel(output, scan).computeStats() + + assert(stats.sizeInBytes === BigInt(64)) + assert(stats.rowCount.isEmpty) + assert(stats.attributeStats.isEmpty) + } + } + + test("DataSourceV2ScanRelation.computeStats infers size-only estimates from row count") { + val idAttr = AttributeReference("id", IntegerType)() + val output = Seq(idAttr) + val scan = new Scan with SupportsReportStatistics { + override def readSchema(): StructType = StructType(Seq(StructField("id", IntegerType))) + override def estimateStatistics(): V2Statistics = new V2Statistics { + override def sizeInBytes(): OptionalLong = OptionalLong.empty() + override def numRows(): OptionalLong = OptionalLong.of(5L) + } + } + + withSQLConf( + SQLConf.CBO_ENABLED.key -> "false", + SQLConf.PLAN_STATS_ENABLED.key -> "false") { + val stats = scanRel(output, scan).computeStats() + + assert(stats.sizeInBytes === EstimationUtils.getSizePerRow(output) * BigInt(5)) + assert(stats.rowCount.isEmpty) + assert(stats.attributeStats.isEmpty) + } + } + + test("DataSourceV2ScanRelation.computeStats uses default size without CBO for empty stats") { + val output = Seq(AttributeReference("id", IntegerType)()) + val scan = new Scan with SupportsReportStatistics { + override def readSchema(): StructType = StructType(Seq(StructField("id", IntegerType))) + override def estimateStatistics(): V2Statistics = new V2Statistics { + override def sizeInBytes(): OptionalLong = OptionalLong.empty() + override def numRows(): OptionalLong = OptionalLong.empty() + } + } + + withSQLConf( + SQLConf.CBO_ENABLED.key -> "false", + SQLConf.PLAN_STATS_ENABLED.key -> "false", + SQLConf.DEFAULT_SIZE_IN_BYTES.key -> "12345") { + val stats = scanRel(output, scan).computeStats() + + assert(stats.sizeInBytes === BigInt(12345)) + assert(stats.rowCount.isEmpty) + } + } + + test("DataSourceV2ScanRelation.computeStats uses default size without reported stats") { + val output = Seq(AttributeReference("id", IntegerType)()) + val scan = new Scan { + override def readSchema(): StructType = StructType(Seq(StructField("id", IntegerType))) + } + + withSQLConf( + SQLConf.CBO_ENABLED.key -> "false", + SQLConf.PLAN_STATS_ENABLED.key -> "false", + SQLConf.DEFAULT_SIZE_IN_BYTES.key -> "12345") { + val stats = scanRel(output, scan).computeStats() + + assert(stats.sizeInBytes === BigInt(12345)) + assert(stats.rowCount.isEmpty) + } + + withSQLConf( + SQLConf.CBO_ENABLED.key -> "true", + SQLConf.DEFAULT_SIZE_IN_BYTES.key -> "12345") { + val stats = scanRel(output, scan).computeStats() + + assert(stats.sizeInBytes === BigInt(12345)) + assert(stats.rowCount.isEmpty) + } + } + + test("DataSourceV2ScanRelation.computeStats uses default size for empty scan stats") { + val output = Seq(AttributeReference("id", IntegerType)()) + val scan = new Scan with SupportsReportStatistics { + override def readSchema(): StructType = StructType(Seq(StructField("id", IntegerType))) + override def estimateStatistics(): V2Statistics = new V2Statistics { + override def sizeInBytes(): OptionalLong = OptionalLong.empty() + override def numRows(): OptionalLong = OptionalLong.empty() + } + } + + withSQLConf( + SQLConf.CBO_ENABLED.key -> "true", + SQLConf.DEFAULT_SIZE_IN_BYTES.key -> "12345") { + val stats = scanRel(output, scan).computeStats() + + assert(stats.sizeInBytes === BigInt(12345)) + assert(stats.rowCount.isEmpty) + } + } + + test("DataSourceV2ScanRelation.computeStats uses default size for null scan stats") { + val output = Seq(AttributeReference("id", IntegerType)()) + val nullStatsScan = new Scan with SupportsReportStatistics { + override def readSchema(): StructType = StructType(Seq(StructField("id", IntegerType))) + override def estimateStatistics(): V2Statistics = null + } + val nullColumnStatsScan = new Scan with SupportsReportStatistics { + override def readSchema(): StructType = StructType(Seq(StructField("id", IntegerType))) + override def estimateStatistics(): V2Statistics = new V2Statistics { + override def sizeInBytes(): OptionalLong = OptionalLong.empty() + override def numRows(): OptionalLong = OptionalLong.empty() + override def columnStats(): java.util.Map[NamedReference, ColumnStatistics] = null + } + } + + withSQLConf( + SQLConf.CBO_ENABLED.key -> "true", + SQLConf.DEFAULT_SIZE_IN_BYTES.key -> "12345") { + Seq(nullStatsScan, nullColumnStatsScan).foreach { scan => + val stats = scanRel(output, scan).computeStats() + + assert(stats.sizeInBytes === BigInt(12345)) + assert(stats.rowCount.isEmpty) + } + } + + withSQLConf( + SQLConf.CBO_ENABLED.key -> "false", + SQLConf.PLAN_STATS_ENABLED.key -> "false", + SQLConf.DEFAULT_SIZE_IN_BYTES.key -> "12345") { + Seq(nullStatsScan, nullColumnStatsScan).foreach { scan => + val stats = scanRel(output, scan).computeStats() + + assert(stats.sizeInBytes === BigInt(12345)) + assert(stats.rowCount.isEmpty) + } + } + } + test("create strips leaked internal metadata but preserves column IDs") { // A column carrying both a column ID (surfaced on purpose) and every internal metadata key // (listed in INTERNAL_METADATA_KEYS), simulating a v2 source that leaks internal metadata. @@ -147,3 +441,12 @@ class DataSourceV2RelationSuite extends SparkFunSuite { assert(field.metadata.contains(FIELD_ID_METADATA_KEY)) } } + +private class FakeTableWithSchema( + tableSchema: StructType = StructType(Seq(StructField("id", IntegerType)))) + extends Table { + + override def name(): String = "fake" + override def schema(): StructType = tableSchema + override def capabilities(): java.util.Set[TableCapability] = java.util.Set.of() +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala index 14d9f754f95c5..ef1f95cbd932e 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala @@ -34,7 +34,7 @@ import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes import org.apache.spark.sql.connector.expressions.{SortOrder => V2SortOrder} import org.apache.spark.sql.connector.expressions.aggregate.{Aggregation, Avg, Count, CountStar, Max, Min, Sum} import org.apache.spark.sql.connector.expressions.filter.Predicate -import org.apache.spark.sql.connector.read.{Scan, ScanBuilder, SupportsPushDownAggregates, SupportsPushDownFilters, SupportsPushDownJoin, SupportsPushDownRequiredColumns, SupportsPushDownVariantExtractions, V1Scan, VariantExtraction} +import org.apache.spark.sql.connector.read.{Scan, ScanBuilder, Statistics => V2Statistics, SupportsPushDownAggregates, SupportsPushDownFilters, SupportsPushDownJoin, SupportsPushDownRequiredColumns, SupportsPushDownVariantExtractions, SupportsReportStatistics, V1Scan, VariantExtraction} import org.apache.spark.sql.execution.datasources.{DataSourceStrategy, VariantInRelation, VariantMetadata} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.internal.connector.VariantExtractionImpl @@ -960,14 +960,21 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan] with PredicateHelper { val withFilter = finalFilters.foldLeft[LogicalPlan](scanRelation)((plan, cond) => { Filter(cond, plan) }) + // Best effort: column pruning can make fully-pushed filters unavailable in the scan output. + // `remappedPushedFilters` already drops those filters, so Spark post-pushdown adjustment can + // only re-add predicates that still reference the pruned scan output. + val withPostPushdownAdjustmentFilters = + withSparkPostPushdownAdjustmentFilters(withFilter, remappedPushedFilters) - if (withFilter.output != project) { + if (withPostPushdownAdjustmentFilters.output != project) { val newProjects = normalizedProjects .map(projectionFunc) .asInstanceOf[Seq[NamedExpression]] - Project(restoreOriginalOutputNames(newProjects, project.map(_.name)), withFilter) + Project( + restoreOriginalOutputNames(newProjects, project.map(_.name)), + withPostPushdownAdjustmentFilters) } else { - withFilter + withPostPushdownAdjustmentFilters } } @@ -1161,6 +1168,34 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan] with PredicateHelper { sHolder.joinedRelationsPushedDownOperators, optRelationName) } + private def withSparkPostPushdownAdjustmentFilters( + plan: LogicalPlan, + pushedFilters: Seq[Expression]): LogicalPlan = { + pushedFilters.reduceLeftOption(And) match { + case None => plan + case Some(pushedCondition) => + def shouldAddPushedFilter(scanRelation: DataSourceV2ScanRelation): Boolean = { + scanRelation.scan match { + case s: SupportsReportStatistics => !s.reflectsFullyPushedDownFilters() + case _ => false + } + } + + def addToScan(plan: LogicalPlan): LogicalPlan = plan match { + case Filter(condition, scanRelation: DataSourceV2ScanRelation) + if shouldAddPushedFilter(scanRelation) => + Filter(And(condition, pushedCondition), scanRelation) + case Filter(condition, child) => + Filter(condition, addToScan(child)) + case scanRelation: DataSourceV2ScanRelation if shouldAddPushedFilter(scanRelation) => + Filter(pushedCondition, scanRelation) + case other => other + } + + addToScan(plan) + } + } + } case class ScanBuilderHolder( @@ -1199,6 +1234,30 @@ case class ScanBuilderHolder( case class V1ScanWrapper( v1Scan: V1Scan, handledFilters: Seq[sources.Filter], - pushedDownOperators: PushedDownOperators) extends Scan { + pushedDownOperators: PushedDownOperators) extends Scan with SupportsReportStatistics { override def readSchema(): StructType = v1Scan.readSchema() + + override def estimateStatistics(): V2Statistics = { + v1Scan match { + case r: SupportsReportStatistics => r.estimateStatistics() + case _ => new V2Statistics { + override def sizeInBytes(): java.util.OptionalLong = java.util.OptionalLong.empty() + override def numRows(): java.util.OptionalLong = java.util.OptionalLong.empty() + } + } + } + + override def estimateSizeInBytes(): java.util.OptionalLong = { + v1Scan match { + case r: SupportsReportStatistics => r.estimateSizeInBytes() + case _ => java.util.OptionalLong.empty() + } + } + + override def reflectsFullyPushedDownFilters(): Boolean = { + v1Scan match { + case r: SupportsReportStatistics => r.reflectsFullyPushedDownFilters() + case _ => true + } + } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2Suite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2Suite.scala index c6f3c1c3886c6..b5d27dc3e4ace 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2Suite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2Suite.scala @@ -19,6 +19,7 @@ package org.apache.spark.sql.connector import java.io.File import java.util +import java.util.Optional import java.util.OptionalLong import scala.jdk.CollectionConverters._ @@ -26,22 +27,26 @@ import scala.jdk.CollectionConverters._ import test.org.apache.spark.sql.connector._ import org.apache.spark.SparkUnsupportedOperationException -import org.apache.spark.sql.{AnalysisException, DataFrame, Row} +import org.apache.spark.sql.{AnalysisException, DataFrame, Row, SQLContext} import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.{ AttributeReference, Expression => CatalystExpression, GreaterThan => CatalystGreaterThan, - Literal => CatalystLiteral, ScalarSubquery} + LessThan => CatalystLessThan, Literal => CatalystLiteral, ScalarSubquery} import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter => LogicalFilter, Project} +import org.apache.spark.sql.catalyst.plans.logical.statsEstimation.EstimationUtils import org.apache.spark.sql.connector.catalog.{PartitionInternalRow, SupportsRead, Table, TableCapability, TableProvider} import org.apache.spark.sql.connector.catalog.TableCapability._ import org.apache.spark.sql.connector.expressions.{Expression, FieldReference, Literal, NamedReference, NullOrdering, SortDirection, SortOrder, Transform} import org.apache.spark.sql.connector.expressions.filter.Predicate import org.apache.spark.sql.connector.read._ import org.apache.spark.sql.connector.read.Scan.ColumnarSupportMode +import org.apache.spark.sql.connector.read.colstats.ColumnStatistics import org.apache.spark.sql.connector.read.partitioning.{KeyGroupedPartitioning, Partitioning, UnknownPartitioning} import org.apache.spark.sql.execution.SortExec import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper -import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, DataSourceV2Relation, DataSourceV2ScanRelation, V2ScanPartitioningAndOrdering} +import org.apache.spark.sql.execution.datasources.v2.{ + BatchScanExec, DataSourceV2Relation, DataSourceV2ScanRelation, PushedDownOperators, + V1ScanWrapper, V2ScanPartitioningAndOrdering} import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Implicits._ import org.apache.spark.sql.execution.exchange.{Exchange, ShuffleExchangeExec} import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector @@ -49,7 +54,7 @@ import org.apache.spark.sql.expressions.Window import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.internal.connector.SupportsPushDownCatalystFilters -import org.apache.spark.sql.sources.{Filter, GreaterThan} +import org.apache.spark.sql.sources.{BaseRelation, Filter, GreaterThan, TableScan} import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.sql.types.{IntegerType, StructField, StructType} import org.apache.spark.sql.util.CaseInsensitiveStringMap @@ -387,16 +392,30 @@ class DataSourceV2Suite extends SharedSparkSession with AdaptiveSparkPlanHelper Seq(classOf[ReportStatisticsDataSource], classOf[JavaReportStatisticsDataSource]).foreach { cls => withClue(cls.getName) { - val df = spark.read.format(cls.getName).load() - val logical = df.queryExecution.optimizedPlan.collect { - case d: DataSourceV2ScanRelation => d - }.head - - val statics = logical.computeStats() - assert(statics.rowCount.isDefined && statics.rowCount.get === 10, - "Row count statics should be reported by data source") - assert(statics.sizeInBytes === 80, - "Size in bytes statics should be reported by data source") + def scanStats: org.apache.spark.sql.catalyst.plans.logical.Statistics = { + val df = spark.read.format(cls.getName).load() + df.queryExecution.optimizedPlan.collect { + case d: DataSourceV2ScanRelation => d + }.head.computeStats() + } + + withSQLConf( + SQLConf.CBO_ENABLED.key -> "false", + SQLConf.PLAN_STATS_ENABLED.key -> "false") { + val statics = scanStats + assert(statics.rowCount.isEmpty, + "Row count statics should not be reported when Spark only needs size") + assert(statics.sizeInBytes === 80, + "Size in bytes statics should be reported by data source") + } + + withSQLConf(SQLConf.CBO_ENABLED.key -> "true") { + val statics = scanStats + assert(statics.rowCount.isDefined && statics.rowCount.get === 10, + "Row count statics should be reported by data source") + assert(statics.sizeInBytes === 80, + "Size in bytes statics should be reported by data source") + } } } } @@ -1168,6 +1187,225 @@ class DataSourceV2Suite extends SharedSparkSession with AdaptiveSparkPlanHelper }.head } + private def hasIGt3(condition: CatalystExpression): Boolean = { + condition.exists { + case CatalystGreaterThan(attr: AttributeReference, CatalystLiteral(value: Int, _)) => + attr.name == "i" && value == 3 + case _ => false + } + } + + private def hasJLtNeg5(condition: CatalystExpression): Boolean = { + condition.exists { + case CatalystLessThan(attr: AttributeReference, CatalystLiteral(value: Int, _)) => + attr.name == "j" && value == -5 + case _ => false + } + } + + test("Spark post-pushdown adjustments re-add fully pushed predicates") { + val df = spark.read.format( + classOf[AdvancedDataSourceV2WithSparkPostPushdownAdjustments].getName).load() + val q = df.filter($"i" > 3) + checkAnswer(q, (4 until 10).map(i => Row(i, -i))) + + val plan = q.queryExecution.optimizedPlan + assert(plan.collect { case f: LogicalFilter => f }.exists(f => hasIGt3(f.condition)), + s"Expected i > 3 in a post-scan Filter when Spark owns post-pushdown adjustments:\n$plan") + val scan = getScanRelation(q) + assert(scan.pushedFilters.exists(hasIGt3), + "scan.pushedFilters should still record the pushed predicate") + } + + test("Spark post-pushdown adjustments use delegated stats end-to-end") { + withSQLConf(SQLConf.CBO_ENABLED.key -> "true") { + val df = spark.read.format( + classOf[AdvancedDataSourceV2WithSparkPostPushdownAdjustments].getName).load() + val q = df.filter($"i" > 3) + checkAnswer(q, (4 until 10).map(i => Row(i, -i))) + + val scan = getScanRelation(q) + val stats = scan.stats + val expectedSize = EstimationUtils.getOutputSize( + scan.output, BigInt(10), stats.attributeStats) + assert(stats.rowCount.contains(BigInt(10)), + "fake connector should delegate numRows through scan stats") + assert(stats.sizeInBytes === expectedSize, + "fake connector should use Spark's projection-aware scan-delegated stats size") + assert(q.queryExecution.optimizedPlan.stats.rowCount.contains(BigInt(7)), + "re-added pushed predicate should adjust plan row count from the scan's pre-filter stats") + } + } + + test("scan stats alone do not imply Spark post-pushdown adjustments") { + withSQLConf(SQLConf.CBO_ENABLED.key -> "true") { + val df = spark.read.format(classOf[AdvancedDataSourceV2WithScanStats].getName).load() + val q = df.filter($"i" > 3) + checkAnswer(q, (4 until 10).map(i => Row(i, -i))) + + val plan = q.queryExecution.optimizedPlan + assert(!plan.collect { case f: LogicalFilter => f }.exists(f => hasIGt3(f.condition)), + s"i > 3 should be stripped from the post-scan Filter by default:\n$plan") + + val scan = getScanRelation(q) + val stats = scan.stats + assert(scan.scan.asInstanceOf[SupportsReportStatistics].reflectsFullyPushedDownFilters(), + "this fake connector reports scan stats as post-pushdown by default") + assert(stats.rowCount.contains(BigInt(2)), + "connector should use scan rowCount") + assert(stats.sizeInBytes === BigInt(32), + "connector should use scan sizeInBytes") + } + } + + test("V1ScanWrapper delegates optional reported statistics") { + val v1Scan = new V1Scan with SupportsReportStatistics { + override def readSchema(): StructType = TestingV2Source.schema + override def toV1TableScan[T <: BaseRelation with TableScan](context: SQLContext): T = + throw new UnsupportedOperationException("not used") + override def estimateStatistics(): Statistics = new Statistics { + override def sizeInBytes(): OptionalLong = OptionalLong.of(128L) + override def numRows(): OptionalLong = OptionalLong.of(8L) + } + override def estimateSizeInBytes(): OptionalLong = OptionalLong.of(64L) + override def reflectsFullyPushedDownFilters(): Boolean = false + } + + val wrapper = V1ScanWrapper(v1Scan, Nil, + PushedDownOperators(None, None, None, None, Nil, Nil, Nil, None)) + + val stats = wrapper.estimateStatistics() + assert(stats.sizeInBytes().getAsLong === 128L) + assert(stats.numRows().getAsLong === 8L) + assert(wrapper.estimateSizeInBytes().getAsLong === 64L) + assert(!wrapper.reflectsFullyPushedDownFilters()) + } + + test("V1ScanWrapper uses empty/default statistics when V1 scan does not report them") { + val v1Scan = new V1Scan { + override def readSchema(): StructType = TestingV2Source.schema + override def toV1TableScan[T <: BaseRelation with TableScan](context: SQLContext): T = + throw new UnsupportedOperationException("not used") + } + + val wrapper = V1ScanWrapper(v1Scan, Nil, + PushedDownOperators(None, None, None, None, Nil, Nil, Nil, None)) + + val stats = wrapper.estimateStatistics() + assert(!stats.sizeInBytes().isPresent) + assert(!stats.numRows().isPresent) + assert(!wrapper.estimateSizeInBytes().isPresent) + assert(wrapper.reflectsFullyPushedDownFilters()) + } + + test("Spark post-pushdown adjustments are not added for scans without reported stats") { + withSQLConf(SQLConf.CONSTRAINT_PROPAGATION_ENABLED.key -> "false") { + val df = spark.read.format(classOf[AdvancedDataSourceV2].getName).load() + val q = df.filter($"i" > 3) + checkAnswer(q, (4 until 10).map(i => Row(i, -i))) + + val plan = q.queryExecution.optimizedPlan + assert(!plan.collect { case f: LogicalFilter => f }.exists(f => hasIGt3(f.condition)), + s"i > 3 should not be re-added without SupportsReportStatistics opt-in:\n$plan") + + val scan = getScanRelation(q) + assert(scan.pushedFilters.exists(hasIGt3), + "scan.pushedFilters should still record the pushed predicate") + } + } + + test("Spark post-pushdown adjustments merge pushed and unpushed predicates") { + val df = spark.read.format( + classOf[AdvancedDataSourceV2WithSparkPostPushdownAdjustments].getName).load() + val q = df.filter($"i" > 3 && $"j" < -5) + checkAnswer(q, (6 until 10).map(i => Row(i, -i))) + + val optimizedPlan = q.queryExecution.optimizedPlan + val directScanFilters = optimizedPlan.collect { + case LogicalFilter(condition, _: DataSourceV2ScanRelation) => condition + } + assert(directScanFilters.length == 1, + s"pushed and unpushed predicates should be merged into one Filter above the scan:\n" + + optimizedPlan) + assert(hasIGt3(directScanFilters.head)) + assert(hasJLtNeg5(directScanFilters.head)) + } + + test("Spark post-pushdown adjustments preserve projection when re-adding predicates") { + withSQLConf(SQLConf.CONSTRAINT_PROPAGATION_ENABLED.key -> "false") { + val df = spark.read.format( + classOf[AdvancedDataSourceV2WithSparkPostPushdownAdjustments].getName).load() + val q = df.filter($"i" > 3).select($"j") + checkAnswer(q, (4 until 10).map(i => Row(-i))) + + val optimizedPlan = q.queryExecution.optimizedPlan + val directScanFilters = optimizedPlan.collect { + case LogicalFilter(condition, _: DataSourceV2ScanRelation) => condition + } + assert(directScanFilters.length == 1, + s"pushed predicate should be re-added directly above the scan before projection:\n" + + optimizedPlan) + assert(hasIGt3(directScanFilters.head)) + assert(optimizedPlan.output.map(_.name) == Seq("j")) + + val scan = getScanRelation(q) + assert(scan.output.exists(_.name == "i"), + "scan output should retain the pushed-filter column needed for the re-added predicate") + assert(scan.output.exists(_.name == "j"), + "scan output should retain the projected column") + } + } + + test("Spark post-pushdown adjustments skip pushed predicates pruned from scan output") { + withSQLConf(SQLConf.CONSTRAINT_PROPAGATION_ENABLED.key -> "false") { + val df = spark.read.format( + classOf[AdvancedDataSourceV2WithBestEffortSparkPostPushdownAdjustments].getName).load() + val q = df.filter($"i" > 3).select($"j") + checkAnswer(q, (4 until 10).map(i => Row(-i))) + + val scan = getScanRelation(q) + assert(!scan.output.exists(_.name == "i"), + "column i should be pruned from the scan output") + assert(!scan.scan.asInstanceOf[SupportsReportStatistics].reflectsFullyPushedDownFilters(), + "fake connector should request Spark post-pushdown adjustments") + assert(scan.pushedFilters.isEmpty, + "pushedFilters should drop filters whose references were pruned") + + val optimizedPlan = q.queryExecution.optimizedPlan + assert(!optimizedPlan.collect { case f: LogicalFilter => f }.exists(f => + hasIGt3(f.condition)), + s"pruned pushed predicate should not be re-added above the scan:\n$optimizedPlan") + } + } + + test("Spark post-pushdown adjustments re-add pushed predicates below stacked residual filters") { + withSQLConf(SQLConf.CONSTRAINT_PROPAGATION_ENABLED.key -> "false") { + // Non-deterministic, always-true residual filters: they cannot be pushed down and cannot be + // combined with each other, so two Filter nodes stay stacked above the scan. This forces + // addToScan to recurse through the outer Filter to re-add the fully pushed i > 3. + val alwaysTrue = udf((_: Int) => true).asNondeterministic() + val df = spark.read.format( + classOf[AdvancedDataSourceV2WithSparkPostPushdownAdjustments].getName).load() + val q = df.filter($"i" > 3).filter(alwaysTrue($"i")).filter(alwaysTrue($"j")) + checkAnswer(q, (4 until 10).map(i => Row(i, -i))) + + val optimizedPlan = q.queryExecution.optimizedPlan + val filters = optimizedPlan.collect { case f: LogicalFilter => f } + assert(filters.length == 2, + s"the two non-deterministic residual filters should stay stacked above the scan:\n" + + optimizedPlan) + val directScanFilters = optimizedPlan.collect { + case LogicalFilter(condition, _: DataSourceV2ScanRelation) => condition + } + assert(directScanFilters.length == 1, + s"pushed predicate should be re-added into the filter directly above the scan:\n" + + optimizedPlan) + assert(hasIGt3(directScanFilters.head), + s"addToScan should recurse through the outer residual filter to re-add i > 3:\n" + + optimizedPlan) + } + } + test("pushedFilters are set for fully pushed filters") { val df = spark.read.format(classOf[AdvancedDataSourceV2].getName).load() // AdvancedDataSourceV2 only supports pushing GreaterThan on column "i". @@ -1610,6 +1848,89 @@ class AdvancedBatch(val filters: Array[Filter], val requiredSchema: StructType) } } +class AdvancedDataSourceV2WithSparkPostPushdownAdjustments extends TestingV2Source { + + override def getTable(options: CaseInsensitiveStringMap): Table = new SimpleBatchTable { + override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = { + new AdvancedScanBuilderWithSparkPostPushdownAdjustments() + } + } +} + +class AdvancedScanBuilderWithSparkPostPushdownAdjustments + extends AdvancedScanBuilder with SupportsReportStatistics { + + override def reflectsFullyPushedDownFilters(): Boolean = false + + override def pruneColumns(requiredSchema: StructType): Unit = { + this.requiredSchema = schemaWithPushedFilterColumns(requiredSchema) + } + + override def estimateStatistics(): Statistics = new Statistics { + override def sizeInBytes(): OptionalLong = OptionalLong.empty() + override def numRows(): OptionalLong = OptionalLong.of(10L) + override def columnStats(): java.util.Map[NamedReference, ColumnStatistics] = { + val stats = new java.util.HashMap[NamedReference, ColumnStatistics]() + stats.put(FieldReference.column("i"), new ColumnStatistics { + override def distinctCount(): OptionalLong = OptionalLong.of(10L) + override def min(): Optional[AnyRef] = Optional.of(Int.box(0)) + override def max(): Optional[AnyRef] = Optional.of(Int.box(9)) + }) + stats + } + } + + private def schemaWithPushedFilterColumns(prunedSchema: StructType): StructType = { + val existingFields = prunedSchema.fieldNames.toSet + val pushedFilterFields = filters.collect { + case GreaterThan(columnName: String, _) => columnName + }.distinct + .filterNot(existingFields.contains) + .flatMap { columnName => + TestingV2Source.schema.fields.find(_.name == columnName) + } + StructType(prunedSchema.fields ++ pushedFilterFields) + } +} + +class AdvancedDataSourceV2WithBestEffortSparkPostPushdownAdjustments extends TestingV2Source { + + override def getTable(options: CaseInsensitiveStringMap): Table = new SimpleBatchTable { + override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = { + new AdvancedScanBuilderWithBestEffortSparkPostPushdownAdjustments() + } + } +} + +class AdvancedScanBuilderWithBestEffortSparkPostPushdownAdjustments + extends AdvancedScanBuilder with SupportsReportStatistics { + + override def reflectsFullyPushedDownFilters(): Boolean = false + + override def estimateStatistics(): Statistics = new Statistics { + override def sizeInBytes(): OptionalLong = OptionalLong.empty() + override def numRows(): OptionalLong = OptionalLong.of(10L) + } +} + +class AdvancedDataSourceV2WithScanStats extends TestingV2Source { + + override def getTable(options: CaseInsensitiveStringMap): Table = new SimpleBatchTable { + override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = { + new AdvancedScanBuilderWithReportedStatistics() + } + } +} + +class AdvancedScanBuilderWithReportedStatistics + extends AdvancedScanBuilder with SupportsReportStatistics { + + override def estimateStatistics(): Statistics = new Statistics { + override def sizeInBytes(): OptionalLong = OptionalLong.of(32L) + override def numRows(): OptionalLong = OptionalLong.of(2L) + } +} + class AdvancedDataSourceV2WithV2Filter extends TestingV2Source { override def getTable(options: CaseInsensitiveStringMap): Table = new SimpleBatchTable {