[lake/iceberg] Implement SupportsRecordBatchWrite for IcebergLakeWriter

fluss-filesystems/fluss-fs-s3/src/main/java/org/apache/fluss/fs/s3/S3FileSystemPlugin.java

@@ -56,6 +56,13 @@ public class S3FileSystemPlugin implements FileSystemPlugin {
 
     private static final String ROLE_ARN_KEY = "fs.s3a.assumed.role.arn";
 
+    /**
+     * The Hadoop S3A AssumedRoleCredentialProvider class that uses the configured role ARN to
+     * assume an IAM role for S3 access.
+     */
+    private static final String ASSUMED_ROLE_CREDENTIAL_PROVIDER =
+            "org.apache.hadoop.fs.s3a.auth.AssumedRoleCredentialProvider";
+
     private static final String[][] MIRRORED_CONFIG_KEYS = {
         {"fs.s3a.access-key", "fs.s3a.access.key"},
         {"fs.s3a.secret-key", "fs.s3a.secret.key"},
@@ -163,10 +170,16 @@ private void setCredentialProvider(org.apache.hadoop.conf.Configuration hadoopCo
         }
 
         if (hasStaticKeys || hasRoleArn) {
-            LOG.info(
-                    hasStaticKeys
-                            ? "Using provided static credentials."
-                            : "Using default AWS credential chain with AssumeRole.");
+            if (hasRoleArn && !hasStaticKeys) {
+                // When only role ARN is configured, use the AssumedRoleCredentialProvider
+                // to assume the role for all S3 operations (reads/writes).
+                hadoopConfig.set(PROVIDER_CONFIG_NAME, ASSUMED_ROLE_CREDENTIAL_PROVIDER);
+                LOG.info(
+                        "Using AssumedRoleCredentialProvider with role ARN for S3 access: {}",
+                        hadoopConfig.get(ROLE_ARN_KEY));
+            } else {
+                LOG.info("Using provided static credentials.");
+            }
             return;
         }
 

fluss-filesystems/fluss-fs-s3/src/test/java/org/apache/fluss/fs/s3/S3FileSystemPluginTest.java

@@ -62,10 +62,33 @@ void testServerModeWithRoleArnOnly() {
         org.apache.hadoop.conf.Configuration hadoopConfig =
                 plugin.buildHadoopConfiguration(flussConfig);
 
+        // When only role ARN is configured, AssumedRoleCredentialProvider should be used
         String providers = hadoopConfig.get(PROVIDER_CONFIG, "");
+        assertThat(providers)
+                .isEqualTo("org.apache.hadoop.fs.s3a.auth.AssumedRoleCredentialProvider");
         assertThat(providers).doesNotContain(DynamicTemporaryAWSCredentialsProvider.NAME);
     }
 
+    @Test
+    void testServerModeWithStaticKeysAndRoleArn() {
+        // When both static keys and role ARN are provided, static keys should take precedence
+        Configuration flussConfig = new Configuration();
+        flussConfig.setString("fs.s3a.access.key", "testAccessKey");
+        flussConfig.setString("fs.s3a.secret.key", "testSecretKey");
+        flussConfig.setString(
+                "fs.s3a.assumed.role.arn", "arn:aws:iam::123456789012:role/test-role");
+
+        S3FileSystemPlugin plugin = new S3FileSystemPlugin();
+        org.apache.hadoop.conf.Configuration hadoopConfig =
+                plugin.buildHadoopConfiguration(flussConfig);
+
+        // Static keys take precedence, AssumedRoleCredentialProvider should NOT be set
+        String providers = hadoopConfig.get(PROVIDER_CONFIG, "");
+        assertThat(providers).doesNotContain(DynamicTemporaryAWSCredentialsProvider.NAME);
+        assertThat(providers)
+                .doesNotContain("org.apache.hadoop.fs.s3a.auth.AssumedRoleCredentialProvider");
+    }
+
     @Test
     void testServerModeWithConfiguredCredentialProvider() {
         Configuration flussConfig = new Configuration();

fluss-lake/fluss-lake-iceberg/pom.xml

@@ -82,6 +82,24 @@
             <scope>provided</scope>
         </dependency>
 
+        <!-- Arrow dependencies are provided because they are only needed
+             at runtime when the tiering service writes Arrow record batches to Iceberg.
+             The Flink tiering plugin bundles these jars; fluss-lake-iceberg itself is a
+             lightweight module that should not pull them transitively. -->
+        <dependency>
+            <groupId>org.apache.arrow</groupId>
+            <artifactId>arrow-vector</artifactId>
+            <version>${arrow.version}</version>
+            <scope>provided</scope>
+        </dependency>
+
+        <dependency>
+            <groupId>org.apache.arrow</groupId>
+            <artifactId>arrow-memory-netty</artifactId>
+            <version>${arrow.version}</version>
+            <scope>provided</scope>
+        </dependency>
+
         <!-- test dependency -->
         <dependency>
             <groupId>org.apache.fluss</groupId>

fluss-lake/fluss-lake-iceberg/src/main/java/org/apache/fluss/lake/iceberg/tiering/IcebergLakeWriter.java

@@ -17,15 +17,20 @@
 
 package org.apache.fluss.lake.iceberg.tiering;
 
+import org.apache.fluss.lake.batch.ArrowRecordBatch;
+import org.apache.fluss.lake.batch.RecordBatch;
 import org.apache.fluss.lake.iceberg.maintenance.IcebergRewriteDataFiles;
 import org.apache.fluss.lake.iceberg.maintenance.RewriteDataFileResult;
 import org.apache.fluss.lake.iceberg.tiering.writer.AppendOnlyTaskWriter;
 import org.apache.fluss.lake.iceberg.tiering.writer.DeltaTaskWriter;
+import org.apache.fluss.lake.iceberg.tiering.writer.IcebergArrowBatchHelper;
 import org.apache.fluss.lake.iceberg.tiering.writer.TaskWriterFactory;
 import org.apache.fluss.lake.writer.LakeWriter;
+import org.apache.fluss.lake.writer.SupportsRecordBatchWrite;
 import org.apache.fluss.lake.writer.WriterInitContext;
 import org.apache.fluss.metadata.TablePath;
 import org.apache.fluss.record.LogRecord;
+import org.apache.fluss.types.RowType;
 import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
 
 import org.apache.iceberg.Table;
@@ -52,14 +57,17 @@
 import static org.apache.fluss.lake.iceberg.utils.IcebergConversions.toIceberg;
 
 /** Implementation of {@link LakeWriter} for Iceberg. */
-public class IcebergLakeWriter implements LakeWriter<IcebergWriteResult> {
+public class IcebergLakeWriter implements LakeWriter<IcebergWriteResult>, SupportsRecordBatchWrite {
 
     protected static final Logger LOG = LoggerFactory.getLogger(IcebergLakeWriter.class);
 
     private final Catalog icebergCatalog;
     private final Table icebergTable;
     private final RecordWriter recordWriter;
+    private final boolean isAppendOnly;
+    private final RowType flussRowType;
 
+    @Nullable private AutoCloseable arrowBatchHelper;
     @Nullable private final ExecutorService compactionExecutor;
     @Nullable private CompletableFuture<RewriteDataFileResult> compactionFuture;
 
@@ -72,6 +80,11 @@ public IcebergLakeWriter(
 
         // Create a record writer
         this.recordWriter = createRecordWriter(writerInitContext);
+        this.flussRowType = writerInitContext.tableInfo().getRowType();
+
+        List<Integer> equalityFieldIds =
+                new ArrayList<>(icebergTable.schema().identifierFieldIds());
+        this.isAppendOnly = equalityFieldIds.isEmpty();
 
         if (writerInitContext.tableInfo().getTableConfig().isDataLakeAutoCompaction()) {
             this.compactionExecutor =
@@ -108,6 +121,36 @@ public void write(LogRecord record) throws IOException {
         }
     }
 
+    @Override
+    public void write(RecordBatch recordBatch) throws IOException {
+        if (!(recordBatch instanceof ArrowRecordBatch)) {
+            throw new IllegalArgumentException(
+                    "IcebergLakeWriter only supports ArrowRecordBatch, but got "
+                            + recordBatch.getClass().getSimpleName());
+        }
+        if (!isAppendOnly) {
+            throw new IllegalStateException(
+                    "Arrow record batch writing is only supported for append-only tables.");
+        }
+        try {
+            IcebergArrowBatchHelper helper;
+            if (arrowBatchHelper == null) {
+                helper =
+                        new IcebergArrowBatchHelper(
+                                recordWriter.taskWriter,
+                                icebergTable.schema(),
+                                flussRowType,
+                                recordWriter.bucket);
+                arrowBatchHelper = helper;
+            } else {
+                helper = (IcebergArrowBatchHelper) arrowBatchHelper;
+            }
+            helper.writeArrowBatch(((ArrowRecordBatch) recordBatch).getArrowBatchData());
+        } catch (Exception e) {
+            throw new IOException("Failed to write Arrow record batch to Iceberg.", e);
+        }
+    }
+
     @Override
     public IcebergWriteResult complete() throws IOException {
         try {
@@ -137,6 +180,11 @@ public void close() throws IOException {
                 }
             }
 
+            if (arrowBatchHelper != null) {
+                arrowBatchHelper.close();
+                arrowBatchHelper = null;
+            }
+
             if (recordWriter != null) {
                 recordWriter.close();
             }