This is an automated email from the ASF dual-hosted git repository.
claudevdm pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 34a00c7781f AddFiles: extract bounded async task plumbing and Parquet
footer reads (#39896)
34a00c7781f is described below
commit 34a00c7781ff8333de79f6554747e039ea8f137e
Author: claudevdm <[email protected]>
AuthorDate: Thu Aug 27 14:28:35 2026 -0400
AddFiles: extract bounded async task plumbing and Parquet footer reads
(#39896)
* AddFiles: extract bounded async task plumbing and Parquet footer reads
ConvertToDataFile ran its per-file work (footer read, metrics, partition
resolution) on a private thread pool.
This change moves that queue into BoundedAsyncTasks and the footer read
into ParquetFooters so the upcoming schema pre-pass (a second DoFn that
reads every Parquet footer) can share both instead of copying them.
* trigger tests
* ParquetFooterTest
---------
Co-authored-by: Claude <[email protected]>
---
.../IO_Iceberg_Integration_Tests.json | 2 +-
.../org/apache/beam/sdk/io/iceberg/AddFiles.java | 103 +++-------
.../beam/sdk/io/iceberg/BoundedAsyncTasks.java | 113 +++++++++++
.../apache/beam/sdk/io/iceberg/ParquetFooters.java | 51 +++++
.../apache/beam/sdk/io/iceberg/AddFilesTest.java | 4 +-
.../beam/sdk/io/iceberg/BoundedAsyncTasksTest.java | 220 +++++++++++++++++++++
.../beam/sdk/io/iceberg/ParquetFootersTest.java | 113 +++++++++++
7 files changed, 527 insertions(+), 79 deletions(-)
diff --git a/.github/trigger_files/IO_Iceberg_Integration_Tests.json
b/.github/trigger_files/IO_Iceberg_Integration_Tests.json
index 34a6e02150e..5d04b2c0a8c 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": 4
+ "modification": 5
}
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 103df541e50..4645be9bfde 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
@@ -27,31 +27,20 @@ import static
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Pr
import java.io.FileNotFoundException;
import java.io.IOException;
import java.nio.ByteBuffer;
-import java.nio.channels.SeekableByteChannel;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
-import java.util.Iterator;
-import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.UUID;
import java.util.concurrent.Callable;
-import java.util.concurrent.ExecutionException;
-import java.util.concurrent.ExecutorService;
-import java.util.concurrent.Executors;
-import java.util.concurrent.Future;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import org.apache.beam.sdk.coders.KvCoder;
import org.apache.beam.sdk.coders.VarIntCoder;
import org.apache.beam.sdk.coders.VarLongCoder;
-import org.apache.beam.sdk.io.Compression;
-import org.apache.beam.sdk.io.FileSystems;
-import org.apache.beam.sdk.io.fs.ResourceId;
-import org.apache.beam.sdk.io.parquet.ParquetIO.ReadFiles.BeamParquetInputFile;
import org.apache.beam.sdk.metrics.Counter;
import org.apache.beam.sdk.schemas.Schema;
import org.apache.beam.sdk.schemas.SchemaCoder;
@@ -75,8 +64,6 @@ import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.TupleTagList;
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.collect.Iterables;
-import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.hash.Hasher;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.hash.Hashing;
import org.apache.iceberg.AppendFiles;
@@ -111,7 +98,6 @@ import org.apache.iceberg.parquet.ParquetUtil;
import org.apache.iceberg.transforms.Transform;
import org.apache.iceberg.types.Conversions;
import org.apache.iceberg.types.Type;
-import org.apache.parquet.hadoop.ParquetFileReader;
import org.apache.parquet.hadoop.metadata.FileMetaData;
import org.apache.parquet.hadoop.metadata.ParquetMetadata;
import org.apache.parquet.schema.MessageType;
@@ -251,15 +237,15 @@ public class AddFiles extends
PTransform<PCollection<String>, PCollectionRowTupl
* <p><b>Asynchronous Bundle Processing:</b> Because file I/O, catalog
lookups, and metadata
* inference can be highly latency-bound, this DoFn implements an
asynchronous processing pattern
* to maximize throughput. By default, Beam processes elements in a bundle
sequentially. To avoid
- * bottlenecking the pipeline, we use an internal {@link ExecutorService} to
process multiple
+ * bottlenecking the pipeline, we use an internal {@link BoundedAsyncTasks}
to process multiple
* files concurrently within a single DoFn instance.
*
* <p><b>Lifecycle & Thread Safety:</b>
*
* <ul>
* <li><b>{@link ProcessElement}:</b> Submits the heavy lifting (format
inference, metrics
- * collection, and partition resolution) to a background thread pool
and stores the
- * resulting {@link Future}.
+ * collection, and partition resolution) to a background thread pool
and emits results as
+ * they complete.
* <li><b>{@link FinishBundle}:</b> Blocks and awaits the completion of
all futures in the
* current bundle. It safely emits the successfully parsed {@link
DataFile}s, or error rows,
* back to the runner on the main thread, as {@link
MultiOutputReceiver} is not thread-safe.
@@ -274,8 +260,7 @@ public class AddFiles extends
PTransform<PCollection<String>, PCollectionRowTupl
private final @Nullable List<String> partitionFields;
private final @Nullable List<String> sortFields;
private final @Nullable Map<String, String> tableProps;
- private transient @MonotonicNonNull ExecutorService executor;
- private transient @MonotonicNonNull LinkedList<Future<ProcessResult>>
activeTasks;
+ private transient @MonotonicNonNull BoundedAsyncTasks<ProcessResult> tasks;
private transient volatile @MonotonicNonNull Table table;
// Number of parallel threads processing incoming files
@@ -329,21 +314,23 @@ public class AddFiles extends
PTransform<PCollection<String>, PCollectionRowTupl
@Setup
public void setup() {
- executor = Executors.newFixedThreadPool(THREAD_POOL_SIZE);
+ tasks = new BoundedAsyncTasks<>(THREAD_POOL_SIZE, MAX_IN_FLIGHT_TASKS);
+ }
+
+ /** Clears anything left behind if the runner reuses this instance after a
failed bundle. */
+ @StartBundle
+ public void startBundle() {
+
+ checkStateNotNull(tasks).cancelAll();
}
@Teardown
public void teardown() {
- if (executor != null) {
- executor.shutdownNow();
+ if (tasks != null) {
+ tasks.shutdown();
}
}
- @StartBundle
- public void startBundle() {
- activeTasks = Lists.newLinkedList();
- }
-
@ProcessElement
public void process(
@Element String filePath,
@@ -351,28 +338,9 @@ public class AddFiles extends
PTransform<PCollection<String>, PCollectionRowTupl
BoundedWindow window,
PaneInfo paneInfo,
MultiOutputReceiver output)
- throws IOException, InterruptedException, ExecutionException {
- LinkedList<Future<ProcessResult>> activeTasks =
checkStateNotNull(this.activeTasks);
-
- // start draining finished tasks, but don't block
- Iterator<Future<ProcessResult>> iterator = activeTasks.iterator();
- while (iterator.hasNext()) {
- Future<ProcessResult> future = iterator.next();
- if (future.isDone()) {
- outputResult(future.get(), output);
- iterator.remove();
- }
- }
-
- // if we have too many active tasks, wait until some finish
- while (activeTasks.size() >= MAX_IN_FLIGHT_TASKS) {
- Future<ProcessResult> oldestTask = activeTasks.removeFirst();
- outputResult(oldestTask.get(), output); // .get() blocks until the
task completes
- }
-
- // create a new task for the current element and add to queue
+ throws Exception {
Callable<ProcessResult> task = createProcessTask(filePath, timestamp,
window, paneInfo);
- activeTasks.add(checkStateNotNull(executor).submit(task));
+ checkStateNotNull(tasks).submit(task, result -> outputResult(result,
output));
}
private void outputResult(ProcessResult result, MultiOutputReceiver
output) {
@@ -398,18 +366,16 @@ public class AddFiles extends
PTransform<PCollection<String>, PCollectionRowTupl
@FinishBundle
public void finishBundle(FinishBundleContext context) throws Exception {
- // Block and wait for threads to finish their work
- int numErrors = 0;
- for (Future<ProcessResult> future : checkStateNotNull(activeTasks)) {
- ProcessResult result = future.get();
- if (result.errorRow != null) {
- context.output(ERRORS, result.errorRow, result.timestamp,
result.window);
- numErrors++;
- } else if (result.dataFile != null) {
- context.output(DATA_FILES, result.dataFile, result.timestamp,
result.window);
- }
+ checkStateNotNull(tasks).awaitAll(result -> outputAtFinish(result,
context));
+ }
+
+ private static void outputAtFinish(ProcessResult result,
FinishBundleContext context) {
+ if (result.errorRow != null) {
+ 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);
}
- numErrorFiles.inc(numErrors);
}
private Callable<ProcessResult> createProcessTask(
@@ -449,7 +415,7 @@ public class AddFiles extends
PTransform<PCollection<String>, PCollectionRowTupl
@Nullable ParquetMetadata parquetFooter = null;
if (format.equals(FileFormat.PARQUET)) {
try {
- parquetFooter = readParquetFooter(filePath);
+ parquetFooter = ParquetFooters.read(filePath);
} catch (Exception e) {
return errorResult(filePath, errorMessage(e), timestamp, window,
paneInfo);
}
@@ -561,7 +527,7 @@ public class AddFiles extends
PTransform<PCollection<String>, PCollectionRowTupl
throws IOException {
Preconditions.checkArgument(
format.equals(FileFormat.PARQUET), "Table creation is only supported
for Parquet files.");
- MessageType messageType =
readParquetFooter(filePath).getFileMetaData().getSchema();
+ MessageType messageType =
ParquetFooters.read(filePath).getFileMetaData().getSchema();
return ParquetSchemaUtil.convert(messageType);
}
@@ -872,12 +838,6 @@ public class AddFiles extends
PTransform<PCollection<String>, PCollectionRowTupl
}
}
- static ParquetMetadata readParquetFooter(String filePath) throws IOException
{
- try (ParquetFileReader reader =
ParquetFileReader.open(getParquetInputFile(filePath))) {
- return reader.getFooter();
- }
- }
-
/**
* Some exceptions carry a null message (bare EOFException, NPE); the
error-routing path must
* never throw on one.
@@ -911,15 +871,6 @@ public class AddFiles extends
PTransform<PCollection<String>, PCollectionRowTupl
return new ParquetMetadata(newFileMeta, footer.getBlocks());
}
- static org.apache.parquet.io.InputFile getParquetInputFile(String filePath)
throws IOException {
- ResourceId resourceId =
-
Iterables.getOnlyElement(FileSystems.match(filePath).metadata()).resourceId();
- Compression compression =
Compression.detect(checkStateNotNull(resourceId.getFilename()));
- SeekableByteChannel channel =
- (SeekableByteChannel)
compression.readDecompressed(FileSystems.open(resourceId));
- return new BeamParquetInputFile(channel);
- }
-
static class UnknownFormatException extends IllegalArgumentException {}
static class UnknownPartitionException extends IllegalStateException {
diff --git
a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/BoundedAsyncTasks.java
b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/BoundedAsyncTasks.java
new file mode 100644
index 00000000000..3349c9e07fe
--- /dev/null
+++
b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/BoundedAsyncTasks.java
@@ -0,0 +1,113 @@
+/*
+ * 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 java.util.ArrayDeque;
+import java.util.Deque;
+import java.util.Iterator;
+import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.function.Consumer;
+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.util.concurrent.ThreadFactoryBuilder;
+
+/**
+ * Runs tasks on a fixed thread pool while bounding how many are in flight.
Results are handed to
+ * {@code onDone} on the caller's thread (from {@link #submit} and {@link
#awaitAll}), so a DoFn can
+ * emit them without a thread-safe output.
+ */
+class BoundedAsyncTasks<T> {
+ private final ExecutorService executor;
+ private final int maxInFlight;
+ private final Deque<Future<T>> active = new ArrayDeque<>();
+
+ BoundedAsyncTasks(int threads, int maxInFlight) {
+ Preconditions.checkArgument(threads > 0, "threads must be positive, got:
%s", threads);
+ Preconditions.checkArgument(
+ maxInFlight > 0, "maxInFlight must be positive, got: %s", maxInFlight);
+ this.executor =
+ Executors.newFixedThreadPool(
+ threads,
+ new ThreadFactoryBuilder()
+ .setDaemon(true)
+ .setNameFormat("iceberg-async-task-%d")
+ .build());
+ this.maxInFlight = maxInFlight;
+ }
+
+ /**
+ * Submits a task, first delivering any finished results. Blocks while
{@code maxInFlight} tasks
+ * are outstanding. If a task failed, its exception is rethrown and every
other outstanding task
+ * is cancelled.
+ */
+ void submit(Callable<T> task, Consumer<T> onDone) throws Exception {
+ try {
+ drainFinished(onDone);
+ while (active.size() >= maxInFlight) {
+ Future<T> oldest = active.removeFirst();
+ onDone.accept(oldest.get()); // blocks until the oldest task completes
+ }
+ active.add(executor.submit(task));
+ } catch (Exception e) {
+ cancelAll();
+ throw e;
+ }
+ }
+
+ /**
+ * Delivers every outstanding result. Finished tasks drained during
execution may have been
+ * delivered out of submission order; remaining tasks are delivered in queue
order. The queue is
+ * empty afterwards.
+ */
+ void awaitAll(Consumer<T> onDone) throws Exception {
+ try {
+ while (!active.isEmpty()) {
+ Future<T> oldest = active.removeFirst();
+ onDone.accept(oldest.get());
+ }
+ } finally {
+ cancelAll();
+ }
+ }
+
+ /** Cancels and forgets every outstanding task. Results already delivered
are unaffected. */
+ void cancelAll() {
+ for (Future<T> future : active) {
+ future.cancel(true);
+ }
+ active.clear();
+ }
+
+ void shutdown() {
+ cancelAll();
+ executor.shutdownNow();
+ }
+
+ private void drainFinished(Consumer<T> onDone) throws Exception {
+ Iterator<Future<T>> iterator = active.iterator();
+ while (iterator.hasNext()) {
+ Future<T> future = iterator.next();
+ if (future.isDone()) {
+ iterator.remove();
+ onDone.accept(future.get());
+ }
+ }
+ }
+}
diff --git
a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/ParquetFooters.java
b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/ParquetFooters.java
new file mode 100644
index 00000000000..3ea62ac6efc
--- /dev/null
+++
b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/ParquetFooters.java
@@ -0,0 +1,51 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.io.iceberg;
+
+import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull;
+
+import java.io.IOException;
+import java.nio.channels.SeekableByteChannel;
+import org.apache.beam.sdk.io.Compression;
+import org.apache.beam.sdk.io.FileSystems;
+import org.apache.beam.sdk.io.fs.ResourceId;
+import org.apache.beam.sdk.io.parquet.ParquetIO.ReadFiles.BeamParquetInputFile;
+import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables;
+import org.apache.parquet.hadoop.ParquetFileReader;
+import org.apache.parquet.hadoop.metadata.ParquetMetadata;
+import org.apache.parquet.io.InputFile;
+
+/** Reads Parquet footers through Beam's {@link FileSystems}, so no table or
FileIO is needed. */
+final class ParquetFooters {
+ private ParquetFooters() {}
+
+ static ParquetMetadata read(String filePath) throws IOException {
+ try (ParquetFileReader reader =
ParquetFileReader.open(inputFile(filePath))) {
+ return reader.getFooter();
+ }
+ }
+
+ private static InputFile inputFile(String filePath) throws IOException {
+ ResourceId resourceId =
+
Iterables.getOnlyElement(FileSystems.match(filePath).metadata()).resourceId();
+ Compression compression =
Compression.detect(checkStateNotNull(resourceId.getFilename()));
+ SeekableByteChannel channel =
+ (SeekableByteChannel)
compression.readDecompressed(FileSystems.open(resourceId));
+ return new BeamParquetInputFile(channel);
+ }
+}
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 40c038cc7f5..96da59f4505 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
@@ -807,7 +807,7 @@ public class AddFilesTest {
writer.close();
InputFile file = table.io().newInputFile(fileName);
- ParquetMetadata footer = AddFiles.readParquetFooter(fileName);
+ ParquetMetadata footer = ParquetFooters.read(fileName);
Metrics metrics =
AddFiles.getFileMetrics(
file, FileFormat.PARQUET, metricsConfig,
MappingUtil.create(icebergSchema), footer);
@@ -873,7 +873,7 @@ public class AddFilesTest {
writer.close();
InputFile file = table.io().newInputFile(fileName);
- ParquetMetadata footer = AddFiles.readParquetFooter(fileName);
+ ParquetMetadata footer = ParquetFooters.read(fileName);
Metrics metrics =
AddFiles.getFileMetrics(
file, FileFormat.PARQUET, metricsConfig,
MappingUtil.create(icebergSchema), footer);
diff --git
a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/BoundedAsyncTasksTest.java
b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/BoundedAsyncTasksTest.java
new file mode 100644
index 00000000000..09471c1f414
--- /dev/null
+++
b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/BoundedAsyncTasksTest.java
@@ -0,0 +1,220 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.io.iceberg;
+
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.containsInAnyOrder;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import org.junit.After;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class BoundedAsyncTasksTest {
+ private final BoundedAsyncTasks<String> tasks = new BoundedAsyncTasks<>(2,
4);
+
+ @After
+ public void teardown() {
+ tasks.shutdown();
+ }
+
+ @Test
+ public void testConstructorValidatesArguments() {
+ assertThrows(IllegalArgumentException.class, () -> new
BoundedAsyncTasks<>(0, 4));
+ assertThrows(IllegalArgumentException.class, () -> new
BoundedAsyncTasks<>(2, 0));
+ }
+
+ @Test
+ public void testAllResultsAreDelivered() throws Exception {
+ List<String> delivered = new ArrayList<>();
+ for (int i = 0; i < 10; i++) {
+ String value = "t" + i;
+ tasks.submit(() -> value, delivered::add);
+ }
+ tasks.awaitAll(delivered::add);
+ List<String> expected = new ArrayList<>();
+ for (int i = 0; i < 10; i++) {
+ expected.add("t" + i);
+ }
+ assertThat(delivered, containsInAnyOrder(expected.toArray(new String[0])));
+ }
+
+ @Test
+ public void testOutOfOrderCompletionDrainsFinishedTasks() throws Exception {
+ CountDownLatch slowTaskHold = new CountDownLatch(1);
+ CountDownLatch fastTaskDone = new CountDownLatch(1);
+ List<String> delivered = new ArrayList<>();
+
+ // Task 0: slow, waiting on latch
+ tasks.submit(
+ () -> {
+ slowTaskHold.await();
+ return "slow";
+ },
+ delivered::add);
+
+ // Task 1: fast, signals when completed
+ tasks.submit(
+ () -> {
+ fastTaskDone.countDown();
+ return "fast";
+ },
+ delivered::add);
+
+ // Wait until fast task has finished running in the background
+ fastTaskDone.await();
+
+ // Trigger a drain by submitting another task
+ tasks.submit(() -> "noop", delivered::add);
+
+ // "fast" should have been drained while "slow" is still blocked
+ assertTrue(delivered.contains("fast"));
+ assertFalse(delivered.contains("slow"));
+
+ slowTaskHold.countDown();
+ tasks.awaitAll(delivered::add);
+ assertTrue(delivered.contains("slow"));
+ }
+
+ @Test
+ public void testSubmitBlocksAtMaxInFlight() throws Exception {
+ CountDownLatch release = new CountDownLatch(1);
+ CountDownLatch fifthExecuted = new CountDownLatch(1);
+ CountDownLatch fifthReturned = new CountDownLatch(1);
+ List<String> delivered = Collections.synchronizedList(new ArrayList<>());
+
+ for (int i = 0; i < 4; i++) {
+ tasks.submit(
+ () -> {
+ release.await();
+ return "blocked";
+ },
+ delivered::add);
+ }
+
+ // The fifth submit must block because 4 tasks are already outstanding.
+ Thread fifthSubmitter =
+ new Thread(
+ () -> {
+ try {
+ tasks.submit(
+ () -> {
+ fifthExecuted.countDown();
+ return "fifth";
+ },
+ delivered::add);
+ fifthReturned.countDown();
+ } catch (Exception e) {
+ throw new RuntimeException(e);
+ }
+ });
+ fifthSubmitter.start();
+
+ // Verify submit() is blocked and has not completed while 4 tasks are in
flight
+ assertFalse(
+ "submit() must block while maxInFlight tasks are outstanding",
+ fifthReturned.await(50, TimeUnit.MILLISECONDS));
+ assertEquals(
+ "fifth task must not execute before oldest task completes", 1,
fifthExecuted.getCount());
+
+ // Release the blocked tasks; fifth submit should unblock and complete
+ release.countDown();
+ fifthReturned.await();
+ fifthExecuted.await();
+ fifthSubmitter.join();
+
+ tasks.awaitAll(delivered::add);
+ assertEquals(5, delivered.size());
+ assertTrue(delivered.contains("fifth"));
+ }
+
+ @Test
+ public void testAwaitAllFailureCancelsRemainingTasks() throws Exception {
+ CountDownLatch release = new CountDownLatch(1);
+ List<String> delivered = new ArrayList<>();
+ tasks.submit(
+ () -> {
+ release.await();
+ return "slow";
+ },
+ delivered::add);
+ tasks.submit(
+ () -> {
+ release.await();
+ throw new IllegalStateException("boom");
+ },
+ delivered::add);
+ // Submit a trailing task that would remain in the queue if awaitAll does
not clean up on
+ // failure
+ tasks.submit(() -> "trailing", delivered::add);
+
+ // awaitAll delivers in order: the slow task first, then the failure
propagates.
+ release.countDown();
+ ExecutionException failure =
+ assertThrows(ExecutionException.class, () ->
tasks.awaitAll(delivered::add));
+ assertTrue(failure.getCause() instanceof IllegalStateException);
+ assertEquals(Arrays.asList("slow"), delivered);
+
+ // Subsequent submissions on this instance should only see their own
results, not "trailing"
+ List<String> freshBatch = new ArrayList<>();
+ tasks.submit(() -> "fresh", freshBatch::add);
+ tasks.awaitAll(freshBatch::add);
+ assertEquals(Arrays.asList("fresh"), freshBatch);
+ }
+
+ @Test
+ public void testFailureOnSubmitCancelsOutstandingTasks() throws Exception {
+ CountDownLatch release = new CountDownLatch(1);
+ List<String> delivered = new ArrayList<>();
+ // The failure surfaces from whichever submit first drains the failed task.
+ boolean failed = false;
+ try {
+ tasks.submit(
+ () -> {
+ throw new IllegalStateException("boom");
+ },
+ delivered::add);
+ for (int i = 0; i < 5; i++) {
+ tasks.submit(
+ () -> {
+ release.await();
+ return "never";
+ },
+ delivered::add);
+ }
+ } catch (ExecutionException e) {
+ failed = true;
+ }
+ assertTrue(failed);
+ release.countDown();
+ tasks.awaitAll(delivered::add);
+ assertTrue(delivered.isEmpty());
+ }
+}
diff --git
a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/ParquetFootersTest.java
b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/ParquetFootersTest.java
new file mode 100644
index 00000000000..631adfb6431
--- /dev/null
+++
b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/ParquetFootersTest.java
@@ -0,0 +1,113 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.io.iceberg;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
+
+import java.io.File;
+import java.io.FileNotFoundException;
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Schema;
+import org.apache.iceberg.data.GenericRecord;
+import org.apache.iceberg.data.Record;
+import org.apache.iceberg.data.parquet.GenericParquetWriter;
+import org.apache.iceberg.io.DataWriter;
+import org.apache.iceberg.parquet.Parquet;
+import org.apache.iceberg.types.Types;
+import org.apache.parquet.hadoop.metadata.ParquetMetadata;
+import org.apache.parquet.schema.MessageType;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class ParquetFootersTest {
+ @Rule public TemporaryFolder temp = new TemporaryFolder();
+
+ private static final Schema SCHEMA =
+ new Schema(
+ Types.NestedField.required(1, "id", Types.IntegerType.get()),
+ Types.NestedField.optional(2, "name", Types.StringType.get()));
+
+ @Test
+ public void testReadsSchemaAndRowCount() throws IOException {
+ String path = writeParquet("data.parquet", 3);
+
+ ParquetMetadata footer = ParquetFooters.read(path);
+
+ MessageType schema = footer.getFileMetaData().getSchema();
+ assertEquals(2, schema.getFieldCount());
+ assertEquals("id", schema.getFieldName(0));
+ assertEquals("name", schema.getFieldName(1));
+ long rows = 0;
+ for (org.apache.parquet.hadoop.metadata.BlockMetaData block :
footer.getBlocks()) {
+ rows += block.getRowCount();
+ }
+ assertEquals(3, rows);
+ }
+
+ @Test
+ public void testMissingFileThrowsFileNotFound() {
+ String path = new File(temp.getRoot(),
"missing.parquet").getAbsolutePath();
+ assertThrows(FileNotFoundException.class, () -> ParquetFooters.read(path));
+ }
+
+ @Test
+ public void testEmptyFileThrows() throws IOException {
+ File file = temp.newFile("empty.parquet");
+ assertThrows(RuntimeException.class, () ->
ParquetFooters.read(file.getAbsolutePath()));
+ }
+
+ @Test
+ public void testNonParquetFileThrows() throws IOException {
+ File file = temp.newFile("garbage.parquet");
+ Files.write(file.toPath(), "this is not a parquet
file".getBytes(StandardCharsets.UTF_8));
+ assertThrows(RuntimeException.class, () ->
ParquetFooters.read(file.getAbsolutePath()));
+ }
+
+ @Test
+ public void testTruncatedFileThrows() throws IOException {
+ String path = writeParquet("full.parquet", 3);
+ byte[] bytes = Files.readAllBytes(new File(path).toPath());
+ File truncated = temp.newFile("truncated.parquet");
+ // Drop the trailing footer-length + magic bytes so the footer cannot be
located.
+ Files.write(truncated.toPath(), java.util.Arrays.copyOf(bytes,
bytes.length - 8));
+ assertThrows(RuntimeException.class, () ->
ParquetFooters.read(truncated.getAbsolutePath()));
+ }
+
+ private String writeParquet(String name, int rows) throws IOException {
+ String path = new File(temp.getRoot(), name).getAbsolutePath();
+ DataWriter<Record> writer =
+ Parquet.writeData(org.apache.iceberg.Files.localOutput(path))
+ .schema(SCHEMA)
+ .withSpec(PartitionSpec.unpartitioned())
+ .createWriterFunc(GenericParquetWriter::create)
+ .build();
+ for (int i = 0; i < rows; i++) {
+ writer.write(GenericRecord.create(SCHEMA).copy("id", i, "name", "n" +
i));
+ }
+ writer.close();
+ return path;
+ }
+}
