forked from delta-io/delta
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[Spark] Add read support for RowId (delta-io#2856)
<!-- Thanks for sending a pull request! Here are some tips for you: 1. If this is your first time, please read our contributor guidelines: https://github.com/delta-io/delta/blob/master/CONTRIBUTING.md 2. If the PR is unfinished, add '[WIP]' in your PR title, e.g., '[WIP] Your PR title ...'. 3. Be sure to keep the PR description updated to reflect all changes. 4. Please write your PR title to summarize what this PR proposes. 5. If possible, provide a concise example to reproduce the issue for a faster review. 6. If applicable, include the corresponding issue number in the PR title and link it in the body. --> #### Which Delta project/connector is this regarding? <!-- Please add the component selected below to the beginning of the pull request title For example: [Spark] Title of my pull request --> - [x] Spark - [ ] Standalone - [ ] Flink - [ ] Kernel - [ ] Other (fill in here) ## Description 1. Add the Analyzer Rule `GenerateRowIds` to generate default Row IDs. 2. Add the `row_id` field to the `_metadata` column for Delta tables, allowing us to read the `row_id` from the file metadata after it is stored. <!-- - Describe what this PR changes. - Describe why we need the change. If this PR resolves an issue be sure to include "Resolves #XXX" to correctly link and close the issue upon merge. --> ## How was this patch tested? Added UTs. <!-- If tests were added, say they were added here. Please make sure to test the changes thoroughly including negative and positive cases if possible. If the changes were tested in any way other than unit tests, please clarify how you tested step by step (ideally copy and paste-able, so that other reviewers can test and check, and descendants can verify in the future). If the changes were not tested, please explain why. --> ## Does this PR introduce _any_ user-facing changes? <!-- If yes, please clarify the previous behavior and the change this PR proposes - provide the console output, description and/or an example to show the behavior difference if possible. If possible, please also clarify if this is a user-facing change compared to the released Delta Lake versions or within the unreleased branches such as master. If no, write 'No'. --> No.
- Loading branch information
1 parent
b0ab2e6
commit b132344
Showing
15 changed files
with
915 additions
and
14 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
139 changes: 139 additions & 0 deletions
139
spark/src/main/scala/org/apache/spark/sql/delta/GenerateRowIDs.scala
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,139 @@ | ||
| /* | ||
| * Copyright (2021) The Delta Lake Project Authors. | ||
| * | ||
| * Licensed 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.delta | ||
|
|
||
| import scala.collection.mutable | ||
|
|
||
| import org.apache.spark.sql.catalyst.expressions._ | ||
| import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, Project} | ||
| import org.apache.spark.sql.catalyst.rules.Rule | ||
| import org.apache.spark.sql.catalyst.trees.TreePattern.PLAN_EXPRESSION | ||
| import org.apache.spark.sql.execution.datasources.{FileFormat, HadoopFsRelation, LogicalRelation} | ||
| import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat | ||
| import org.apache.spark.sql.types.StructType | ||
|
|
||
| /** | ||
| * This rule adds a Project on top of Delta tables that support the Row tracking table feature to | ||
| * provide a default generated Row ID for rows that don't have them materialized in the data file. | ||
| */ | ||
| object GenerateRowIDs extends Rule[LogicalPlan] { | ||
|
|
||
| /** | ||
| * Matcher for a scan on a Delta table that has Row tracking enabled. | ||
| */ | ||
| private object DeltaScanWithRowTrackingEnabled { | ||
| def unapply(plan: LogicalPlan): Option[LogicalRelation] = plan match { | ||
| case scan @ LogicalRelation(relation: HadoopFsRelation, _, _, _) => | ||
| relation.fileFormat match { | ||
| case format: DeltaParquetFileFormat | ||
| if RowTracking.isEnabled(format.protocol, format.metadata) => Some(scan) | ||
| case _ => None | ||
| } | ||
| case _ => None | ||
| } | ||
| } | ||
|
|
||
| override def apply(plan: LogicalPlan): LogicalPlan = plan.transformUpWithNewOutput { | ||
| case DeltaScanWithRowTrackingEnabled(scan) => | ||
| // While Row IDs are non-nullable, we'll use the Row ID attributes to read | ||
| // the materialized values from now on, which can be null. We make | ||
| // the materialized Row ID attributes nullable in the scan here. | ||
|
|
||
| // Update nullability in the scan `metadataOutput` by updating the delta file format. | ||
| val baseRelation = scan.relation.asInstanceOf[HadoopFsRelation] | ||
| val newFileFormat = baseRelation.fileFormat match { | ||
| case format: DeltaParquetFileFormat => | ||
| format.copy(nullableRowTrackingFields = true) | ||
| } | ||
| val newBaseRelation = baseRelation.copy(fileFormat = newFileFormat)(baseRelation.sparkSession) | ||
|
|
||
| // Update the output metadata column's data type (now with nullable row tracking fields). | ||
| val newOutput = scan.output.map { | ||
| case MetadataAttributeWithLogicalName(metadata, FileFormat.METADATA_NAME) => | ||
| metadata.withDataType(newFileFormat.createFileMetadataCol().dataType) | ||
| case other => other | ||
| } | ||
| val newScan = scan.copy(relation = newBaseRelation, output = newOutput) | ||
| newScan.copyTagsFrom(scan) | ||
|
|
||
| // Add projection with row tracking column expressions. | ||
| val updatedAttributes = mutable.Buffer.empty[(Attribute, Attribute)] | ||
| val projectList = newOutput.map { | ||
| case MetadataAttributeWithLogicalName(metadata, FileFormat.METADATA_NAME) => | ||
| val updatedMetadata = metadataWithRowTrackingColumnsProjection(metadata) | ||
| updatedAttributes += metadata -> updatedMetadata.toAttribute | ||
| updatedMetadata | ||
| case other => other | ||
| } | ||
| Project(projectList = projectList, child = newScan) -> updatedAttributes.toSeq | ||
| case o => | ||
| val newPlan = o.transformExpressionsWithPruning(_.containsPattern(PLAN_EXPRESSION)) { | ||
| // Recurse into subquery plans. Similar to how [[transformUpWithSubqueries]] works except | ||
| // that it allows us to still use [[transformUpWithNewOutput]] on subquery plans to | ||
| // correctly update references to the metadata attribute when going up the plan. | ||
| // Get around type erasure by explicitly checking the plan type and removing warning. | ||
| case planExpression: PlanExpression[LogicalPlan @unchecked] | ||
| if planExpression.plan.isInstanceOf[LogicalPlan] => | ||
| planExpression.withNewPlan(apply(planExpression.plan)) | ||
| } | ||
| newPlan -> Nil | ||
| } | ||
|
|
||
| /** | ||
| * Expression that reads the Row IDs from the materialized Row ID column if the value is | ||
| * present and returns the default generated Row ID using the file's base Row ID and current row | ||
| * index if not: | ||
| * coalesce(_metadata.row_id, _metadata.base_row_id + _metadata.row_index). | ||
| */ | ||
| private def rowIdExpr(metadata: AttributeReference): Expression = { | ||
| Coalesce(Seq( | ||
| getField(metadata, RowId.ROW_ID), | ||
| Add( | ||
| getField(metadata, RowId.BASE_ROW_ID), | ||
| getField(metadata, ParquetFileFormat.ROW_INDEX)))) | ||
| } | ||
|
|
||
| /** | ||
| * Extract a field from the metadata column. | ||
| */ | ||
| private def getField(metadata: AttributeReference, name: String): GetStructField = { | ||
| ExtractValue(metadata, Literal(name), conf.resolver) match { | ||
| case field: GetStructField => field | ||
| case _ => | ||
| throw new IllegalStateException(s"The metadata column '${metadata.name}' is not a struct.") | ||
| } | ||
| } | ||
|
|
||
| /** | ||
| * Create a new metadata struct where the Row ID values are populated using | ||
| * the materialized values if present, or the default Row ID values if not. | ||
| */ | ||
| private def metadataWithRowTrackingColumnsProjection(metadata: AttributeReference) | ||
| : NamedExpression = { | ||
| val metadataFields = metadata.dataType.asInstanceOf[StructType].map { | ||
| case field if field.name == RowId.ROW_ID => | ||
| field -> rowIdExpr(metadata) | ||
| case field => | ||
| field -> getField(metadata, field.name) | ||
| }.flatMap { case (oldField, newExpr) => | ||
| // Propagate the type metadata from the old fields to the new fields. | ||
| val newField = Alias(newExpr, oldField.name)(explicitMetadata = Some(oldField.metadata)) | ||
| Seq(Literal(oldField.name), newField) | ||
| } | ||
| Alias(CreateNamedStruct(metadataFields), metadata.name)() | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.