GCP: Add GCS Single shot upload for smaller files -- gitIssue [17524]

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<String, String> 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<Integer> 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<String> clientLibToken() {
     return Optional.ofNullable(clientLibToken);
   }

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

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 <a href="https://cloud.google.com/storage/docs/streaming">Streaming Transfers</a>
+ * Uploads to GCS.
+ *
+ * <ul>
+ *   <li>If size &lt; {@link GCPProperties#GCS_WRITE_THRESHOLD_BYTES} (default 8MB): single-shot
+ *       {@link Storage#create}
+ *   <li>Otherwise: {@link WriteChannel} output stream
+ * </ul>
  */
 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<BlobWriteOption> 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,15 +151,33 @@ public void close() throws IOException {
 
     super.close();
     closed = true;
-    stream.close();
+
+    if (useWriteChannel) {
+      stream.close();
+      return;
+    }
+
+    // size < threshold → single-shot upload
+    List<BlobTargetOption> 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"})
   @Override
   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);
     }

gcp/src/test/java/org/apache/iceberg/gcp/TestGCPProperties.java

@@ -22,14 +22,36 @@
 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;
 import org.junit.jupiter.api.Test;
 
 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()

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,58 @@ 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);