diff --git a/sdks/java/io/iceberg/build.gradle b/sdks/java/io/iceberg/build.gradle index 983ebd07fefc..c98711d55da1 100644 --- a/sdks/java/io/iceberg/build.gradle +++ b/sdks/java/io/iceberg/build.gradle @@ -64,6 +64,7 @@ dependencies { provided "org.immutables:value:2.8.8" permitUnusedDeclared "org.immutables:value:2.8.8" implementation library.java.vendored_calcite_1_40_0 + implementation library.java.jackson_databind runtimeOnly "org.apache.iceberg:iceberg-gcp:$iceberg_version" runtimeOnly "org.apache.iceberg:iceberg-aws:$iceberg_version" runtimeOnly "org.apache.iceberg:iceberg-aws-bundle:$iceberg_version" diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/maintenance/ExecutedGroup.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/maintenance/ExecutedGroup.java new file mode 100644 index 000000000000..50f1e9464357 --- /dev/null +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/maintenance/ExecutedGroup.java @@ -0,0 +1,108 @@ +/* + * 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 com.google.auto.value.AutoValue; +import java.util.ArrayList; +import java.util.List; +import org.apache.beam.sdk.io.iceberg.SerializableDataFile; +import org.apache.beam.sdk.schemas.AutoValueSchema; +import org.apache.beam.sdk.schemas.annotations.DefaultSchema; +import org.apache.beam.sdk.schemas.annotations.SchemaFieldNumber; + +/** Represents the result of one rewritten subgroup as compact commit descriptors. */ +@AutoValue +@DefaultSchema(AutoValueSchema.class) +public abstract class ExecutedGroup { + + @SchemaFieldNumber("0") + public abstract long getStartingSnapshotId(); + + /** The rewrite operation's id used to stamp the commit. */ + @SchemaFieldNumber("1") + public abstract String getOperationId(); + + /** Index of the planned parent group this subgroup belongs to. */ + @SchemaFieldNumber("2") + public abstract int getParentGroupIndex(); + + /** Total number of subgroups the parent was split into. */ + @SchemaFieldNumber("3") + public abstract int getParentSubgroupCount(); + + /** Total input byte size of this group, for partial-progress failure accounting. */ + @SchemaFieldNumber("4") + public abstract long getTotalInputByteSize(); + + /** Newly written compacted data files to ADD (full metrics). */ + @SchemaFieldNumber("5") + public abstract List getNewFiles(); + + /** Rewritten input data files to DELETE (no metrics). */ + @SchemaFieldNumber("6") + public abstract List getRewrittenDataFiles(); + + /** Dangling deletion vector JSONs to DELETE. */ + @SchemaFieldNumber("7") + public abstract List getDanglingDeleteFileJsons(); + + /** The starting snapshot's sequence number; the floor for the commit's idempotency stamp scan. */ + @SchemaFieldNumber("8") + public abstract long getStartingSequenceNumber(); + + public static Builder builder() { + return new AutoValue_ExecutedGroup.Builder(); + } + + /** + * Locations of every newly written output file across {@code groups}. After a failed commit these + * are orphans; they carry the operation id so a later remove-orphan-files run can find them. + */ + static List newFilePaths(Iterable groups) { + List paths = new ArrayList<>(); + for (ExecutedGroup g : groups) { + for (SerializableDataFile sdf : g.getNewFiles()) { + paths.add(sdf.getPath()); + } + } + return paths; + } + + @AutoValue.Builder + public abstract static class Builder { + public abstract Builder setStartingSnapshotId(long v); + + public abstract Builder setStartingSequenceNumber(long v); + + public abstract Builder setOperationId(String v); + + public abstract Builder setParentGroupIndex(int v); + + public abstract Builder setParentSubgroupCount(int v); + + public abstract Builder setTotalInputByteSize(long v); + + public abstract Builder setNewFiles(List v); + + public abstract Builder setRewrittenDataFiles(List v); + + public abstract Builder setDanglingDeleteFileJsons(List v); + + public abstract ExecutedGroup build(); + } +} diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/maintenance/RangeFileScanTask.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/maintenance/RangeFileScanTask.java new file mode 100644 index 000000000000..c885a9bca38a --- /dev/null +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/maintenance/RangeFileScanTask.java @@ -0,0 +1,83 @@ +/* + * 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 java.util.List; +import org.apache.iceberg.DataFile; +import org.apache.iceberg.DeleteFile; +import org.apache.iceberg.FileScanTask; +import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.expressions.Expression; +import org.apache.iceberg.expressions.Expressions; + +/** + * A worker-side {@link FileScanTask} reconstructed from a {@link TaskDescriptor}. Represents one + * row-group range of a data file, plus the delete files that apply to it. + */ +class RangeFileScanTask implements FileScanTask { + private final DataFile file; + private final List deletes; + private final long start; + private final long length; + private final PartitionSpec spec; + + RangeFileScanTask( + DataFile file, List deletes, long start, long length, PartitionSpec spec) { + this.file = file; + this.deletes = deletes; + this.start = start; + this.length = length; + this.spec = spec; + } + + @Override + public DataFile file() { + return file; + } + + @Override + public List deletes() { + return deletes; + } + + @Override + public long start() { + return start; + } + + @Override + public long length() { + return length; + } + + @Override + public PartitionSpec spec() { + return spec; + } + + @Override + public Expression residual() { + return Expressions.alwaysTrue(); + } + + @Override + public Iterable split(long targetSplitSize) { + throw new UnsupportedOperationException( + "RangeFileScanTask is already a fixed row-group range and cannot be re-split."); + } +} diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/maintenance/RewriteSubGroup.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/maintenance/RewriteSubGroup.java new file mode 100644 index 000000000000..f5c97142ea8c --- /dev/null +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/maintenance/RewriteSubGroup.java @@ -0,0 +1,116 @@ +/* + * 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 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.FileScanTask; +import org.apache.iceberg.PartitionSpec; + +@AutoValue +@DefaultSchema(AutoValueSchema.class) +public abstract class RewriteSubGroup { + static Builder builder() { + return new AutoValue_RewriteSubGroup.Builder(); + } + + @SchemaFieldNumber("0") + abstract int getGlobalIndex(); + + /** Index of the planned parent group this subgroup belongs to; shared by all its subgroups. */ + @SchemaFieldNumber("1") + abstract int getParentGroupIndex(); + + /** Total number of subgroups belonging to the parent. */ + @SchemaFieldNumber("2") + abstract int getParentSubgroupCount(); + + /** The compact per-range descriptors this subgroup rewrites (one per row-group range). */ + @SchemaFieldNumber("3") + abstract List getTaskDescriptors(); + + @SchemaFieldNumber("4") + abstract int getOutputSpecId(); + + @SchemaFieldNumber("5") + abstract long getWriteMaxFileSize(); + + @SchemaFieldNumber("6") + abstract long getTotalInputFileByteSize(); + + @SchemaFieldNumber("7") + abstract long getStartingSnapshotId(); + + /** + * The rewrite operation's id, unique to this pipeline execution. Used to name/tag output files + * and to stamp commits for idempotency. + */ + @SchemaFieldNumber("8") + abstract String getOperationId(); + + /** + * The starting snapshot's own sequence number, captured at planning. It floors the commit's + * idempotency stamp scan, and still bounds the walk if that snapshot has since been expired. + */ + @SchemaFieldNumber("9") + abstract long getStartingSequenceNumber(); + + @AutoValue.Builder + abstract static class Builder { + abstract Builder setGlobalIndex(int globalIndex); + + abstract Builder setParentGroupIndex(int parentGroupIndex); + + abstract Builder setParentSubgroupCount(int parentSubgroupCount); + + abstract Builder setTaskDescriptors(List taskDescriptors); + + /** + * Builds compact per-range descriptors from planned range tasks and records the group's total + * input byte size (the summed range lengths). + */ + Builder setFileScanTasks(List tasks, Map specs) { + long byteSize = 0; + List taskDescriptors = new ArrayList<>(tasks.size()); + for (FileScanTask task : tasks) { + byteSize += task.length(); + taskDescriptors.add(TaskDescriptor.from(task, specs)); + } + return setTotalInputFileByteSize(byteSize).setTaskDescriptors(taskDescriptors); + } + + abstract Builder setOutputSpecId(int outputSpecId); + + abstract Builder setWriteMaxFileSize(long writeMaxFileSize); + + abstract Builder setTotalInputFileByteSize(long byteSize); + + abstract Builder setStartingSnapshotId(long startingSnapshotId); + + abstract Builder setStartingSequenceNumber(long startingSequenceNumber); + + abstract Builder setOperationId(String operationId); + + abstract RewriteSubGroup build(); + } +} diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/maintenance/TaskDescriptor.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/maintenance/TaskDescriptor.java new file mode 100644 index 000000000000..73d9b37a1d3d --- /dev/null +++ b/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 getDeleteFileJsons(); + + static Builder builder() { + return new AutoValue_TaskDescriptor.Builder(); + } + + /** Builds a compact descriptor from one planned range task. */ + static TaskDescriptor from(FileScanTask task, Map 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 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 specs) { + DataFile file = + (DataFile) + JsonUtil.parse(getDataFileJson(), node -> ContentFileParser.fromJson(node, specs)); + List deletes = new ArrayList<>(getDeleteFileJsons().size()); + for (String deleteJson : getDeleteFileJsons()) { + deletes.add( + (DeleteFile) JsonUtil.parse(deleteJson, node -> ContentFileParser.fromJson(node, specs))); + } + PartitionSpec spec = + checkStateNotNull( + specs.get(getSpecId()), + "Spec id %s not found in table specs %s", + getSpecId(), + specs.keySet()); + return new RangeFileScanTask(file, deletes, getStart(), getLength(), spec); + } + + @AutoValue.Builder + abstract static class Builder { + abstract Builder setDataFileJson(String dataFileJson); + + abstract Builder setSpecId(int specId); + + abstract Builder setStart(long start); + + abstract Builder setLength(long length); + + abstract Builder setDataSequenceNumber(long dataSequenceNumber); + + abstract Builder setDeleteFileJsons(List deleteFileJsons); + + abstract TaskDescriptor build(); + } +} diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/maintenance/WriterFactory.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/maintenance/WriterFactory.java new file mode 100644 index 000000000000..fcc5aee033ee --- /dev/null +++ b/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 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 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 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 { + + private final PartitionKey partitionKey; + private final InternalRecordWrapper recordWrapper; + private final StructLikeSet openPartitions; + + RecordPartitionedFanoutWriter( + PartitionSpec spec, + FileFormat format, + FileAppenderFactory 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", + maxOpenFanoutWriters)); + } + openPartitions.add(partitionKey.copy()); + openFanoutWriters.inc(); + } + return partitionKey; + } + } +} diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/maintenance/package-info.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/maintenance/package-info.java new file mode 100644 index 000000000000..5e45f4b3fd11 --- /dev/null +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/maintenance/package-info.java @@ -0,0 +1,20 @@ +/* + * 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. + */ + +/** Iceberg table maintenance actions (e.g. rewrite data files / bin-pack compaction). */ +package org.apache.beam.sdk.io.iceberg.maintenance; diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/TestDataWarehouse.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/TestDataWarehouse.java index 2e711219349c..61f96f0e684d 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/TestDataWarehouse.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/TestDataWarehouse.java @@ -162,6 +162,33 @@ public DataFile writeRecords( .build(); } + /** + * Writes {@code records} to an unpartitioned Parquet file applying {@code writerProperties} (for + * example a tiny {@code write.parquet.row-group-size-bytes} to force several row groups) and + * records the resulting split offsets on the returned {@link DataFile} so it is splittable by row + * group. {@code withMetrics} does not carry split offsets, so they must be set explicitly. + */ + public DataFile writeRecords( + String filename, Schema schema, List records, Map writerProperties) + throws IOException { + Path path = new Path(location, filename); + FileAppender appender = + Parquet.write(fromPath(path, hadoopConf)) + .createWriterFunc(GenericParquetWriter::create) + .schema(schema) + .setAll(writerProperties) + .overwrite() + .build(); + appender.addAll(records); + appender.close(); + + return DataFiles.builder(PartitionSpec.unpartitioned()) + .withInputFile(HadoopInputFile.fromPath(path, hadoopConf)) + .withMetrics(appender.metrics()) + .withSplitOffsets(appender.splitOffsets()) + .build(); + } + public Table createTable(TableIdentifier tableId, Schema schema) { return createTable(tableId, schema, null, null); } diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/maintenance/TaskDescriptorTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/maintenance/TaskDescriptorTest.java new file mode 100644 index 000000000000..2d498a3d1421 --- /dev/null +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/maintenance/TaskDescriptorTest.java @@ -0,0 +1,310 @@ +/* + * 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.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; + +import java.io.ByteArrayOutputStream; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.List; +import org.apache.beam.sdk.io.iceberg.TestDataWarehouse; +import org.apache.beam.sdk.io.iceberg.TestFixtures; +import org.apache.beam.sdk.schemas.SchemaCoder; +import org.apache.beam.sdk.schemas.SchemaRegistry; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; +import org.apache.iceberg.DataFile; +import org.apache.iceberg.DeleteFile; +import org.apache.iceberg.FileFormat; +import org.apache.iceberg.FileScanTask; +import org.apache.iceberg.RowDelta; +import org.apache.iceberg.ScanTaskParser; +import org.apache.iceberg.Schema; +import org.apache.iceberg.Table; +import org.apache.iceberg.catalog.TableIdentifier; +import org.apache.iceberg.data.GenericAppenderFactory; +import org.apache.iceberg.data.GenericRecord; +import org.apache.iceberg.data.Record; +import org.apache.iceberg.deletes.BaseDVFileWriter; +import org.apache.iceberg.deletes.DVFileWriter; +import org.apache.iceberg.deletes.PositionDelete; +import org.apache.iceberg.deletes.PositionDeleteWriter; +import org.apache.iceberg.encryption.EncryptedOutputFile; +import org.apache.iceberg.io.CloseableIterable; +import org.apache.iceberg.io.OutputFileFactory; +import org.apache.iceberg.types.Types; +import org.junit.ClassRule; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Round-trip and payload-size tests for {@link TaskDescriptor}. */ +@RunWith(JUnit4.class) +public class TaskDescriptorTest { + @ClassRule public static final TemporaryFolder TEMPORARY_FOLDER = new TemporaryFolder(); + + @Rule public TestDataWarehouse warehouse = new TestDataWarehouse(TEMPORARY_FOLDER, "default"); + + @Test + public void roundTripPreservesFileStartLengthSequenceAndDeletes() throws Exception { + // Golden round-trip: descriptor -> task must reproduce the file identity, range, data sequence + // number, and the applying delete files. A v2 table + a positional delete gives a non-null + // sequence number and one delete on the whole-file task. + TableIdentifier id = TableIdentifier.of("default", "td_" + System.nanoTime()); + Table table = + warehouse.createTable( + id, TestFixtures.SCHEMA, null, ImmutableMap.of("format-version", "2")); + DataFile dataFile = + warehouse.writeRecords( + "d_" + System.nanoTime() + ".parquet", table.schema(), rows(1, 2, 3)); + table.newAppend().appendFile(dataFile).commit(); + table.refresh(); + addPositionalDelete(table, dataFile, 0L); + + FileScanTask task; + try (CloseableIterable it = table.newScan().planFiles()) { + task = it.iterator().next(); + } + assertEquals("fixture: the task must carry the positional delete", 1, task.deletes().size()); + + TaskDescriptor descriptor = TaskDescriptor.from(task, table.specs()); + FileScanTask reconstructed = descriptor.toScanTask(table.specs()); + + assertEquals(task.file().location(), reconstructed.file().location()); + assertEquals(task.file().format(), reconstructed.file().format()); + assertEquals(task.file().recordCount(), reconstructed.file().recordCount()); + assertEquals(task.file().fileSizeInBytes(), reconstructed.file().fileSizeInBytes()); + assertEquals(task.file().specId(), reconstructed.file().specId()); + assertEquals(task.start(), reconstructed.start()); + assertEquals(task.length(), reconstructed.length()); + assertEquals(task.file().dataSequenceNumber().longValue(), descriptor.getDataSequenceNumber()); + assertEquals(1, reconstructed.deletes().size()); + assertEquals(task.deletes().get(0).location(), reconstructed.deletes().get(0).location()); + } + + @Test + public void roundTripPreservesDeletionVectorTopLevelFields() throws Exception { + // A deletion vector's contentOffset, contentSizeInBytes, and referencedDataFile locate its + // blob inside the Puffin file, so they must survive the ContentFileParser JSON round-trip. + // Assert them field-by-field on a REAL v3 DV, not just that a delete round-trips by location. + TableIdentifier id = TableIdentifier.of("default", "tddv_" + System.nanoTime()); + Table table = + warehouse.createTable( + id, TestFixtures.SCHEMA, null, ImmutableMap.of("format-version", "3")); + DataFile dataFile = + warehouse.writeRecords( + "dv_" + System.nanoTime() + ".parquet", table.schema(), rows(1, 2, 3)); + table.newAppend().appendFile(dataFile).commit(); + table.refresh(); + addDeletionVector(table, dataFile, 0L); + + FileScanTask task; + try (CloseableIterable it = table.newScan().planFiles()) { + task = it.iterator().next(); + } + assertEquals("fixture: the task must carry the deletion vector", 1, task.deletes().size()); + DeleteFile originalDv = task.deletes().get(0); + // Sanity: a real DV carries a Puffin blob offset/size and references exactly this data file. + assertNotNull( + "fixture must be a deletion vector with a content offset", originalDv.contentOffset()); + assertNotNull("fixture DV must carry a content size", originalDv.contentSizeInBytes()); + assertEquals(dataFile.location(), originalDv.referencedDataFile()); + + TaskDescriptor descriptor = TaskDescriptor.from(task, table.specs()); + DeleteFile reconstructedDv = descriptor.toScanTask(table.specs()).deletes().get(0); + assertEquals( + "contentOffset must round-trip", + originalDv.contentOffset(), + reconstructedDv.contentOffset()); + assertEquals( + "contentSizeInBytes must round-trip", + originalDv.contentSizeInBytes(), + reconstructedDv.contentSizeInBytes()); + assertEquals( + "referencedDataFile must round-trip", + originalDv.referencedDataFile(), + reconstructedDv.referencedDataFile()); + assertEquals("DV location must round-trip", originalDv.location(), reconstructedDv.location()); + } + + @Test + public void roundTripPreservesRangeStartAndLength() throws Exception { + // A start>0 row-group range must round-trip its exact start/length (the descriptor carries them + // as scalars, not re-derived). + TableIdentifier id = TableIdentifier.of("default", "tdr_" + System.nanoTime()); + Table table = warehouse.createTable(id, TestFixtures.SCHEMA); + List rows = new ArrayList<>(); + for (int i = 0; i < 800; i++) { + Record r = GenericRecord.create(TestFixtures.SCHEMA); + r.setField("id", (long) i); + r.setField("data", "row-" + i + "-padding-0123456789abcdef0123456789abcdef"); + rows.add(r); + } + DataFile dataFile = + warehouse.writeRecords( + "mrg_" + System.nanoTime() + ".parquet", + table.schema(), + rows, + ImmutableMap.builder() + .put("write.parquet.row-group-size-bytes", "8192") + .put("parquet.enable.dictionary", "false") + .put("write.parquet.page-size-bytes", "1024") + .put("write.parquet.row-group-check-max-record-count", "100") + .put("write.parquet.compression-codec", "uncompressed") + .build()); + table.newAppend().appendFile(dataFile).commit(); + table.refresh(); + + List ranges; + try (CloseableIterable it = table.newScan().planFiles()) { + ranges = Lists.newArrayList(it.iterator().next().split(1L)); + } + FileScanTask ranged = + ranges.stream() + .filter(t -> t.start() > 0) + .findFirst() + .orElseThrow(() -> new AssertionError("fixture must produce a start>0 range")); + + TaskDescriptor descriptor = TaskDescriptor.from(ranged, table.specs()); + FileScanTask reconstructed = descriptor.toScanTask(table.specs()); + assertEquals(ranged.start(), reconstructed.start()); + assertEquals(ranged.length(), reconstructed.length()); + assertEquals(ranged.file().location(), reconstructed.file().location()); + // A compaction read is unfiltered: the config filter selects files, not rows. + assertEquals(org.apache.iceberg.expressions.Expressions.alwaysTrue(), reconstructed.residual()); + } + + @Test + public void compactDescriptorPayloadShrinksVsFullScanTaskJson() throws Exception { + // The compact descriptor payload must be far smaller than embedding a full ScanTaskParser JSON + // (table schema + spec + residual) per range: a wide 50-column schema makes the shrink >=10x. + // The coder-encoded group is compared against the summed per-task JSON as the ceiling. + Schema wide = wideSchema(50); + TableIdentifier id = TableIdentifier.of("default", "c1size_" + System.nanoTime()); + Table table = warehouse.createTable(id, wide); + int numFiles = 30; + org.apache.iceberg.AppendFiles append = table.newAppend(); + for (int f = 0; f < numFiles; f++) { + Record r = GenericRecord.create(wide); + r.setField("id", (long) f); + append.appendFile( + warehouse.writeRecords( + "w" + f + "_" + System.nanoTime() + ".parquet", wide, Lists.newArrayList(r))); + } + append.commit(); + table.refresh(); + + List tasks; + try (CloseableIterable it = table.newScan().planFiles()) { + tasks = Lists.newArrayList(it); + } + assertEquals(numFiles, tasks.size()); + + long oldBytes = 0; + for (FileScanTask t : tasks) { + oldBytes += ScanTaskParser.toJson(t).getBytes(StandardCharsets.UTF_8).length; + } + + RewriteSubGroup group = + RewriteSubGroup.builder() + .setGlobalIndex(1) + .setParentGroupIndex(0) + .setParentSubgroupCount(1) + .setFileScanTasks(tasks, table.specs()) + .setOutputSpecId(table.spec().specId()) + .setWriteMaxFileSize(Long.MAX_VALUE) + .setStartingSnapshotId(table.currentSnapshot().snapshotId()) + .setStartingSequenceNumber(table.currentSnapshot().sequenceNumber()) + .setOperationId("op-test") + .build(); + SchemaCoder coder = + SchemaRegistry.createDefault().getSchemaCoder(RewriteSubGroup.class); + ByteArrayOutputStream out = new ByteArrayOutputStream(); + coder.encode(group, out); + long newBytes = out.size(); + + assertTrue( + "compact descriptor payload (" + + newBytes + + " B) must be >=10x smaller than the old per-task ScanTaskParser JSON (" + + oldBytes + + " B)", + newBytes * 10 <= oldBytes); + } + + private static Schema wideSchema(int columns) { + List fields = new ArrayList<>(); + fields.add(Types.NestedField.required(1, "id", Types.LongType.get())); + for (int i = 1; i < columns; i++) { + fields.add(Types.NestedField.optional(i + 1, "col_" + i, Types.StringType.get())); + } + return new Schema(fields); + } + + private static List rows(long... ids) { + List recs = new ArrayList<>(); + for (long id : ids) { + Record r = GenericRecord.create(TestFixtures.SCHEMA); + r.setField("id", id); + r.setField("data", "row-" + id); + recs.add(r); + } + return recs; + } + + private void addPositionalDelete(Table table, DataFile dataFile, long position) throws Exception { + GenericAppenderFactory appenderFactory = + new GenericAppenderFactory(table.schema(), table.spec()); + OutputFileFactory fileFactory = + OutputFileFactory.builderFor(table, 1, 1L).format(FileFormat.PARQUET).build(); + EncryptedOutputFile outputFile = fileFactory.newOutputFile(); + PositionDeleteWriter writer = + appenderFactory.newPosDeleteWriter(outputFile, FileFormat.PARQUET, null); + PositionDelete positionDelete = PositionDelete.create(); + try { + positionDelete.set(dataFile.location().toString(), position, null); + writer.write(positionDelete); + } finally { + writer.close(); + } + DeleteFile deleteFile = writer.toDeleteFile(); + table.newRowDelta().addDeletes(deleteFile).commit(); + table.refresh(); + } + + /** Writes a v3 deletion vector deleting {@code position} in {@code dataFile} and commits it. */ + private void addDeletionVector(Table table, DataFile dataFile, long position) throws Exception { + OutputFileFactory fileFactory = + OutputFileFactory.builderFor(table, 3, 3L).format(FileFormat.PUFFIN).build(); + DVFileWriter writer = new BaseDVFileWriter(fileFactory, path -> null); + try { + writer.delete(dataFile.location().toString(), position, table.spec(), null); + } finally { + writer.close(); + } + RowDelta rowDelta = table.newRowDelta(); + writer.result().deleteFiles().forEach(rowDelta::addDeletes); + rowDelta.commit(); + table.refresh(); + } +} diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/maintenance/WriterFactoryTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/maintenance/WriterFactoryTest.java new file mode 100644 index 000000000000..c4de4daaefb8 --- /dev/null +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/maintenance/WriterFactoryTest.java @@ -0,0 +1,181 @@ +/* + * 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.junit.Assert.assertEquals; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +import java.util.ArrayList; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import java.util.concurrent.ThreadLocalRandom; +import org.apache.beam.sdk.io.iceberg.TestDataWarehouse; +import org.apache.beam.sdk.io.iceberg.TestFixtures; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; +import org.apache.iceberg.DataFile; +import org.apache.iceberg.FileFormat; +import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.Schema; +import org.apache.iceberg.Table; +import org.apache.iceberg.catalog.TableIdentifier; +import org.apache.iceberg.data.GenericRecord; +import org.apache.iceberg.data.Record; +import org.apache.iceberg.io.TaskWriter; +import org.apache.iceberg.types.Types; +import org.junit.ClassRule; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Tests for {@link WriterFactory}. */ +@RunWith(JUnit4.class) +public class WriterFactoryTest { + @ClassRule public static final TemporaryFolder TEMPORARY_FOLDER = new TemporaryFolder(); + + @Rule public TestDataWarehouse warehouse = new TestDataWarehouse(TEMPORARY_FOLDER, "default"); + + private static final Schema SHARDED_SCHEMA = + new Schema( + Types.NestedField.required(1, "id", Types.LongType.get()), + Types.NestedField.required(2, "shard", Types.IntegerType.get())); + + private WriterFactory factoryFor(Table table, PartitionSpec spec, long targetFileSize) { + return factoryFor(table, spec, targetFileSize, ImmutableMap.of()); + } + + private WriterFactory factoryFor( + Table table, PartitionSpec spec, long targetFileSize, java.util.Map props) { + WriterFactory wf = + new WriterFactory( + FileFormat.PARQUET, + targetFileSize, + ThreadLocalRandom.current().nextLong(), + 1, + "op-test", + spec, + props, + false); + wf.init(table); + return wf; + } + + private static Record shardedRow(long id, int shard) { + Record r = GenericRecord.create(SHARDED_SCHEMA); + r.setField("id", id); + r.setField("shard", shard); + return r; + } + + @Test + public void unpartitionedWriteProducesOneFileWithAllRecords() throws Exception { + TableIdentifier id = TableIdentifier.of("default", "wf_unp_" + System.nanoTime()); + Table table = warehouse.createTable(id, TestFixtures.SCHEMA); + + TaskWriter writer = factoryFor(table, table.spec(), Long.MAX_VALUE).create(); + for (int i = 0; i < 5; i++) { + Record r = GenericRecord.create(TestFixtures.SCHEMA); + r.setField("id", (long) i); + r.setField("data", "row-" + i); + writer.write(r); + } + DataFile[] files = writer.dataFiles(); + + assertEquals("an unpartitioned write under target rolls into one file", 1, files.length); + assertEquals(5L, files[0].recordCount()); + } + + @Test + public void partitionedWriteRoutesEachRecordToItsOwnPartition() throws Exception { + // The fanout writer keys every record individually, so rows of different partitions must land + // in different files, each registered under the partition its rows belong to. + TableIdentifier id = TableIdentifier.of("default", "wf_part_" + System.nanoTime()); + PartitionSpec spec = PartitionSpec.builderFor(SHARDED_SCHEMA).identity("shard").build(); + Table table = warehouse.createTable(id, SHARDED_SCHEMA, spec); + + TaskWriter writer = factoryFor(table, spec, Long.MAX_VALUE).create(); + for (int shard = 0; shard < 3; shard++) { + for (int i = 0; i < 2; i++) { + writer.write(shardedRow(shard * 10L + i, shard)); + } + } + DataFile[] files = writer.dataFiles(); + + assertEquals("one output file per partition", 3, files.length); + Set registeredShards = new HashSet<>(); + for (DataFile f : files) { + registeredShards.add(f.partition().get(0, Integer.class)); + assertEquals("each partition's rows stay together", 2L, f.recordCount()); + } + assertEquals(new HashSet<>(java.util.Arrays.asList(0, 1, 2)), registeredShards); + } + + @Test + public void fanoutBeyondOpenWriterCapFailsWithGuidance() throws Exception { + // A repartitioning subgroup can fan out to many partitions, each holding an open appender with + // its own row-group buffers. The cap turns that OOM into a diagnosable failure. + TableIdentifier id = TableIdentifier.of("default", "wf_cap_" + System.nanoTime()); + PartitionSpec spec = PartitionSpec.builderFor(SHARDED_SCHEMA).identity("shard").build(); + Table table = warehouse.createTable(id, SHARDED_SCHEMA, spec); + + int originalCap = WriterFactory.maxOpenFanoutWriters; + WriterFactory.maxOpenFanoutWriters = 1; + try { + TaskWriter writer = factoryFor(table, spec, Long.MAX_VALUE).create(); + writer.write(shardedRow(0L, 0)); + IllegalStateException ex = + assertThrows(IllegalStateException.class, () -> writer.write(shardedRow(1L, 1))); + assertTrue( + "the message must guide the operator: " + ex.getMessage(), + ex.getMessage().contains("simultaneously-open writers")); + writer.abort(); + } finally { + WriterFactory.maxOpenFanoutWriters = originalCap; + } + } + + @Test + public void writesRollOnceTheTargetFileSizeIsExceeded() throws Exception { + TableIdentifier id = TableIdentifier.of("default", "wf_roll_" + System.nanoTime()); + Table table = warehouse.createTable(id, TestFixtures.SCHEMA); + + TaskWriter writer = factoryFor(table, table.spec(), 1L).create(); + List rows = new ArrayList<>(); + // Iceberg only evaluates the size every 1000th row + for (int i = 0; i < 2500; i++) { + Record r = GenericRecord.create(TestFixtures.SCHEMA); + r.setField("id", (long) i); + r.setField("data", "row-" + i); + rows.add(r); + } + for (Record r : rows) { + writer.write(r); + } + DataFile[] files = writer.dataFiles(); + + assertTrue("a 1-byte target must roll into several files", files.length > 1); + long total = 0; + for (DataFile f : files) { + total += f.recordCount(); + } + assertEquals("no record may be lost across the roll", 2500L, total); + } +}