vinishjail97 commented on code in PR #992: URL: https://github.com/apache/incubator-xtable/pull/992#discussion_r4427525867
########## xtable-spark-extensions/src/main/scala/org/apache/xtable/spark/extensions/optimizer/IndexPruningRule.scala: ########## @@ -0,0 +1,362 @@ +/* + * 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, Snapshot, 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.hudi.index.HudiBackedSecondaryIndexUtils.SyncedVersion +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. The rule fails the query + * when the indexed column was renamed after the index was synced. + */ +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)) Review Comment: Added `!table.table().isInstanceOf[BaseMetadataTable]` to the guard in 6d6a2574. Before this, a metadata table fell back only because the lookup threw a `ClassCastException`; `metadataTablesDoNotUseIndex` checks that no lookup runs. ########## xtable-spark-extensions/src/main/scala/org/apache/xtable/spark/extensions/optimizer/IndexPruningRule.scala: ########## @@ -0,0 +1,362 @@ +/* + * 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, Snapshot, 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.hudi.index.HudiBackedSecondaryIndexUtils.SyncedVersion +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. The rule fails the query + * when the indexed column was renamed after the index was synced. + */ +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 e: IndexedColumnRenamedException => throw e + 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()) + version = index.syncedVersion + if isCurrent(table, column, version, snapshot.snapshotId()) + _ = checkNotRenamed(table, column, snapshot) + // 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, rollback or rebuild that completes between the check above and the end of the + // lookup can point the index at other files than the candidates, even at the same snapshot. + // Each of them completes a new action on the index, so an equal version after the lookup + // shows that the lookup read the checked version. + if isUnchanged(table, column, version, index.syncedVersion) + } 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) } + } + + /** + * Fails when the column names another field than in the schema of the synced snapshot. A rename + * keeps the snapshot, so the index stays current, but it keys the column by name and then holds + * the values of another field. + */ + private def checkNotRenamed(table: Table, column: String, snapshot: Snapshot): Unit = { + val syncedField = + Option(snapshot.schemaId()) Review Comment: Done in 6d6a2574, the rule now skips the index when the synced snapshot has no schema id, because it cannot detect a rename then. ########## xtable-spark-extensions/src/main/scala/org/apache/xtable/spark/extensions/optimizer/StagedScans.scala: ########## @@ -0,0 +1,104 @@ +/* + * 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.UUID + +import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan +import org.apache.spark.sql.execution.QueryExecution +import org.apache.spark.sql.execution.datasources.v2.{ + DataSourceV2Relation, + DataSourceV2ScanRelation +} +import org.apache.spark.sql.util.QueryExecutionListener + +import org.apache.iceberg.{FileScanTask, Table} +import org.apache.iceberg.spark.ScanTaskSetManager + +import scala.collection.JavaConverters._ +import scala.util.Try + +/** + * Stages the file scan tasks that an index lookup selects. Iceberg reads a staged task set when a + * scan has the {@code scan-task-set-id} option. Staged tasks live in a process wide registry until + * they are removed, so this object removes them when the query that uses them ends. Iceberg reads + * the tasks of a scan only when Spark plans its partitions, so a set is never removed before its + * query ends. Plans that are only explained never end and leave their sets staged, so this object + * stages no new set once too many are live, and the scans that would use them read the table as + * usual. + */ +object StagedScans { + val ScanTaskSetIdOption = "scan-task-set-id" + + val MaxLiveSets = 1000 + + private val liveSets = new java.util.HashMap[String, Table]() + + /** Stages the tasks and returns the id of their set, or nothing when too many sets are live. */ + def stage(table: Table, tasks: Seq[FileScanTask]): Option[String] = { + val setId = UUID.randomUUID().toString + val reserved = liveSets.synchronized { + liveSets.size() < MaxLiveSets && liveSets.put(setId, table) == null + } + if (reserved) { + ScanTaskSetManager.get().stageTasks(table, setId, tasks.asJava) + Some(setId) + } else { + None + } + } + + def release(setId: String): Unit = { + val table = liveSets.synchronized(liveSets.remove(setId)) + if (table != null) { + removeTasks(table, setId) + } + } + + def liveSetIds: Set[String] = liveSets.synchronized(liveSets.keySet().asScala.toSet) + + /** Releases the task sets that the scans of a plan use. */ + def releaseAll(plan: LogicalPlan): Unit = { + def setIdOf(relation: DataSourceV2Relation): Option[String] = + Option(relation.options.get(ScanTaskSetIdOption)) + plan + .collectWithSubqueries { + case scan: DataSourceV2ScanRelation => setIdOf(scan.relation) + case relation: DataSourceV2Relation => setIdOf(relation) + } + .flatten + .foreach(release) + } + + private def removeTasks(table: Table, setId: String): Unit = + ScanTaskSetManager.get().removeTasks[FileScanTask](table, setId) +} + +/** Releases the staged task sets of a query once it ends. */ +class StagedScanCleanupListener extends QueryExecutionListener { + override def onSuccess(funcName: String, qe: QueryExecution, durationNs: Long): Unit = + release(qe) + + override def onFailure(funcName: String, qe: QueryExecution, exception: Exception): Unit = + release(qe) + + // a query that failed during optimization has no optimized plan, and nothing was staged for it + private def release(qe: QueryExecution): Unit = + Try(qe.optimizedPlan).foreach(StagedScans.releaseAll) Review Comment: Done in 6d6a2574, the listener releases the task sets only when `qe.tracker.phases` contains `OPTIMIZATION`, so a failed query does not run the optimizer again. ########## xtable-spark-extensions/src/main/scala/org/apache/xtable/spark/extensions/optimizer/IndexPruningRule.scala: ########## @@ -0,0 +1,362 @@ +/* + * 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, Snapshot, 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.hudi.index.HudiBackedSecondaryIndexUtils.SyncedVersion +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. The rule fails the query + * when the indexed column was renamed after the index was synced. + */ +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 e: IndexedColumnRenamedException => throw e + 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()) + version = index.syncedVersion + if isCurrent(table, column, version, snapshot.snapshotId()) + _ = checkNotRenamed(table, column, snapshot) + // 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, rollback or rebuild that completes between the check above and the end of the + // lookup can point the index at other files than the candidates, even at the same snapshot. + // Each of them completes a new action on the index, so an equal version after the lookup + // shows that the lookup read the checked version. + if isUnchanged(table, column, version, index.syncedVersion) Review Comment: Done in 6d6a2574, the relation records the key filters it declined in a `TreeNodeTag`, so the lookup runs once per query. Spark clones the plan for each query, so the tag lasts for one query only; `declinedLookupRunsOnce` covers it. -- 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]
