diff --git a/.github/trigger_files/IO_Iceberg_Integration_Tests.json b/.github/trigger_files/IO_Iceberg_Integration_Tests.json index 7ab7bcd9a9c6..37dd25bf9029 100644 --- a/.github/trigger_files/IO_Iceberg_Integration_Tests.json +++ b/.github/trigger_files/IO_Iceberg_Integration_Tests.json @@ -1,4 +1,4 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run.", - "modification": 2 + "modification": 3 } diff --git a/CHANGES.md b/CHANGES.md index 04ac961fcd93..abc3c02a36ad 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -99,6 +99,7 @@ * (Python) `Regex.all_matches` now returns every group of the regex, instead of failing with a `TypeError` when the regex has no groups and dropping nested or unmatched trailing groups ([#40562](https://github.com/apache/beam/issues/40562)). * (Python) `--dataflow_endpoint` now defaults to `''` instead of `https://dataflow.googleapis.com`. When unset, the Google Cloud client library resolves the Dataflow endpoint and honors `GOOGLE_CLOUD_UNIVERSE_DOMAIN`, so Dataflow jobs can be submitted in non-default universes. Code that reads `GoogleCloudOptions.dataflow_endpoint` should handle an empty value ([#40569](https://github.com/apache/beam/issues/40569)). * (Python) The `xgboost` extras now install `xgboost>=2.0.0,<4` and `datatable>=1.1.0,<2`, moving both the lower and upper bounds for both. The previous minimum `xgboost` and `datatable` versions (`1.6.0` and `1.0.0` respectively) are no longer allowed ([#40694](https://github.com/apache/beam/issues/40694)). +* (Java) [IcebergIO] AddFiles (`IcebergAddFiles` in YAML) now carries the destination table alongside each file through the registration chain, so that one transform can register files into many tables. The coders of its internal `PCollection`s change, which makes this an update incompatible change for streaming pipelines that use it ([#40817](https://github.com/apache/beam/pull/40817)). ## Deprecations diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/AddFiles.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/AddFiles.java index 00cf77af1dca..c97a20c6221b 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/AddFiles.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/AddFiles.java @@ -38,11 +38,12 @@ import java.util.UUID; import java.util.concurrent.Callable; import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; import java.util.concurrent.atomic.AtomicBoolean; import java.util.stream.Collectors; import org.apache.beam.sdk.coders.KvCoder; import org.apache.beam.sdk.coders.RowCoder; -import org.apache.beam.sdk.coders.VarIntCoder; +import org.apache.beam.sdk.coders.StringUtf8Coder; import org.apache.beam.sdk.coders.VarLongCoder; import org.apache.beam.sdk.io.iceberg.SchemaEvolutionConfig.UnverifiableFileHandling; import org.apache.beam.sdk.metrics.Counter; @@ -58,6 +59,7 @@ import org.apache.beam.sdk.transforms.GroupIntoBatches; import org.apache.beam.sdk.transforms.PTransform; import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.Values; import org.apache.beam.sdk.transforms.Wait; import org.apache.beam.sdk.transforms.WithKeys; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; @@ -71,6 +73,7 @@ import org.apache.beam.sdk.values.Row; import org.apache.beam.sdk.values.TupleTag; import org.apache.beam.sdk.values.TupleTagList; +import org.apache.beam.sdk.values.TypeDescriptors; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Strings; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.hash.Hasher; @@ -284,13 +287,17 @@ public PCollectionRowTuple expand(PCollection input) { paths = windowed.apply("WaitForSchemaCommit", Wait.on(committed)); } - PCollectionTuple dataFiles = + PCollection> keyedPaths = paths.apply( + "KeyByTable", + WithKeys.of(tableIdentifier).withKeyType(TypeDescriptors.strings())); + + PCollectionTuple dataFiles = + keyedPaths.apply( "ConvertToDataFiles", ParDo.of( new ConvertToDataFile( catalogConfig, - tableIdentifier, locationPrefix, partitionFields, sortFields, @@ -298,20 +305,37 @@ public PCollectionRowTuple expand(PCollection input) { evolution)) .withOutputTags(DATA_FILES, TupleTagList.of(ERRORS))); SchemaCoder sdfCoder; + SchemaCoder manifestKeyCoder; try { - sdfCoder = SchemaRegistry.createDefault().getSchemaCoder(SerializableDataFile.class); + SchemaRegistry schemas = SchemaRegistry.createDefault(); + sdfCoder = schemas.getSchemaCoder(SerializableDataFile.class); + manifestKeyCoder = schemas.getSchemaCoder(ManifestKey.class); } catch (Exception e) { throw new RuntimeException(e); } - PCollection> keyedFiles = + PCollection> keyedFiles = dataFiles .get(DATA_FILES) - .setCoder(sdfCoder) - .apply("AddSpecIdKey", WithKeys.of(SerializableDataFile::getPartitionSpecId)) - .setCoder(KvCoder.of(VarIntCoder.of(), sdfCoder)); - - GroupIntoBatches batchDataFiles = + .setCoder(KvCoder.of(StringUtf8Coder.of(), sdfCoder)) + .apply( + "AddTableAndSpecIdKey", + ParDo.of( + new DoFn< + KV, KV>() { + @ProcessElement + public void process( + @Element KV element, + OutputReceiver> out) { + SerializableDataFile file = element.getValue(); + out.output( + KV.of( + ManifestKey.of(element.getKey(), file.getPartitionSpecId()), file)); + } + })) + .setCoder(KvCoder.of(manifestKeyCoder, sdfCoder)); + + GroupIntoBatches batchDataFiles = GroupIntoBatches.ofSize(manifestFileSize); GroupIntoBatches batchManifestFiles = GroupIntoBatches.ofSize(DEFAULT_MAX_MANIFESTS_PER_SNAPSHOT); @@ -322,29 +346,38 @@ public PCollectionRowTuple expand(PCollection input) { batchManifestFiles.withMaxBufferingDuration(checkStateNotNull(intervalTrigger)); } - PCollection>> groupedFiles = + PCollection>> groupedFiles = keyedFiles.apply("GroupDataFilesIntoBatches", batchDataFiles); PCollection> manifests = - groupedFiles.apply( - "CreateManifests", ParDo.of(new CreateManifests(catalogConfig, tableIdentifier))); + groupedFiles.apply("CreateManifests", ParDo.of(new CreateManifests(catalogConfig))); PCollection snapshots = manifests .apply("GatherManifests", batchManifestFiles) - .apply( - "CommitManifests", - ParDo.of(new CommitManifestFilesDoFn(catalogConfig, tableIdentifier))) + .apply("CommitManifests", ParDo.of(new CommitManifestFilesDoFn(catalogConfig))) .setRowSchema(SnapshotInfo.getSchema()); return PCollectionRowTuple.of( OUTPUT_TAG, snapshots, ERROR_TAG, dataFiles.get(ERRORS).setRowSchema(ERROR_SCHEMA)); } + /** The footer read is keyed, so the single-table path supplies its one table as the key. */ + private PCollection> fileSchemas( + PCollection> keyed) { + return keyed + .apply("ReadFooterSchema", ParDo.of(new ReadFooterSchema(evolution, locationPrefix))) + .setCoder(KvCoder.of(StringUtf8Coder.of(), CollectDistinctSchemas.groupCoder())); + } + private PCollection> distinctSchemas( PCollection windowed) { - return windowed - .apply("ReadFooterSchema", ParDo.of(new ReadFooterSchema(evolution, locationPrefix))) + PCollection> keyed = + windowed.apply( + "KeyBySingleTable", + WithKeys.of(tableIdentifier).withKeyType(TypeDescriptors.strings())); + return fileSchemas(keyed) + .apply("DropTableKey", Values.create()) .setCoder(CollectDistinctSchemas.groupCoder()) .apply( "CollectDistinctSchemas", @@ -413,21 +446,37 @@ AddFiles withSchemaCommitter(CommitSchemaUnion.Committer committer) { * back to the runner on the main thread, as {@link MultiOutputReceiver} is not thread-safe. * */ - static class ConvertToDataFile extends DoFn { + static class ConvertToDataFile + extends DoFn, KV> { private final IcebergCatalogConfig catalogConfig; - private final String identifier; public static final TupleTag ERRORS = new TupleTag<>(); - public static final TupleTag DATA_FILES = new TupleTag<>(); + public static final TupleTag> DATA_FILES = new TupleTag<>(); private final @Nullable String prefix; private final @Nullable List partitionFields; private final @Nullable List sortFields; private final @Nullable Map tableProps; private final SchemaEvolutionConfig evolution; private transient @MonotonicNonNull BoundedAsyncTasks tasks; - private transient volatile @MonotonicNonNull Table table; - private transient volatile boolean tableMissing; private transient @MonotonicNonNull Set warned; - private final AtomicBoolean refreshedThisBundle = new AtomicBoolean(); + private transient @MonotonicNonNull ConcurrentMap states; + + /** One table's cached handle and flags; the single-table fields became a map of these. */ + private static final class TableState { + volatile @Nullable Table table; + volatile boolean missing; + final AtomicBoolean refreshedThisBundle = new AtomicBoolean(); + } + + private TableState stateFor(String tableName) { + ConcurrentMap map = checkStateNotNull(states); + TableState existing = map.get(tableName); + if (existing != null) { + return existing; + } + TableState created = new TableState(); + TableState raced = map.putIfAbsent(tableName, created); + return raced != null ? raced : created; + } // Number of parallel threads processing incoming files private static final int THREAD_POOL_SIZE = 10; @@ -435,14 +484,12 @@ static class ConvertToDataFile extends DoFn { public ConvertToDataFile( IcebergCatalogConfig catalogConfig, - String identifier, @Nullable String prefix, @Nullable List partitionFields, @Nullable List sortFields, @Nullable Map tableProps) { this( catalogConfig, - identifier, prefix, partitionFields, sortFields, @@ -452,14 +499,12 @@ public ConvertToDataFile( public ConvertToDataFile( IcebergCatalogConfig catalogConfig, - String identifier, @Nullable String prefix, @Nullable List partitionFields, @Nullable List sortFields, @Nullable Map tableProps, SchemaEvolutionConfig evolution) { this.catalogConfig = catalogConfig; - this.identifier = identifier; this.prefix = prefix; this.partitionFields = partitionFields; this.sortFields = sortFields; @@ -516,6 +561,7 @@ static Verdict unverified(Unverified what) { } private static class ProcessResult { + final String table; final @Nullable SerializableDataFile dataFile; final @Nullable Row errorRow; @@ -527,6 +573,7 @@ private static class ProcessResult { final PaneInfo paneInfo; ProcessResult( + String table, @Nullable SerializableDataFile dataFile, @Nullable Row errorRow, @Nullable Unverified unverified, @@ -538,6 +585,7 @@ private static class ProcessResult { "Expected only one of dataFile or errorRow, but got both:%n\tfile: %s%n\terror: %s", dataFile != null ? dataFile.getPath() : null, errorRow); + this.table = table; this.dataFile = dataFile; this.errorRow = errorRow; this.unverified = unverified; @@ -551,13 +599,16 @@ private static class ProcessResult { public void setup() { tasks = new BoundedAsyncTasks<>(THREAD_POOL_SIZE, MAX_IN_FLIGHT_TASKS); warned = ConcurrentHashMap.newKeySet(); + states = new ConcurrentHashMap<>(); } /** Clears anything left behind if the runner reuses this instance after a failed bundle. */ @StartBundle public void startBundle() { checkStateNotNull(tasks).cancelAll(); - refreshedThisBundle.set(false); + for (TableState state : checkStateNotNull(states).values()) { + state.refreshedThisBundle.set(false); + } } @Teardown @@ -569,13 +620,14 @@ public void teardown() { @ProcessElement public void process( - @Element String filePath, + @Element KV element, @Timestamp Instant timestamp, BoundedWindow window, PaneInfo paneInfo, MultiOutputReceiver output) throws Exception { - Callable task = createProcessTask(filePath, timestamp, window, paneInfo); + Callable task = + createProcessTask(element.getKey(), element.getValue(), timestamp, window, paneInfo); checkStateNotNull(tasks).submit(task, result -> outputResult(result, output)); } @@ -593,7 +645,7 @@ private void outputResult(ProcessResult result, MultiOutputReceiver output) { output .get(DATA_FILES) .outputWindowedValue( - result.dataFile, + KV.of(result.table, result.dataFile), result.timestamp, Collections.singleton(result.window), result.paneInfo); @@ -611,7 +663,8 @@ private static void outputAtFinish(ProcessResult result, FinishBundleContext con context.output(ERRORS, result.errorRow, result.timestamp, result.window); numErrorFiles.inc(); } else if (result.dataFile != null) { - context.output(DATA_FILES, result.dataFile, result.timestamp, result.window); + context.output( + DATA_FILES, KV.of(result.table, result.dataFile), result.timestamp, result.window); countUnverified(result); } } @@ -625,9 +678,14 @@ private static void countUnverified(ProcessResult result) { } private Callable createProcessTask( - String filePath, Instant timestamp, BoundedWindow window, PaneInfo paneInfo) { + String tableName, + String filePath, + Instant timestamp, + BoundedWindow window, + PaneInfo paneInfo) { return () -> { + TableState state = stateFor(tableName); FileFormat format; @Nullable ParquetMetadata parquetFooter = null; if (hasNoExtension(filePath)) { @@ -635,52 +693,65 @@ private Callable createProcessTask( parquetFooter = ParquetFooters.read(filePath); } catch (Exception e) { return errorResult( - filePath, NOT_PARQUET_ERROR + errorMessage(e), timestamp, window, paneInfo); + tableName, + filePath, + NOT_PARQUET_ERROR + errorMessage(e), + timestamp, + window, + paneInfo); } format = FileFormat.PARQUET; } else { try { format = inferFormat(filePath); } catch (UnknownFormatException e) { - return errorResult(filePath, UNKNOWN_FORMAT_ERROR, timestamp, window, paneInfo); + return errorResult( + tableName, filePath, UNKNOWN_FORMAT_ERROR, timestamp, window, paneInfo); } } // ---- Infrastructure phase. Failures propagate so the runner retries the bundle; // per-file error rows here would silently drop in-flight files on a transient blip. // Only conditions that are properties of the file go to the error output. - if (tableMissing) { + if (state.missing) { return errorResult( - filePath, MISSING_TABLE_ERROR + identifier, timestamp, window, paneInfo); + tableName, filePath, MISSING_TABLE_ERROR + tableName, timestamp, window, paneInfo); } - if (table == null) { - synchronized (this) { - if (table == null) { + if (state.table == null) { + synchronized (state) { + if (state.table == null) { try { - table = getOrCreateTable(filePath, format); + state.table = getOrCreateTable(tableName, filePath, format); } catch (FileNotFoundException e) { - return errorResult(filePath, errorMessage(e), timestamp, window, paneInfo); + return errorResult( + tableName, filePath, errorMessage(e), timestamp, window, paneInfo); } catch (NoSuchTableException e) { - tableMissing = true; + state.missing = true; return errorResult( - filePath, MISSING_TABLE_ERROR + identifier, timestamp, window, paneInfo); + tableName, + filePath, + MISSING_TABLE_ERROR + tableName, + timestamp, + window, + paneInfo); } } } } + Table table = checkStateNotNull(state.table); // Check if the file path contains the provided prefix if (table.spec().isPartitioned() && !Strings.isNullOrEmpty(prefix) && !filePath.startsWith(checkStateNotNull(prefix))) { - return errorResult(filePath, PREFIX_ERROR, timestamp, window, paneInfo); + return errorResult(tableName, filePath, PREFIX_ERROR, timestamp, window, paneInfo); } - if (table.schema().columns().isEmpty() && firstTime("empty schema")) { + if (table.schema().columns().isEmpty() && firstTime(tableName, "empty schema")) { LOG.warn( "Table {} has no columns: files register with no readable columns and no stats." + " Enable schema evolution to infer the schema from the files.", - identifier); + tableName); } // ---- Per-file phase: every failure below is one error row, never a failed bundle. @@ -688,15 +759,15 @@ private Callable createProcessTask( try { parquetFooter = ParquetFooters.read(filePath); } catch (Exception e) { - return errorResult(filePath, errorMessage(e), timestamp, window, paneInfo); + return errorResult(tableName, filePath, errorMessage(e), timestamp, window, paneInfo); } } Verdict verdict = Verdict.OK; if (evolution.isEnabled()) { - verdict = verify(filePath, format, parquetFooter); + verdict = verify(table, tableName, state, filePath, format, parquetFooter); } if (verdict.error != null) { - return errorResult(filePath, verdict.error, timestamp, window, paneInfo); + return errorResult(tableName, filePath, verdict.error, timestamp, window, paneInfo); } // What readers resolve a file without field ids with, so its metrics and partition agree @@ -712,7 +783,7 @@ private Callable createProcessTask( resolvedFooter = ParquetFieldIds.resolve(parquetFooter, table, mapping); } catch (ParquetFieldIds.ConflictException e) { return errorResult( - filePath, FIELD_ID_ERROR + e.getMessage(), timestamp, window, paneInfo); + tableName, filePath, FIELD_ID_ERROR + e.getMessage(), timestamp, window, paneInfo); } } @@ -729,7 +800,7 @@ private Callable createProcessTask( table.schema(), resolvedFooter); } catch (Exception e) { - return errorResult(filePath, errorMessage(e), timestamp, window, paneInfo); + return errorResult(tableName, filePath, errorMessage(e), timestamp, window, paneInfo); } // Figure out which partition this DataFile should go to @@ -739,7 +810,7 @@ private Callable createProcessTask( if (partitioned && !Strings.isNullOrEmpty(prefix)) { // option 1: use directory structure to determine partition // Note: we don't validate the DataFile content here - partitionPath = getPartitionFromFilePath(filePath); + partitionPath = getPartitionFromFilePath(table, filePath); } else if (partitioned) { try { // option 2: examine DataFile min/max statistics to determine partition @@ -747,12 +818,22 @@ private Callable createProcessTask( getPartitionFromMetrics(metrics, inputFile, table, mapping, resolvedFooter); } catch (UnknownPartitionException e) { return errorResult( - filePath, UNKNOWN_PARTITION_ERROR + e.getMessage(), timestamp, window, paneInfo); + tableName, + filePath, + UNKNOWN_PARTITION_ERROR + e.getMessage(), + timestamp, + window, + paneInfo); } catch (RuntimeException e) { // e.g. a bound that does not decode as the table's type: one error row, not a failed // bundle return errorResult( - filePath, UNKNOWN_PARTITION_ERROR + errorMessage(e), timestamp, window, paneInfo); + tableName, + filePath, + UNKNOWN_PARTITION_ERROR + errorMessage(e), + timestamp, + window, + paneInfo); } } @@ -771,6 +852,7 @@ private Callable createProcessTask( } DataFile df = builder.build(); return new ProcessResult( + tableName, SerializableDataFile.from(df, table.spec()), null, verdict.unverified, @@ -779,7 +861,7 @@ private Callable createProcessTask( paneInfo); } catch (Exception e) { // getLength is a per-file read (e.g. the file was deleted mid-flight). - return errorResult(filePath, errorMessage(e), timestamp, window, paneInfo); + return errorResult(tableName, filePath, errorMessage(e), timestamp, window, paneInfo); } }; } @@ -788,19 +870,25 @@ private Callable createProcessTask( * The checks the options promise, in order: a format the checks can read, a convertible schema, * coverage by the table, pins. The first failure is the verdict. */ - private Verdict verify(String filePath, FileFormat format, @Nullable ParquetMetadata footer) { + private Verdict verify( + Table table, + String tableName, + TableState state, + String filePath, + FileFormat format, + @Nullable ParquetMetadata footer) { if (!format.equals(FileFormat.PARQUET)) { if (evolution.getUnverifiableFileHandling() == UnverifiableFileHandling.REJECT) { return Verdict.error(UNCHECKED_FORMAT_ERROR + format.name()); } - if (firstTime("unchecked format")) { + if (firstTime(tableName, "unchecked format")) { LOG.warn( "Registering {} files in table {} unchecked (UnverifiableFileHandling.ACCEPT):" + " coverage and pin checks read only Parquet, so a required column such a file" + " lacks or holds nulls in fails reads of the table, not registration. First" + " file: {}", format, - identifier, + tableName, filePath); } return Verdict.unverified(Unverified.FORMAT); @@ -812,12 +900,15 @@ private Verdict verify(String filePath, FileFormat format, @Nullable ParquetMeta } catch (Exception e) { return Verdict.error(UNREADABLE_SCHEMA_ERROR + errorMessage(e)); } - Set identityPartitionColumnsFromPath = identityPartitionColumnsFromPath(filePath); - @Nullable String uncovered = uncoveredReason(fileSchema, identityPartitionColumnsFromPath); + Set identityPartitionColumnsFromPath = + identityPartitionColumnsFromPath(table, filePath); + @Nullable String uncovered = + uncoveredReason(table, state, fileSchema, identityPartitionColumnsFromPath); if (uncovered != null) { return Verdict.error(uncovered); } - return checkPins(filePath, fileSchema, parquetFooter, identityPartitionColumnsFromPath); + return checkPins( + table, tableName, filePath, fileSchema, parquetFooter, identityPartitionColumnsFromPath); } /** @@ -825,8 +916,7 @@ private Verdict verify(String filePath, FileFormat format, @Nullable ParquetMeta * for a file under {@code dt=2024-01-15/} when the table is partitioned by {@code dt}. Empty * without a prefix. */ - private Set identityPartitionColumnsFromPath(String filePath) { - Table table = checkStateNotNull(this.table); + private Set identityPartitionColumnsFromPath(Table table, String filePath) { if (Strings.isNullOrEmpty(prefix) || table.spec().isUnpartitioned()) { return Collections.emptySet(); } @@ -840,16 +930,18 @@ private Set identityPartitionColumnsFromPath(String filePath) { * was cached) and report the remaining delta. Never changes the schema. */ private @Nullable String uncoveredReason( - org.apache.iceberg.Schema fileSchema, Set identityPartitionColumnsFromPath) { - Table table = checkStateNotNull(this.table); + Table table, + TableState state, + org.apache.iceberg.Schema fileSchema, + Set identityPartitionColumnsFromPath) { SchemaDelta delta = SchemaDelta.classify(table, fileSchema, identityPartitionColumnsFromPath); if (delta.isEmpty()) { return null; } // a commit can land after the table was cached; one refresh per bundle is enough to see it, // and a bundle full of routed files must not load the table once per file - if (refreshedThisBundle.compareAndSet(false, true)) { - synchronized (this) { + if (state.refreshedThisBundle.compareAndSet(false, true)) { + synchronized (state) { table.refresh(); } } @@ -874,11 +966,12 @@ private Set identityPartitionColumnsFromPath(String filePath) { * goes on, so a pin the footer does count nulls for still fails the file. */ private Verdict checkPins( + Table table, + String tableName, String filePath, org.apache.iceberg.Schema fileSchema, ParquetMetadata footer, Set identityPartitionColumnsFromPath) { - Table table = checkStateNotNull(this.table); List unproven = new ArrayList<>(); for (String pinned : evolution.getRequiredColumns()) { org.apache.iceberg.types.Types.@Nullable NestedField tableField = @@ -910,12 +1003,12 @@ private Verdict checkPins( if (unproven.isEmpty()) { return Verdict.OK; } - if (firstTime("unproven pins")) { + if (firstTime(tableName, "unproven pins")) { LOG.warn( "Registering files in table {} whose footer has no null count statistics for pinned" + " column(s) {} on trust (UnverifiableFileHandling.ACCEPT): nulls there fail" + " reads of the table, not registration. First file: {}", - identifier, + tableName, unproven, filePath); } @@ -923,13 +1016,19 @@ private Verdict checkPins( } /** Once per instance per key: at volume a line per file would drown the log. */ - private boolean firstTime(String key) { - return checkStateNotNull(warned).add(key); + private boolean firstTime(String tableName, String key) { + return checkStateNotNull(warned).add(tableName + ":" + key); } private static ProcessResult errorResult( - String filePath, String message, Instant timestamp, BoundedWindow window, PaneInfo pane) { + String tableName, + String filePath, + String message, + Instant timestamp, + BoundedWindow window, + PaneInfo pane) { return new ProcessResult( + tableName, null, Row.withSchema(ERROR_SCHEMA).addValues(filePath, message).build(), null, @@ -942,8 +1041,9 @@ static T transformValue(Transform transform, Type type, ByteBuffer return transform.bind(type).apply(Conversions.fromByteBuffer(type, bytes)); } - private Table getOrCreateTable(String filePath, FileFormat format) throws IOException { - TableIdentifier tableId = IcebergUtils.parseTableIdentifier(identifier); + private Table getOrCreateTable(String tableName, String filePath, FileFormat format) + throws IOException { + TableIdentifier tableId = IcebergUtils.parseTableIdentifier(tableName); try { return catalogConfig.catalog().loadTable(tableId); } catch (NoSuchTableException e) { @@ -973,7 +1073,7 @@ private Table getOrCreateTable(String filePath, FileFormat format) throws IOExce .create(); } catch (AlreadyExistsException e2) { // if table already exists, just load it - return catalogConfig.catalog().loadTable(IcebergUtils.parseTableIdentifier(identifier)); + return catalogConfig.catalog().loadTable(tableId); } } } @@ -991,8 +1091,8 @@ private static org.apache.iceberg.Schema getSchema(String filePath, FileFormat f return ParquetSchemaUtil.convert(messageType); } - private String getPartitionFromFilePath(String filePath) { - if (checkStateNotNull(table).spec().isUnpartitioned()) { + private String getPartitionFromFilePath(Table table, String filePath) { + if (table.spec().isUnpartitioned()) { return ""; } return PartitionPaths.relative(filePath, checkStateNotNull(prefix)); @@ -1188,29 +1288,34 @@ private static boolean lacksColumn(ParquetFieldIds.@Nullable Resolved footer, in * downstream {@link CommitManifestFilesDoFn}. */ static class CreateManifests - extends DoFn>, KV> { + extends DoFn>, KV> { private final IcebergCatalogConfig catalogConfig; - private final String identifier; - private transient @MonotonicNonNull Table table; + private transient @MonotonicNonNull Map tables; - public CreateManifests(IcebergCatalogConfig catalogConfig, String identifier) { + public CreateManifests(IcebergCatalogConfig catalogConfig) { this.catalogConfig = catalogConfig; - this.identifier = identifier; } @ProcessElement public void process( - @Element KV> batch, + @Element KV> batch, OutputReceiver> output) throws IOException { if (!batch.getValue().iterator().hasNext()) { return; } + String identifier = batch.getKey().getTable(); + if (tables == null) { + tables = new HashMap<>(); + } + Table table = tables.get(identifier); if (table == null) { table = catalogConfig.catalog().loadTable(IcebergUtils.parseTableIdentifier(identifier)); + tables.put(identifier, table); } - PartitionSpec spec = checkStateNotNull(table.specs().get(batch.getKey())); + PartitionSpec spec = + checkStateNotNull(table.specs().get(batch.getKey().getPartitionSpecId())); String manifestPath = String.format( @@ -1258,17 +1363,15 @@ public void process( */ static class CommitManifestFilesDoFn extends DoFn>, Row> { private final IcebergCatalogConfig catalogConfig; - private final String identifier; - private transient @MonotonicNonNull Table table = null; + private transient @MonotonicNonNull Map tables = null; private static final String COMMIT_ID_KEY = "beam.add-files-commit-id"; @StateId("lastCommitTimestamp") private final StateSpec> lastCommitTimestamp = StateSpecs.value(VarLongCoder.of()); - public CommitManifestFilesDoFn(IcebergCatalogConfig catalogConfig, String identifier) { + public CommitManifestFilesDoFn(IcebergCatalogConfig catalogConfig) { this.catalogConfig = catalogConfig; - this.identifier = identifier; } private static void ensureNameMappingPresent(Table table) { @@ -1303,13 +1406,19 @@ public void process( manifests.add(ManifestFiles.decode(bytes)); } String commitId = commitHash(manifests); + String identifier = batch.getKey(); + if (tables == null) { + tables = new HashMap<>(); + } + Table table = tables.get(identifier); if (table == null) { table = catalogConfig.catalog().loadTable(IcebergUtils.parseTableIdentifier(identifier)); + tables.put(identifier, table); } table.refresh(); ensureNameMappingPresent(table); - if (shouldSkip(commitId, lastCommitTimestamp.read())) { + if (shouldSkip(table, commitId, lastCommitTimestamp.read())) { return; } @@ -1349,11 +1458,10 @@ private String commitHash(Iterable files) { * Performs a look-back through Iceberg table history to determine if this specific batch of * files has already been successfully committed. */ - private boolean shouldSkip(String commitUID, @Nullable Long lastCommitTimestamp) { + private boolean shouldSkip(Table table, String commitUID, @Nullable Long lastCommitTimestamp) { if (lastCommitTimestamp == null) { return false; } - Table table = checkStateNotNull(this.table); // check past snapshots to see if they contain the commit ID @Nullable Snapshot current = table.currentSnapshot(); diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/ManifestKey.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/ManifestKey.java new file mode 100644 index 000000000000..90a8d62942df --- /dev/null +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/ManifestKey.java @@ -0,0 +1,45 @@ +/* + * 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 com.google.auto.value.AutoValue; +import org.apache.beam.sdk.schemas.AutoValueSchema; +import org.apache.beam.sdk.schemas.annotations.DefaultSchema; +import org.apache.beam.sdk.schemas.annotations.SchemaFieldNumber; + +/** + * As much as one manifest file can hold: entries for a single table, all written against a single + * partition spec. Data files are batched by this before a manifest is written for each batch. + * + *

