claudevdm commented on code in PR #39706: URL: https://github.com/apache/beam/pull/39706#discussion_r3785358871
########## sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/maintenance/WriterFactory.java: ########## @@ -0,0 +1,173 @@ +/* + * 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.maintenance; + +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; + +import java.util.Map; +import org.apache.beam.sdk.metrics.Counter; +import org.apache.beam.sdk.metrics.Metrics; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; +import org.apache.iceberg.FileFormat; +import org.apache.iceberg.MetadataColumns; +import org.apache.iceberg.PartitionKey; +import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.Schema; +import org.apache.iceberg.Table; +import org.apache.iceberg.data.GenericAppenderFactory; +import org.apache.iceberg.data.InternalRecordWrapper; +import org.apache.iceberg.data.Record; +import org.apache.iceberg.io.FileAppenderFactory; +import org.apache.iceberg.io.FileIO; +import org.apache.iceberg.io.OutputFileFactory; +import org.apache.iceberg.io.PartitionedFanoutWriter; +import org.apache.iceberg.io.TaskWriter; +import org.apache.iceberg.io.UnpartitionedWriter; +import org.apache.iceberg.util.StructLikeSet; +import org.checkerframework.checker.nullness.qual.MonotonicNonNull; + +/** Builds the {@link TaskWriter} for one rewrite subgroup. */ +class WriterFactory { + @VisibleForTesting static int maxOpenFanoutWriters = 100; + // Number of output partitions opened while writing (one appender each). + private static final Counter openFanoutWriters = + Metrics.counter(WriterFactory.class, "openFanoutWriters"); + + private final long targetFileSizeBytes; + private final String operationId; + private final long attemptId; + private final int globalIndex; + private final PartitionSpec outputSpec; + private final FileFormat format; + private final Map<String, String> writeProperties; + private final boolean preserveRowLineage; + private @MonotonicNonNull OutputFileFactory outputFileFactory; + private @MonotonicNonNull Table table; + + /** + * @param attemptId unique id minted per rewrite attempt. + * @param globalIndex the rewrite group's global index. + * @param outputSpec the spec the planner chose for the rewritten files; may differ from the + * table's current default when {@code output-spec-id} is set or the spec has evolved. + * @param writeProperties write properties that override the table's for the rewrite operation. + * @param preserveRowLineage for v3 row-lineage tables, carry each record's {@code _row_id} / + * {@code _last_updated_sequence_number} metadata columns through the rewrite. + */ + WriterFactory( + FileFormat format, + long targetFileSizeBytes, + long attemptId, + int globalIndex, + String operationId, + PartitionSpec outputSpec, + Map<String, String> writeProperties, + boolean preserveRowLineage) { + this.format = format; + this.targetFileSizeBytes = targetFileSizeBytes; + this.operationId = operationId; + this.attemptId = attemptId; + this.globalIndex = globalIndex; + this.outputSpec = outputSpec; + this.writeProperties = writeProperties; + this.preserveRowLineage = preserveRowLineage; + } + + void init(Table table) { + if (outputFileFactory == null) { + this.table = table; + + outputFileFactory = + OutputFileFactory.builderFor(table, globalIndex, attemptId) + .format(format) + .ioSupplier(table::io) + .defaultSpec(outputSpec) + .operationId(operationId) + .build(); + } + } + + TaskWriter<Record> create() { + Table table = checkStateNotNull(this.table); + Schema writeSchema = + preserveRowLineage ? MetadataColumns.schemaWithRowLineage(table.schema()) : table.schema(); + GenericAppenderFactory appenderFactory = new GenericAppenderFactory(writeSchema, outputSpec); + + // The rewrite's write properties override the table's. + appenderFactory.setAll(table.properties()); + appenderFactory.setAll(writeProperties); + + if (outputSpec.isUnpartitioned()) { + return new UnpartitionedWriter<>( + outputSpec, + format, + appenderFactory, + checkStateNotNull(outputFileFactory), + table.io(), + targetFileSizeBytes); + } else { + return new RecordPartitionedFanoutWriter( + outputSpec, + format, + appenderFactory, + checkStateNotNull(outputFileFactory), + table.io(), + targetFileSizeBytes, + table.schema()); + } + } + + private static class RecordPartitionedFanoutWriter extends PartitionedFanoutWriter<Record> { + + private final PartitionKey partitionKey; + private final InternalRecordWrapper recordWrapper; + private final StructLikeSet openPartitions; + + RecordPartitionedFanoutWriter( + PartitionSpec spec, + FileFormat format, + FileAppenderFactory<Record> appenderFactory, + OutputFileFactory fileFactory, + FileIO io, + long targetFileSize, + Schema schema) { + super(spec, format, appenderFactory, fileFactory, io, targetFileSize); + this.partitionKey = new PartitionKey(spec, schema); + this.openPartitions = StructLikeSet.create(spec.partitionType()); + this.recordWrapper = new InternalRecordWrapper(schema.asStruct()); + } + + @Override + protected PartitionKey partition(Record row) { + // Cap simultaneously-open appenders so a runaway fan-out fails fast instead of OOMing. + partitionKey.partition(recordWrapper.wrap(row)); + if (!openPartitions.contains(partitionKey)) { + if (openPartitions.size() >= maxOpenFanoutWriters) { + throw new IllegalStateException( + String.format( + "Repartitioning compaction fanned out to more than %d simultaneously-open writers on one " + + "subgroup. Compact with the table's current spec (so each subgroup stays within one " + + "partition), or raise worker memory", Review Comment: Should we make `maxOpenFanoutWriters` a configurable option? How will increasing worker memory help if maxOpenFanoutWriters is always 100 and not changeable? ########## sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/maintenance/TaskDescriptor.java: ########## @@ -0,0 +1,134 @@ +/* + * 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.maintenance; + +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; + +import com.google.auto.value.AutoValue; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import org.apache.beam.sdk.schemas.AutoValueSchema; +import org.apache.beam.sdk.schemas.annotations.DefaultSchema; +import org.apache.beam.sdk.schemas.annotations.SchemaFieldNumber; +import org.apache.iceberg.ContentFileParser; +import org.apache.iceberg.DataFile; +import org.apache.iceberg.DeleteFile; +import org.apache.iceberg.FileScanTask; +import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.util.JsonUtil; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** + * A lightweight serializable descriptor of a {@link FileScanTask}, dropping the table schema, + * partition spec and residual that the full task JSON carries, and the data file's column metrics. + */ +@AutoValue +@DefaultSchema(AutoValueSchema.class) +public abstract class TaskDescriptor { + @SchemaFieldNumber("0") + public abstract String getDataFileJson(); + + @SchemaFieldNumber("1") + public abstract int getSpecId(); + + @SchemaFieldNumber("2") + public abstract long getStart(); + + @SchemaFieldNumber("3") + public abstract long getLength(); + + /** + * The input file's data sequence number, carried alongside the file JSON: v3 row lineage derives + * {@code _last_updated_sequence_number} from it on rewrite. + */ + @SchemaFieldNumber("4") + public abstract long getDataSequenceNumber(); + + @SchemaFieldNumber("5") + public abstract List<String> getDeleteFileJsons(); + + static Builder builder() { + return new AutoValue_TaskDescriptor.Builder(); + } + + /** Builds a compact descriptor from one planned range task. */ + static TaskDescriptor from(FileScanTask task, Map<Integer, PartitionSpec> specs) { + PartitionSpec dataSpec = + checkStateNotNull( + specs.get(task.file().specId()), + "Data file spec id %s not found in table specs %s", + task.file().specId(), + specs.keySet()); + List<String> deleteJsons = new ArrayList<>(task.deletes().size()); + for (DeleteFile delete : task.deletes()) { + PartitionSpec deleteSpec = + checkStateNotNull( + specs.get(delete.specId()), + "Delete file spec id %s not found in table specs %s", + delete.specId(), + specs.keySet()); + deleteJsons.add(ContentFileParser.toJson(delete, deleteSpec)); + } + @Nullable Long seq = task.file().dataSequenceNumber(); + return builder() + .setDataFileJson(ContentFileParser.toJson(task.file().copyWithoutStats(), dataSpec)) + .setSpecId(task.file().specId()) + .setStart(task.start()) + .setLength(task.length()) + .setDataSequenceNumber(seq != null ? seq : 0L) + .setDeleteFileJsons(deleteJsons) + .build(); + } + + /** Rebuilds the worker-side {@link FileScanTask} for this range. */ + FileScanTask toScanTask(Map<Integer, PartitionSpec> specs) { + DataFile file = Review Comment: Should we setDataSequenceNumber? ########## sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/maintenance/TaskDescriptor.java: ########## @@ -0,0 +1,134 @@ +/* + * 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.maintenance; + +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; + +import com.google.auto.value.AutoValue; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import org.apache.beam.sdk.schemas.AutoValueSchema; +import org.apache.beam.sdk.schemas.annotations.DefaultSchema; +import org.apache.beam.sdk.schemas.annotations.SchemaFieldNumber; +import org.apache.iceberg.ContentFileParser; +import org.apache.iceberg.DataFile; +import org.apache.iceberg.DeleteFile; +import org.apache.iceberg.FileScanTask; +import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.util.JsonUtil; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** + * A lightweight serializable descriptor of a {@link FileScanTask}, dropping the table schema, + * partition spec and residual that the full task JSON carries, and the data file's column metrics. + */ +@AutoValue +@DefaultSchema(AutoValueSchema.class) +public abstract class TaskDescriptor { + @SchemaFieldNumber("0") + public abstract String getDataFileJson(); + + @SchemaFieldNumber("1") + public abstract int getSpecId(); + + @SchemaFieldNumber("2") + public abstract long getStart(); + + @SchemaFieldNumber("3") + public abstract long getLength(); + + /** + * The input file's data sequence number, carried alongside the file JSON: v3 row lineage derives + * {@code _last_updated_sequence_number} from it on rewrite. + */ + @SchemaFieldNumber("4") + public abstract long getDataSequenceNumber(); + + @SchemaFieldNumber("5") + public abstract List<String> getDeleteFileJsons(); + + static Builder builder() { + return new AutoValue_TaskDescriptor.Builder(); + } + + /** Builds a compact descriptor from one planned range task. */ + static TaskDescriptor from(FileScanTask task, Map<Integer, PartitionSpec> specs) { + PartitionSpec dataSpec = + checkStateNotNull( + specs.get(task.file().specId()), + "Data file spec id %s not found in table specs %s", + task.file().specId(), + specs.keySet()); + List<String> deleteJsons = new ArrayList<>(task.deletes().size()); + for (DeleteFile delete : task.deletes()) { + PartitionSpec deleteSpec = + checkStateNotNull( + specs.get(delete.specId()), + "Delete file spec id %s not found in table specs %s", + delete.specId(), + specs.keySet()); + deleteJsons.add(ContentFileParser.toJson(delete, deleteSpec)); + } + @Nullable Long seq = task.file().dataSequenceNumber(); + return builder() + .setDataFileJson(ContentFileParser.toJson(task.file().copyWithoutStats(), dataSpec)) + .setSpecId(task.file().specId()) + .setStart(task.start()) + .setLength(task.length()) + .setDataSequenceNumber(seq != null ? seq : 0L) Review Comment: Does null mean the same as 0? Should we fail if null, or make this nallable? ########## sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/maintenance/WriterFactory.java: ########## @@ -0,0 +1,173 @@ +/* + * 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.maintenance; + +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; + +import java.util.Map; +import org.apache.beam.sdk.metrics.Counter; +import org.apache.beam.sdk.metrics.Metrics; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; +import org.apache.iceberg.FileFormat; +import org.apache.iceberg.MetadataColumns; +import org.apache.iceberg.PartitionKey; +import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.Schema; +import org.apache.iceberg.Table; +import org.apache.iceberg.data.GenericAppenderFactory; +import org.apache.iceberg.data.InternalRecordWrapper; +import org.apache.iceberg.data.Record; +import org.apache.iceberg.io.FileAppenderFactory; +import org.apache.iceberg.io.FileIO; +import org.apache.iceberg.io.OutputFileFactory; +import org.apache.iceberg.io.PartitionedFanoutWriter; +import org.apache.iceberg.io.TaskWriter; +import org.apache.iceberg.io.UnpartitionedWriter; +import org.apache.iceberg.util.StructLikeSet; +import org.checkerframework.checker.nullness.qual.MonotonicNonNull; + +/** Builds the {@link TaskWriter} for one rewrite subgroup. */ +class WriterFactory { + @VisibleForTesting static int maxOpenFanoutWriters = 100; + // Number of output partitions opened while writing (one appender each). + private static final Counter openFanoutWriters = + Metrics.counter(WriterFactory.class, "openFanoutWriters"); + + private final long targetFileSizeBytes; + private final String operationId; + private final long attemptId; + private final int globalIndex; + private final PartitionSpec outputSpec; + private final FileFormat format; + private final Map<String, String> writeProperties; + private final boolean preserveRowLineage; + private @MonotonicNonNull OutputFileFactory outputFileFactory; + private @MonotonicNonNull Table table; + + /** + * @param attemptId unique id minted per rewrite attempt. + * @param globalIndex the rewrite group's global index. + * @param outputSpec the spec the planner chose for the rewritten files; may differ from the + * table's current default when {@code output-spec-id} is set or the spec has evolved. + * @param writeProperties write properties that override the table's for the rewrite operation. + * @param preserveRowLineage for v3 row-lineage tables, carry each record's {@code _row_id} / + * {@code _last_updated_sequence_number} metadata columns through the rewrite. + */ + WriterFactory( + FileFormat format, + long targetFileSizeBytes, + long attemptId, + int globalIndex, + String operationId, + PartitionSpec outputSpec, + Map<String, String> writeProperties, + boolean preserveRowLineage) { + this.format = format; + this.targetFileSizeBytes = targetFileSizeBytes; + this.operationId = operationId; + this.attemptId = attemptId; + this.globalIndex = globalIndex; + this.outputSpec = outputSpec; + this.writeProperties = writeProperties; + this.preserveRowLineage = preserveRowLineage; + } + + void init(Table table) { + if (outputFileFactory == null) { + this.table = table; + + outputFileFactory = + OutputFileFactory.builderFor(table, globalIndex, attemptId) + .format(format) + .ioSupplier(table::io) + .defaultSpec(outputSpec) + .operationId(operationId) + .build(); + } + } + + TaskWriter<Record> create() { + Table table = checkStateNotNull(this.table); + Schema writeSchema = + preserveRowLineage ? MetadataColumns.schemaWithRowLineage(table.schema()) : table.schema(); + GenericAppenderFactory appenderFactory = new GenericAppenderFactory(writeSchema, outputSpec); + + // The rewrite's write properties override the table's. + appenderFactory.setAll(table.properties()); + appenderFactory.setAll(writeProperties); + + if (outputSpec.isUnpartitioned()) { + return new UnpartitionedWriter<>( + outputSpec, + format, + appenderFactory, + checkStateNotNull(outputFileFactory), + table.io(), + targetFileSizeBytes); + } else { + return new RecordPartitionedFanoutWriter( + outputSpec, + format, + appenderFactory, + checkStateNotNull(outputFileFactory), + table.io(), + targetFileSizeBytes, + table.schema()); + } + } + + private static class RecordPartitionedFanoutWriter extends PartitionedFanoutWriter<Record> { + + private final PartitionKey partitionKey; + private final InternalRecordWrapper recordWrapper; + private final StructLikeSet openPartitions; + + RecordPartitionedFanoutWriter( + PartitionSpec spec, + FileFormat format, + FileAppenderFactory<Record> appenderFactory, + OutputFileFactory fileFactory, + FileIO io, + long targetFileSize, + Schema schema) { + super(spec, format, appenderFactory, fileFactory, io, targetFileSize); + this.partitionKey = new PartitionKey(spec, schema); + this.openPartitions = StructLikeSet.create(spec.partitionType()); + this.recordWrapper = new InternalRecordWrapper(schema.asStruct()); + } + + @Override + protected PartitionKey partition(Record row) { + // Cap simultaneously-open appenders so a runaway fan-out fails fast instead of OOMing. + partitionKey.partition(recordWrapper.wrap(row)); Review Comment: Is all of this necessary if the partition did not change? Wont all rows map to the same partition as a result of planning if there was not repartitioning and in that case we can create a fixed partition writer? -- 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]
