From f347aa29d7768a5571c4ef6026c8a2f3ad2bd8eb Mon Sep 17 00:00:00 2001 From: Claire McGinty Date: Wed, 5 Aug 2026 16:35:55 -0400 Subject: [PATCH] (IcebergIO) bugfix: wire writeProperties through table create request, not DataWriteBuilder --- .../beam/sdk/io/iceberg/RecordWriter.java | 24 +--- .../sdk/io/iceberg/RecordWriterManager.java | 8 +- .../iceberg/WritePartitionedRowsToFiles.java | 8 +- .../sdk/io/iceberg/IcebergIOWriteTest.java | 106 ++++++++++++++++++ .../io/iceberg/RecordWriterManagerTest.java | 69 ------------ 5 files changed, 120 insertions(+), 95 deletions(-) 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 c3b63b2a336f..fd3d5d63327c 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,7 +18,6 @@ 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; @@ -35,7 +34,6 @@ 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; @@ -56,22 +54,11 @@ class RecordWriter { catalog.loadTable(destination.getTableIdentifier()), destination.getFileFormat(), filename, - partitionKey, - null); + partitionKey); } 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 writeProperties) - throws IOException { this.table = table; this.fileFormat = fileFormat; @@ -104,17 +91,14 @@ class RecordWriter { .build(); break; case PARQUET: - Parquet.DataWriteBuilder parquetBuilder = + icebergDataWriter = Parquet.writeData(outputFile) .forTable(table) .createWriterFunc(GenericParquetWriter::create) .withPartition(partitionKey) .withKeyMetadata(keyMetadata) - .overwrite(); - if (writeProperties != null && !writeProperties.isEmpty()) { - parquetBuilder.setAll(writeProperties); - } - icebergDataWriter = parquetBuilder.build(); + .overwrite() + .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 6893c743f431..014475714050 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,8 +202,7 @@ private RecordWriter createWriter(PartitionKey partitionKey) { table, icebergDestination.getFileFormat(), filePrefix + "_" + stateToken + "_" + recordIndex, - partitionKey, - writeProperties); + partitionKey); openWriters++; return writer; } catch (IOException e) { @@ -311,8 +310,11 @@ private Table loadOrCreateTable(IcebergDestination destination, Schema dataSchem SortOrder sortOrder = createConfig != null ? createConfig.getSortOrder() : SortOrder.unsorted(); Map tableProperties = createConfig != null && createConfig.getTableProperties() != null - ? createConfig.getTableProperties() + ? Maps.newHashMap(createConfig.getTableProperties()) : Maps.newHashMap(); + if (writeProperties != null) { + tableProperties.putAll(writeProperties); + } // Create namespace if it does not exist yet if (!namespace.isEmpty() && catalog instanceof SupportsNamespaces) { 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 d1a08980fa9d..dc912ff1e4e1 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 @@ -141,8 +141,7 @@ public void processElement( .addExtension(String.format("%s-%s", filePrefix, UUID.randomUUID())); RecordWriter writer = - new RecordWriter( - table, destination.getFileFormat(), fileName, partitionData, writeProperties); + new RecordWriter(table, destination.getFileFormat(), fileName, partitionData); try { for (Row row : element.getValue()) { Record record = IcebergUtils.beamRowToIcebergRecord(table.schema(), row); @@ -194,8 +193,11 @@ private Table loadOrCreateTable( createConfig != null ? createConfig.getSortOrder() : SortOrder.unsorted(); Map tableProperties = createConfig != null && createConfig.getTableProperties() != null - ? createConfig.getTableProperties() + ? Maps.newHashMap(createConfig.getTableProperties()) : Maps.newHashMap(); + if (writeProperties != null) { + tableProperties.putAll(writeProperties); + } // Create namespace if it does not exist yet if (!namespace.isEmpty() && catalog instanceof SupportsNamespaces) { diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergIOWriteTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergIOWriteTest.java index 384e11b761bb..2d3243d8b9e6 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergIOWriteTest.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergIOWriteTest.java @@ -88,9 +88,13 @@ import org.apache.iceberg.data.IcebergGenerics; import org.apache.iceberg.data.Record; import org.apache.iceberg.data.parquet.GenericParquetWriter; +import org.apache.iceberg.io.CloseableIterable; import org.apache.iceberg.io.DataWriter; import org.apache.iceberg.io.OutputFile; import org.apache.iceberg.parquet.Parquet; +import org.apache.parquet.hadoop.ParquetFileReader; +import org.apache.parquet.hadoop.metadata.BlockMetaData; +import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData; import org.hamcrest.Matchers; import org.joda.time.Duration; import org.joda.time.Instant; @@ -829,4 +833,106 @@ public void testCreateTableWithPartitionSpecAndSortOrder() { List writtenRecords = ImmutableList.copyOf(IcebergGenerics.read(table).build()); assertThat(writtenRecords, Matchers.containsInAnyOrder(TestFixtures.FILE1SNAPSHOT1.toArray())); } + + @Test + public void testWriteWithParquetProperties() throws Exception { + TableIdentifier tableId = + TableIdentifier.of( + "default", "parquet_props_" + Long.toString(UUID.randomUUID().hashCode(), 16)); + + Schema beamSchema = IcebergUtils.icebergSchemaToBeamSchema(TestFixtures.SCHEMA); + + Map catalogProps = + ImmutableMap.builder() + .put("type", CatalogUtil.ICEBERG_CATALOG_TYPE_HADOOP) + .put("warehouse", warehouse.location) + .build(); + + IcebergCatalogConfig catalog = + IcebergCatalogConfig.builder() + .setCatalogName("name") + .setCatalogProperties(catalogProps) + .build(); + + testPipeline + .apply("Records To Add", Create.of(TestFixtures.asRows(TestFixtures.FILE1SNAPSHOT1))) + .setRowSchema(beamSchema) + .apply( + "Append To Table", + writeTransform(catalog, tableId) + .withWriteProperties( + ImmutableMap.of("write.parquet.bloom-filter-enabled.column.data", "true"))); + + testPipeline.run().waitUntilFinish(); + + Table table = warehouse.loadTable(tableId); + + List writtenRecords = ImmutableList.copyOf(IcebergGenerics.read(table).build()); + assertThat(writtenRecords, Matchers.containsInAnyOrder(TestFixtures.FILE1SNAPSHOT1.toArray())); + + // verify bloom filter is present on 'data' column in written parquet files + try (CloseableIterable tasks = table.newScan().planFiles()) { + for (org.apache.iceberg.FileScanTask task : tasks) { + String path = task.file().path().toString(); + try (ParquetFileReader reader = + ParquetFileReader.open( + org.apache.parquet.hadoop.util.HadoopInputFile.fromPath( + new org.apache.hadoop.fs.Path(path), + new org.apache.hadoop.conf.Configuration()))) { + for (BlockMetaData block : reader.getFooter().getBlocks()) { + for (ColumnChunkMetaData col : block.getColumns()) { + boolean hasBloom = col.getBloomFilterOffset() > 0; + if (col.getPath().toDotString().equals("data")) { + assertTrue("Expected bloom filter on column 'data', but none was found", hasBloom); + } else { + assertFalse( + "Expected no bloom filter on column '" + col.getPath().toDotString() + "'", + hasBloom); + } + } + } + } + } + } + } + + @Test + public void testWriteWithTableProperties() throws Exception { + TableIdentifier tableId = + TableIdentifier.of( + "default", "table_props_" + Long.toString(UUID.randomUUID().hashCode(), 16)); + + Schema beamSchema = IcebergUtils.icebergSchemaToBeamSchema(TestFixtures.SCHEMA); + + Map catalogProps = + ImmutableMap.builder() + .put("type", CatalogUtil.ICEBERG_CATALOG_TYPE_HADOOP) + .put("warehouse", warehouse.location) + .build(); + + IcebergCatalogConfig catalog = + IcebergCatalogConfig.builder() + .setCatalogName("name") + .setCatalogProperties(catalogProps) + .build(); + + testPipeline + .apply("Records To Add", Create.of(TestFixtures.asRows(TestFixtures.FILE1SNAPSHOT1))) + .setRowSchema(beamSchema) + .apply( + "Append To Table", + writeTransform(catalog, tableId) + .withWriteProperties( + ImmutableMap.of("write.data.path", warehouse.location + "/custom_data_path"))); + + testPipeline.run().waitUntilFinish(); + + Table table = warehouse.loadTable(tableId); + + List writtenRecords = ImmutableList.copyOf(IcebergGenerics.read(table).build()); + assertThat(writtenRecords, Matchers.containsInAnyOrder(TestFixtures.FILE1SNAPSHOT1.toArray())); + + assertEquals( + warehouse.location + "/custom_data_path", table.properties().get("write.data.path")); + } } 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 821fb2ac7b24..390e8d87af28 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,8 +53,6 @@ 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; @@ -79,10 +77,6 @@ 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; @@ -1306,67 +1300,4 @@ public void testFileIOSurvivesAcrossBundles() throws IOException { 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 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 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 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 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); - } - } - } - } - } }