vinishjail97 commented on code in PR #992: URL: https://github.com/apache/incubator-xtable/pull/992#discussion_r4395574705
########## xtable-spark-extensions/src/main/scala/org/apache/xtable/spark/extensions/optimizer/IndexPruningRule.scala: ########## @@ -0,0 +1,328 @@ +/* + * 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.xtable.spark.extensions.optimizer + +import java.util.concurrent.atomic.AtomicBoolean + +import org.apache.hadoop.fs.Path +import org.apache.spark.sql.{Encoders, SparkSession} +import org.apache.spark.sql.catalyst.expressions.{ + And, + Attribute, + AttributeReference, + EqualTo, + Expression, + In, + InSet, + Literal +} +import org.apache.spark.sql.catalyst.plans.logical.{Filter, LocalRelation, LogicalPlan} +import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.util.CaseInsensitiveStringMap + +import org.apache.iceberg.{FileScanTask, Table} +import org.apache.iceberg.expressions.{Expression => IcebergExpression, Expressions} +import org.apache.iceberg.spark.SparkReadOptions +import org.apache.iceberg.spark.source.SparkTable +import org.apache.iceberg.types.Type.TypeID +import org.apache.iceberg.types.Types.NestedField + +import scala.collection.JavaConverters._ +import scala.util.control.NonFatal + +import org.apache.xtable.hudi.index.HudiBackedSecondaryIndexReader +import org.apache.xtable.spark.extensions.{XTableIndex, XTableIndexConf} + +/** + * Uses an XTable index to read only the files that hold the keys of a filter such as + * {@code col = 'x'} or {@code col IN ('x', 'y')} on an indexed column of an Iceberg table. The rule + * runs for SQL queries and DataFrame reads alike. + * + * The rule looks the keys up in the index, keeps the files that Iceberg's own partition and min/max + * pruning also keeps, and stages them as the task set of the scan. The filter stays in the plan, so + * the query returns the same rows as without the rule. The rule leaves the plan unchanged when the + * index is behind the table, when the query reads an older snapshot, a branch, a tag or the changes + * between snapshots, or when the scan is too small for a lookup to pay off. + */ +case class IndexPruningRule(spark: SparkSession) extends Rule[LogicalPlan] { + private val listenerRegistered = new AtomicBoolean(false) + + override def apply(plan: LogicalPlan): LogicalPlan = { + if (!XTableIndexConf.pruningEnabled(conf)) { + return plan + } + plan.transformDown { + case filter @ Filter( + condition, + relation @ DataSourceV2Relation(table: SparkTable, _, _, _, options)) + if !IndexPruningRule.SnapshotReadOptions.exists(options.containsKey) && + table.snapshotId() == null && table.branch() == null => + try { + prune(filter, condition, relation, table).getOrElse(filter) + } catch { + case NonFatal(e) => + logWarning(s"Not using the XTable index for ${table.name()}", e) + filter + } + } + } + + private def prune( + filter: Filter, + condition: Expression, + relation: DataSourceV2Relation, + sparkTable: SparkTable): Option[LogicalPlan] = { + val keyFilters = IndexPruningRule.findKeyFilters(condition, relation.output) + if (keyFilters.isEmpty) { + return None + } + val table = sparkTable.table() + val index = + XTableIndex.of(spark, table, relation.catalog.map(_.name()), relation.identifier) + for { + (column, keys) <- findIndexedFilter(table, index, keyFilters) + if withinKeyLimit(table, column, keys) + snapshot <- Option(table.currentSnapshot()) + if isCurrent(table, column, index.lastSyncedSnapshot, snapshot.snapshotId()) + // the files a scan reads after Iceberg's partition and min/max pruning + candidates = planFiles(table, snapshot.snapshotId(), column, keys) + if isLargeEnough(table, column, candidates) + indexedFiles = lookup(index, column, keys) + // A sync that completes between the check above and the lookup can point the index at files + // of a newer snapshot than the candidates. The index stays at the checked snapshot until the + // next sync commits, so an unchanged snapshot after the lookup shows that it read that one. + if isUnchanged(table, column, index.lastSyncedSnapshot, snapshot.snapshotId()) + } yield { + val tasks = candidates.filter(task => + indexedFiles.contains(IndexPruningRule.normalizePath(task.file().location()))) + if (tasks.isEmpty) { + logInfo( + s"Using the index of ${table.name()} on $column for ${keys.size} keys: no file of " + + s"${candidates.size} candidate files holds them") + filter.copy(child = LocalRelation(relation.output, Nil, isStreaming = false)) + } else { + StagedScans.stage(table, tasks) match { + case Some(setId) => + logInfo( + s"Using the index of ${table.name()} on $column for ${keys.size} keys: reading " + + s"${tasks.size} of ${candidates.size} candidate files") + registerCleanupListener() + val stagedOptions = new java.util.HashMap[String, String](options(relation)) + stagedOptions.put(StagedScans.ScanTaskSetIdOption, setId) + filter.copy(child = + relation.copy(options = new CaseInsensitiveStringMap(stagedOptions))) + case None => + logWarning( + s"Not using the index of ${table.name()} on $column: too many staged scans are " + + "live, which happens when plans are optimized but not run") + filter + } + } + } + } + + /** + * Returns the first filter on an indexed column, with the name of the column in the table + * schema. Hudi records the indexed column under that name. + */ + private def findIndexedFilter( + table: Table, + index: XTableIndex, + keyFilters: Seq[(String, Seq[Any])]): Option[(String, Seq[Any])] = { + val schema = table.schema() + keyFilters.iterator + .flatMap { case (attribute, keys) => + val field: NestedField = + if (conf.caseSensitiveAnalysis) schema.findField(attribute) + else schema.caseInsensitiveFindField(attribute) + Option(field) + .filter(field => XTableIndex.SupportedTypes.contains(field.`type`().typeId())) + .map(field => field.name() -> keys) + } + .find { case (column, _) => index.isIndexed(column) } Review Comment: It's better to fail fast when an indexed column is renamed instead of adding more checks like these? I think there would be more edge cases. -- 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]
