-
Notifications
You must be signed in to change notification settings - Fork 4.6k
[#39723] Implement model for Iceberg side input cache #39724
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: master
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,219 @@ | ||
| /* | ||
| * 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.beam.sdk.io.iceberg; | ||
|
|
||
| import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; | ||
|
|
||
| import com.google.auto.value.AutoValue; | ||
| import java.io.Serializable; | ||
| import java.util.Map; | ||
| import org.apache.beam.sdk.schemas.AutoValueSchema; | ||
| import org.apache.beam.sdk.schemas.NoSuchSchemaException; | ||
| import org.apache.beam.sdk.schemas.SchemaCoder; | ||
| import org.apache.beam.sdk.schemas.SchemaRegistry; | ||
| import org.apache.beam.sdk.schemas.annotations.DefaultSchema; | ||
| import org.apache.beam.sdk.schemas.annotations.SchemaFieldNumber; | ||
| import org.apache.beam.sdk.schemas.annotations.SchemaIgnore; | ||
| import org.apache.iceberg.PartitionSpec; | ||
| import org.apache.iceberg.PartitionSpecParser; | ||
| import org.apache.iceberg.Schema; | ||
| import org.apache.iceberg.SchemaParser; | ||
| import org.apache.iceberg.SortOrder; | ||
| import org.apache.iceberg.SortOrderParser; | ||
| import org.apache.iceberg.Table; | ||
| import org.apache.iceberg.catalog.TableIdentifier; | ||
| import org.checkerframework.checker.nullness.qual.MonotonicNonNull; | ||
|
|
||
| /** | ||
| * A serializable, lightweight representation of an Iceberg {@link Table}'s declarative metadata. | ||
| * | ||
| * <p>Captures the table's schema, partition spec, sort order, location, properties, and identifier. | ||
| * Suitable for broadcasting across worker nodes via Beam's side-input mechanism. | ||
| */ | ||
| @DefaultSchema(AutoValueSchema.class) | ||
| @AutoValue | ||
| public abstract class SerializableTableSpec implements Serializable { | ||
|
|
||
| @SchemaFieldNumber("0") | ||
| public abstract String getTableIdentifierString(); | ||
|
|
||
| @SchemaFieldNumber("1") | ||
| public abstract String getName(); | ||
|
|
||
| @SchemaFieldNumber("2") | ||
| public abstract String getLocation(); | ||
|
|
||
| @SchemaFieldNumber("3") | ||
| public abstract int getSpecId(); | ||
|
|
||
| @SchemaFieldNumber("4") | ||
| public abstract String getSchemaJson(); | ||
|
|
||
| @SchemaFieldNumber("5") | ||
| public abstract String getPartitionSpecJson(); | ||
|
|
||
| @SchemaFieldNumber("6") | ||
| public abstract String getSortOrderJson(); | ||
|
|
||
| @SchemaFieldNumber("7") | ||
| public abstract Map<String, String> getProperties(); | ||
|
|
||
| private transient volatile @MonotonicNonNull Schema cachedSchema; | ||
| private transient volatile @MonotonicNonNull PartitionSpec cachedPartitionSpec; | ||
| private transient volatile @MonotonicNonNull SortOrder cachedSortOrder; | ||
| private transient volatile @MonotonicNonNull TableIdentifier cachedTableIdentifier; | ||
|
|
||
| private static volatile @MonotonicNonNull SchemaCoder<SerializableTableSpec> cachedCoder; | ||
|
|
||
| @SchemaIgnore | ||
| public Schema getSchema() { | ||
| Schema local = cachedSchema; | ||
| if (local == null) { | ||
| synchronized (this) { | ||
| local = cachedSchema; | ||
| if (local == null) { | ||
| cachedSchema = local = SchemaParser.fromJson(getSchemaJson()); | ||
| } | ||
| } | ||
| } | ||
| return local; | ||
| } | ||
|
|
||
| @SchemaIgnore | ||
| public PartitionSpec getPartitionSpec() { | ||
| PartitionSpec local = cachedPartitionSpec; | ||
| if (local == null) { | ||
| synchronized (this) { | ||
| local = cachedPartitionSpec; | ||
| if (local == null) { | ||
| cachedPartitionSpec = | ||
| local = PartitionSpecParser.fromJson(getSchema(), getPartitionSpecJson()); | ||
| } | ||
| } | ||
| } | ||
| return local; | ||
| } | ||
|
|
||
| @SchemaIgnore | ||
| public SortOrder getSortOrder() { | ||
| SortOrder local = cachedSortOrder; | ||
| if (local == null) { | ||
| synchronized (this) { | ||
| local = cachedSortOrder; | ||
| if (local == null) { | ||
| cachedSortOrder = local = SortOrderParser.fromJson(getSchema(), getSortOrderJson()); | ||
| } | ||
| } | ||
| } | ||
| return local; | ||
| } | ||
|
|
||
| @SchemaIgnore | ||
| public TableIdentifier getTableIdentifier() { | ||
| TableIdentifier local = cachedTableIdentifier; | ||
| if (local == null) { | ||
| synchronized (this) { | ||
| local = cachedTableIdentifier; | ||
| if (local == null) { | ||
| cachedTableIdentifier = | ||
| local = IcebergUtils.parseTableIdentifier(getTableIdentifierString()); | ||
| } | ||
| } | ||
| } | ||
| return local; | ||
| } | ||
|
|
||
| public static Builder builder() { | ||
| return new AutoValue_SerializableTableSpec.Builder(); | ||
| } | ||
|
|
||
| public abstract Builder toBuilder(); | ||
|
|
||
| @AutoValue.Builder | ||
| public abstract static class Builder { | ||
| public abstract Builder setTableIdentifierString(String tableIdentifierString); | ||
|
|
||
| public abstract Builder setName(String name); | ||
|
|
||
| public abstract Builder setLocation(String location); | ||
|
|
||
| public abstract Builder setSpecId(int specId); | ||
|
|
||
| public abstract Builder setSchemaJson(String schemaJson); | ||
|
|
||
| public abstract Builder setPartitionSpecJson(String partitionSpecJson); | ||
|
|
||
| public abstract Builder setSortOrderJson(String sortOrderJson); | ||
|
|
||
| public abstract Builder setProperties(Map<String, String> properties); | ||
|
|
||
| public abstract SerializableTableSpec build(); | ||
| } | ||
|
|
||
| /** | ||
| * Constructs a {@link SerializableTableSpec} from a {@link Table}, using {@link Table#name()} as | ||
| * the table identifier string. | ||
| * | ||
| * <p>Note: When possible, prefer {@link #fromTable(TableIdentifier, Table)} to avoid catalog name | ||
| * prefix ambiguities in {@link Table#name()}. | ||
| */ | ||
| public static SerializableTableSpec fromTable(Table table) { | ||
| return fromTable(table.name(), table); | ||
| } | ||
|
|
||
| /** | ||
| * Constructs a {@link SerializableTableSpec} from a {@link TableIdentifier} and a {@link Table}. | ||
| */ | ||
| public static SerializableTableSpec fromTable(TableIdentifier tableIdentifier, Table table) { | ||
| return fromTable(IcebergUtils.tableIdentifierToString(tableIdentifier), table); | ||
| } | ||
|
|
||
| /** | ||
| * Constructs a {@link SerializableTableSpec} from an explicit table identifier string and a | ||
| * {@link Table}. | ||
| */ | ||
| public static SerializableTableSpec fromTable(String tableIdentifierString, Table table) { | ||
| return builder() | ||
| .setTableIdentifierString(tableIdentifierString) | ||
| .setName(table.name()) | ||
| .setLocation(table.location()) | ||
| .setSpecId(table.spec().specId()) | ||
| .setSchemaJson(SchemaParser.toJson(table.schema())) | ||
| .setPartitionSpecJson(PartitionSpecParser.toJson(table.spec())) | ||
| .setSortOrderJson(SortOrderParser.toJson(table.sortOrder())) | ||
| .setProperties(table.properties()) | ||
| .build(); | ||
| } | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Also include fileIO json here. Iceberg has a FileIOParser that works similarly to the other properties
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I think we can also serialize properties to reconstruct the EncryptionManager, instead of making it an arg for SideInputTable. Mostly because we won't have access to the EncryptionManager on the worker side to properly create SideInputTable When loading the real table, we can serialize and store the encrypted keys in (AFAICT, we can assume the table implements Then to reconstruct it in We'd probably need to pass |
||
|
|
||
| /** Returns the cached {@link SchemaCoder} for {@link SerializableTableSpec}. */ | ||
| public static SchemaCoder<SerializableTableSpec> getCoder() { | ||
| if (cachedCoder == null) { | ||
| synchronized (SerializableTableSpec.class) { | ||
| if (cachedCoder == null) { | ||
| try { | ||
| cachedCoder = | ||
| SchemaRegistry.createDefault().getSchemaCoder(SerializableTableSpec.class); | ||
| } catch (NoSuchSchemaException e) { | ||
| throw new RuntimeException(e); | ||
| } | ||
| } | ||
| } | ||
| } | ||
| return checkStateNotNull(cachedCoder); | ||
| } | ||
| } | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think we should keep track of the full map of schemas / partition specs / sort orders. There are instances where we're not always working with the most recent table metadata, so we'll need to have access to historical schema/spec/sort orders.
Might make sense to add 3 int fields to indicate the "current" schemaId, specId, orderId, and use them to return the respective metadata property.