vinishjail97 commented on code in PR #992: URL: https://github.com/apache/incubator-xtable/pull/992#discussion_r4395185020
########## 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: [P1] Validate the index generation across lookup This extends the existing finding with a local reproduction against `b57e1b83`. The second snapshot check still accepts a different index generation after rollback and rebuild. 1. The first check accepts S1, and planning collects S1 candidate files. 2. A rewrite preserves ID 7, and sync publishes an S2 index before lookup. 3. Lookup collects S2 file paths. 4. Iceberg rolls back to S1. Deleting `.hoodie` and running FULL sync rebuilds an S1 index. 5. The second check accepts S1 again. The S1 candidates and S2 lookup paths have an empty intersection. The query returned zero rows instead of ID 7. The reproduction used actual Iceberg writes and Hudi syncs. A temporary callback after lookup collection controlled the sequence. The callback changed timing, but kept the pruning decisions unchanged. Would you validate the index generation and completed commit identity, including rebuilds, or pin lookup to that generation? The source snapshot ID alone does not identify the index generation. I consider this blocking because the query silently removes a matching row. ########## 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: [P1] Validate indexed field identity across schema changes I reproduced this on unmodified commit `b57e1b83` with Spark 3.4.2, Iceberg 1.9.2, and Hudi 1.2.1. 1. Build the index on `email`. 2. In one schema commit, rename `email` to `old_email` and `level` to `email`. 3. Run `SELECT id, email FROM table WHERE email = 'level_1'`. The query returned 13 rows without pruning and zero rows with pruning. The schema commit kept the same snapshot ID, so both freshness checks passed. The rule matched the old index by its reused column name. Would you validate the indexed Iceberg field ID and schema compatibility before pruning? A regression test should compare actual rows with pruning enabled and disabled. I consider this blocking because schema changes can silently remove matching rows. -- 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]
