vinishjail97 commented on code in PR #992: URL: https://github.com/apache/incubator-xtable/pull/992#discussion_r4395438430
########## xtable-spark-extensions/src/main/scala/org/apache/xtable/spark/extensions/optimizer/IndexPruningRule.scala: ########## @@ -0,0 +1,258 @@ +/* + * 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.source.SparkTable +import org.apache.iceberg.types.Type.TypeID + +import scala.collection.JavaConverters._ +import scala.util.control.NonFatal + +import org.apache.xtable.hudi.index.HudiBackedSecondaryIndexReader +import org.apache.xtable.spark.extensions.{IndexDefinition, 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 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 or a branch, 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 !options.containsKey(StagedScans.ScanTaskSetIdOption) && + 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 definitions = IndexDefinition.fromProperties(sparkTable.properties()) + val table = sparkTable.table() + for { + (definition, keys) <- IndexPruningRule + .findIndexedKeys(condition, relation.output, definitions, conf.resolver) + if withinKeyLimit(table, definition, keys) + snapshot <- Option(table.currentSnapshot()) + index = IndexDefinition.newIndex( + spark, + table, + relation.catalog.map(_.name()), + relation.identifier, + definition) + reader = index.reader + if isCurrent(table, definition, reader.getLastSyncedSourceIdentifier, snapshot.snapshotId()) Review Comment: Fixed in 3e1ccc0c. The rule now compares the index version (latest completion time on the Hudi timeline plus the synced snapshot) before and after the lookup, so a rollback and rebuild to the same snapshot falls back. `rebuildDuringLookupKeepsTheRows` reproduces the sequence and returns no rows without the fix. ########## 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: Fixed in 3e1ccc0c. The rule skips the index when any retained schema maps the column name to another field id or type, or renames the field. `swappedColumnNamesDoNotUseIndex` compares rows with pruning on and off. -- 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]