A schema type rather than a {@code KV} of the two. The field numbers pin each field's place in + * the encoding, so a third field can be added after these two without moving them. + */ +@DefaultSchema(AutoValueSchema.class) +@AutoValue +abstract class ManifestKey { + + static ManifestKey of(String table, int partitionSpecId) { + return new AutoValue_ManifestKey(table, partitionSpecId); + } + + @SchemaFieldNumber("0") + abstract String getTable(); + + @SchemaFieldNumber("1") + abstract int getPartitionSpecId(); +} diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/ReadFooterSchema.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/ReadFooterSchema.java index 37310b8d7f71..9f7eafb74552 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/ReadFooterSchema.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/ReadFooterSchema.java @@ -28,6 +28,7 @@ import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; import org.apache.beam.sdk.transforms.windowing.PaneInfo; +import org.apache.beam.sdk.values.KV; import org.apache.iceberg.FileFormat; import org.apache.parquet.hadoop.metadata.ParquetMetadata; import org.checkerframework.checker.nullness.qual.MonotonicNonNull; @@ -40,7 +41,8 @@ * Emits one {@link CollectDistinctSchemas.SchemaGroup} of one file per readable Parquet file: its * declared schema and its null-free columns. Unreadable or non-Parquet files contribute nothing. */ -class ReadFooterSchema extends DoFn { +class ReadFooterSchema + extends DoFn, KV> { private static final Logger LOG = LoggerFactory.getLogger(ReadFooterSchema.class); static final int DEFAULT_THREAD_POOL_SIZE = 10; @@ -94,6 +96,7 @@ class ReadFooterSchema extends DoFn /** Counters are updated on the processing thread: metrics touched from the executor are lost. */ private static class ReadResult { + final String table; final CollectDistinctSchemas.@Nullable SchemaGroup schema; final boolean footerError; final Instant timestamp; @@ -101,11 +104,13 @@ private static class ReadResult { final PaneInfo paneInfo; ReadResult( + String table, CollectDistinctSchemas.@Nullable SchemaGroup schema, boolean footerError, Instant timestamp, BoundedWindow window, PaneInfo paneInfo) { + this.table = table; this.schema = schema; this.footerError = footerError; this.timestamp = timestamp; @@ -134,14 +139,15 @@ public void teardown() { @ProcessElement public void process( - @Element String filePath, + @Element KV element, @Timestamp Instant timestamp, BoundedWindow window, PaneInfo paneInfo, - OutputReceiver output) + OutputReceiver> output) throws Exception { numFilesRead.inc(); - Callable task = createReadTask(filePath, timestamp, window, paneInfo); + Callable task = + createReadTask(element.getKey(), element.getValue(), timestamp, window, paneInfo); checkStateNotNull(tasks).submit(task, result -> outputResult(result, output)); } @@ -153,16 +159,19 @@ public void finishBundle(FinishBundleContext context) throws Exception { private static void outputAtFinish(ReadResult result, FinishBundleContext context) { count(result); if (result.schema != null) { - context.output(result.schema, result.timestamp, result.window); + context.output(KV.of(result.table, result.schema), result.timestamp, result.window); } } private static void outputResult( - ReadResult result, OutputReceiver output) { + ReadResult result, OutputReceiver> output) { count(result); if (result.schema != null) { output.outputWindowedValue( - result.schema, result.timestamp, Collections.singleton(result.window), result.paneInfo); + KV.of(result.table, result.schema), + result.timestamp, + Collections.singleton(result.window), + result.paneInfo); } } @@ -176,30 +185,30 @@ private static void count(ReadResult result) { } private Callable createReadTask( - String filePath, Instant timestamp, BoundedWindow window, PaneInfo paneInfo) { + String table, String filePath, Instant timestamp, BoundedWindow window, PaneInfo paneInfo) { return () -> { ParquetMetadata footer; if (AddFiles.hasNoExtension(filePath)) { try { footer = ParquetFooters.read(filePath); } catch (Exception e) { - return unknownFormat(timestamp, window, paneInfo); + return unknownFormat(table, timestamp, window, paneInfo); } } else { FileFormat format; try { format = AddFiles.inferFormat(filePath); } catch (AddFiles.UnknownFormatException e) { - return unknownFormat(timestamp, window, paneInfo); + return unknownFormat(table, timestamp, window, paneInfo); } if (!format.equals(FileFormat.PARQUET)) { return new ReadResult( - dryRunMarker(UNCHECKED_FORMAT_KEY), false, timestamp, window, paneInfo); + table, dryRunMarker(UNCHECKED_FORMAT_KEY), false, timestamp, window, paneInfo); } try { footer = ParquetFooters.read(filePath); } catch (Exception e) { - return unreadableFooter(filePath, e, timestamp, window, paneInfo); + return unreadableFooter(table, filePath, e, timestamp, window, paneInfo); } } try { @@ -207,33 +216,44 @@ private Callable createReadTask( new ArrayList<>(PartitionPaths.partitionKeysInPath(filePath, locationPrefix)); CollectDistinctSchemas.SchemaGroup group = FileSchemas.schemaGroup(footer).withPartitionKeysInPath(partitionKeysInPath); - return new ReadResult(group, false, timestamp, window, paneInfo); + return new ReadResult(table, group, false, timestamp, window, paneInfo); } catch (Exception e) { - return unconvertibleSchema(filePath, e, timestamp, window, paneInfo); + return unconvertibleSchema(table, filePath, e, timestamp, window, paneInfo); } }; } - private ReadResult unknownFormat(Instant timestamp, BoundedWindow window, PaneInfo paneInfo) { - return new ReadResult(dryRunMarker(UNREADABLE_KEY), false, timestamp, window, paneInfo); + private ReadResult unknownFormat( + String table, Instant timestamp, BoundedWindow window, PaneInfo paneInfo) { + return new ReadResult(table, dryRunMarker(UNREADABLE_KEY), false, timestamp, window, paneInfo); } private ReadResult unreadableFooter( - String filePath, Exception e, Instant timestamp, BoundedWindow window, PaneInfo paneInfo) { + String table, + String filePath, + Exception e, + Instant timestamp, + BoundedWindow window, + PaneInfo paneInfo) { LOG.warn( "Could not read the footer of {}; the file will not contribute to schema inference: {}", filePath, AddFiles.errorMessage(e)); - return new ReadResult(dryRunMarker(UNREADABLE_KEY), true, timestamp, window, paneInfo); + return new ReadResult(table, dryRunMarker(UNREADABLE_KEY), true, timestamp, window, paneInfo); } private ReadResult unconvertibleSchema( - String filePath, Exception e, Instant timestamp, BoundedWindow window, PaneInfo paneInfo) { + String table, + String filePath, + Exception e, + Instant timestamp, + BoundedWindow window, + PaneInfo paneInfo) { LOG.warn( "Could not convert the schema of {}; the file will not contribute to schema inference: {}", filePath, AddFiles.errorMessage(e)); - return new ReadResult(dryRunMarker(UNREADABLE_KEY), false, timestamp, window, paneInfo); + return new ReadResult(table, dryRunMarker(UNREADABLE_KEY), false, timestamp, window, paneInfo); } /** A dry run counts files that contribute no schema; a real run handles them at registration. */ diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/AddFilesTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/AddFilesTest.java index 47aa8290f724..64ea7e873828 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/AddFilesTest.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/AddFilesTest.java @@ -52,12 +52,15 @@ import java.util.stream.Collectors; import org.apache.beam.sdk.Pipeline; import org.apache.beam.sdk.PipelineResult; +import org.apache.beam.sdk.coders.KvCoder; import org.apache.beam.sdk.coders.StringUtf8Coder; import org.apache.beam.sdk.io.iceberg.SchemaEvolutionConfig.UnverifiableFileHandling; import org.apache.beam.sdk.metrics.MetricNameFilter; import org.apache.beam.sdk.metrics.MetricResult; import org.apache.beam.sdk.metrics.MetricsFilter; import org.apache.beam.sdk.runners.TransformHierarchy; +import org.apache.beam.sdk.schemas.NoSuchSchemaException; +import org.apache.beam.sdk.schemas.SchemaRegistry; import org.apache.beam.sdk.testing.ExpectedLogs; import org.apache.beam.sdk.testing.PAssert; import org.apache.beam.sdk.testing.TestPipeline; @@ -65,6 +68,7 @@ import org.apache.beam.sdk.transforms.Count; import org.apache.beam.sdk.transforms.Create; import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.WithKeys; import org.apache.beam.sdk.transforms.windowing.FixedWindows; import org.apache.beam.sdk.transforms.windowing.Window; import org.apache.beam.sdk.values.PCollection; @@ -73,6 +77,7 @@ import org.apache.beam.sdk.values.Row; import org.apache.beam.sdk.values.TimestampedValue; import org.apache.beam.sdk.values.TupleTagList; +import org.apache.beam.sdk.values.TypeDescriptors; 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.Iterables; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; @@ -1857,14 +1862,27 @@ private PCollectionTuple convert(SchemaEvolutionConfig config, String... files) PCollectionTuple out = pipeline .apply("Create Input", Create.of(Arrays.asList(files))) + .apply( + "KeyByTable", + WithKeys.of(tableId.toString()) + .withKeyType(TypeDescriptors.strings())) .apply( ParDo.of( new AddFiles.ConvertToDataFile( - catalogConfig, tableId.toString(), null, null, null, null, config)) + catalogConfig, null, null, null, null, config)) .withOutputTags( AddFiles.ConvertToDataFile.DATA_FILES, TupleTagList.of(AddFiles.ConvertToDataFile.ERRORS))); out.get(AddFiles.ConvertToDataFile.ERRORS).setRowSchema(AddFiles.ERROR_SCHEMA); + try { + out.get(AddFiles.ConvertToDataFile.DATA_FILES) + .setCoder( + KvCoder.of( + StringUtf8Coder.of(), + SchemaRegistry.createDefault().getSchemaCoder(SerializableDataFile.class))); + } catch (NoSuchSchemaException e) { + throw new RuntimeException(e); + } return out; } diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/ReadFooterSchemaTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/ReadFooterSchemaTest.java index ea5a6e768685..55b0e5fa7c71 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/ReadFooterSchemaTest.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/ReadFooterSchemaTest.java @@ -32,6 +32,8 @@ import java.util.List; import org.apache.avro.SchemaBuilder; import org.apache.beam.sdk.PipelineResult; +import org.apache.beam.sdk.coders.KvCoder; +import org.apache.beam.sdk.coders.StringUtf8Coder; import org.apache.beam.sdk.metrics.MetricNameFilter; import org.apache.beam.sdk.metrics.MetricQueryResults; import org.apache.beam.sdk.metrics.MetricResult; @@ -40,7 +42,10 @@ import org.apache.beam.sdk.testing.TestPipeline; import org.apache.beam.sdk.transforms.Create; import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.Values; +import org.apache.beam.sdk.transforms.WithKeys; import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.TypeDescriptors; import org.apache.hadoop.fs.Path; import org.apache.iceberg.PartitionSpec; import org.apache.iceberg.Schema; @@ -341,9 +346,15 @@ private static void writeLegacyRepeated(File file) throws IOException { } private PCollection run(String... paths) { + // the read is keyed by destination table; these tests are about the schema, so one key return pipeline .apply(Create.of(Arrays.asList(paths))) + .apply( + "KeyByTable", + WithKeys.of("db.t").withKeyType(TypeDescriptors.strings())) .apply(ParDo.of(new ReadFooterSchema(SchemaEvolutionConfig.disabled()))) + .setCoder(KvCoder.of(StringUtf8Coder.of(), CollectDistinctSchemas.groupCoder())) + .apply("DropKey", Values.create()) .setCoder(CollectDistinctSchemas.groupCoder()); } @@ -351,7 +362,12 @@ private PCollection runWithPrefix( String prefix, String... paths) { return pipeline .apply(Create.of(Arrays.asList(paths))) + .apply( + "KeyByTable", + WithKeys.of("db.t").withKeyType(TypeDescriptors.strings())) .apply(ParDo.of(new ReadFooterSchema(SchemaEvolutionConfig.disabled(), prefix))) + .setCoder(KvCoder.of(StringUtf8Coder.of(), CollectDistinctSchemas.groupCoder())) + .apply("DropKey", Values.create()) .setCoder(CollectDistinctSchemas.groupCoder()); }