feat: DH-18149: Simplified passing s3 instructions for iceberg writing (deephaven#6668)

If the S3 instructions are not passed when constructing the 
table writer, it will automatically inherit the configuration from
the catalog properties.

Author: malhotrashivam

extensions/iceberg/s3/src/test/java/io/deephaven/iceberg/util/TableParquetWriterOptionsTest.java

@@ -5,6 +5,8 @@
 
 import io.deephaven.engine.table.ColumnDefinition;
 import io.deephaven.engine.table.TableDefinition;
+import io.deephaven.extensions.s3.S3Instructions;
+import io.deephaven.parquet.table.CompletedParquetWrite;
 import io.deephaven.parquet.table.ParquetInstructions;
 import org.junit.jupiter.api.Test;
 
@@ -144,8 +146,12 @@ void toParquetInstructionTest() {
                 ColumnDefinition.ofInt("PC2").withPartitioning(),
                 ColumnDefinition.ofLong("I"));
         final Map<Integer, String> fieldIdToName = Map.of(2, "field2", 3, "field3");
+        final S3Instructions s3Instructions = S3Instructions.builder().regionName("test-region").build();
+        final ParquetInstructions.OnWriteCompleted onWriteCompleted =
+                (final CompletedParquetWrite completedParquetWrite) -> {
+                    /* Do nothing */ };
         final ParquetInstructions parquetInstructions = writeInstructions.toParquetInstructions(
-                null, definition, fieldIdToName);
+                onWriteCompleted, definition, fieldIdToName, s3Instructions);
 
         assertThat(parquetInstructions.getCompressionCodecName()).isEqualTo("GZIP");
         assertThat(parquetInstructions.getMaximumDictionaryKeys()).isEqualTo(100);
@@ -154,7 +160,8 @@ void toParquetInstructionTest() {
         assertThat(parquetInstructions.getFieldId("field1")).isEmpty();
         assertThat(parquetInstructions.getFieldId("field2")).hasValue(2);
         assertThat(parquetInstructions.getFieldId("field3")).hasValue(3);
-        assertThat(parquetInstructions.onWriteCompleted()).isEmpty();
+        assertThat(parquetInstructions.onWriteCompleted()).hasValue(onWriteCompleted);
         assertThat(parquetInstructions.getTableDefinition()).hasValue(definition);
+        assertThat(parquetInstructions.getSpecialInstructions()).isEqualTo(s3Instructions);
     }
 }

extensions/iceberg/src/main/java/io/deephaven/iceberg/base/IcebergUtils.java

@@ -131,15 +131,15 @@ public static <T> Stream<T> toStream(final org.apache.iceberg.io.CloseableIterab
         });
     }
 
