uros-b commented on code in PR #57439: URL: https://github.com/apache/spark/pull/57439#discussion_r3637550486
########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GeneratedColumnPartitionFilters.scala: ########## @@ -0,0 +1,342 @@ +/* + * 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.execution.datasources.v2 + +import java.util.Locale + +import org.apache.spark.SparkException +import org.apache.spark.internal.Logging +import org.apache.spark.sql.catalyst.expressions.{Alias, AttributeReference, Cast, DateFormatClass, DayOfMonth, EqualTo, Expression, GetStructField, GreaterThan, GreaterThanOrEqual, Hour, IntegerLiteral, IsNull, LessThan, LessThanOrEqual, Literal, Month, StringLiteral, Substring, TruncDate, TruncTimestamp, Year} +import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, Project} +import org.apache.spark.sql.catalyst.types.DataTypeUtils.{fromAttributes, toAttributes} +import org.apache.spark.sql.catalyst.util.{quoteIfNeeded, CaseInsensitiveMap, GeneratedColumn} +import org.apache.spark.sql.classic.SparkSession +import org.apache.spark.sql.connector.expressions.IdentityTransform +import org.apache.spark.sql.types.{DataType, DateType, StringType, StructField, StructType, TimestampType} + +/** + * Derives partition filters from data filters when a DataSource V2 table is partitioned by a + * generated column. For example, if a table has a partition column + * `date GENERATED ALWAYS AS (CAST(eventTime AS DATE))`, a data filter `eventTime < '2021-01-01'` + * can be used to derive the partition filter `date <= DATE '2021-01-01'`, which lets the connector + * prune partitions without scanning all of them. + * + * Used by [[V2ScanRelationPushDown]]. The derived filters are only used to help the data source + * prune data; because each derived filter is implied by the original data filter it is derived + * from, applying it never removes valid rows. + */ +object GeneratedColumnPartitionFilters extends Logging { + + private val DATE_FORMAT_YEAR_MONTH = "yyyy-MM" + private val DATE_FORMAT_YEAR_MONTH_DAY = "yyyy-MM-dd" + private val DATE_FORMAT_YEAR_MONTH_DAY_HOUR = "yyyy-MM-dd-HH" + + /** + * Extracts the base (source) column referenced by a data filter. Handles both top level columns + * and nested fields, returning the field path and the column data type. + */ + private object ExtractBaseColumn { + def unapply(e: Expression): Option[(Seq[String], DataType)] = e match { + case a: AttributeReference => + Some((Seq(a.name), a.dataType)) + case g: GetStructField => g.child match { + case ExtractBaseColumn(nameParts, _) => + Some((nameParts :+ g.extractFieldName, g.dataType)) + case _ => None + } + case _ => None + } + } + + private def createFieldPath(nameParts: Seq[String]): String = { + nameParts.map(quoteIfNeeded).mkString(".") + } + + /** + * Returns the schema of the table's top level, identity-partitioned columns, preserving field + * metadata (which carries the generation expression). Only such columns can be generated + * partition columns for this optimization; the schema is empty if the table has none. + */ + private[v2] def identityPartitionSchema(relation: DataSourceV2Relation): StructType = { + val partitionColumnNames = relation.table.partitioning().collect { + case t: IdentityTransform if t.ref.fieldNames().length == 1 => t.ref.fieldNames().head + }.toSet + StructType(relation.output.collect { + case a if partitionColumnNames.contains(a.name) => + StructField(a.name, a.dataType, a.nullable, a.metadata) + }) + } + + /** + * The main entry point used by [[V2ScanRelationPushDown]]. Given a relation and its data filters, + * returns additional partition filters derived from the relation's generated partition columns. + * The returned filters are analyzed (resolved and type coerced) against the relation output and + * reference only the partition columns, so they are safe to add to the list of filters that are + * pushed down to the data source. + */ + def generatePartitionFilters( + spark: SparkSession, + relation: DataSourceV2Relation, + dataFilters: Seq[Expression]): Seq[Expression] = { + if (dataFilters.isEmpty) { + return Nil + } + + val output = relation.output + val partitionSchema = identityPartitionSchema(relation) + if (partitionSchema.isEmpty) { + return Nil + } + // Cheap guard: skip the (more expensive) analysis unless a partition column is generated. + if (!partitionSchema.exists(GeneratedColumn.isGeneratedColumn)) { + return Nil + } + + val fullSchema = fromAttributes(output) + + val rawExpressions = getOptimizablePartitionExpressions(spark, fullSchema, partitionSchema) + if (rawExpressions.isEmpty) { + return Nil + } + val optimizablePartitionExpressions = + if (spark.sessionState.conf.caseSensitiveAnalysis) { + rawExpressions + } else { + CaseInsensitiveMap(rawExpressions) + } + + // Put the column on the left and the literal on the right. + def preprocess(filter: Expression): Expression = filter match { + case LessThan(lit: Literal, e: Expression) => GreaterThan(e, lit) + case LessThanOrEqual(lit: Literal, e: Expression) => GreaterThanOrEqual(e, lit) + case EqualTo(lit: Literal, e: Expression) => EqualTo(e, lit) + case GreaterThan(lit: Literal, e: Expression) => LessThan(e, lit) + case GreaterThanOrEqual(lit: Literal, e: Expression) => LessThanOrEqual(e, lit) + case e => e + } + + def toPartitionFilter( + nameParts: Seq[String], + func: OptimizablePartitionExpression => Option[Expression]): Seq[Expression] = { + optimizablePartitionExpressions.get(createFieldPath(nameParts)).toSeq.flatMap { exprs => + exprs.flatMap(expr => func(expr)) + } + } + + val partitionFilters = dataFilters.flatMap { filter => + preprocess(filter) match { + case LessThan(ExtractBaseColumn(nameParts, _), lit: Literal) => + toPartitionFilter(nameParts, _.lessThan(lit)) + case LessThanOrEqual(ExtractBaseColumn(nameParts, _), lit: Literal) => + toPartitionFilter(nameParts, _.lessThanOrEqual(lit)) + case EqualTo(ExtractBaseColumn(nameParts, _), lit: Literal) => + toPartitionFilter(nameParts, _.equalTo(lit)) + case GreaterThan(ExtractBaseColumn(nameParts, _), lit: Literal) => + toPartitionFilter(nameParts, _.greaterThan(lit)) + case GreaterThanOrEqual(ExtractBaseColumn(nameParts, _), lit: Literal) => + toPartitionFilter(nameParts, _.greaterThanOrEqual(lit)) + case IsNull(ExtractBaseColumn(nameParts, _)) => + toPartitionFilter(nameParts, _.isNull()) + case _ => Nil + } + } + + resolveAndCoerce(spark, partitionFilters, output) + } + + /** + * Analyzes the derived partition filters against the relation output so that partition column + * references are resolved and the expressions are type coerced. The derived filters reference + * only partition columns that exist in `output`, so they are always expected to resolve; a + * failure indicates a bug and is surfaced as an internal error. + */ + private def resolveAndCoerce( Review Comment: Note: a pure optimization can hard-fail query planning. The generation-expression string comes from table metadata; sqlParser.parseExpression can raise ParseException, analyzer.execute can raise AnalysisException, and resolveAndCoerce throws SparkException.internalError on any expression shape it does not expect. Since skipping derivation never changes results, an unforeseen or malformed generation expression turns a previously-successful scan into a hard error. Wrapping the parse/analyze/resolve path in a catch that logs and falls back to Nil (derive nothing) is the safe design for an optimization. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
