From c8078b880a3c30e5dd4ce27afdabc78109dfa8e5 Mon Sep 17 00:00:00 2001 From: Manideep Addepalli <42642602+manideep4597@users.noreply.github.com> Date: Fri, 7 Aug 2026 11:17:38 +0000 Subject: [PATCH 1/3] GCS Single shot upload for smaller files --- .../org/apache/iceberg/gcp/GCPProperties.java | 25 ++++++ .../apache/iceberg/gcp/gcs/GCSOutputFile.java | 8 +- .../iceberg/gcp/gcs/GCSOutputStream.java | 60 ++++++++++++--- .../apache/iceberg/gcp/TestGCPProperties.java | 24 ++++++ .../iceberg/gcp/gcs/TestGCSOutputStream.java | 76 +++++++++++++++++-- 5 files changed, 171 insertions(+), 22 deletions(-) diff --git a/gcp/src/main/java/org/apache/iceberg/gcp/GCPProperties.java b/gcp/src/main/java/org/apache/iceberg/gcp/GCPProperties.java index 2702ce565d4e..87e9219b1d5b 100644 --- a/gcp/src/main/java/org/apache/iceberg/gcp/GCPProperties.java +++ b/gcp/src/main/java/org/apache/iceberg/gcp/GCPProperties.java @@ -47,6 +47,14 @@ public class GCPProperties implements Serializable { public static final String GCS_CHANNEL_READ_CHUNK_SIZE = "gcs.channel.read.chunk-size-bytes"; public static final String GCS_CHANNEL_WRITE_CHUNK_SIZE = "gcs.channel.write.chunk-size-bytes"; + /** + * If object size is less than this value (default: 8MB), use single-shot {@code storage.create}; + * otherwise use a WriteChannel output stream. Override with {@code gcs.write.threshold-bytes}. + */ + public static final String GCS_WRITE_THRESHOLD_BYTES = "gcs.write.threshold-bytes"; + + public static final long GCS_WRITE_THRESHOLD_BYTES_DEFAULT = 8L * 1024 * 1024; + public static final String GCS_OAUTH2_TOKEN = "gcs.oauth2.token"; public static final String GCS_OAUTH2_TOKEN_EXPIRES_AT = "gcs.oauth2.token-expires-at"; // Boolean to explicitly configure "no authentication" for testing purposes using a GCS emulator @@ -91,6 +99,7 @@ public class GCPProperties implements Serializable { private Integer gcsChannelReadChunkSize; private Integer gcsChannelWriteChunkSize; + private long gcsWriteThresholdBytes = GCS_WRITE_THRESHOLD_BYTES_DEFAULT; private boolean gcsNoAuth; private String gcsOAuth2Token; @@ -161,6 +170,15 @@ public GCPProperties(Map properties) { gcsChannelWriteChunkSize = Integer.parseInt(properties.get(GCS_CHANNEL_WRITE_CHUNK_SIZE)); } + gcsWriteThresholdBytes = + PropertyUtil.propertyAsLong( + properties, GCS_WRITE_THRESHOLD_BYTES, GCS_WRITE_THRESHOLD_BYTES_DEFAULT); + Preconditions.checkArgument( + gcsWriteThresholdBytes >= 0, + "Property %s must be >= 0: %s", + GCS_WRITE_THRESHOLD_BYTES, + gcsWriteThresholdBytes); + gcsOAuth2Token = properties.get(GCS_OAUTH2_TOKEN); if (properties.containsKey(GCS_OAUTH2_TOKEN_EXPIRES_AT)) { gcsOAuth2TokenExpiresAt = @@ -207,6 +225,13 @@ public Optional channelWriteChunkSize() { return Optional.ofNullable(gcsChannelWriteChunkSize); } + /** + * Returns the max size in bytes for single-shot uploads. See {@link #GCS_WRITE_THRESHOLD_BYTES}. + */ + public long writeThresholdBytes() { + return gcsWriteThresholdBytes; + } + public Optional clientLibToken() { return Optional.ofNullable(clientLibToken); } diff --git a/gcp/src/main/java/org/apache/iceberg/gcp/gcs/GCSOutputFile.java b/gcp/src/main/java/org/apache/iceberg/gcp/gcs/GCSOutputFile.java index 18bf53b791fb..db3f03e45738 100644 --- a/gcp/src/main/java/org/apache/iceberg/gcp/gcs/GCSOutputFile.java +++ b/gcp/src/main/java/org/apache/iceberg/gcp/gcs/GCSOutputFile.java @@ -20,8 +20,6 @@ import com.google.cloud.storage.BlobId; import com.google.cloud.storage.Storage; -import java.io.IOException; -import java.io.UncheckedIOException; import org.apache.iceberg.exceptions.AlreadyExistsException; import org.apache.iceberg.gcp.GCPProperties; import org.apache.iceberg.io.InputFile; @@ -69,11 +67,7 @@ public PositionOutputStream create() { @Override public PositionOutputStream createOrOverwrite() { - try { - return new GCSOutputStream(storage(), blobId(), gcpProperties(), metrics()); - } catch (IOException e) { - throw new UncheckedIOException("Failed to create output stream for location: " + uri(), e); - } + return new GCSOutputStream(storage(), blobId(), gcpProperties(), metrics()); } @Override diff --git a/gcp/src/main/java/org/apache/iceberg/gcp/gcs/GCSOutputStream.java b/gcp/src/main/java/org/apache/iceberg/gcp/gcs/GCSOutputStream.java index 3fb8aa4801ae..01ae308007a5 100644 --- a/gcp/src/main/java/org/apache/iceberg/gcp/gcs/GCSOutputStream.java +++ b/gcp/src/main/java/org/apache/iceberg/gcp/gcs/GCSOutputStream.java @@ -22,7 +22,9 @@ import com.google.cloud.storage.BlobId; import com.google.cloud.storage.BlobInfo; import com.google.cloud.storage.Storage; +import com.google.cloud.storage.Storage.BlobTargetOption; import com.google.cloud.storage.Storage.BlobWriteOption; +import java.io.ByteArrayOutputStream; import java.io.IOException; import java.io.OutputStream; import java.nio.channels.Channels; @@ -40,8 +42,13 @@ import org.slf4j.LoggerFactory; /** - * The GCSOutputStream leverages native streaming channels from the GCS API for streaming uploads. - * See Streaming Transfers + * Uploads to GCS. + * + *
    + *
  • If size < {@link GCPProperties#GCS_WRITE_THRESHOLD_BYTES} (default 8MB): single-shot + * {@link Storage#create} + *
  • Otherwise: {@link WriteChannel} output stream + *
*/ class GCSOutputStream extends PositionOutputStream { private static final Logger LOG = LoggerFactory.getLogger(GCSOutputStream.class); @@ -50,8 +57,11 @@ class GCSOutputStream extends PositionOutputStream { private final Storage storage; private final BlobId blobId; private final GCPProperties gcpProperties; + private final long writeThreshold; + private final ByteArrayOutputStream buffer = new ByteArrayOutputStream(); private OutputStream stream; + private boolean useWriteChannel = false; private final Counter writeBytes; private final Counter writeOperations; @@ -60,18 +70,19 @@ class GCSOutputStream extends PositionOutputStream { private boolean closed = false; GCSOutputStream( - Storage storage, BlobId blobId, GCPProperties gcpProperties, MetricsContext metrics) - throws IOException { + Storage storage, BlobId blobId, GCPProperties gcpProperties, MetricsContext metrics) { this.storage = storage; this.blobId = blobId; this.gcpProperties = gcpProperties; + this.writeThreshold = gcpProperties.writeThresholdBytes(); createStack = Thread.currentThread().getStackTrace(); this.writeBytes = metrics.counter(FileIOMetricsContext.WRITE_BYTES, Unit.BYTES); this.writeOperations = metrics.counter(FileIOMetricsContext.WRITE_OPERATIONS); - openStream(); + // Buffer in memory until size reaches the threshold. + this.stream = buffer; } @Override @@ -90,6 +101,10 @@ public void write(int b) throws IOException { pos += 1; writeBytes.increment(); writeOperations.increment(); + + if (!useWriteChannel && pos >= writeThreshold) { + openStream(); + } } @Override @@ -98,9 +113,14 @@ public void write(byte[] b, int off, int len) throws IOException { pos += len; writeBytes.increment(len); writeOperations.increment(); + + if (!useWriteChannel && pos >= writeThreshold) { + openStream(); + } } - private void openStream() { + /** Switch from in-memory buffer to a WriteChannel once size >= threshold. */ + private void openStream() throws IOException { List writeOptions = Lists.newArrayList(); gcpProperties @@ -116,7 +136,11 @@ private void openStream() { gcpProperties.channelWriteChunkSize().ifPresent(channel::setChunkSize); - stream = Channels.newOutputStream(channel); + OutputStream channelStream = Channels.newOutputStream(channel); + buffer.writeTo(channelStream); + buffer.reset(); + stream = channelStream; + useWriteChannel = true; } @Override @@ -127,7 +151,25 @@ public void close() throws IOException { super.close(); closed = true; - stream.close(); + + if (useWriteChannel) { + stream.close(); + return; + } + + // size < threshold → single-shot upload + List targetOptions = Lists.newArrayList(); + gcpProperties + .encryptionKey() + .ifPresent(key -> targetOptions.add(BlobTargetOption.encryptionKey(key))); + gcpProperties + .userProject() + .ifPresent(userProject -> targetOptions.add(BlobTargetOption.userProject(userProject))); + + storage.create( + BlobInfo.newBuilder(blobId).build(), + buffer.toByteArray(), + targetOptions.toArray(new BlobTargetOption[0])); } @SuppressWarnings({"checkstyle:NoFinalizer", "Finalize", "deprecation"}) @@ -135,7 +177,7 @@ public void close() throws IOException { protected void finalize() throws Throwable { super.finalize(); if (!closed) { - close(); // releasing resources is more important than printing the warning + close(); String trace = Joiner.on("\n\t").join(Arrays.copyOfRange(createStack, 1, createStack.length)); LOG.warn("Unclosed output stream created by:\n\t{}", trace); } diff --git a/gcp/src/test/java/org/apache/iceberg/gcp/TestGCPProperties.java b/gcp/src/test/java/org/apache/iceberg/gcp/TestGCPProperties.java index 0ec2183fc355..4adf3f844dbf 100644 --- a/gcp/src/test/java/org/apache/iceberg/gcp/TestGCPProperties.java +++ b/gcp/src/test/java/org/apache/iceberg/gcp/TestGCPProperties.java @@ -22,7 +22,10 @@ import static org.apache.iceberg.gcp.GCPProperties.GCS_OAUTH2_REFRESH_CREDENTIALS_ENABLED; import static org.apache.iceberg.gcp.GCPProperties.GCS_OAUTH2_REFRESH_CREDENTIALS_ENDPOINT; import static org.apache.iceberg.gcp.GCPProperties.GCS_OAUTH2_TOKEN; +import static org.apache.iceberg.gcp.GCPProperties.GCS_WRITE_THRESHOLD_BYTES; +import static org.apache.iceberg.gcp.GCPProperties.GCS_WRITE_THRESHOLD_BYTES_DEFAULT; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatIllegalArgumentException; import static org.assertj.core.api.Assertions.assertThatIllegalStateException; import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap; @@ -30,6 +33,27 @@ public class TestGCPProperties { + @Test + public void testWriteThresholdDefault() { + assertThat(new GCPProperties().writeThresholdBytes()) + .isEqualTo(GCS_WRITE_THRESHOLD_BYTES_DEFAULT); + } + + @Test + public void testWriteThresholdOverride() { + GCPProperties props = + new GCPProperties(ImmutableMap.of(GCS_WRITE_THRESHOLD_BYTES, "1048576")); + assertThat(props.writeThresholdBytes()).isEqualTo(1_048_576L); + } + + @Test + public void testWriteThresholdNegativeRejected() { + assertThatIllegalArgumentException() + .isThrownBy( + () -> new GCPProperties(ImmutableMap.of(GCS_WRITE_THRESHOLD_BYTES, "-1"))) + .withMessageContaining(GCS_WRITE_THRESHOLD_BYTES); + } + @Test public void testOAuthWithNoAuth() { assertThatIllegalStateException() diff --git a/gcp/src/test/java/org/apache/iceberg/gcp/gcs/TestGCSOutputStream.java b/gcp/src/test/java/org/apache/iceberg/gcp/gcs/TestGCSOutputStream.java index ebf56ffb7eb4..d9d887343131 100644 --- a/gcp/src/test/java/org/apache/iceberg/gcp/gcs/TestGCSOutputStream.java +++ b/gcp/src/test/java/org/apache/iceberg/gcp/gcs/TestGCSOutputStream.java @@ -18,18 +18,31 @@ */ package org.apache.iceberg.gcp.gcs; +import static org.apache.iceberg.gcp.GCPProperties.GCS_WRITE_THRESHOLD_BYTES; import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; +import com.google.cloud.WriteChannel; +import com.google.cloud.storage.Blob; import com.google.cloud.storage.BlobId; +import com.google.cloud.storage.BlobInfo; import com.google.cloud.storage.Storage; +import com.google.cloud.storage.Storage.BlobTargetOption; +import com.google.cloud.storage.Storage.BlobWriteOption; import com.google.cloud.storage.contrib.nio.testing.LocalStorageHelper; import java.io.IOException; import java.io.UncheckedIOException; +import java.nio.ByteBuffer; import java.util.Random; import java.util.UUID; import java.util.stream.Stream; import org.apache.iceberg.gcp.GCPProperties; import org.apache.iceberg.metrics.MetricsContext; +import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap; import org.junit.jupiter.api.Test; public class TestGCSOutputStream { @@ -45,11 +58,12 @@ public void testWrite() { Stream.of(true, false) .forEach( arrayWrite -> { - // Test small file write - writeAndVerify(storage, randomBlobId(), randomData(1024), arrayWrite); + // Below default threshold → single-shot create + writeAndVerify(storage, randomBlobId(), randomData(1024), arrayWrite, properties); - // Test large file - writeAndVerify(storage, randomBlobId(), randomData(10 * 1024 * 1024), arrayWrite); + // At/above default 8 MiB threshold → WriteChannel + writeAndVerify( + storage, randomBlobId(), randomData(8 * 1024 * 1024), arrayWrite, properties); }); } @@ -61,9 +75,59 @@ public void testMultipleClose() throws IOException { stream.close(); } - private void writeAndVerify(Storage client, BlobId uri, byte[] data, boolean arrayWrite) { + @Test + public void testSingleShotBelowThresholdUsesCreate() throws IOException { + Storage mockStorage = mock(Storage.class); + when(mockStorage.create(any(BlobInfo.class), any(byte[].class), any(BlobTargetOption[].class))) + .thenReturn(mock(Blob.class)); + + GCPProperties props = + new GCPProperties(ImmutableMap.of(GCS_WRITE_THRESHOLD_BYTES, "1024")); + byte[] data = randomData(1023); + + try (GCSOutputStream stream = + new GCSOutputStream(mockStorage, randomBlobId(), props, MetricsContext.nullMetrics())) { + stream.write(data); + } + + verify(mockStorage).create(any(BlobInfo.class), any(byte[].class), any(BlobTargetOption[].class)); + verify(mockStorage, never()).writer(any(BlobInfo.class), any(BlobWriteOption[].class)); + } + + @Test + public void testAtThresholdUsesWriteChannel() throws IOException { + Storage mockStorage = mock(Storage.class); + WriteChannel mockChannel = mock(WriteChannel.class); + when(mockStorage.writer(any(BlobInfo.class), any(BlobWriteOption[].class))) + .thenReturn(mockChannel); + when(mockChannel.write(any(ByteBuffer.class))) + .thenAnswer( + invocation -> { + ByteBuffer buf = invocation.getArgument(0); + int remaining = buf.remaining(); + buf.position(buf.limit()); + return remaining; + }); + + // S3-style: pos >= threshold switches off single-shot + GCPProperties props = + new GCPProperties(ImmutableMap.of(GCS_WRITE_THRESHOLD_BYTES, "1024")); + byte[] data = randomData(1024); + + try (GCSOutputStream stream = + new GCSOutputStream(mockStorage, randomBlobId(), props, MetricsContext.nullMetrics())) { + stream.write(data); + } + + verify(mockStorage).writer(any(BlobInfo.class), any(BlobWriteOption[].class)); + verify(mockStorage, never()) + .create(any(BlobInfo.class), any(byte[].class), any(BlobTargetOption[].class)); + } + + private void writeAndVerify( + Storage client, BlobId uri, byte[] data, boolean arrayWrite, GCPProperties props) { try (GCSOutputStream stream = - new GCSOutputStream(client, uri, properties, MetricsContext.nullMetrics())) { + new GCSOutputStream(client, uri, props, MetricsContext.nullMetrics())) { if (arrayWrite) { stream.write(data); assertThat(stream.getPos()).isEqualTo(data.length); From 45c15935da99b167a4dd454ec02418087b12cb37 Mon Sep 17 00:00:00 2001 From: Manideep Addepalli <42642602+manideep4597@users.noreply.github.com> Date: Fri, 7 Aug 2026 17:00:51 +0530 Subject: [PATCH 2/3] Added new props TestGCPProperties.java --- .../test/java/org/apache/iceberg/gcp/TestGCPProperties.java | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/gcp/src/test/java/org/apache/iceberg/gcp/TestGCPProperties.java b/gcp/src/test/java/org/apache/iceberg/gcp/TestGCPProperties.java index 4adf3f844dbf..113ebbde3e64 100644 --- a/gcp/src/test/java/org/apache/iceberg/gcp/TestGCPProperties.java +++ b/gcp/src/test/java/org/apache/iceberg/gcp/TestGCPProperties.java @@ -41,16 +41,14 @@ public void testWriteThresholdDefault() { @Test public void testWriteThresholdOverride() { - GCPProperties props = - new GCPProperties(ImmutableMap.of(GCS_WRITE_THRESHOLD_BYTES, "1048576")); + GCPProperties props = new GCPProperties(ImmutableMap.of(GCS_WRITE_THRESHOLD_BYTES, "1048576")); assertThat(props.writeThresholdBytes()).isEqualTo(1_048_576L); } @Test public void testWriteThresholdNegativeRejected() { assertThatIllegalArgumentException() - .isThrownBy( - () -> new GCPProperties(ImmutableMap.of(GCS_WRITE_THRESHOLD_BYTES, "-1"))) + .isThrownBy(() -> new GCPProperties(ImmutableMap.of(GCS_WRITE_THRESHOLD_BYTES, "-1"))) .withMessageContaining(GCS_WRITE_THRESHOLD_BYTES); } From 1be1b125e7974d6cc72119559b5b533f822655e9 Mon Sep 17 00:00:00 2001 From: Manideep Addepalli <42642602+manideep4597@users.noreply.github.com> Date: Fri, 7 Aug 2026 17:01:28 +0530 Subject: [PATCH 3/3] added tests to TestGCSOutputStream --- .../org/apache/iceberg/gcp/gcs/TestGCSOutputStream.java | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/gcp/src/test/java/org/apache/iceberg/gcp/gcs/TestGCSOutputStream.java b/gcp/src/test/java/org/apache/iceberg/gcp/gcs/TestGCSOutputStream.java index d9d887343131..26513a27773d 100644 --- a/gcp/src/test/java/org/apache/iceberg/gcp/gcs/TestGCSOutputStream.java +++ b/gcp/src/test/java/org/apache/iceberg/gcp/gcs/TestGCSOutputStream.java @@ -81,8 +81,7 @@ public void testSingleShotBelowThresholdUsesCreate() throws IOException { when(mockStorage.create(any(BlobInfo.class), any(byte[].class), any(BlobTargetOption[].class))) .thenReturn(mock(Blob.class)); - GCPProperties props = - new GCPProperties(ImmutableMap.of(GCS_WRITE_THRESHOLD_BYTES, "1024")); + GCPProperties props = new GCPProperties(ImmutableMap.of(GCS_WRITE_THRESHOLD_BYTES, "1024")); byte[] data = randomData(1023); try (GCSOutputStream stream = @@ -90,7 +89,8 @@ public void testSingleShotBelowThresholdUsesCreate() throws IOException { stream.write(data); } - verify(mockStorage).create(any(BlobInfo.class), any(byte[].class), any(BlobTargetOption[].class)); + verify(mockStorage) + .create(any(BlobInfo.class), any(byte[].class), any(BlobTargetOption[].class)); verify(mockStorage, never()).writer(any(BlobInfo.class), any(BlobWriteOption[].class)); } @@ -110,8 +110,7 @@ public void testAtThresholdUsesWriteChannel() throws IOException { }); // S3-style: pos >= threshold switches off single-shot - GCPProperties props = - new GCPProperties(ImmutableMap.of(GCS_WRITE_THRESHOLD_BYTES, "1024")); + GCPProperties props = new GCPProperties(ImmutableMap.of(GCS_WRITE_THRESHOLD_BYTES, "1024")); byte[] data = randomData(1024); try (GCSOutputStream stream =