-    private static String path(String path, FileIO io) {
+    private static String path(@NotNull final String path, @NotNull final FileIO io) {
         return io instanceof RelativeFileIO ? ((RelativeFileIO) io).absoluteLocation(path) : path;
     }
 
-    public static URI locationUri(Table table) {
+    public static URI locationUri(@NotNull final Table table) {
         return FileUtils.convertToURI(path(table.location(), table.io()), true);
     }
 
-    public static URI dataFileUri(Table table, DataFile dataFile) {
+    public static URI dataFileUri(@NotNull final Table table, @NotNull final DataFile dataFile) {
         return FileUtils.convertToURI(path(dataFile.path().toString(), table.io()), false);
     }
 

extensions/iceberg/src/main/java/io/deephaven/iceberg/internal/DataInstructionsProviderLoader.java

@@ -4,6 +4,7 @@
 package io.deephaven.iceberg.internal;
 
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -82,6 +83,7 @@ private DataInstructionsProviderLoader(
      * @param uriScheme The URI scheme
      * @return A data instructions object for the given URI scheme or null if one cannot be found
      */
+    @Nullable
     public Object load(@NotNull final String uriScheme) {
         for (final DataInstructionsProviderPlugin plugin : providers) {
             final Object pluginInstructions = plugin.createInstructions(uriScheme, properties);

extensions/iceberg/src/main/java/io/deephaven/iceberg/layout/IcebergBaseLayout.java

@@ -31,7 +31,6 @@
 
 import static io.deephaven.iceberg.base.IcebergUtils.allManifestFiles;
 import static io.deephaven.iceberg.base.IcebergUtils.dataFileUri;
-import static io.deephaven.iceberg.base.IcebergUtils.locationUri;
 
 public abstract class IcebergBaseLayout implements TableLocationKeyFinder<IcebergTableLocationKey> {
     /**
@@ -71,11 +70,6 @@ public abstract class IcebergBaseLayout implements TableLocationKeyFinder<Iceber
      */
     final TableDefinition tableDef;
 
-    /**
-     * The URI scheme from the Table {@link Table#location() location}.
-     */
-    private final String uriScheme;
-
     /**
      * The {@link Snapshot} from which to discover data files.
      */
@@ -123,6 +117,8 @@ protected IcebergTableLocationKey locationKey(
     /**
      * @param tableAdapter The {@link IcebergTableAdapter} that will be used to access the table.
      * @param instructions The instructions for customizations while reading.
+     * @param dataInstructionsProvider The provider for special instructions, to be used if special instructions not
+     *        provided in the {@code instructions}.
      */
     public IcebergBaseLayout(
             @NotNull final IcebergTableAdapter tableAdapter,
@@ -147,7 +143,8 @@ public IcebergBaseLayout(
 
         this.snapshot = tableAdapter.getSnapshot(instructions);
         this.tableDef = tableAdapter.definition(instructions);
-        this.uriScheme = locationUri(tableAdapter.icebergTable()).getScheme();
+
+        final String uriScheme = tableAdapter.locationUri().getScheme();
         // Add the data instructions if provided as part of the IcebergReadInstructions, or else attempt to create
         // data instructions from the properties collection and URI scheme.
         final Object specialInstructions = instructions.dataInstructions()

extensions/iceberg/src/main/java/io/deephaven/iceberg/layout/IcebergFlatLayout.java

@@ -22,6 +22,8 @@ public final class IcebergFlatLayout extends IcebergBaseLayout {
     /**
      * @param tableAdapter The {@link IcebergTableAdapter} that will be used to access the table.
      * @param instructions The instructions for customizations while reading.
+     * @param dataInstructionsProvider The provider for special instructions, to be used if special instructions not
+     *        provided in the {@code instructions}.
      */
     public IcebergFlatLayout(
             @NotNull final IcebergTableAdapter tableAdapter,

extensions/iceberg/src/main/java/io/deephaven/iceberg/layout/IcebergKeyValuePartitionedLayout.java

@@ -42,6 +42,8 @@ private IdentityPartitioningColData(String name, Class<?> type, int index) {
      * @param tableAdapter The {@link IcebergTableAdapter} that will be used to access the table.
      * @param partitionSpec The Iceberg {@link PartitionSpec partition spec} for the table.
      * @param instructions The instructions for customizations while reading.
+     * @param dataInstructionsProvider The provider for special instructions, to be used if special instructions not
+     *        provided in the {@code instructions}.
      */
     public IcebergKeyValuePartitionedLayout(
             @NotNull final IcebergTableAdapter tableAdapter,

extensions/iceberg/src/main/java/io/deephaven/iceberg/util/IcebergReadInstructions.java

@@ -37,7 +37,8 @@ public static Builder builder() {
 
     /**
      * The data instructions to use for reading the Iceberg data files (might be S3Instructions or other cloud
-     * provider-specific instructions).
+     * provider-specific instructions). If not provided, data instructions will be derived from the properties of the
+     * catalog.
      */
     public abstract Optional<Object> dataInstructions();
 

extensions/iceberg/src/main/java/io/deephaven/iceberg/util/IcebergTableAdapter.java

@@ -19,6 +19,7 @@
 import io.deephaven.engine.table.impl.sources.regioned.RegionedTableComponentFactoryImpl;
 import io.deephaven.engine.updategraph.UpdateSourceRegistrar;
 import io.deephaven.engine.util.TableTools;
+import io.deephaven.iceberg.base.IcebergUtils;
 import io.deephaven.iceberg.base.IcebergUtils.SpecAndSchema;
 import io.deephaven.iceberg.internal.DataInstructionsProviderLoader;
 import io.deephaven.iceberg.layout.*;
@@ -38,6 +39,7 @@
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
+import java.net.URI;
 import java.time.Instant;
 import java.util.*;
 import java.util.stream.Collectors;
@@ -61,6 +63,8 @@ public class IcebergTableAdapter {
     private final TableIdentifier tableIdentifier;
     private final DataInstructionsProviderLoader dataInstructionsProviderLoader;
 
+    private final URI locationUri;
+
     public IcebergTableAdapter(
             final Catalog catalog,
             final TableIdentifier tableIdentifier,
@@ -70,6 +74,7 @@ public IcebergTableAdapter(
         this.table = table;
         this.tableIdentifier = tableIdentifier;
         this.dataInstructionsProviderLoader = dataInstructionsProviderLoader;
+        this.locationUri = IcebergUtils.locationUri(table);
     }
 
     /**
@@ -584,6 +589,13 @@ private static TableDefinition fromSchema(
      * @return A new instance of {@link IcebergTableWriter} configured with the provided options.
      */
     public IcebergTableWriter tableWriter(final TableWriterOptions tableWriterOptions) {
-        return new IcebergTableWriter(tableWriterOptions, this);
+        return new IcebergTableWriter(tableWriterOptions, this, dataInstructionsProviderLoader);
+    }
+
+    /**
+     * Get the location URI of the Iceberg table.
+     */
+    public URI locationUri() {
+        return locationUri;
     }
 }

extensions/iceberg/src/main/java/io/deephaven/iceberg/util/IcebergTableWriter.java

@@ -12,6 +12,7 @@
 import io.deephaven.engine.table.Table;
 import io.deephaven.engine.table.TableDefinition;
 import io.deephaven.engine.table.impl.locations.TableDataException;
+import io.deephaven.iceberg.internal.DataInstructionsProviderLoader;
 import io.deephaven.parquet.table.CompletedParquetWrite;
 import io.deephaven.parquet.table.ParquetInstructions;
 import io.deephaven.parquet.table.ParquetTools;
@@ -101,15 +102,30 @@ public class IcebergTableWriter {
      */
     private final OutputFileFactory outputFileFactory;
 
+    /**
+     * The special instructions to use for writing the Iceberg data files (might be S3Instructions or other cloud
+     * provider-specific instructions).
+     */
+    private final Object specialInstructions;
+
     /**
      * Characters to be used for generating random variable names of length {@link #VARIABLE_NAME_LENGTH}.
      */
     private static final String CHARACTERS = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz";
     private static final int VARIABLE_NAME_LENGTH = 6;
 
+    /**
+     * Create a new Iceberg table writer instance.
+     *
+     * @param tableWriterOptions The options to configure the behavior of this writer instance.
+     * @param tableAdapter The Iceberg table adapter corresponding to the Iceberg table to write to.
+     * @param dataInstructionsProvider The provider for special instructions, to be used if special instructions not
+     *        provided in the {@code tableWriterOptions}.
+     */
     IcebergTableWriter(
             final TableWriterOptions tableWriterOptions,
-            final IcebergTableAdapter tableAdapter) {
+            final IcebergTableAdapter tableAdapter,
+            final DataInstructionsProviderLoader dataInstructionsProvider) {
         this.tableWriterOptions = verifyWriterOptions(tableWriterOptions);
         this.table = tableAdapter.icebergTable();
 
@@ -131,6 +147,11 @@ public class IcebergTableWriter {
         outputFileFactory = OutputFileFactory.builderFor(table, 0, 0)
                 .format(FileFormat.PARQUET)
                 .build();
+
+        final String uriScheme = tableAdapter.locationUri().getScheme();
+        this.specialInstructions = tableWriterOptions.dataInstructions()
+                .orElseGet(() -> dataInstructionsProvider.load(uriScheme));
+
     }
 
     private static TableParquetWriterOptions verifyWriterOptions(
@@ -459,7 +480,7 @@ private List<CompletedParquetWrite> writeParquet(
         final List<CompletedParquetWrite> parquetFilesWritten = new ArrayList<>(dhTables.size());
         final ParquetInstructions.OnWriteCompleted onWriteCompleted = parquetFilesWritten::add;
         final ParquetInstructions parquetInstructions = tableWriterOptions.toParquetInstructions(
-                onWriteCompleted, tableDefinition, fieldIdToColumnName);
+                onWriteCompleted, tableDefinition, fieldIdToColumnName, specialInstructions);
 
         // Write the data to parquet files
         for (int idx = 0; idx < dhTables.size(); idx++) {

extensions/iceberg/src/main/java/io/deephaven/iceberg/util/TableParquetWriterOptions.java

@@ -8,6 +8,7 @@
 import io.deephaven.parquet.table.ParquetInstructions;
 import org.immutables.value.Value;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 
 import java.util.Map;
 
@@ -71,14 +72,18 @@ public int targetPageSize() {
      * @param onWriteCompleted The callback to be invoked after writing the parquet file.
      * @param tableDefinition The table definition to be populated inside the parquet file's schema
      * @param fieldIdToName Mapping of field id to field name, to be populated inside the parquet file's schema
+     * @param specialInstructions Special instructions to be passed to the parquet writer
      */
     ParquetInstructions toParquetInstructions(
             @NotNull final ParquetInstructions.OnWriteCompleted onWriteCompleted,
             @NotNull final TableDefinition tableDefinition,
-            @NotNull final Map<Integer, String> fieldIdToName) {
+            @NotNull final Map<Integer, String> fieldIdToName,
+            @Nullable final Object specialInstructions) {
         final ParquetInstructions.Builder builder = new ParquetInstructions.Builder();
 
-        dataInstructions().ifPresent(builder::setSpecialInstructions);
+        if (specialInstructions != null) {
+            builder.setSpecialInstructions(specialInstructions);
+        }
 
         // Add parquet writing specific instructions.
         builder.setTableDefinition(tableDefinition);