This is an automated email from the ASF dual-hosted git repository.
ahmedabu98 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 668c6a38f37 (IcebergIO) bugfix: propagate config properties to
RecordWriter (#39250)
668c6a38f37 is described below
commit 668c6a38f37191e9d5e92be939d9da7e18542a3c
Author: Claire McGinty <[email protected]>
AuthorDate: Thu Jul 9 23:11:47 2026 -0400
(IcebergIO) bugfix: propagate config properties to RecordWriter (#39250)
* (IcebergIO) bugfix: propagate config properties to RecordWriter
* Refactor: add new writeProperties option to
IcebergWriteSchemaTransformProvider
---
.../org/apache/beam/sdk/io/iceberg/IcebergIO.java | 15 ++++-
.../IcebergWriteSchemaTransformProvider.java | 12 ++++
.../apache/beam/sdk/io/iceberg/RecordWriter.java | 24 ++++++--
.../beam/sdk/io/iceberg/RecordWriterManager.java | 14 ++++-
.../sdk/io/iceberg/WriteDirectRowsToFiles.java | 15 +++--
.../sdk/io/iceberg/WriteGroupedRowsToFiles.java | 17 ++++--
.../io/iceberg/WritePartitionedRowsToFiles.java | 17 ++++--
.../beam/sdk/io/iceberg/WriteToDestinations.java | 29 +++++++--
.../beam/sdk/io/iceberg/WriteToPartitions.java | 9 ++-
.../sdk/io/iceberg/WriteUngroupedRowsToFiles.java | 16 +++--
.../sdk/io/iceberg/RecordWriterManagerTest.java | 69 ++++++++++++++++++++++
11 files changed, 207 insertions(+), 30 deletions(-)
diff --git
a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergIO.java
b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergIO.java
index 5c5f934ea20..04feea5037d 100644
---
a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergIO.java
+++
b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergIO.java
@@ -22,6 +22,7 @@ import static
org.apache.beam.sdk.util.Preconditions.checkStateNotNull;
import com.google.auto.value.AutoValue;
import java.util.Arrays;
import java.util.List;
+import java.util.Map;
import org.apache.beam.sdk.annotations.Internal;
import org.apache.beam.sdk.io.Read;
import org.apache.beam.sdk.schemas.Schema;
@@ -406,6 +407,8 @@ public class IcebergIO {
abstract boolean getAutoSharding();
+ abstract @Nullable Map<String, String> getWriteProperties();
+
abstract Builder toBuilder();
@AutoValue.Builder
@@ -424,6 +427,8 @@ public class IcebergIO {
abstract Builder setAutoSharding(boolean autoSharding);
+ abstract Builder setWriteProperties(Map<String, String> writeProperties);
+
abstract WriteRows build();
}
@@ -474,6 +479,10 @@ public class IcebergIO {
return toBuilder().setAutoSharding(true).build();
}
+ public WriteRows withWriteProperties(Map<String, String> writeProperties) {
+ return toBuilder().setWriteProperties(writeProperties).build();
+ }
+
@Override
public IcebergWriteResult expand(PCollection<Row> input) {
List<?> allToArgs = Arrays.asList(getTableIdentifier(),
getDynamicDestinations());
@@ -509,7 +518,8 @@ public class IcebergIO {
getCatalogConfig(),
destinations,
getTriggeringFrequency(),
- getDirectWriteByteLimit()));
+ getDirectWriteByteLimit(),
+ getWriteProperties()));
case HASH:
return input
.apply(
@@ -521,7 +531,8 @@ public class IcebergIO {
getCatalogConfig(),
destinations,
getTriggeringFrequency(),
- getAutoSharding()));
+ getAutoSharding(),
+ getWriteProperties()));
default:
throw new UnsupportedOperationException(
"Unsupported distribution mode: " + getDistributionMode());
diff --git
a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergWriteSchemaTransformProvider.java
b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergWriteSchemaTransformProvider.java
index 8db4fb77a8e..0ae9d5fb0ec 100644
---
a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergWriteSchemaTransformProvider.java
+++
b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergWriteSchemaTransformProvider.java
@@ -158,6 +158,11 @@ public class IcebergWriteSchemaTransformProvider
+ "during high-throughput writes. Only available with 'hash'
distribution mode.")
public abstract @Nullable Boolean getAutosharding();
+ @SchemaFieldDescription(
+ "Properties applied to the underlying file writer (e.g. Parquet write
properties like "
+ + "'write.parquet.bloom-filter-enabled.column.<col>').")
+ public abstract @Nullable Map<String, String> getWriteProperties();
+
@AutoValue.Builder
public abstract static class Builder {
public abstract Builder setTable(String table);
@@ -188,6 +193,8 @@ public class IcebergWriteSchemaTransformProvider
public abstract Builder setAutosharding(Boolean autosharding);
+ public abstract Builder setWriteProperties(Map<String, String>
writeProperties);
+
public abstract Configuration build();
}
@@ -279,6 +286,11 @@ public class IcebergWriteSchemaTransformProvider
writeTransform = writeTransform.withAutosharding();
}
+ @Nullable Map<String, String> writeProperties =
configuration.getWriteProperties();
+ if (writeProperties != null && !writeProperties.isEmpty()) {
+ writeTransform = writeTransform.withWriteProperties(writeProperties);
+ }
+
// TODO: support dynamic destinations
IcebergWriteResult result = rows.apply(writeTransform);
diff --git
a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/RecordWriter.java
b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/RecordWriter.java
index fd3d5d63327..c3b63b2a336 100644
---
a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/RecordWriter.java
+++
b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/RecordWriter.java
@@ -18,6 +18,7 @@
package org.apache.beam.sdk.io.iceberg;
import java.io.IOException;
+import java.util.Map;
import org.apache.beam.sdk.metrics.Counter;
import org.apache.beam.sdk.metrics.Metrics;
import org.apache.iceberg.DataFile;
@@ -34,6 +35,7 @@ import org.apache.iceberg.encryption.EncryptionKeyMetadata;
import org.apache.iceberg.io.DataWriter;
import org.apache.iceberg.io.OutputFile;
import org.apache.iceberg.parquet.Parquet;
+import org.checkerframework.checker.nullness.qual.Nullable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -54,11 +56,22 @@ class RecordWriter {
catalog.loadTable(destination.getTableIdentifier()),
destination.getFileFormat(),
filename,
- partitionKey);
+ partitionKey,
+ null);
}
RecordWriter(Table table, FileFormat fileFormat, String filename, StructLike
partitionKey)
throws IOException {
+ this(table, fileFormat, filename, partitionKey, null);
+ }
+
+ RecordWriter(
+ Table table,
+ FileFormat fileFormat,
+ String filename,
+ StructLike partitionKey,
+ @Nullable Map<String, String> writeProperties)
+ throws IOException {
this.table = table;
this.fileFormat = fileFormat;
@@ -91,14 +104,17 @@ class RecordWriter {
.build();
break;
case PARQUET:
- icebergDataWriter =
+ Parquet.DataWriteBuilder parquetBuilder =
Parquet.writeData(outputFile)
.forTable(table)
.createWriterFunc(GenericParquetWriter::create)
.withPartition(partitionKey)
.withKeyMetadata(keyMetadata)
- .overwrite()
- .build();
+ .overwrite();
+ if (writeProperties != null && !writeProperties.isEmpty()) {
+ parquetBuilder.setAll(writeProperties);
+ }
+ icebergDataWriter = parquetBuilder.build();
break;
case ORC:
throw new UnsupportedOperationException("ORC file format not currently
supported.");
diff --git
a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/RecordWriterManager.java
b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/RecordWriterManager.java
index d5439ad8383..6893c743f43 100644
---
a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/RecordWriterManager.java
+++
b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/RecordWriterManager.java
@@ -202,7 +202,8 @@ class RecordWriterManager implements AutoCloseable {
table,
icebergDestination.getFileFormat(),
filePrefix + "_" + stateToken + "_" + recordIndex,
- partitionKey);
+ partitionKey,
+ writeProperties);
openWriters++;
return writer;
} catch (IOException e) {
@@ -253,6 +254,7 @@ class RecordWriterManager implements AutoCloseable {
private final String filePrefix;
private final long maxFileSize;
private final int maxNumWriters;
+ private final @Nullable Map<String, String> writeProperties;
@VisibleForTesting int openWriters = 0;
@VisibleForTesting
@@ -265,10 +267,20 @@ class RecordWriterManager implements AutoCloseable {
RecordWriterManager(
IcebergCatalogConfig catalogConfig, String filePrefix, long maxFileSize,
int maxNumWriters) {
+ this(catalogConfig, filePrefix, maxFileSize, maxNumWriters, null);
+ }
+
+ RecordWriterManager(
+ IcebergCatalogConfig catalogConfig,
+ String filePrefix,
+ long maxFileSize,
+ int maxNumWriters,
+ @Nullable Map<String, String> writeProperties) {
this.catalogConfig = catalogConfig;
this.filePrefix = filePrefix;
this.maxFileSize = maxFileSize;
this.maxNumWriters = maxNumWriters;
+ this.writeProperties = writeProperties;
}
/**
diff --git
a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WriteDirectRowsToFiles.java
b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WriteDirectRowsToFiles.java
index 5cf095dc3c3..e03085e6be7 100644
---
a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WriteDirectRowsToFiles.java
+++
b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WriteDirectRowsToFiles.java
@@ -40,16 +40,19 @@ class WriteDirectRowsToFiles
private final IcebergCatalogConfig catalogConfig;
private final String filePrefix;
private final long maxBytesPerFile;
+ private final @Nullable Map<String, String> writeProperties;
WriteDirectRowsToFiles(
IcebergCatalogConfig catalogConfig,
DynamicDestinations dynamicDestinations,
String filePrefix,
- long maxBytesPerFile) {
+ long maxBytesPerFile,
+ @Nullable Map<String, String> writeProperties) {
this.catalogConfig = catalogConfig;
this.dynamicDestinations = dynamicDestinations;
this.filePrefix = filePrefix;
this.maxBytesPerFile = maxBytesPerFile;
+ this.writeProperties = writeProperties;
}
@Override
@@ -57,7 +60,7 @@ class WriteDirectRowsToFiles
return input.apply(
ParDo.of(
new WriteDirectRowsToFilesDoFn(
- catalogConfig, dynamicDestinations, maxBytesPerFile,
filePrefix)));
+ catalogConfig, dynamicDestinations, maxBytesPerFile,
filePrefix, writeProperties)));
}
private static class WriteDirectRowsToFilesDoFn extends DoFn<KV<String,
Row>, FileWriteResult> {
@@ -66,24 +69,28 @@ class WriteDirectRowsToFiles
private final IcebergCatalogConfig catalogConfig;
private final String filePrefix;
private final long maxFileSize;
+ private final @Nullable Map<String, String> writeProperties;
private transient @Nullable RecordWriterManager recordWriterManager;
WriteDirectRowsToFilesDoFn(
IcebergCatalogConfig catalogConfig,
DynamicDestinations dynamicDestinations,
long maxFileSize,
- String filePrefix) {
+ String filePrefix,
+ @Nullable Map<String, String> writeProperties) {
this.catalogConfig = catalogConfig;
this.dynamicDestinations = dynamicDestinations;
this.filePrefix = filePrefix;
this.maxFileSize = maxFileSize;
+ this.writeProperties = writeProperties;
this.recordWriterManager = null;
}
@StartBundle
public void startBundle() {
recordWriterManager =
- new RecordWriterManager(catalogConfig, filePrefix, maxFileSize,
Integer.MAX_VALUE);
+ new RecordWriterManager(
+ catalogConfig, filePrefix, maxFileSize, Integer.MAX_VALUE,
writeProperties);
}
@ProcessElement
diff --git
a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WriteGroupedRowsToFiles.java
b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WriteGroupedRowsToFiles.java
index b16496240d1..e74715a7eeb 100644
---
a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WriteGroupedRowsToFiles.java
+++
b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WriteGroupedRowsToFiles.java
@@ -18,6 +18,7 @@
package org.apache.beam.sdk.io.iceberg;
import java.util.List;
+import java.util.Map;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.ParDo;
@@ -30,6 +31,7 @@ import org.apache.beam.sdk.values.Row;
import org.apache.beam.sdk.values.WindowedValue;
import org.apache.beam.sdk.values.WindowedValues;
import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions;
+import org.checkerframework.checker.nullness.qual.Nullable;
class WriteGroupedRowsToFiles
extends PTransform<
@@ -39,16 +41,19 @@ class WriteGroupedRowsToFiles
private final DynamicDestinations dynamicDestinations;
private final IcebergCatalogConfig catalogConfig;
private final String filePrefix;
+ private final @Nullable Map<String, String> writeProperties;
WriteGroupedRowsToFiles(
IcebergCatalogConfig catalogConfig,
DynamicDestinations dynamicDestinations,
String filePrefix,
- long maxBytesPerFile) {
+ long maxBytesPerFile,
+ @Nullable Map<String, String> writeProperties) {
this.catalogConfig = catalogConfig;
this.dynamicDestinations = dynamicDestinations;
this.filePrefix = filePrefix;
this.maxBytesPerFile = maxBytesPerFile;
+ this.writeProperties = writeProperties;
}
@Override
@@ -57,7 +62,7 @@ class WriteGroupedRowsToFiles
return input.apply(
ParDo.of(
new WriteGroupedRowsToFilesDoFn(
- catalogConfig, dynamicDestinations, maxBytesPerFile,
filePrefix)));
+ catalogConfig, dynamicDestinations, maxBytesPerFile,
filePrefix, writeProperties)));
}
private static class WriteGroupedRowsToFilesDoFn
@@ -67,16 +72,19 @@ class WriteGroupedRowsToFiles
private final IcebergCatalogConfig catalogConfig;
private final String filePrefix;
private final long maxFileSize;
+ private final @Nullable Map<String, String> writeProperties;
WriteGroupedRowsToFilesDoFn(
IcebergCatalogConfig catalogConfig,
DynamicDestinations dynamicDestinations,
long maxFileSize,
- String filePrefix) {
+ String filePrefix,
+ @Nullable Map<String, String> writeProperties) {
this.catalogConfig = catalogConfig;
this.dynamicDestinations = dynamicDestinations;
this.filePrefix = filePrefix;
this.maxFileSize = maxFileSize;
+ this.writeProperties = writeProperties;
}
@ProcessElement
@@ -93,7 +101,8 @@ class WriteGroupedRowsToFiles
WindowedValues.of(destination, window.maxTimestamp(), window,
paneInfo);
RecordWriterManager writer;
try (RecordWriterManager openWriter =
- new RecordWriterManager(catalogConfig, filePrefix, maxFileSize,
Integer.MAX_VALUE)) {
+ new RecordWriterManager(
+ catalogConfig, filePrefix, maxFileSize, Integer.MAX_VALUE,
writeProperties)) {
writer = openWriter;
for (Row e : element.getValue()) {
writer.write(windowedDestination, e);
diff --git
a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WritePartitionedRowsToFiles.java
b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WritePartitionedRowsToFiles.java
index 8b4ae0863f7..84e7d05b124 100644
---
a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WritePartitionedRowsToFiles.java
+++
b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WritePartitionedRowsToFiles.java
@@ -59,14 +59,17 @@ class WritePartitionedRowsToFiles
private final DynamicDestinations dynamicDestinations;
private final IcebergCatalogConfig catalogConfig;
private final String filePrefix;
+ private final @Nullable Map<String, String> writeProperties;
WritePartitionedRowsToFiles(
IcebergCatalogConfig catalogConfig,
DynamicDestinations dynamicDestinations,
- String filePrefix) {
+ String filePrefix,
+ @Nullable Map<String, String> writeProperties) {
this.catalogConfig = catalogConfig;
this.dynamicDestinations = dynamicDestinations;
this.filePrefix = filePrefix;
+ this.writeProperties = writeProperties;
}
@Override
@@ -78,7 +81,9 @@ class WritePartitionedRowsToFiles
.getElemCoder())
.getSchema();
return input.apply(
- ParDo.of(new WriteDoFn(catalogConfig, dynamicDestinations, filePrefix,
dataSchema)));
+ ParDo.of(
+ new WriteDoFn(
+ catalogConfig, dynamicDestinations, filePrefix, dataSchema,
writeProperties)));
}
private static class WriteDoFn extends DoFn<KV<Row, Iterable<Row>>,
FileWriteResult> {
@@ -87,6 +92,7 @@ class WritePartitionedRowsToFiles
private final IcebergCatalogConfig catalogConfig;
private final String filePrefix;
private final Schema dataSchema;
+ private final @Nullable Map<String, String> writeProperties;
private transient @MonotonicNonNull Map<TableIdentifier, Integer> specIds;
private transient @MonotonicNonNull Map<TableIdentifier, Map<String,
PartitionField>>
partitionFieldMaps;
@@ -95,11 +101,13 @@ class WritePartitionedRowsToFiles
IcebergCatalogConfig catalogConfig,
DynamicDestinations dynamicDestinations,
String filePrefix,
- Schema dataSchema) {
+ Schema dataSchema,
+ @Nullable Map<String, String> writeProperties) {
this.catalogConfig = catalogConfig;
this.dynamicDestinations = dynamicDestinations;
this.filePrefix = filePrefix;
this.dataSchema = dataSchema;
+ this.writeProperties = writeProperties;
}
@Setup
@@ -132,7 +140,8 @@ class WritePartitionedRowsToFiles
.addExtension(String.format("%s-%s", filePrefix,
UUID.randomUUID()));
RecordWriter writer =
- new RecordWriter(table, destination.getFileFormat(), fileName,
partitionData);
+ new RecordWriter(
+ table, destination.getFileFormat(), fileName, partitionData,
writeProperties);
try {
for (Row row : element.getValue()) {
Record record = IcebergUtils.beamRowToIcebergRecord(table.schema(),
row);
diff --git
a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WriteToDestinations.java
b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WriteToDestinations.java
index bea84fc826b..684ef350a20 100644
---
a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WriteToDestinations.java
+++
b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WriteToDestinations.java
@@ -59,16 +59,19 @@ class WriteToDestinations extends
PTransform<PCollection<KV<String, Row>>, Icebe
private final @Nullable Duration triggeringFrequency;
private final String filePrefix;
private final @Nullable Integer directWriteByteLimit;
+ private final @Nullable Map<String, String> writeProperties;
WriteToDestinations(
IcebergCatalogConfig catalogConfig,
DynamicDestinations dynamicDestinations,
@Nullable Duration triggeringFrequency,
- @Nullable Integer directWriteByteLimit) {
+ @Nullable Integer directWriteByteLimit,
+ @Nullable Map<String, String> writeProperties) {
this.dynamicDestinations = dynamicDestinations;
this.catalogConfig = catalogConfig;
this.triggeringFrequency = triggeringFrequency;
this.directWriteByteLimit = directWriteByteLimit;
+ this.writeProperties = writeProperties;
// single unique prefix per write transform
this.filePrefix = UUID.randomUUID().toString();
}
@@ -112,7 +115,11 @@ class WriteToDestinations extends
PTransform<PCollection<KV<String, Row>>, Icebe
return groupedRecords.apply(
"WriteGroupedRows",
new WriteGroupedRowsToFiles(
- catalogConfig, dynamicDestinations, filePrefix,
DEFAULT_MAX_BYTES_PER_FILE));
+ catalogConfig,
+ dynamicDestinations,
+ filePrefix,
+ DEFAULT_MAX_BYTES_PER_FILE,
+ writeProperties));
}
private PCollection<FileWriteResult>
applyUserTriggering(PCollection<FileWriteResult> input) {
@@ -157,7 +164,11 @@ class WriteToDestinations extends
PTransform<PCollection<KV<String, Row>>, Icebe
largeBatches.apply(
"WriteDirectRowsToFiles",
new WriteDirectRowsToFiles(
- catalogConfig, dynamicDestinations, filePrefix,
DEFAULT_MAX_BYTES_PER_FILE));
+ catalogConfig,
+ dynamicDestinations,
+ filePrefix,
+ DEFAULT_MAX_BYTES_PER_FILE,
+ writeProperties));
PCollection<FileWriteResult> groupedFileWrites =
groupAndWriteRecords(smallBatches);
@@ -189,7 +200,11 @@ class WriteToDestinations extends
PTransform<PCollection<KV<String, Row>>, Icebe
input.apply(
"Fast-path write rows",
new WriteUngroupedRowsToFiles(
- catalogConfig, dynamicDestinations, filePrefix,
DEFAULT_MAX_BYTES_PER_FILE));
+ catalogConfig,
+ dynamicDestinations,
+ filePrefix,
+ DEFAULT_MAX_BYTES_PER_FILE,
+ writeProperties));
// Then write the rest by shuffling on the destination
PCollection<FileWriteResult> writeGroupedResult =
@@ -199,7 +214,11 @@ class WriteToDestinations extends
PTransform<PCollection<KV<String, Row>>, Icebe
.apply(
"Write remaining rows to files",
new WriteGroupedRowsToFiles(
- catalogConfig, dynamicDestinations, filePrefix,
DEFAULT_MAX_BYTES_PER_FILE));
+ catalogConfig,
+ dynamicDestinations,
+ filePrefix,
+ DEFAULT_MAX_BYTES_PER_FILE,
+ writeProperties));
return PCollectionList.of(writeUngroupedResult.getWrittenFiles())
.and(writeGroupedResult)
diff --git
a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WriteToPartitions.java
b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WriteToPartitions.java
index 310fa1bede4..84b9c6a9c0f 100644
---
a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WriteToPartitions.java
+++
b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WriteToPartitions.java
@@ -22,6 +22,7 @@ import static
org.apache.beam.sdk.values.TypeDescriptors.iterables;
import static org.apache.beam.sdk.values.TypeDescriptors.kvs;
import static org.apache.beam.sdk.values.TypeDescriptors.rows;
+import java.util.Map;
import java.util.UUID;
import org.apache.beam.sdk.coders.IterableCoder;
import org.apache.beam.sdk.coders.KvCoder;
@@ -46,18 +47,21 @@ class WriteToPartitions extends
PTransform<PCollection<KV<Row, Row>>, IcebergWri
private final @Nullable Duration triggeringFrequency;
private final String filePrefix;
private final boolean autoSharding;
+ private final @Nullable Map<String, String> writeProperties;
WriteToPartitions(
IcebergCatalogConfig catalogConfig,
DynamicDestinations dynamicDestinations,
@Nullable Duration triggeringFrequency,
- boolean autoSharding) {
+ boolean autoSharding,
+ @Nullable Map<String, String> writeProperties) {
this.dynamicDestinations = dynamicDestinations;
this.catalogConfig = catalogConfig;
this.triggeringFrequency = triggeringFrequency;
// single unique prefix per write transform
this.filePrefix = UUID.randomUUID().toString();
this.autoSharding = autoSharding;
+ this.writeProperties = writeProperties;
}
private PCollection<KV<Row, Iterable<Row>>>
groupByPartition(PCollection<KV<Row, Row>> input) {
@@ -95,7 +99,8 @@ class WriteToPartitions extends
PTransform<PCollection<KV<Row, Row>>, IcebergWri
PCollection<FileWriteResult> writtenFiles =
groupedRows.apply(
- new WritePartitionedRowsToFiles(catalogConfig,
dynamicDestinations, filePrefix));
+ new WritePartitionedRowsToFiles(
+ catalogConfig, dynamicDestinations, filePrefix,
writeProperties));
if (IcebergUtils.isUnbounded(input) && triggeringFrequency != null) {
writtenFiles =
diff --git
a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WriteUngroupedRowsToFiles.java
b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WriteUngroupedRowsToFiles.java
index ff9a98c6200..7c780e6395d 100644
---
a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WriteUngroupedRowsToFiles.java
+++
b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WriteUngroupedRowsToFiles.java
@@ -72,16 +72,19 @@ class WriteUngroupedRowsToFiles
private final DynamicDestinations dynamicDestinations;
private final IcebergCatalogConfig catalogConfig;
private final long maxBytesPerFile;
+ private final @Nullable Map<String, String> writeProperties;
WriteUngroupedRowsToFiles(
IcebergCatalogConfig catalogConfig,
DynamicDestinations dynamicDestinations,
String filePrefix,
- long maxBytesPerFile) {
+ long maxBytesPerFile,
+ @Nullable Map<String, String> writeProperties) {
this.catalogConfig = catalogConfig;
this.dynamicDestinations = dynamicDestinations;
this.filePrefix = filePrefix;
this.maxBytesPerFile = maxBytesPerFile;
+ this.writeProperties = writeProperties;
}
@Override
@@ -95,7 +98,8 @@ class WriteUngroupedRowsToFiles
dynamicDestinations,
filePrefix,
DEFAULT_MAX_WRITERS_PER_BUNDLE,
- maxBytesPerFile))
+ maxBytesPerFile,
+ writeProperties))
.withOutputTags(
WRITTEN_FILES_TAG,
TupleTagList.of(ImmutableList.of(WRITTEN_ROWS_TAG,
SPILLED_ROWS_TAG))));
@@ -191,6 +195,7 @@ class WriteUngroupedRowsToFiles
private final long maxFileSize;
private final DynamicDestinations dynamicDestinations;
private final IcebergCatalogConfig catalogConfig;
+ private final @Nullable Map<String, String> writeProperties;
private transient @Nullable RecordWriterManager recordWriterManager;
private int spilledShardNumber;
@@ -199,18 +204,21 @@ class WriteUngroupedRowsToFiles
DynamicDestinations dynamicDestinations,
String filename,
int maximumWritersPerBundle,
- long maxFileSize) {
+ long maxFileSize,
+ @Nullable Map<String, String> writeProperties) {
this.catalogConfig = catalogConfig;
this.dynamicDestinations = dynamicDestinations;
this.filename = filename;
this.maxWritersPerBundle = maximumWritersPerBundle;
this.maxFileSize = maxFileSize;
+ this.writeProperties = writeProperties;
}
@StartBundle
public void startBundle() {
recordWriterManager =
- new RecordWriterManager(catalogConfig, filename, maxFileSize,
maxWritersPerBundle);
+ new RecordWriterManager(
+ catalogConfig, filename, maxFileSize, maxWritersPerBundle,
writeProperties);
this.spilledShardNumber =
ThreadLocalRandom.current().nextInt(SPILLED_RECORD_SHARDING_FACTOR);
}
diff --git
a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/RecordWriterManagerTest.java
b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/RecordWriterManagerTest.java
index 390e8d87af2..821fb2ac7b2 100644
---
a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/RecordWriterManagerTest.java
+++
b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/RecordWriterManagerTest.java
@@ -53,6 +53,8 @@ import org.apache.beam.sdk.values.WindowedValue;
import org.apache.beam.sdk.values.WindowedValues;
import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
import org.apache.commons.lang3.RandomStringUtils;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.Path;
import org.apache.iceberg.AppendFiles;
import org.apache.iceberg.DataFile;
import org.apache.iceberg.FileFormat;
@@ -77,6 +79,10 @@ import org.apache.iceberg.types.Conversions;
import org.apache.iceberg.types.Type;
import org.apache.iceberg.types.Types;
import org.apache.iceberg.util.DateTimeUtil;
+import org.apache.parquet.hadoop.ParquetFileReader;
+import org.apache.parquet.hadoop.metadata.BlockMetaData;
+import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData;
+import org.apache.parquet.hadoop.util.HadoopInputFile;
import org.checkerframework.checker.nullness.qual.Nullable;
import org.joda.time.DateTime;
import org.joda.time.DateTimeZone;
@@ -1300,4 +1306,67 @@ public class RecordWriterManagerTest {
assertTrue(
"Bundle 2 should produce data files",
bundle2.getSerializableDataFiles().containsKey(dest));
}
+
+ @Test
+ public void testWritePropertiesAppliedToParquetFiles() throws IOException {
+ Schema bloomSchema =
+
Schema.builder().addInt32Field("colWithBf").addInt32Field("colWithoutBf").build();
+ org.apache.iceberg.Schema icebergBloomSchema =
+ IcebergUtils.beamSchemaToIcebergSchema(bloomSchema);
+
+ TableIdentifier tableId = TableIdentifier.of("default",
"test_write_properties");
+ warehouse.createTable(tableId, icebergBloomSchema);
+
+ Map<String, String> writeProperties =
+ ImmutableMap.of(
+ "write.parquet.bloom-filter-enabled.column.colWithBf", "true",
+ "write.parquet.bloom-filter-enabled.column.colWithoutBf", "false");
+
+ IcebergDestination destination =
+ IcebergDestination.builder()
+ .setTableIdentifier(tableId)
+ .setFileFormat(FileFormat.PARQUET)
+ .build();
+ WindowedValue<IcebergDestination> dest =
WindowedValues.valueInGlobalWindow(destination);
+
+ RecordWriterManager writerManager =
+ new RecordWriterManager(catalogConfig, "test_bloom", Long.MAX_VALUE,
3, writeProperties);
+ for (int i = 0; i < 10; i++) {
+ Row row = Row.withSchema(bloomSchema).addValues(i, 100 + i).build();
+ assertTrue(writerManager.write(dest, row));
+ }
+ writerManager.close();
+
+ List<SerializableDataFile> dataFiles =
writerManager.getSerializableDataFiles().get(dest);
+ assertEquals(1, dataFiles.size());
+
+ String dataFilePath = dataFiles.get(0).getPath();
+ assertNotNull(dataFilePath);
+
+ try (ParquetFileReader reader =
+ ParquetFileReader.open(
+ HadoopInputFile.fromPath(new Path(dataFilePath), new
Configuration()))) {
+ List<BlockMetaData> blocks = reader.getFooter().getBlocks();
+ assertFalse("Parquet file should have at least one row group",
blocks.isEmpty());
+
+ for (int i = 0; i < blocks.size(); i++) {
+ BlockMetaData block = blocks.get(i);
+ assertEquals("Each row group should have 2 columns", 2,
block.getColumns().size());
+
+ for (ColumnChunkMetaData col : block.getColumns()) {
+ boolean hasBloomFilter = col.getBloomFilterOffset() > 0;
+ String colName = col.getPath().toDotString();
+ if (colName.equals("colWithBf")) {
+ assertTrue(
+ "Column 'colWithBf' in row group " + i + " should have a bloom
filter",
+ hasBloomFilter);
+ } else if (colName.equals("colWithoutBf")) {
+ assertFalse(
+ "Column 'colWithoutBf' in row group " + i + " should not have
a bloom filter",
+ hasBloomFilter);
+ }
+ }
+ }
+ }
+ }
}
