diff --git a/.github/trigger_files/IO_Iceberg_Integration_Tests.json b/.github/trigger_files/IO_Iceberg_Integration_Tests.json index 7ab7bcd9a9c6..37dd25bf9029 100644 --- a/.github/trigger_files/IO_Iceberg_Integration_Tests.json +++ b/.github/trigger_files/IO_Iceberg_Integration_Tests.json @@ -1,4 +1,4 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run.", - "modification": 2 + "modification": 3 } diff --git a/CHANGES.md b/CHANGES.md index 1c85993ca9c6..1aa8d446c0a6 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -149,6 +149,7 @@ * (Java/Python) The SQL `SqlTransformSchemaTransformProvider` configuration schema (`schematransform:org.apache.beam:sql_transform:v1`) now uses only portable types so it can be discovered from Python: `dialect` is a string, `tableproviders` is an array of strings, and the unused `parameters` field was removed. Invalid dialects and unknown table providers are rejected with a clear error ([#34613](https://github.com/apache/beam/issues/34613)). * (Python) An unknown pipeline option is no longer given the value of a later argument that does not directly follow it, which made `Pipeline.run()` and `get_all_options(retain_unknown_options=True)` fail with `expected one argument` for flags like `--unknown_flag --job_name=foo bar` ([#22434](https://github.com/apache/beam/issues/22434)). * (Python) Fixed `WriteToJson` silently rounding floats to 10 decimal places: it now defaults pandas' `double_precision` to 15, so values such as `0.1234567890123` and `1e-11` are written correctly. Floats that need 16-17 significant digits may still be rounded ([#41102](https://github.com/apache/beam/issues/41102)). +* (Java) Makes some fixes to IcebergIO's hash distribution mode, which sometimes wrote incorrect partition values. ([#41012](https://github.com/apache/beam/issues/41012)). * Fixed X (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). ## Security Fixes diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/BeamRowWrapper.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/BeamRowWrapper.java index 4a918d297282..0fab06d3bf54 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/BeamRowWrapper.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/BeamRowWrapper.java @@ -35,6 +35,7 @@ import org.apache.beam.sdk.schemas.logicaltypes.FixedPrecisionNumeric; import org.apache.beam.sdk.schemas.logicaltypes.MicrosInstant; import org.apache.beam.sdk.schemas.logicaltypes.Time; +import org.apache.beam.sdk.schemas.logicaltypes.Timestamp; import org.apache.beam.sdk.values.Row; import org.apache.iceberg.StructLike; import org.apache.iceberg.types.Type; @@ -158,7 +159,8 @@ private interface PositionalGetter { BeamRowWrapper nestedWrapper = new BeamRowWrapper(beamSchema, structType); return (row, pos) -> nestedWrapper.wrap(row.getRow(pos)); case LOGICAL_TYPE: - if (beamType.isLogicalType(MicrosInstant.IDENTIFIER)) { + if (beamType.isLogicalType(MicrosInstant.IDENTIFIER) + || beamType.isLogicalType(Timestamp.IDENTIFIER)) { return (row, pos) -> { Instant instant = checkStateNotNull(row.getLogicalTypeValue(pos, Instant.class)); return TimeUnit.SECONDS.toMicros(instant.getEpochSecond()) + instant.getNano() / 1000; 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 63d3c80173ca..44524f8b3d52 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 @@ -21,11 +21,6 @@ import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState; import java.io.IOException; -import java.time.LocalDateTime; -import java.time.YearMonth; -import java.time.ZoneOffset; -import java.time.format.DateTimeFormatter; -import java.time.temporal.ChronoUnit; import java.util.ArrayList; import java.util.List; import java.util.Map; @@ -36,7 +31,6 @@ import org.apache.beam.sdk.values.Row; import org.apache.beam.sdk.values.WindowedValue; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Splitter; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.Cache; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.CacheBuilder; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.RemovalNotification; @@ -44,7 +38,6 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps; import org.apache.iceberg.DataFile; import org.apache.iceberg.ManifestFile; -import org.apache.iceberg.PartitionField; import org.apache.iceberg.PartitionKey; import org.apache.iceberg.PartitionSpec; import org.apache.iceberg.SortOrder; @@ -57,7 +50,6 @@ import org.apache.iceberg.data.Record; import org.apache.iceberg.exceptions.AlreadyExistsException; import org.apache.iceberg.exceptions.NoSuchTableException; -import org.apache.iceberg.transforms.Transforms; import org.checkerframework.checker.nullness.qual.Nullable; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -104,7 +96,7 @@ final class DestinationState { private final String stateToken = UUID.randomUUID().toString(); final Cache writers; private final List dataFiles = Lists.newArrayList(); - @VisibleForTesting final Map writerCounts = Maps.newHashMap(); + @VisibleForTesting final Map writerCounts = Maps.newHashMap(); private final List exceptions = Lists.newArrayList(); private final InternalRecordWrapper wrapper; // wrapper that facilitates partitioning @@ -191,13 +183,13 @@ private RecordWriter createWriter(PartitionKey partitionKey) { // keep track of how many writers we opened for each destination-partition path // use this as a prefix to differentiate the new path. // this avoids overwriting a data file written by a previous writer in this destination state. - int recordIndex = writerCounts.merge(partitionKey, 1, Integer::sum); + int writerIndex = writerCounts.merge(partitionKey.toPath(), 1, Integer::sum); try { RecordWriter writer = new RecordWriter( table, icebergDestination.getFileFormat(), - filePrefix + "_" + stateToken + "_" + recordIndex, + filePrefix + "_" + stateToken + "_" + writerIndex, partitionKey, writeProperties); openWriters++; @@ -212,40 +204,6 @@ private RecordWriter createWriter(PartitionKey partitionKey) { } } - /** - * Returns an equivalent partition path that is made up of partition data. Needed to reconstruct a - * {@link DataFile}. - */ - @VisibleForTesting - static String getPartitionDataPath( - String partitionPath, Map partitionFieldMap) { - if (partitionPath.isEmpty() || partitionFieldMap.isEmpty()) { - return partitionPath; - } - List resolved = new ArrayList<>(); - for (String partition : Splitter.on('/').splitToList(partitionPath)) { - List nameAndValue = Splitter.on('=').splitToList(partition); - String name = nameAndValue.get(0); - String value = nameAndValue.get(1); - String transformName = - Preconditions.checkArgumentNotNull(partitionFieldMap.get(name)).transform().toString(); - if (Transforms.month().toString().equals(transformName)) { - long months = - ChronoUnit.MONTHS.between(EPOCH, YearMonth.parse(value).atDay(1).atStartOfDay()); - value = String.valueOf(months); - } else if (Transforms.hour().toString().equals(transformName)) { - long hour = ChronoUnit.HOURS.between(EPOCH, LocalDateTime.parse(value, HOUR_FORMATTER)); - value = String.valueOf(hour); - } - resolved.add(name + "=" + value); - } - return String.join("/", resolved); - } - - private static final DateTimeFormatter HOUR_FORMATTER = - DateTimeFormatter.ofPattern("yyyy-MM-dd-HH"); - private static final LocalDateTime EPOCH = LocalDateTime.ofEpochSecond(0, 0, ZoneOffset.UTC); - private final IcebergCatalogConfig catalogConfig; private final String filePrefix; private final long maxFileSize; 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 fbb6cb08aca4..25a70ec24615 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 @@ -18,46 +18,25 @@ package org.apache.beam.sdk.io.iceberg; import static org.apache.beam.sdk.io.iceberg.AssignDestinationsAndPartitions.DESTINATION; -import static org.apache.beam.sdk.io.iceberg.AssignDestinationsAndPartitions.PARTITION; -import static org.apache.beam.sdk.io.iceberg.RecordWriterManager.getPartitionDataPath; +import static org.apache.beam.sdk.io.iceberg.WriteToPartitions.DEFAULT_BYTES_PER_FILE; import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; import java.util.Map; -import java.util.UUID; -import org.apache.beam.sdk.coders.IterableCoder; -import org.apache.beam.sdk.coders.KvCoder; -import org.apache.beam.sdk.coders.RowCoder; -import org.apache.beam.sdk.schemas.Schema; import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.PTransform; import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.transforms.windowing.PaneInfo; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.PCollection; import org.apache.beam.sdk.values.PCollectionView; import org.apache.beam.sdk.values.Row; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps; -import org.apache.iceberg.DataFiles; -import org.apache.iceberg.PartitionField; -import org.apache.iceberg.PartitionKey; -import org.apache.iceberg.PartitionSpec; -import org.apache.iceberg.SortOrder; -import org.apache.iceberg.StructLike; -import org.apache.iceberg.Table; -import org.apache.iceberg.catalog.Catalog; -import org.apache.iceberg.catalog.Namespace; -import org.apache.iceberg.catalog.SupportsNamespaces; -import org.apache.iceberg.catalog.TableIdentifier; -import org.apache.iceberg.data.Record; -import org.apache.iceberg.exceptions.AlreadyExistsException; -import org.apache.iceberg.exceptions.NoSuchTableException; -import org.checkerframework.checker.nullness.qual.MonotonicNonNull; +import org.apache.beam.sdk.values.WindowedValue; +import org.apache.beam.sdk.values.WindowedValues; import org.checkerframework.checker.nullness.qual.Nullable; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; class WritePartitionedRowsToFiles extends PTransform>>, PCollection> { - private static final Logger LOG = LoggerFactory.getLogger(WritePartitionedRowsToFiles.class); private final DynamicDestinations dynamicDestinations; private final IcebergCatalogConfig catalogConfig; private final String filePrefix; @@ -87,21 +66,10 @@ class WritePartitionedRowsToFiles @Override public PCollection expand(PCollection>> input) { - Schema dataSchema = - ((RowCoder) - ((IterableCoder) - ((KvCoder>) input.getCoder()).getValueCoder()) - .getElemCoder()) - .getSchema(); ParDo.SingleOutput>, FileWriteResult> parDo = ParDo.of( new WriteDoFn( - catalogConfig, - dynamicDestinations, - filePrefix, - dataSchema, - writeProperties, - metadataView)); + catalogConfig, dynamicDestinations, filePrefix, writeProperties, metadataView)); if (metadataView != null) { parDo = parDo.withSideInputs(metadataView); } @@ -113,184 +81,60 @@ private static class WriteDoFn extends DoFn>, FileWriteRes private final DynamicDestinations dynamicDestinations; private final IcebergCatalogConfig catalogConfig; private final String filePrefix; - private final Schema dataSchema; private final @Nullable Map writeProperties; private final @Nullable PCollectionView> metadataView; - private transient @MonotonicNonNull Map specIds; - private transient @MonotonicNonNull Map> - partitionFieldMaps; WriteDoFn( IcebergCatalogConfig catalogConfig, DynamicDestinations dynamicDestinations, String filePrefix, - Schema dataSchema, - @Nullable Map writeProperties) { - this(catalogConfig, dynamicDestinations, filePrefix, dataSchema, writeProperties, null); - } - - WriteDoFn( - IcebergCatalogConfig catalogConfig, - DynamicDestinations dynamicDestinations, - String filePrefix, - Schema dataSchema, @Nullable Map writeProperties, @Nullable PCollectionView> metadataView) { this.catalogConfig = catalogConfig; this.dynamicDestinations = dynamicDestinations; this.filePrefix = filePrefix; - this.dataSchema = dataSchema; this.writeProperties = writeProperties; this.metadataView = metadataView; } - @Setup - public void setup() { - partitionFieldMaps = Maps.newHashMap(); - specIds = Maps.newHashMap(); - } - @ProcessElement public void processElement( ProcessContext c, @Element KV> element, + BoundedWindow window, + PaneInfo paneInfo, OutputReceiver out) throws Exception { String tableIdentifier = checkStateNotNull(element.getKey().getString(DESTINATION)); - String partitionPath = checkStateNotNull(element.getKey().getString(PARTITION)); IcebergDestination destination = dynamicDestinations.instantiateDestination(tableIdentifier); + WindowedValue windowedDestination = + WindowedValues.of(destination, window.maxTimestamp(), window, paneInfo); Map sideInputs = metadataView != null ? c.sideInput(metadataView) : null; - Table table = getOrCreateTable(destination, dataSchema, sideInputs); - partitionPath = - getPartitionDataPath( - partitionPath, getPartitionFieldMap(destination.getTableIdentifier(), table)); - - StructLike partitionData = - table.spec().isPartitioned() - ? DataFiles.data(table.spec(), partitionPath) - : new PartitionKey(table.spec(), table.schema()); - String fileName = - destination - .getFileFormat() - .addExtension(String.format("%s-%s", filePrefix, UUID.randomUUID())); - - RecordWriter writer = - new RecordWriter( - table, destination.getFileFormat(), fileName, partitionData, writeProperties); - try { + RecordWriterManager writer; + try (RecordWriterManager openWriter = + new RecordWriterManager( + catalogConfig, + filePrefix, + DEFAULT_BYTES_PER_FILE, + Integer.MAX_VALUE, + writeProperties, + sideInputs)) { + writer = openWriter; for (Row row : element.getValue()) { - Record record = IcebergUtils.beamRowToIcebergRecord(table.schema(), row); - writer.write(record); + writer.write(windowedDestination, row); } - } finally { - writer.close(); } - // Serialize against the file's own spec - SerializableDataFile sdf = SerializableDataFile.from(writer.getDataFile(), table.specs()); - out.output( - FileWriteResult.builder() - .setTableIdentifier(destination.getTableIdentifier()) - .setSerializableDataFile(sdf) - .build()); - } - - private Map getPartitionFieldMap( - TableIdentifier identifier, Table table) { - @Nullable Integer specId = checkStateNotNull(specIds).get(identifier); - if (specId != null && specId == table.spec().specId()) { - return checkStateNotNull(checkStateNotNull(partitionFieldMaps).get(identifier)); - } - Map partitionFieldMap = Maps.newHashMap(); - for (PartitionField partitionField : table.spec().fields()) { - partitionFieldMap.put(partitionField.name(), partitionField); - } - checkStateNotNull(specIds).put(identifier, table.spec().specId()); - checkStateNotNull(partitionFieldMaps).put(identifier, partitionFieldMap); - return partitionFieldMap; - } - - Table getOrCreateTable( - IcebergDestination destination, - Schema dataSchema, - @Nullable Map sideInputTableSpecs) { - TableIdentifier identifier = destination.getTableIdentifier(); - String tableIdString = IcebergUtils.tableIdentifierToString(identifier); - if (sideInputTableSpecs != null && sideInputTableSpecs.containsKey(tableIdString)) { - SerializableTableSpec spec = sideInputTableSpecs.get(tableIdString); - if (spec != null) { - return new SideInputTable(spec, catalogConfig); - } - } - return TableCache.getAndRefreshIfStale( - catalogConfig, - identifier, - () -> loadOrCreateTable(catalogConfig.catalog(), destination, dataSchema)); - } - - Table getOrCreateTable(IcebergDestination destination, Schema dataSchema) { - return getOrCreateTable(destination, dataSchema, null); - } - - private Table loadOrCreateTable( - Catalog catalog, IcebergDestination destination, Schema dataSchema) { - TableIdentifier identifier = destination.getTableIdentifier(); - Namespace namespace = identifier.namespace(); - @Nullable IcebergTableCreateConfig createConfig = destination.getTableCreateConfig(); - PartitionSpec partitionSpec = - createConfig != null ? createConfig.getPartitionSpec() : PartitionSpec.unpartitioned(); - SortOrder sortOrder = - createConfig != null ? createConfig.getSortOrder() : SortOrder.unsorted(); - Map tableProperties = - createConfig != null && createConfig.getTableProperties() != null - ? createConfig.getTableProperties() - : Maps.newHashMap(); - - // Create namespace if it does not exist yet - if (!namespace.isEmpty() && catalog instanceof SupportsNamespaces) { - SupportsNamespaces supportsNamespaces = (SupportsNamespaces) catalog; - if (!supportsNamespaces.namespaceExists(namespace)) { - try { - supportsNamespaces.createNamespace(namespace); - LOG.info("Created new namespace '{}'.", namespace); - } catch (AlreadyExistsException ignored) { - // race condition: another worker already created this namespace - LOG.info("Namespace '{}' already exists.", namespace); - } - } - } - - // If table exists, just load it - // Note: the implementation of catalog.tableExists() will load the table to check its - // existence. We don't use it here to avoid double loadTable() calls. - try { - return catalog.loadTable(identifier); - } catch (NoSuchTableException e) { // Otherwise, create the table - org.apache.iceberg.Schema tableSchema = IcebergUtils.beamSchemaToIcebergSchema(dataSchema); - try { - Table table = - catalog - .buildTable(identifier, tableSchema) - .withPartitionSpec(partitionSpec) - .withSortOrder(sortOrder) - .withProperties(tableProperties) - .create(); - LOG.info( - "Created Iceberg table '{}' with schema: {}\n" - + ", partition spec: {}, sort order: {}, table properties: {}", - identifier, - tableSchema, - partitionSpec, - sortOrder, - tableProperties); - return table; - } catch (AlreadyExistsException ignored) { - // race condition: another worker already created this table - return catalog.loadTable(identifier); - } + for (SerializableDataFile dataFile : + checkStateNotNull(writer.getSerializableDataFiles().get(windowedDestination))) { + out.output( + FileWriteResult.builder() + .setTableIdentifier(destination.getTableIdentifier()) + .setSerializableDataFile(dataFile) + .build()); } } } 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 7e9d56df98a9..8edb5682e4cc 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 @@ -42,7 +42,7 @@ import org.joda.time.Duration; class WriteToPartitions extends PTransform>, IcebergWriteResult> { - private static final long DEFAULT_BYTES_PER_FILE = (1L << 29); // 512mb + static final long DEFAULT_BYTES_PER_FILE = (1L << 29); // 512mb private final IcebergCatalogConfig catalogConfig; private final DynamicDestinations dynamicDestinations; private final @Nullable Duration triggeringFrequency; diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/BeamRowWrapperTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/BeamRowWrapperTest.java index f361bc94c15b..e8885e198b48 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/BeamRowWrapperTest.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/BeamRowWrapperTest.java @@ -37,6 +37,7 @@ import org.apache.beam.sdk.schemas.logicaltypes.FixedPrecisionNumeric; import org.apache.beam.sdk.schemas.logicaltypes.MicrosInstant; import org.apache.beam.sdk.schemas.logicaltypes.Time; +import org.apache.beam.sdk.schemas.logicaltypes.Timestamp; import org.apache.beam.sdk.values.Row; import org.apache.iceberg.StructLike; import org.apache.iceberg.types.Types; @@ -63,6 +64,7 @@ public class BeamRowWrapperTest { .addLogicalTypeField("fixed_numeric_field", FixedPrecisionNumeric.of(10, 2)) .addRowField("row_field", NESTED_BEAM_SCHEMA) .addInt32Field("pass_through_field") + .addLogicalTypeField("timestamp_field", Timestamp.MICROS) .build(); private static final Types.StructType ICEBERG_STRUCT = Types.StructType.of( @@ -83,7 +85,8 @@ public class BeamRowWrapperTest { "row_field", Types.StructType.of( Types.NestedField.required(1, "nested_int", Types.IntegerType.get()))), - Types.NestedField.required(14, "pass_through_field", Types.IntegerType.get())); + Types.NestedField.required(14, "pass_through_field", Types.IntegerType.get()), + Types.NestedField.required(15, "timestamp_field", Types.TimestampType.withZone())); private static final UUID TEST_UUID = UUID.randomUUID(); private static final Instant TEST_MICROS_INSTANT = Instant.ofEpochSecond(1_609_459_200L, 123_456_000); @@ -107,14 +110,15 @@ public class BeamRowWrapperTest { LocalTime.now(ZoneId.systemDefault()), new BigDecimal("567.89"), NESTED_ROW, - 888) + 888, + TEST_MICROS_INSTANT) .build(); private static final BeamRowWrapper WRAPPER = new BeamRowWrapper(BEAM_SCHEMA, ICEBERG_STRUCT).wrap(ROW); @Test public void testSize() { - assertEquals("Size should match the schema field count", 14, WRAPPER.size()); + assertEquals("Size should match the schema field count", 15, WRAPPER.size()); } @Test @@ -185,6 +189,11 @@ public void testMicrosInstantLogicalTypeConversion() { assertEquals(1_609_459_200_123_456L, (long) WRAPPER.get(7, Long.class)); } + @Test + public void testTimestampLogicalTypeConversion() { + assertEquals(1_609_459_200_123_456L, (long) WRAPPER.get(14, Long.class)); + } + @Test public void testDateTimeLogicalTypeConversion() { long expectedDateTime = 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..6a736930a695 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 @@ -30,6 +30,8 @@ import static org.junit.Assume.assumeTrue; import java.io.Serializable; +import java.time.LocalDate; +import java.time.LocalDateTime; import java.util.Arrays; import java.util.List; import java.util.Map; @@ -47,6 +49,7 @@ import org.apache.beam.sdk.metrics.MetricsFilter; import org.apache.beam.sdk.schemas.Schema; import org.apache.beam.sdk.schemas.SchemaCoder; +import org.apache.beam.sdk.schemas.logicaltypes.SqlTypes; import org.apache.beam.sdk.testing.PAssert; import org.apache.beam.sdk.testing.TestPipeline; import org.apache.beam.sdk.testing.TestStream; @@ -76,7 +79,9 @@ import org.apache.iceberg.DataFile; import org.apache.iceberg.DistributionMode; import org.apache.iceberg.FileFormat; +import org.apache.iceberg.FileScanTask; import org.apache.iceberg.NullOrder; +import org.apache.iceberg.PartitionKey; import org.apache.iceberg.PartitionSpec; import org.apache.iceberg.SortDirection; import org.apache.iceberg.SortOrder; @@ -86,11 +91,14 @@ import org.apache.iceberg.catalog.TableIdentifier; import org.apache.iceberg.data.GenericRecord; import org.apache.iceberg.data.IcebergGenerics; +import org.apache.iceberg.data.InternalRecordWrapper; 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.iceberg.util.StructLikeSet; import org.hamcrest.Matchers; import org.joda.time.Duration; import org.joda.time.Instant; @@ -470,6 +478,74 @@ public void testStreamingWrite() { assertThat(writtenRecords, Matchers.containsInAnyOrder(TestFixtures.FILE1SNAPSHOT1.toArray())); } + /** A file's partition must come from its rows, because a partition path can't be read back. */ + @Test + public void testPartitionValues() throws Exception { + TableIdentifier tableId = + TableIdentifier.of("default", "table" + Long.toString(UUID.randomUUID().hashCode(), 16)); + Schema schema = + Schema.builder() + .addLogicalTypeField("date", SqlTypes.DATE) + .addNullableStringField("str") + .addNullableInt32Field("num") + .addLogicalTypeField("ts", SqlTypes.DATETIME) + .addByteArrayField("bin") + .build(); + org.apache.iceberg.Schema icebergSchema = IcebergUtils.beamSchemaToIcebergSchema(schema); + Table table = + warehouse.createTable( + tableId, + icebergSchema, + PartitionSpec.builderFor(icebergSchema) + .year("date") + .identity("str") + .identity("num") + .identity("ts") + .identity("bin") + .build()); + + LocalDate date = LocalDate.of(2024, 3, 1); + LocalDateTime ts = LocalDateTime.of(2024, 3, 1, 5, 0); + byte[] bin = {1, 2, 3}; + List rows = + ImmutableList.of( + Row.withSchema(schema).addValues(date, "a b/c", 5, ts, bin).build(), + Row.withSchema(schema).addValues(date, null, null, ts, bin).build(), + // has the same partition path as the null string above + Row.withSchema(schema).addValues(date, "null", null, ts, bin).build()); + + 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(Create.of(rows).withRowSchema(schema)) + .apply(writeTransform(catalog, tableId)); + testPipeline.run().waitUntilFinish(); + + StructLikeSet expected = StructLikeSet.create(table.spec().partitionType()); + PartitionKey partitionKey = new PartitionKey(table.spec(), table.schema()); + InternalRecordWrapper wrapper = new InternalRecordWrapper(table.schema().asStruct()); + for (Row row : rows) { + partitionKey.partition(wrapper.wrap(beamRowToIcebergRecord(table.schema(), row))); + expected.add(partitionKey.copy()); + } + + StructLikeSet written = StructLikeSet.create(table.spec().partitionType()); + table.refresh(); + try (CloseableIterable tasks = table.newScan().planFiles()) { + tasks.forEach(task -> written.add(task.file().partition())); + } + assertEquals(expected, written); + } + @Test public void testHashDistribution() { assumeTrue(distributionMode.equals(HASH_WITH_AUTOSHARDING)); 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 98597ba4c4f3..bf8e9455c32f 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 @@ -349,10 +349,10 @@ public void testRespectMaxFileSize() throws IOException { assertEquals(1, writerManager.openWriters); partitionKey.partition(IcebergUtils.beamRowToIcebergRecord(ICEBERG_SCHEMA, row)); - Map writerCounts = + Map writerCounts = writerManager.destinations.get(windowedDestination).writerCounts; // this is our first writer - assertEquals(1, writerCounts.get(partitionKey).intValue()); + assertEquals(1, writerCounts.get(partitionKey.toPath()).intValue()); // row partition:: [aaa, true]. // existing partition. use existing writer @@ -364,7 +364,7 @@ public void testRespectMaxFileSize() throws IOException { assertTrue(writeSuccess); assertEquals(1, writerManager.openWriters); // check that we still use our first writer - assertEquals(1, writerCounts.get(partitionKey).intValue()); + assertEquals(1, writerCounts.get(partitionKey.toPath()).intValue()); // row partition:: [aaa, true]. // writer has reached max file size. create a new writer @@ -372,7 +372,7 @@ public void testRespectMaxFileSize() throws IOException { writeSuccess = writerManager.write(windowedDestination, row); assertTrue(writeSuccess); // check that we have opened and are using a second writer - assertEquals(2, writerCounts.get(partitionKey).intValue()); + assertEquals(2, writerCounts.get(partitionKey.toPath()).intValue()); // check that only one writer is open (we have closed the first writer) assertEquals(1, writerManager.openWriters);