diff --git a/fluss-filesystems/fluss-fs-s3/src/main/java/org/apache/fluss/fs/s3/S3FileSystemPlugin.java b/fluss-filesystems/fluss-fs-s3/src/main/java/org/apache/fluss/fs/s3/S3FileSystemPlugin.java index 7e293f5bd62..d9e261f6928 100644 --- a/fluss-filesystems/fluss-fs-s3/src/main/java/org/apache/fluss/fs/s3/S3FileSystemPlugin.java +++ b/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; } diff --git a/fluss-filesystems/fluss-fs-s3/src/test/java/org/apache/fluss/fs/s3/S3FileSystemPluginTest.java b/fluss-filesystems/fluss-fs-s3/src/test/java/org/apache/fluss/fs/s3/S3FileSystemPluginTest.java index d9463485512..4584bdeaa21 100644 --- a/fluss-filesystems/fluss-fs-s3/src/test/java/org/apache/fluss/fs/s3/S3FileSystemPluginTest.java +++ b/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(); diff --git a/fluss-lake/fluss-lake-iceberg/pom.xml b/fluss-lake/fluss-lake-iceberg/pom.xml index ab3ced21f9b..a4ea74490fb 100644 --- a/fluss-lake/fluss-lake-iceberg/pom.xml +++ b/fluss-lake/fluss-lake-iceberg/pom.xml @@ -82,6 +82,24 @@ provided + + + org.apache.arrow + arrow-vector + ${arrow.version} + provided + + + + org.apache.arrow + arrow-memory-netty + ${arrow.version} + provided + + org.apache.fluss diff --git a/fluss-lake/fluss-lake-iceberg/src/main/java/org/apache/fluss/lake/iceberg/tiering/IcebergLakeWriter.java b/fluss-lake/fluss-lake-iceberg/src/main/java/org/apache/fluss/lake/iceberg/tiering/IcebergLakeWriter.java index 93f60ed4809..e9c6320a759 100644 --- a/fluss-lake/fluss-lake-iceberg/src/main/java/org/apache/fluss/lake/iceberg/tiering/IcebergLakeWriter.java +++ b/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 { +public class IcebergLakeWriter implements LakeWriter, 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 compactionFuture; @@ -72,6 +80,11 @@ public IcebergLakeWriter( // Create a record writer this.recordWriter = createRecordWriter(writerInitContext); + this.flussRowType = writerInitContext.tableInfo().getRowType(); + + List 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(); } diff --git a/fluss-lake/fluss-lake-iceberg/src/main/java/org/apache/fluss/lake/iceberg/tiering/writer/IcebergArrowBatchHelper.java b/fluss-lake/fluss-lake-iceberg/src/main/java/org/apache/fluss/lake/iceberg/tiering/writer/IcebergArrowBatchHelper.java new file mode 100644 index 00000000000..856755a22cb --- /dev/null +++ b/fluss-lake/fluss-lake-iceberg/src/main/java/org/apache/fluss/lake/iceberg/tiering/writer/IcebergArrowBatchHelper.java @@ -0,0 +1,255 @@ +/* + * 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.fluss.lake.iceberg.tiering.writer; + +import org.apache.fluss.record.ArrowBatchData; +import org.apache.fluss.types.ArrayType; +import org.apache.fluss.types.BigIntType; +import org.apache.fluss.types.BinaryType; +import org.apache.fluss.types.BooleanType; +import org.apache.fluss.types.BytesType; +import org.apache.fluss.types.CharType; +import org.apache.fluss.types.DataType; +import org.apache.fluss.types.DateType; +import org.apache.fluss.types.DecimalType; +import org.apache.fluss.types.DoubleType; +import org.apache.fluss.types.FloatType; +import org.apache.fluss.types.IntType; +import org.apache.fluss.types.LocalZonedTimestampType; +import org.apache.fluss.types.MapType; +import org.apache.fluss.types.RowType; +import org.apache.fluss.types.SmallIntType; +import org.apache.fluss.types.StringType; +import org.apache.fluss.types.TimeType; +import org.apache.fluss.types.TimestampType; +import org.apache.fluss.types.TinyIntType; + +import org.apache.arrow.vector.BigIntVector; +import org.apache.arrow.vector.BitVector; +import org.apache.arrow.vector.DateDayVector; +import org.apache.arrow.vector.DecimalVector; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.Float4Vector; +import org.apache.arrow.vector.Float8Vector; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.SmallIntVector; +import org.apache.arrow.vector.TimeMicroVector; +import org.apache.arrow.vector.TimeMilliVector; +import org.apache.arrow.vector.TimeStampMicroVector; +import org.apache.arrow.vector.TimeStampMilliVector; +import org.apache.arrow.vector.TinyIntVector; +import org.apache.arrow.vector.VarBinaryVector; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.iceberg.Schema; +import org.apache.iceberg.data.GenericRecord; +import org.apache.iceberg.data.Record; +import org.apache.iceberg.io.TaskWriter; +import org.apache.iceberg.types.Types; + +import java.io.IOException; +import java.math.BigDecimal; +import java.nio.ByteBuffer; +import java.time.Instant; +import java.time.LocalDate; +import java.time.LocalDateTime; +import java.time.LocalTime; +import java.time.OffsetDateTime; +import java.time.ZoneOffset; +import java.util.List; + +/** + * Helper class that reads Arrow vectors and writes Iceberg records through the existing TaskWriter. + * This avoids the per-record LogRecord deserialization overhead by reading directly from Arrow + * columnar memory. + * + *

This class is lazily loaded to avoid classloading Arrow classes when Arrow is not on the + * classpath. + */ +public class IcebergArrowBatchHelper implements AutoCloseable { + + private final TaskWriter taskWriter; + private final Schema icebergSchema; + private final RowType flussRowType; + private final int bucket; + + public IcebergArrowBatchHelper( + TaskWriter taskWriter, Schema icebergSchema, RowType flussRowType, int bucket) { + this.taskWriter = taskWriter; + this.icebergSchema = icebergSchema; + this.flussRowType = flussRowType; + this.bucket = bucket; + } + + /** + * Writes an Arrow batch to Iceberg by reading vectors directly and producing GenericRecord + * objects. System columns (__bucket, __offset, __timestamp) are computed from batch metadata. + */ + public void writeArrowBatch(ArrowBatchData arrowBatchData) throws IOException { + VectorSchemaRoot root = arrowBatchData.getVectorSchemaRoot(); + int rowCount = root.getRowCount(); + long baseOffset = arrowBatchData.getBaseLogOffset(); + long timestamp = arrowBatchData.getTimestamp(); + + List icebergFields = icebergSchema.columns(); + int userFieldCount = flussRowType.getFieldCount(); + + // Pre-fetch vectors for user columns + FieldVector[] vectors = new FieldVector[userFieldCount]; + for (int col = 0; col < userFieldCount; col++) { + vectors[col] = root.getVector(col); + } + + // Create field extractors based on Fluss types + ArrowFieldExtractor[] extractors = new ArrowFieldExtractor[userFieldCount]; + for (int col = 0; col < userFieldCount; col++) { + extractors[col] = createExtractor(flussRowType.getTypeAt(col), vectors[col]); + } + + // Write records batch-style: iterate rows, build GenericRecord from vectors + for (int row = 0; row < rowCount; row++) { + GenericRecord record = GenericRecord.create(icebergSchema); + + // Set user columns from Arrow vectors + for (int col = 0; col < userFieldCount; col++) { + if (vectors[col].isNull(row)) { + record.set(col, null); + } else { + record.set(col, extractors[col].extract(row)); + } + } + + // Set system columns + record.set(userFieldCount, bucket); // __bucket + record.set(userFieldCount + 1, baseOffset + row); // __offset + record.set( + userFieldCount + 2, + OffsetDateTime.ofInstant( + Instant.ofEpochMilli(timestamp), ZoneOffset.UTC)); // __timestamp + + taskWriter.write(record); + } + } + + @Override + public void close() { + // Nothing to close; the taskWriter is managed by the caller + } + + // --- Arrow field extractors --- + + @FunctionalInterface + private interface ArrowFieldExtractor { + Object extract(int rowIndex); + } + + private ArrowFieldExtractor createExtractor(DataType flussType, FieldVector vector) { + if (flussType instanceof BooleanType) { + BitVector v = (BitVector) vector; + return row -> v.get(row) == 1; + } else if (flussType instanceof TinyIntType) { + TinyIntVector v = (TinyIntVector) vector; + return row -> (int) v.get(row); + } else if (flussType instanceof SmallIntType) { + SmallIntVector v = (SmallIntVector) vector; + return row -> (int) v.get(row); + } else if (flussType instanceof IntType) { + IntVector v = (IntVector) vector; + return row -> v.get(row); + } else if (flussType instanceof BigIntType) { + BigIntVector v = (BigIntVector) vector; + return row -> v.get(row); + } else if (flussType instanceof FloatType) { + Float4Vector v = (Float4Vector) vector; + return row -> v.get(row); + } else if (flussType instanceof DoubleType) { + Float8Vector v = (Float8Vector) vector; + return row -> v.get(row); + } else if (flussType instanceof StringType || flussType instanceof CharType) { + VarCharVector v = (VarCharVector) vector; + return row -> new String(v.get(row), java.nio.charset.StandardCharsets.UTF_8); + } else if (flussType instanceof BytesType || flussType instanceof BinaryType) { + VarBinaryVector v = (VarBinaryVector) vector; + return row -> ByteBuffer.wrap(v.get(row)); + } else if (flussType instanceof DecimalType) { + DecimalVector v = (DecimalVector) vector; + return row -> { + BigDecimal decimal = v.getObject(row); + return decimal; + }; + } else if (flussType instanceof DateType) { + DateDayVector v = (DateDayVector) vector; + return row -> LocalDate.ofEpochDay(v.get(row)); + } else if (flussType instanceof TimeType) { + // Fluss stores time as millis-of-day in INT + if (vector instanceof TimeMilliVector) { + TimeMilliVector v = (TimeMilliVector) vector; + return row -> LocalTime.ofNanoOfDay((long) v.get(row) * 1_000_000L); + } else { + TimeMicroVector v = (TimeMicroVector) vector; + return row -> LocalTime.ofNanoOfDay(v.get(row) * 1_000L); + } + } else if (flussType instanceof TimestampType) { + TimestampType tsType = (TimestampType) flussType; + if (tsType.getPrecision() <= 3) { + TimeStampMilliVector v = (TimeStampMilliVector) vector; + return row -> { + long millis = v.get(row); + return LocalDateTime.ofInstant(Instant.ofEpochMilli(millis), ZoneOffset.UTC); + }; + } else { + TimeStampMicroVector v = (TimeStampMicroVector) vector; + return row -> { + long micros = v.get(row); + long seconds = micros / 1_000_000; + int nanos = (int) ((micros % 1_000_000) * 1_000); + return LocalDateTime.ofInstant( + Instant.ofEpochSecond(seconds, nanos), ZoneOffset.UTC); + }; + } + } else if (flussType instanceof LocalZonedTimestampType) { + LocalZonedTimestampType ltzType = (LocalZonedTimestampType) flussType; + if (ltzType.getPrecision() <= 3) { + TimeStampMilliVector v = (TimeStampMilliVector) vector; + return row -> { + long millis = v.get(row); + return OffsetDateTime.ofInstant(Instant.ofEpochMilli(millis), ZoneOffset.UTC); + }; + } else { + TimeStampMicroVector v = (TimeStampMicroVector) vector; + return row -> { + long micros = v.get(row); + long seconds = micros / 1_000_000; + int nanos = (int) ((micros % 1_000_000) * 1_000); + return OffsetDateTime.ofInstant( + Instant.ofEpochSecond(seconds, nanos), ZoneOffset.UTC); + }; + } + } else if (flussType instanceof ArrayType + || flussType instanceof MapType + || flussType instanceof RowType) { + // Complex types: fall back to null for now. + // TODO: implement nested type extraction from Arrow complex vectors + return row -> null; + } else { + throw new UnsupportedOperationException( + "Unsupported Arrow field extraction for type: " + + flussType.getClass().getSimpleName()); + } + } +}