From 87665ecfaf395d934c8d09e51163cb156de7d75f Mon Sep 17 00:00:00 2001 From: atognolag Date: Fri, 18 Sep 2026 11:42:42 +0000 Subject: [PATCH 1/3] Implement high-performance sorted writing support in IcebergIO --- CHANGES.md | 2 +- .../AssignDestinationsAndPartitions.java | 107 +++- .../apache/beam/sdk/io/iceberg/IcebergIO.java | 259 ++++++++- .../beam/sdk/io/iceberg/IcebergRowSorter.java | 492 ++++++++++++++++++ .../beam/sdk/io/iceberg/IcebergUtils.java | 39 ++ .../IcebergWriteSchemaTransformProvider.java | 14 +- .../io/iceberg/WriteGroupedRowsToFiles.java | 84 ++- .../iceberg/WritePartitionedRowsToFiles.java | 155 ++++-- .../sdk/io/iceberg/WriteToDestinations.java | 34 +- .../sdk/io/iceberg/WriteToPartitions.java | 59 ++- .../io/iceberg/WriteUngroupedRowsToFiles.java | 84 ++- .../IcebergIOSideInputTableCacheTest.java | 80 ++- .../sdk/io/iceberg/IcebergIOWriteTest.java | 294 +++++++++++ .../sdk/io/iceberg/IcebergRowSorterTest.java | 432 +++++++++++++++ .../beam/sdk/io/iceberg/IcebergUtilsTest.java | 13 + ...ebergWriteSchemaTransformProviderTest.java | 54 ++ 16 files changed, 2130 insertions(+), 72 deletions(-) create mode 100644 sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergRowSorter.java create mode 100644 sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergRowSorterTest.java diff --git a/CHANGES.md b/CHANGES.md index c10b7262209c..601ae5f41047 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -67,6 +67,7 @@ * Support for X source added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). * (Java) FirestoreIO: added `RpcQosOptions.Builder.withRampupThrottlingDisabled()`, mirroring the existing DatastoreIO option, to opt out of the 500/50/5 write ramp-up ([#27170](https://github.com/apache/beam/issues/27170)). * Added support for attributing BigQuery API quota and billing to a specific GCP project (quota project): `quota_project_id` parameter in `ReadFromBigQuery` or `--quota_project_id` pipeline option, honored by the Python BigQuery IO transforms (`WriteToBigQuery` with `FILE_LOADS` or `STREAMING_INSERTS`; not yet `STORAGE_WRITE_API`, which is Java-backed), `--bigQueryQuotaProjectId` pipeline option (Java), and `bigqueryio.WithQuotaProject` read/query option (Go) ([#37431](https://github.com/apache/beam/issues/37431)). +* [IcebergIO] Added opt-in sorted write support (`withSortedWrites()` / `sort_on_write`) to write data files in the table's sort order, spilling to local disk via Beam's external sorter for groups that exceed memory. Also added `DistributionMode.RANGE` with a user-supplied `withDistributionFunction(...)`, which shards a partition into non-overlapping key ranges so each shard becomes its own sorted file(s) (Java) ([#38406](https://github.com/apache/beam/issues/38406)). ## New Features / Improvements @@ -135,7 +136,6 @@ * BigQueryIO now supports reading BigQuery Lakehouse runtime catalog (BigLake metastore) Iceberg tables with the Storage Read API, using 4-part `project.catalog.namespace.table` identifiers (or a `TableReference` with a composite `catalog.namespace` dataset id). Previously such references were silently mis-parsed (Java) ([#39597](https://github.com/apache/beam/issues/39597)) . * SolaceIO now supports reading and writing binary and text content data payload (Java) ([#39875](https://github.com/apache/beam/issues/39875)). * ClickHouseIO: support writing `Decimal(P, S)` / `Decimal32/64/128/256` columns (Java) ([#39840](https://github.com/apache/beam/issues/39840)). -* SolaceIO now supports reading and writing user properties (message metadata) (Java) ([#40099](https://github.com/apache/beam/issues/40099)). * [IcebergIO] AddFiles (`IcebergAddFiles` in YAML) can evolve the table schema before registering files, with `schema_evolution_options`, `required_columns`, `incompatible_schema_handling` and `unverifiable_file_handling` (Java/YAML, batch only) ([#40144](https://github.com/apache/beam/issues/40144)). * [IcebergIO] Added batch and streaming CDC writes that applies INSERT/UPDATE_BEFORE/UPDATE_AFTER/DELETE changes to Iceberg V2+ tables by primary key. Invoke with `IcebergIO.writeCdcRows` (Java) or by setting `mode: merge-on-read` on the Managed `ICEBERG` write (Java, Python, YAML) ([#39979](https://github.com/apache/beam/issues/39979)). * [IcebergIO] Added an optional side-input table cache for writes to significantly reduce catalog and table requests for large pipelines. A single worker polls the table and broadcasts it to other workers in the pipeline. Enable with `IcebergIO.writeRows(...).withSideInputTableCache()` (Java) or by setting `use_side_input_table_cache: true` on the Managed `ICEBERG` write (Java, Python) ([#39723](https://github.com/apache/beam/issues/39723)). diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/AssignDestinationsAndPartitions.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/AssignDestinationsAndPartitions.java index 4167ad006244..fa664e7c4aa1 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/AssignDestinationsAndPartitions.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/AssignDestinationsAndPartitions.java @@ -26,6 +26,7 @@ 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.SerializableFunction; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; import org.apache.beam.sdk.transforms.windowing.PaneInfo; import org.apache.beam.sdk.values.KV; @@ -33,6 +34,7 @@ import org.apache.beam.sdk.values.PCollectionView; import org.apache.beam.sdk.values.Row; import org.apache.beam.sdk.values.ValueInSingleWindow; +import org.apache.iceberg.DistributionMode; import org.apache.iceberg.PartitionKey; import org.apache.iceberg.PartitionSpec; import org.apache.iceberg.Schema; @@ -54,34 +56,86 @@ class AssignDestinationsAndPartitions private final DynamicDestinations dynamicDestinations; private final IcebergCatalogConfig catalogConfig; private final @Nullable PCollectionView> metadataView; + private final DistributionMode distributionMode; + private final @Nullable SerializableFunction distributionFunction; static final String DESTINATION = "destination"; static final String PARTITION = "partition"; - + static final String SHARD = "shard"; + + /** + * Grouping key used for {@link DistributionMode#NONE} and {@link DistributionMode#HASH} writes. + * + *

This encoding is load-bearing for pipeline update compatibility: it is the key coder of the + * {@link org.apache.beam.sdk.transforms.GroupIntoBatches} in {@link WriteToPartitions}, so + * changing its shape invalidates the buffered state of in-flight streaming jobs. Do not add + * fields here. + */ static final org.apache.beam.sdk.schemas.Schema OUTPUT_SCHEMA = org.apache.beam.sdk.schemas.Schema.builder() .addStringField(DESTINATION) .addStringField(PARTITION) .build(); + /** + * Grouping key used for {@link DistributionMode#RANGE} writes, which additionally fan each + * partition out into non-overlapping shards so that every shard becomes its own sorted file. + */ + static final org.apache.beam.sdk.schemas.Schema OUTPUT_SCHEMA_WITH_SHARD = + org.apache.beam.sdk.schemas.Schema.builder() + .addStringField(DESTINATION) + .addStringField(PARTITION) + .addInt32Field(SHARD) + .build(); + + /** Returns the grouping key schema that the given distribution mode emits. */ + static org.apache.beam.sdk.schemas.Schema outputSchemaFor(DistributionMode distributionMode) { + return distributionMode == DistributionMode.RANGE ? OUTPUT_SCHEMA_WITH_SHARD : OUTPUT_SCHEMA; + } + public AssignDestinationsAndPartitions( DynamicDestinations dynamicDestinations, IcebergCatalogConfig catalogConfig) { - this(dynamicDestinations, catalogConfig, null); + this(dynamicDestinations, catalogConfig, null, DistributionMode.HASH, null); } public AssignDestinationsAndPartitions( DynamicDestinations dynamicDestinations, IcebergCatalogConfig catalogConfig, @Nullable PCollectionView> metadataView) { + this(dynamicDestinations, catalogConfig, metadataView, DistributionMode.HASH, null); + } + + public AssignDestinationsAndPartitions( + DynamicDestinations dynamicDestinations, + IcebergCatalogConfig catalogConfig, + DistributionMode distributionMode, + @Nullable SerializableFunction distributionFunction) { + this(dynamicDestinations, catalogConfig, null, distributionMode, distributionFunction); + } + + public AssignDestinationsAndPartitions( + DynamicDestinations dynamicDestinations, + IcebergCatalogConfig catalogConfig, + @Nullable PCollectionView> metadataView, + DistributionMode distributionMode, + @Nullable SerializableFunction distributionFunction) { this.dynamicDestinations = dynamicDestinations; this.catalogConfig = catalogConfig; this.metadataView = metadataView; + this.distributionMode = distributionMode; + this.distributionFunction = distributionFunction; } @Override public PCollection> expand(PCollection input) { ParDo.SingleOutput> parDo = - ParDo.of(new AssignDoFn(dynamicDestinations, catalogConfig, metadataView)); + ParDo.of( + new AssignDoFn( + dynamicDestinations, + catalogConfig, + metadataView, + distributionMode, + distributionFunction)); if (metadataView != null) { parDo = parDo.withSideInputs(metadataView); } @@ -89,7 +143,8 @@ public PCollection> expand(PCollection input) { .apply(parDo) .setCoder( KvCoder.of( - RowCoder.of(OUTPUT_SCHEMA), RowCoder.of(dynamicDestinations.getDataSchema()))); + RowCoder.of(outputSchemaFor(distributionMode)), + RowCoder.of(dynamicDestinations.getDataSchema()))); } static class AssignDoFn extends DoFn> { @@ -104,18 +159,39 @@ static class AssignDoFn extends DoFn> { private final DynamicDestinations dynamicDestinations; private final IcebergCatalogConfig catalogConfig; private final @Nullable PCollectionView> metadataView; + private final DistributionMode distributionMode; + private final @Nullable SerializableFunction distributionFunction; AssignDoFn(DynamicDestinations dynamicDestinations, IcebergCatalogConfig catalogConfig) { - this(dynamicDestinations, catalogConfig, null); + this(dynamicDestinations, catalogConfig, null, DistributionMode.HASH, null); } AssignDoFn( DynamicDestinations dynamicDestinations, IcebergCatalogConfig catalogConfig, @Nullable PCollectionView> metadataView) { + this(dynamicDestinations, catalogConfig, metadataView, DistributionMode.HASH, null); + } + + AssignDoFn( + DynamicDestinations dynamicDestinations, + IcebergCatalogConfig catalogConfig, + DistributionMode distributionMode, + @Nullable SerializableFunction distributionFunction) { + this(dynamicDestinations, catalogConfig, null, distributionMode, distributionFunction); + } + + AssignDoFn( + DynamicDestinations dynamicDestinations, + IcebergCatalogConfig catalogConfig, + @Nullable PCollectionView> metadataView, + DistributionMode distributionMode, + @Nullable SerializableFunction distributionFunction) { this.dynamicDestinations = dynamicDestinations; this.catalogConfig = catalogConfig; this.metadataView = metadataView; + this.distributionMode = distributionMode; + this.distributionFunction = distributionFunction; } @Setup @@ -214,10 +290,25 @@ public void processElement( String partitionPath = partitionKey.toPath(); - Row destAndPartition = - Row.withSchema(OUTPUT_SCHEMA).addValues(tableIdentifier, partitionPath).build(); + Row.Builder keyBuilder = + Row.withSchema(outputSchemaFor(distributionMode)) + .addValue(tableIdentifier) + .addValue(partitionPath); + + if (distributionMode == DistributionMode.RANGE) { + SerializableFunction shardFn = + checkStateNotNull( + distributionFunction, + "A distribution function is required when using RANGE distribution mode."); + keyBuilder = + keyBuilder.addValue( + checkStateNotNull( + shardFn.apply(data), + "The RANGE distribution function returned a null shard id. It must return a" + + " non-null Integer for every row.")); + } - out.output(KV.of(destAndPartition, data)); + out.output(KV.of(keyBuilder.build(), data)); } } } 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 d98d9944bf3c..8bf26962e231 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 @@ -30,6 +30,7 @@ import org.apache.beam.sdk.options.StreamingOptions; import org.apache.beam.sdk.schemas.Schema; import org.apache.beam.sdk.transforms.PTransform; +import org.apache.beam.sdk.transforms.SerializableFunction; import org.apache.beam.sdk.transforms.display.DisplayData; import org.apache.beam.sdk.values.PBegin; import org.apache.beam.sdk.values.PCollection; @@ -394,6 +395,7 @@ public static WriteRows writeRows(IcebergCatalogConfig catalog) { .setDistributionMode(DistributionMode.NONE) .setAutoSharding(false) .setUseSideInputTableCache(false) + .setSortedWrites(false) .build(); } @@ -429,8 +431,12 @@ public abstract static class WriteRows extends PTransform, Iceb abstract DistributionMode getDistributionMode(); + abstract @Nullable SerializableFunction getDistributionFunction(); + abstract boolean getAutoSharding(); + abstract boolean getSortedWrites(); + abstract @Nullable Map getWriteProperties(); abstract @Nullable List getPartitionFields(); @@ -461,8 +467,12 @@ abstract static class Builder { abstract Builder setDistributionMode(DistributionMode mode); + abstract Builder setDistributionFunction(SerializableFunction shardFn); + abstract Builder setAutoSharding(boolean autoSharding); + abstract Builder setSortedWrites(boolean sortedWrites); + abstract Builder setWriteProperties(Map writeProperties); abstract Builder setPartitionFields(List partitionFields); @@ -510,23 +520,220 @@ public WriteRows withDirectWriteByteLimit(Integer directWriteByteLimit) { } /** - * Defines distribution of write data. Supported distributions: + * The default distribution mode is {@link DistributionMode#NONE}, which writes without a + * shuffle unless {@link #withSortedWrites()} is enabled on a table with a sort order. + * + *

Warning on HASH mode: Utilizing {@code HASH} distribution mode (with or without + * auto-sharding) can suffer from large unpartitioned or skewed writes if key spaces are not + * uniformly distributed. This can bottleneck workers and produce fragmented layout files. + * + *

Note on RANGE mode: When utilizing {@code RANGE} distribution mode, it is + * recommended that the custom distribution function is designed to produce adequately sized and + * strictly non-overlapping ranges of the sorting column to optimize downstream read + * performance. {@code RANGE} mode automatically enables sorted writes when the target table has + * a {@link org.apache.iceberg.SortOrder}. + * + *

Comparison of Distribution Modes:

* - *
    - *
  1. {@link DistributionMode#NONE}: don't shuffle rows (default) - *
  2. {@link DistributionMode#HASH}: shuffle rows by partition key before writing data - *
+ * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + *
Comparison of Distribution Modes
ModeDescriptionProsCons
{@link DistributionMode#NONE}No partition shuffle is performed by default. When {@link #withSortedWrites()} is enabled on a sorted table, records are grouped by destination shard and sorted prior to writing.Highly lightweight with zero shuffle overhead for unsorted writes.Writers on different workers can write to overlapping min/max key ranges across multiple files. Relies heavily on post-fact compaction or query time merges.
{@link DistributionMode#HASH}Data is shuffled and consolidated by partition key. All records for a partition are routed to a single worker.Consolidates partition files, eliminating cross-worker file overlapping for partition keys. Excellent worker stability.Can suffer from severe data skew if a single partition contains significantly more data than others (hot partitions).
{@link DistributionMode#RANGE}Data is shuffled based on a user-provided shard/bucket function (e.g., hashing/binning continuous keys) and sorted within each shard.Distributes writes for hot partitions across multiple workers. Eliminates skew while keeping file min/max key ranges tight and non-overlapping.Requires providing a custom {@link SerializableFunction} mapping rows to integer shard/bucket IDs.
* - * {@link DistributionMode#RANGE} is not supported yet + *

Recommendation Matrix (Sorting & Partitioning vs. Scale):

+ * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + *
Recommendation Matrix
PartitioningSortingScale / VolumeLatency PriorityRecommended ModeOperational Impact
PartitionedSortedSmallAny{@link DistributionMode#HASH} + {@link #withSortedWrites()}Consolidates partition files and sorts them locally. Avoids file overlaps for small volumes.
PartitionedSortedMedium / LargeLow Write Latency{@link DistributionMode#NONE} + {@link #withSortedWrites()}Shards by destination without partition shuffle. Results in overlapping key ranges across files, which requires downstream compaction.
PartitionedSortedMedium / LargeLow Read Latency{@link DistributionMode#HASH} with auto-sharding + {@link #withSortedWrites()} OR {@link DistributionMode#RANGE}HASH with auto-sharding scales writes for hot partitions but can result in overlapping file ranges requiring query-time sort merges. RANGE sharding distributes hot partitions into sequential, non-overlapping files to optimize reads.
PartitionedUnsortedSmallAny{@link DistributionMode#HASH}Consolidates data files into single partition directories to prevent file fragmentation.
PartitionedUnsortedMedium / LargeAny{@link DistributionMode#HASH} with auto-shardingConsolidates partition files while dynamically balancing hot partition writes across parallel workers.
UnpartitionedSortedSmallAny{@link DistributionMode#NONE} + {@link #withSortedWrites()}Groups by destination shard for fast, low-volume local sorting.
UnpartitionedSortedMedium / LargeLow Write Latency{@link DistributionMode#NONE} + {@link #withSortedWrites()}Parallel worker writes across destination shards. Requires downstream compaction to resolve overlapping file ranges.
UnpartitionedSortedMedium / LargeLow Read Latency{@link DistributionMode#RANGE} (with custom sharding function)Shards continuous keys into non-overlapping worker ranges. Eliminates single-worker bottlenecks and guarantees zero file overlap for fast queries.
UnpartitionedUnsortedAnyAny{@link DistributionMode#NONE}Direct, parallel worker writes with maximum throughput and zero network shuffle overhead.
+ * + *

Code Samples:

+ * + * 
{@code
+     * // 1. Using the default NONE distribution mode
+     * pipeline
+     *     .apply(Create.of(BEAM_ROWS))
+     *     .apply(IcebergIO.writeRows(catalogConfig)
+     *         .to(tableId));
+     *
+     * // 2. Using HASH distribution mode with sorted writes
+     * pipeline
+     *     .apply(Create.of(BEAM_ROWS))
+     *     .apply(IcebergIO.writeRows(catalogConfig)
+     *         .to(tableId)
+     *         .withDistributionMode(DistributionMode.HASH)
+     *         .withSortedWrites());
+     *
+     * // 3. Using RANGE distribution mode with a custom shard/bucket function to avoid data skew
+     * pipeline
+     *     .apply(Create.of(BEAM_ROWS))
+     *     .apply(IcebergIO.writeRows(catalogConfig)
+     *         .to(tableId)
+     *         .withDistributionMode(DistributionMode.RANGE)
+     *         .withDistributionFunction(row -> {
+     *             // Group continuous IDs into parallel, non-overlapping shards of 10,000 IDs each
+     *             long id = row.getInt64("id");
+     *             return (int) (id / 10000);
+     *         }));
+     * }
*/ public WriteRows withDistributionMode(DistributionMode mode) { return toBuilder().setDistributionMode(mode).build(); } + /** + * Sets the custom range-distribution function. + * + *

Only applicable when the distribution mode is set to {@link DistributionMode#RANGE}. The + * function maps a Beam {@link Row} to an Integer representing a shard/bucket ID. + */ + public WriteRows withDistributionFunction(SerializableFunction shardFn) { + return toBuilder().setDistributionFunction(shardFn).build(); + } + + /** + * Enables Beam's dynamic auto-sharding when using {@link DistributionMode#HASH}. + * + *

When enabled, the pipeline uses {@link + * org.apache.beam.sdk.transforms.GroupIntoBatches#withShardedKey()} under the hood. The runner + * (such as Dataflow) dynamically monitors throughput per partition key. If a partition is + * extremely hot, the runner automatically splits it into parallel sub-shards distributed across + * multiple workers to prevent single-worker bottlenecks and out-of-memory (OOM) errors, while + * keeping the number of data files for cold partitions minimal. + * + *

Note that because auto-sharding distributes hot-partition data randomly across worker + * shards, the written data files cannot guarantee non-overlapping key ranges. Downstream + * queries may require read-time sort merges for overlapping file segments until an Iceberg + * compaction job (e.g., `rewriteDataFiles`) is executed. + * + *

Only applicable when using {@link DistributionMode#HASH}. + */ public WriteRows withAutosharding() { return toBuilder().setAutoSharding(true).build(); } + /** + * Enables physically sorting records within each written data file according to the target + * table's {@link org.apache.iceberg.SortOrder}. + * + *

By default, {@link WriteRows} does not sort records unless {@link DistributionMode#RANGE} + * is configured, so existing pipelines writing to tables that have a metadata sort order incur + * no additional sorting or shuffle overhead. + */ + public WriteRows withSortedWrites() { + return withSortedWrites(true); + } + + /** + * Sets whether to physically sort records within each written data file according to the target + * table's {@link org.apache.iceberg.SortOrder}. + */ + public WriteRows withSortedWrites(boolean sortedWrites) { + return toBuilder().setSortedWrites(sortedWrites).build(); + } + /** * Defines properties to be passed to the Iceberg writer itself. Note that these properties are * execution-scoped, meaning that they are applied to a preexisting table and will not mutate @@ -679,6 +886,8 @@ public IcebergWriteResult expand(PCollection input) { metadataView = input.apply("GenerateTableMetadataView", driverBuilder.build().asView()); } + boolean sortedWrites = getSortedWrites() || getDistributionMode() == DistributionMode.RANGE; + switch (getDistributionMode()) { case NONE: Preconditions.checkArgument( @@ -694,13 +903,44 @@ public IcebergWriteResult expand(PCollection input) { getTriggeringFrequency(), getDirectWriteByteLimit(), getWriteProperties(), - metadataView)); + metadataView, + sortedWrites)); case HASH: return input .apply( "AssignDestinationAndPartition", new AssignDestinationsAndPartitions( - destinations, getCatalogConfig(), metadataView)) + destinations, + getCatalogConfig(), + metadataView, + getDistributionMode(), + getDistributionFunction())) + .apply( + "Write Rows to Partitions", + new WriteToPartitions( + getCatalogConfig(), + destinations, + getTriggeringFrequency(), + getAutoSharding(), + getWriteProperties(), + metadataView, + sortedWrites)); + case RANGE: + Preconditions.checkArgument( + !getAutoSharding(), + "Autosharding option is only available with 'hash' distribution mode."); + Preconditions.checkArgument( + getDistributionFunction() != null, + "Must provide a distribution function when using RANGE distribution mode."); + return input + .apply( + "AssignDestinationAndPartitionWithRange", + new AssignDestinationsAndPartitions( + destinations, + getCatalogConfig(), + metadataView, + getDistributionMode(), + getDistributionFunction())) .apply( "Write Rows to Partitions", new WriteToPartitions( @@ -709,7 +949,8 @@ destinations, getCatalogConfig(), metadataView)) getTriggeringFrequency(), getAutoSharding(), getWriteProperties(), - metadataView)); + metadataView, + sortedWrites)); default: throw new UnsupportedOperationException( "Unsupported distribution mode: " + getDistributionMode()); diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergRowSorter.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergRowSorter.java new file mode 100644 index 000000000000..64362ca9f217 --- /dev/null +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergRowSorter.java @@ -0,0 +1,492 @@ +/* + * 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.beam.sdk.io.iceberg; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.math.BigDecimal; +import java.math.BigInteger; +import java.math.RoundingMode; +import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; +import java.nio.file.DirectoryStream; +import java.nio.file.Files; +import java.nio.file.Path; +import java.time.Instant; +import java.time.LocalDate; +import java.time.LocalDateTime; +import java.time.LocalTime; +import java.time.OffsetDateTime; +import java.util.Arrays; +import java.util.Date; +import java.util.Iterator; +import java.util.List; +import java.util.UUID; +import org.apache.beam.sdk.coders.RowCoder; +import org.apache.beam.sdk.extensions.sorter.BufferedExternalSorter; +import org.apache.beam.sdk.extensions.sorter.ExternalSorter; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.Row; +import org.apache.iceberg.NullOrder; +import org.apache.iceberg.Schema; +import org.apache.iceberg.SortDirection; +import org.apache.iceberg.SortField; +import org.apache.iceberg.SortOrder; +import org.apache.iceberg.io.CloseableIterable; +import org.apache.iceberg.transforms.Transform; +import org.apache.iceberg.types.Type; +import org.apache.iceberg.types.Types; +import org.apache.iceberg.util.DateTimeUtil; +import org.apache.iceberg.util.SerializableFunction; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.ReadableInstant; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * A utility class to sort Beam {@link Row}s based on an Iceberg {@link SortOrder}. Leverages {@link + * BufferedExternalSorter} to spill to local disk when elements exceed memory limit. + */ +class IcebergRowSorter { + private static final Logger LOG = LoggerFactory.getLogger(IcebergRowSorter.class); + + /** Width of an encoded decimal key. 16 bytes covers Iceberg's maximum precision of 38 digits. */ + private static final int DECIMAL_KEY_BYTES = 16; + + /** + * Best-effort cleanup of a temporary spill directory. Failures are logged rather than propagated. + */ + private static void deleteTempDir(@Nullable Path tempDir) { + if (tempDir == null) { + return; + } + try (DirectoryStream entries = Files.newDirectoryStream(tempDir)) { + for (Path entry : entries) { + if (!Files.deleteIfExists(entry)) { + LOG.debug("Temporary sort file {} was already removed", entry); + } + } + } catch (IOException | RuntimeException e) { + LOG.debug("Failed to list temporary sort directory {}", tempDir, e); + return; + } + try { + if (!Files.deleteIfExists(tempDir)) { + LOG.debug("Temporary sort directory {} was already removed", tempDir); + } + } catch (IOException | RuntimeException e) { + LOG.debug("Failed to delete temporary sort directory {}", tempDir, e); + } + } + + /** + * Sorts {@code rows} according to {@code sortOrder}, spilling to local disk beyond the in-memory + * buffer. + * + *

The caller must close the returned iterable (ideally via try-with-resources). Closing + * reclaims the temporary spill directory; skipping it leaks that directory on the worker's local + * disk, which matters on long-lived runner harnesses that process many bundles. + * + *

Note that spilled rows are written unencrypted to the worker's local disk, in a directory + * created with owner-only permissions. + */ + @SuppressWarnings({"unchecked", "rawtypes"}) + public static CloseableIterable sortRows( + Iterable rows, + SortOrder sortOrder, + Schema icebergSchema, + org.apache.beam.sdk.schemas.Schema beamSchema) { + + if (sortOrder == null || !sortOrder.isSorted()) { + return CloseableIterable.withNoopClose(rows); + } + + final Path sortTempDir; + try { + sortTempDir = Files.createTempDirectory("iceberg-row-sorter-"); + } catch (IOException e) { + throw new RuntimeException("Failed to create temporary directory for IcebergRowSorter", e); + } + + BufferedExternalSorter.Options sorterOptions = + BufferedExternalSorter.options() + .withTempLocation(sortTempDir.toString()) + .withExternalSorterType(ExternalSorter.Options.SorterType.NATIVE); + BufferedExternalSorter sorter = BufferedExternalSorter.create(sorterOptions); + RowCoder rowCoder = RowCoder.of(beamSchema); + + List fields = sortOrder.fields(); + String[] columnNames = new String[fields.size()]; + Type[] sourceTypes = new Type[fields.size()]; + Type[] resultTypes = new Type[fields.size()]; + @Nullable SerializableFunction[] boundTransforms = + new SerializableFunction[fields.size()]; + for (int i = 0; i < fields.size(); i++) { + SortField field = fields.get(i); + columnNames[i] = icebergSchema.findColumnName(field.sourceId()); + Type sourceType = icebergSchema.findType(field.sourceId()); + sourceTypes[i] = sourceType; + if (field.transform().isIdentity()) { + resultTypes[i] = sourceType; + boundTransforms[i] = null; + } else { + Transform transform = (Transform) field.transform(); + resultTypes[i] = transform.getResultType(sourceType); + boundTransforms[i] = transform.bind(sourceType); + } + } + + // Create reusable ByteArrayOutputStreams for key and value encoding + ByteArrayOutputStream keyBaos = new ByteArrayOutputStream(); + ByteArrayOutputStream valBaos = new ByteArrayOutputStream(); + + try { + for (Row row : rows) { + keyBaos.reset(); + valBaos.reset(); + encodeSortKey( + row, sortOrder, columnNames, sourceTypes, resultTypes, boundTransforms, keyBaos); + byte[] keyBytes = keyBaos.toByteArray(); + + rowCoder.encode(row, valBaos); + byte[] valBytes = valBaos.toByteArray(); + sorter.add(KV.of(keyBytes, valBytes)); + } + + Iterable> sortedKVs = sorter.sort(); + Iterable decodedRows = + new Iterable() { + @Override + public Iterator iterator() { + final Iterator> it = sortedKVs.iterator(); + return new Iterator() { + @Override + public boolean hasNext() { + return it.hasNext(); + } + + @Override + public Row next() { + KV next = it.next(); + try { + return rowCoder.decode(new ByteArrayInputStream(next.getValue())); + } catch (IOException e) { + deleteTempDir(sortTempDir); + throw new RuntimeException("Failed to decode Row during sorting", e); + } + } + }; + } + }; + + return CloseableIterable.combine(decodedRows, () -> deleteTempDir(sortTempDir)); + + } catch (IOException | RuntimeException e) { + deleteTempDir(sortTempDir); + throw new RuntimeException("Failed to sort rows with external sorter", e); + } + } + + @SuppressWarnings({"nullness", "unchecked", "rawtypes"}) + public static void encodeSortKey( + Row row, + SortOrder sortOrder, + String[] columnNames, + ByteArrayOutputStream baos, + Schema icebergSchema, + org.apache.beam.sdk.schemas.Schema beamSchema) + throws IOException { + + List fields = sortOrder.fields(); + Type[] sourceTypes = new Type[fields.size()]; + Type[] resultTypes = new Type[fields.size()]; + @Nullable SerializableFunction[] boundTransforms = + new SerializableFunction[fields.size()]; + for (int i = 0; i < fields.size(); i++) { + SortField field = fields.get(i); + Type sourceType = icebergSchema.findType(field.sourceId()); + sourceTypes[i] = sourceType; + if (field.transform().isIdentity()) { + resultTypes[i] = sourceType; + boundTransforms[i] = null; + } else { + Transform transform = (Transform) field.transform(); + resultTypes[i] = transform.getResultType(sourceType); + boundTransforms[i] = transform.bind(sourceType); + } + } + encodeSortKey(row, sortOrder, columnNames, sourceTypes, resultTypes, boundTransforms, baos); + } + + @SuppressWarnings("nullness") + private static void encodeSortKey( + Row row, + SortOrder sortOrder, + String[] columnNames, + Type[] sourceTypes, + Type[] resultTypes, + @Nullable SerializableFunction[] boundTransforms, + ByteArrayOutputStream baos) + throws IOException { + + List fields = sortOrder.fields(); + + for (int i = 0; i < fields.size(); i++) { + SortField field = fields.get(i); + String colName = columnNames[i]; + + // Normalize into the Iceberg value domain up front so that the encoder only ever deals with + // Iceberg's representation, regardless of which Beam logical type the row carried. + Type sourceType = sourceTypes[i]; + Type valueType = resultTypes[i]; + Object val = IcebergUtils.beamValueToIcebergValue(sourceType, row.getValue(colName)); + + SerializableFunction boundTransform = boundTransforms[i]; + if (boundTransform != null && val != null) { + val = boundTransform.apply(toTransformInput(sourceType, val)); + } + + boolean isNull = (val == null); + boolean isDesc = (field.direction() == SortDirection.DESC); + boolean nullsFirst = (field.nullOrder() == NullOrder.NULLS_FIRST); + + // Determine correct header prefix to fulfill the NullOrder contracts + byte prefixByte; + if (isNull) { + prefixByte = nullsFirst ? (byte) 0x00 : (byte) 0xFF; + } else { + prefixByte = nullsFirst ? (byte) 0x01 : (byte) 0x00; + } + + baos.write(prefixByte); + + if (!isNull) { + writeValue(val, valueType, baos, isDesc); + } + } + } + + /** + * Converts a normalized Iceberg value into the internal representation expected by bound Iceberg + * {@link Transform} functions (e.g. epoch days as {@code Integer} for {@code DATE}, epoch micros + * as {@code Long} for {@code TIME} / {@code TIMESTAMP}). + */ + @SuppressWarnings("JavaUtilDate") + private static Object toTransformInput(Type sourceType, Object val) { + if (val instanceof LocalDate) { + return DateTimeUtil.daysFromDate((LocalDate) val); + } else if (val instanceof LocalTime) { + return DateTimeUtil.microsFromTime((LocalTime) val); + } else if (val instanceof LocalDateTime) { + return DateTimeUtil.microsFromTimestamp((LocalDateTime) val); + } else if (val instanceof OffsetDateTime) { + return DateTimeUtil.microsFromTimestamptz((OffsetDateTime) val); + } else if (val instanceof ReadableInstant) { + return ((ReadableInstant) val).getMillis() * 1000L; + } else if (val instanceof Instant) { + return DateTimeUtil.microsFromInstant((Instant) val); + } else if (val instanceof Date) { + return ((Date) val).getTime() * 1000L; + } else if (val instanceof Short || val instanceof Byte) { + return ((Number) val).intValue(); + } else if (val instanceof byte[]) { + return ByteBuffer.wrap((byte[]) val); + } else if (val instanceof BigDecimal && sourceType instanceof Types.DecimalType) { + int scale = ((Types.DecimalType) sourceType).scale(); + return ((BigDecimal) val).setScale(scale, RoundingMode.UNNECESSARY); + } + return val; + } + + private static void writeInt(int v, ByteArrayOutputStream baos, boolean invert) { + byte b3 = (byte) (v >>> 24); + byte b2 = (byte) (v >>> 16); + byte b1 = (byte) (v >>> 8); + byte b0 = (byte) v; + if (invert) { + baos.write(~b3); + baos.write(~b2); + baos.write(~b1); + baos.write(~b0); + } else { + baos.write(b3); + baos.write(b2); + baos.write(b1); + baos.write(b0); + } + } + + private static void writeLong(long v, ByteArrayOutputStream baos, boolean invert) { + byte b7 = (byte) (v >>> 56); + byte b6 = (byte) (v >>> 48); + byte b5 = (byte) (v >>> 40); + byte b4 = (byte) (v >>> 32); + byte b3 = (byte) (v >>> 24); + byte b2 = (byte) (v >>> 16); + byte b1 = (byte) (v >>> 8); + byte b0 = (byte) v; + if (invert) { + baos.write(~b7); + baos.write(~b6); + baos.write(~b5); + baos.write(~b4); + baos.write(~b3); + baos.write(~b2); + baos.write(~b1); + baos.write(~b0); + } else { + baos.write(b7); + baos.write(b6); + baos.write(b5); + baos.write(b4); + baos.write(b3); + baos.write(b2); + baos.write(b1); + baos.write(b0); + } + } + + @SuppressWarnings("JavaUtilDate") + private static void writeValue( + Object val, @Nullable Type type, ByteArrayOutputStream baos, boolean invert) + throws IOException { + if (val instanceof CharSequence) { + writeString(val.toString(), baos, invert); + } else if (val instanceof Integer || val instanceof Short || val instanceof Byte) { + writeInt(((Number) val).intValue() ^ Integer.MIN_VALUE, baos, invert); + } else if (val instanceof Long) { + writeLong((Long) val ^ Long.MIN_VALUE, baos, invert); + } else if (val instanceof Float) { + int bits = Float.floatToIntBits((Float) val); + bits = (bits >= 0) ? (bits ^ Integer.MIN_VALUE) : ~bits; + writeInt(bits, baos, invert); + } else if (val instanceof Double) { + long bits = Double.doubleToLongBits((Double) val); + bits = (bits >= 0) ? (bits ^ Long.MIN_VALUE) : ~bits; + writeLong(bits, baos, invert); + } else if (val instanceof Boolean) { + byte b = ((Boolean) val) ? (byte) 0x01 : (byte) 0x00; + baos.write(invert ? ~b : b); + } else if (val instanceof BigDecimal) { + writeDecimal((BigDecimal) val, type, baos, invert); + } else if (val instanceof byte[]) { + writeByteArray((byte[]) val, baos, invert); + } else if (val instanceof ByteBuffer) { + writeByteArray(toByteArray((ByteBuffer) val), baos, invert); + } else if (val instanceof LocalDate) { + writeLong(((LocalDate) val).toEpochDay() ^ Long.MIN_VALUE, baos, invert); + } else if (val instanceof LocalTime) { + writeLong(((LocalTime) val).toNanoOfDay() ^ Long.MIN_VALUE, baos, invert); + } else if (val instanceof LocalDateTime) { + writeLong( + DateTimeUtil.microsFromTimestamp((LocalDateTime) val) ^ Long.MIN_VALUE, baos, invert); + } else if (val instanceof OffsetDateTime) { + writeLong( + DateTimeUtil.microsFromTimestamptz((OffsetDateTime) val) ^ Long.MIN_VALUE, baos, invert); + } else if (val instanceof ReadableInstant) { + writeLong(((ReadableInstant) val).getMillis() ^ Long.MIN_VALUE, baos, invert); + } else if (val instanceof Instant) { + writeLong(((Instant) val).toEpochMilli() ^ Long.MIN_VALUE, baos, invert); + } else if (val instanceof Date) { + writeLong(((Date) val).getTime() ^ Long.MIN_VALUE, baos, invert); + } else if (val instanceof UUID) { + // Iceberg orders UUIDs by their 16-byte big-endian representation, so flip the sign bit of + // each half to turn Java's signed longs into an unsigned lexicographic ordering. + UUID uuid = (UUID) val; + writeLong(uuid.getMostSignificantBits() ^ Long.MIN_VALUE, baos, invert); + writeLong(uuid.getLeastSignificantBits() ^ Long.MIN_VALUE, baos, invert); + } else { + throw new UnsupportedOperationException( + "Unsupported type for sorting: " + val.getClass().getName()); + } + } + + /** + * Encodes a decimal as a fixed-width, sign-flipped two's complement big-endian integer. + * + *

Iceberg pins the scale of a decimal column in the schema, so rescaling to that scale makes + * every value in the column directly comparable as a fixed-width unscaled integer. A fixed width + * is required because this encoding carries no terminator; {@value #DECIMAL_KEY_BYTES} bytes + * covers Iceberg's maximum precision of 38 digits. + */ + private static void writeDecimal( + BigDecimal val, @Nullable Type type, ByteArrayOutputStream baos, boolean invert) { + int scale = + (type instanceof Types.DecimalType) ? ((Types.DecimalType) type).scale() : val.scale(); + BigInteger unscaled; + try { + unscaled = val.setScale(scale, RoundingMode.UNNECESSARY).unscaledValue(); + } catch (ArithmeticException e) { + throw new IllegalArgumentException( + String.format( + "Decimal sort value %s does not fit the column's declared scale of %d.", val, scale), + e); + } + + byte[] magnitude = unscaled.toByteArray(); + if (magnitude.length > DECIMAL_KEY_BYTES) { + throw new IllegalArgumentException( + String.format("Decimal sort value %s exceeds Iceberg's maximum decimal precision.", val)); + } + + byte[] encoded = new byte[DECIMAL_KEY_BYTES]; + // Sign-extend so that shorter negative magnitudes stay ordered against longer ones. + Arrays.fill( + encoded, + 0, + DECIMAL_KEY_BYTES - magnitude.length, + (byte) (unscaled.signum() < 0 ? 0xFF : 0x00)); + System.arraycopy(magnitude, 0, encoded, DECIMAL_KEY_BYTES - magnitude.length, magnitude.length); + // Flip the sign bit so two's complement sorts correctly as unsigned bytes. + encoded[0] ^= (byte) 0x80; + + for (byte b : encoded) { + baos.write(invert ? ~b : b); + } + } + + /** Copies the readable region of a buffer without assuming an accessible backing array. */ + private static byte[] toByteArray(ByteBuffer buffer) { + ByteBuffer duplicate = buffer.duplicate(); + byte[] bytes = new byte[duplicate.remaining()]; + duplicate.get(bytes); + return bytes; + } + + private static void writeString(String s, ByteArrayOutputStream baos, boolean invert) + throws IOException { + byte[] bytes = s.getBytes(StandardCharsets.UTF_8); + writeByteArray(bytes, baos, invert); + } + + private static void writeByteArray(byte[] bytes, ByteArrayOutputStream baos, boolean invert) { + for (byte b : bytes) { + if (b == 0x00) { + baos.write(invert ? ~(byte) 0x01 : (byte) 0x01); + baos.write(invert ? ~(byte) 0x01 : (byte) 0x01); + } else if (b == 0x01) { + baos.write(invert ? ~(byte) 0x01 : (byte) 0x01); + baos.write(invert ? ~(byte) 0x02 : (byte) 0x02); + } else { + baos.write(invert ? ~b : b); + } + } + baos.write(invert ? ~(byte) 0x00 : (byte) 0x00); + } +} diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergUtils.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergUtils.java index 8db3b643cf8b..87aa73d716a9 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergUtils.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergUtils.java @@ -336,6 +336,45 @@ public static org.apache.iceberg.Schema beamSchemaToIcebergSchema(final Schema s return new org.apache.iceberg.Schema(fields.toArray(new Types.NestedField[fields.size()])); } + /** + * Converts a Beam field value to its Iceberg-compatible equivalent based on the Iceberg {@link + * Type}. + */ + public static @Nullable Object beamValueToIcebergValue(Type type, @Nullable Object value) { + if (value == null) { + return null; + } + switch (type.typeId()) { + case BOOLEAN: + case INTEGER: + case LONG: + case FLOAT: + case DOUBLE: + case DATE: + case TIME: + case DECIMAL: + case STRING: + return value; + case TIMESTAMP: + Types.TimestampType ts = (Types.TimestampType) type.asPrimitiveType(); + return getIcebergTimestampValue(value, ts.shouldAdjustToUTC()); + case UUID: + if (value instanceof byte[]) { + return UUID.nameUUIDFromBytes((byte[]) value); + } + return value; + case BINARY: + if (value instanceof byte[]) { + return ByteBuffer.wrap((byte[]) value); + } + return value; + case FIXED: + throw new UnsupportedOperationException("Fixed-precision fields are not yet supported."); + default: + return value; + } + } + /** Converts a Beam {@link Row} to an Iceberg {@link Record}. */ public static Record beamRowToIcebergRecord(org.apache.iceberg.Schema schema, Row row) { if (row.getSchema().getFieldCount() != schema.columns().size()) { 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 91d0ed73835e..47c9469f8b1d 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 @@ -188,10 +188,15 @@ public static Builder builder() { + "is a field name or one of the partition transforms (e.g. `bucket(col, 4)`, `day(ts)`). " + "Direction defaults to ascending; null order defaults to nulls-first for ascending and " + "nulls-last for descending. Note: this sets the table's declared sort order as metadata; " - + "it does not cause Beam to physically sort records before writing.\n" + + "to also physically sort records before writing, set `sort_on_write` to true.\n" + "For more information on sort orders, please visit https://iceberg.apache.org/spec/#sort-orders.") public abstract @Nullable List getSortFields(); + @SchemaFieldDescription( + "When true, physically sorts records within each written data file according to the target" + + " table's sort order. Defaults to false.") + public abstract @Nullable Boolean getSortOnWrite(); + @SchemaFieldDescription( "Defines distribution of write data. Supported distributions:" + "\n- none: don't shuffle rows (default)" @@ -310,6 +315,8 @@ public abstract static class Builder { public abstract Builder setSortFields(List sortFields); + public abstract Builder setSortOnWrite(Boolean sortOnWrite); + public abstract Builder setDistributionMode(String mode); public abstract Builder setAutosharding(Boolean autosharding); @@ -530,6 +537,11 @@ private PCollectionRowTuple expandAppend(PCollection rows) { writeTransform = writeTransform.withAutosharding(); } + @Nullable Boolean sortOnWrite = configuration.getSortOnWrite(); + if (sortOnWrite != null && sortOnWrite) { + writeTransform = writeTransform.withSortedWrites(); + } + @Nullable Map writeProperties = configuration.getWriteProperties(); if (writeProperties != null && !writeProperties.isEmpty()) { writeTransform = writeTransform.withWriteProperties(writeProperties); 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 a2ed40c87f99..7b850103b1cf 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 @@ -19,6 +19,7 @@ import java.util.List; import java.util.Map; +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; @@ -32,6 +33,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.base.Preconditions; +import org.apache.iceberg.Table; +import org.apache.iceberg.io.CloseableIterable; import org.checkerframework.checker.nullness.qual.Nullable; class WriteGroupedRowsToFiles @@ -44,6 +47,7 @@ class WriteGroupedRowsToFiles private final String filePrefix; private final @Nullable Map writeProperties; private final @Nullable PCollectionView> metadataView; + private final boolean sortedWrites; WriteGroupedRowsToFiles( IcebergCatalogConfig catalogConfig, @@ -51,7 +55,14 @@ class WriteGroupedRowsToFiles String filePrefix, long maxBytesPerFile, @Nullable Map writeProperties) { - this(catalogConfig, dynamicDestinations, filePrefix, maxBytesPerFile, writeProperties, null); + this( + catalogConfig, + dynamicDestinations, + filePrefix, + maxBytesPerFile, + writeProperties, + null, + false); } WriteGroupedRowsToFiles( @@ -61,17 +72,37 @@ class WriteGroupedRowsToFiles long maxBytesPerFile, @Nullable Map writeProperties, @Nullable PCollectionView> metadataView) { + this( + catalogConfig, + dynamicDestinations, + filePrefix, + maxBytesPerFile, + writeProperties, + metadataView, + false); + } + + WriteGroupedRowsToFiles( + IcebergCatalogConfig catalogConfig, + DynamicDestinations dynamicDestinations, + String filePrefix, + long maxBytesPerFile, + @Nullable Map writeProperties, + @Nullable PCollectionView> metadataView, + boolean sortedWrites) { this.catalogConfig = catalogConfig; this.dynamicDestinations = dynamicDestinations; this.filePrefix = filePrefix; this.maxBytesPerFile = maxBytesPerFile; this.writeProperties = writeProperties; this.metadataView = metadataView; + this.sortedWrites = sortedWrites; } @Override public PCollection expand( PCollection, Iterable>> input) { + Schema dataSchema = dynamicDestinations.getDataSchema(); ParDo.SingleOutput, Iterable>, FileWriteResult> parDo = ParDo.of( new WriteGroupedRowsToFilesDoFn( @@ -79,8 +110,10 @@ public PCollection expand( dynamicDestinations, maxBytesPerFile, filePrefix, + dataSchema, writeProperties, - metadataView)); + metadataView, + sortedWrites)); if (metadataView != null) { parDo = parDo.withSideInputs(metadataView); } @@ -94,16 +127,27 @@ private static class WriteGroupedRowsToFilesDoFn private final IcebergCatalogConfig catalogConfig; private final String filePrefix; private final long maxFileSize; + private final Schema dataSchema; private final @Nullable Map writeProperties; private final @Nullable PCollectionView> metadataView; + private final boolean sortedWrites; WriteGroupedRowsToFilesDoFn( IcebergCatalogConfig catalogConfig, DynamicDestinations dynamicDestinations, long maxFileSize, String filePrefix, + Schema dataSchema, @Nullable Map writeProperties) { - this(catalogConfig, dynamicDestinations, maxFileSize, filePrefix, writeProperties, null); + this( + catalogConfig, + dynamicDestinations, + maxFileSize, + filePrefix, + dataSchema, + writeProperties, + null, + false); } WriteGroupedRowsToFilesDoFn( @@ -111,14 +155,37 @@ private static class WriteGroupedRowsToFilesDoFn DynamicDestinations dynamicDestinations, long maxFileSize, String filePrefix, + Schema dataSchema, @Nullable Map writeProperties, @Nullable PCollectionView> metadataView) { + this( + catalogConfig, + dynamicDestinations, + maxFileSize, + filePrefix, + dataSchema, + writeProperties, + metadataView, + false); + } + + WriteGroupedRowsToFilesDoFn( + IcebergCatalogConfig catalogConfig, + DynamicDestinations dynamicDestinations, + long maxFileSize, + String filePrefix, + Schema dataSchema, + @Nullable Map writeProperties, + @Nullable PCollectionView> metadataView, + boolean sortedWrites) { this.catalogConfig = catalogConfig; this.dynamicDestinations = dynamicDestinations; this.filePrefix = filePrefix; this.maxFileSize = maxFileSize; + this.dataSchema = dataSchema; this.writeProperties = writeProperties; this.metadataView = metadataView; + this.sortedWrites = sortedWrites; } @ProcessElement @@ -145,8 +212,15 @@ public void processElement( writeProperties, sideInputs)) { writer = openWriter; - for (Row e : element.getValue()) { - writer.write(windowedDestination, e); + Table table = sortedWrites ? openWriter.getOrCreateTable(destination, dataSchema) : null; + try (CloseableIterable sortedOrUnsortedRows = + table != null && table.sortOrder().isSorted() + ? IcebergRowSorter.sortRows( + element.getValue(), table.sortOrder(), table.schema(), dataSchema) + : CloseableIterable.withNoopClose(element.getValue())) { + for (Row e : sortedOrUnsortedRows) { + 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 fbb6cb08aca4..fef0ab034d7d 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 @@ -22,11 +22,10 @@ import static org.apache.beam.sdk.io.iceberg.RecordWriterManager.getPartitionDataPath; import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; +import java.util.ArrayList; +import java.util.List; 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; @@ -50,6 +49,7 @@ import org.apache.iceberg.data.Record; import org.apache.iceberg.exceptions.AlreadyExistsException; import org.apache.iceberg.exceptions.NoSuchTableException; +import org.apache.iceberg.io.CloseableIterable; import org.checkerframework.checker.nullness.qual.MonotonicNonNull; import org.checkerframework.checker.nullness.qual.Nullable; import org.slf4j.Logger; @@ -61,15 +61,24 @@ class WritePartitionedRowsToFiles private final DynamicDestinations dynamicDestinations; private final IcebergCatalogConfig catalogConfig; private final String filePrefix; + private final long maxFileSize; private final @Nullable Map writeProperties; private final @Nullable PCollectionView> metadataView; + private final boolean sortedWrites; WritePartitionedRowsToFiles( IcebergCatalogConfig catalogConfig, DynamicDestinations dynamicDestinations, String filePrefix, @Nullable Map writeProperties) { - this(catalogConfig, dynamicDestinations, filePrefix, writeProperties, null); + this( + catalogConfig, + dynamicDestinations, + filePrefix, + WriteToPartitions.DEFAULT_BYTES_PER_FILE, + writeProperties, + null, + false); } WritePartitionedRowsToFiles( @@ -78,30 +87,64 @@ class WritePartitionedRowsToFiles String filePrefix, @Nullable Map writeProperties, @Nullable PCollectionView> metadataView) { + this( + catalogConfig, + dynamicDestinations, + filePrefix, + WriteToPartitions.DEFAULT_BYTES_PER_FILE, + writeProperties, + metadataView, + false); + } + + WritePartitionedRowsToFiles( + IcebergCatalogConfig catalogConfig, + DynamicDestinations dynamicDestinations, + String filePrefix, + @Nullable Map writeProperties, + @Nullable PCollectionView> metadataView, + boolean sortedWrites) { + this( + catalogConfig, + dynamicDestinations, + filePrefix, + WriteToPartitions.DEFAULT_BYTES_PER_FILE, + writeProperties, + metadataView, + sortedWrites); + } + + WritePartitionedRowsToFiles( + IcebergCatalogConfig catalogConfig, + DynamicDestinations dynamicDestinations, + String filePrefix, + long maxFileSize, + @Nullable Map writeProperties, + @Nullable PCollectionView> metadataView, + boolean sortedWrites) { this.catalogConfig = catalogConfig; this.dynamicDestinations = dynamicDestinations; this.filePrefix = filePrefix; + this.maxFileSize = maxFileSize; this.writeProperties = writeProperties; this.metadataView = metadataView; + this.sortedWrites = sortedWrites; } @Override public PCollection expand(PCollection>> input) { - Schema dataSchema = - ((RowCoder) - ((IterableCoder) - ((KvCoder>) input.getCoder()).getValueCoder()) - .getElemCoder()) - .getSchema(); + Schema dataSchema = dynamicDestinations.getDataSchema(); ParDo.SingleOutput>, FileWriteResult> parDo = ParDo.of( new WriteDoFn( catalogConfig, dynamicDestinations, filePrefix, + maxFileSize, dataSchema, writeProperties, - metadataView)); + metadataView, + sortedWrites)); if (metadataView != null) { parDo = parDo.withSideInputs(metadataView); } @@ -113,9 +156,11 @@ private static class WriteDoFn extends DoFn>, FileWriteRes private final DynamicDestinations dynamicDestinations; private final IcebergCatalogConfig catalogConfig; private final String filePrefix; + private final long maxFileSize; private final Schema dataSchema; private final @Nullable Map writeProperties; private final @Nullable PCollectionView> metadataView; + private final boolean sortedWrites; private transient @MonotonicNonNull Map specIds; private transient @MonotonicNonNull Map> partitionFieldMaps; @@ -126,7 +171,15 @@ private static class WriteDoFn extends DoFn>, FileWriteRes String filePrefix, Schema dataSchema, @Nullable Map writeProperties) { - this(catalogConfig, dynamicDestinations, filePrefix, dataSchema, writeProperties, null); + this( + catalogConfig, + dynamicDestinations, + filePrefix, + WriteToPartitions.DEFAULT_BYTES_PER_FILE, + dataSchema, + writeProperties, + null, + false); } WriteDoFn( @@ -136,12 +189,34 @@ private static class WriteDoFn extends DoFn>, FileWriteRes Schema dataSchema, @Nullable Map writeProperties, @Nullable PCollectionView> metadataView) { + this( + catalogConfig, + dynamicDestinations, + filePrefix, + WriteToPartitions.DEFAULT_BYTES_PER_FILE, + dataSchema, + writeProperties, + metadataView, + false); + } + + WriteDoFn( + IcebergCatalogConfig catalogConfig, + DynamicDestinations dynamicDestinations, + String filePrefix, + long maxFileSize, + Schema dataSchema, + @Nullable Map writeProperties, + @Nullable PCollectionView> metadataView, + boolean sortedWrites) { this.catalogConfig = catalogConfig; this.dynamicDestinations = dynamicDestinations; this.filePrefix = filePrefix; + this.maxFileSize = maxFileSize; this.dataSchema = dataSchema; this.writeProperties = writeProperties; this.metadataView = metadataView; + this.sortedWrites = sortedWrites; } @Setup @@ -172,30 +247,50 @@ public void processElement( ? 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 { - for (Row row : element.getValue()) { + List writtenFiles = new ArrayList<>(); + RecordWriter writer = null; + try (CloseableIterable sortedOrUnsortedRows = + sortedWrites && table.sortOrder().isSorted() + ? IcebergRowSorter.sortRows( + element.getValue(), table.sortOrder(), table.schema(), dataSchema) + : CloseableIterable.withNoopClose(element.getValue())) { + for (Row row : sortedOrUnsortedRows) { + if (writer != null && writer.bytesWritten() > maxFileSize) { + RecordWriter finishedWriter = writer; + writer = null; + finishedWriter.close(); + writtenFiles.add( + SerializableDataFile.from(finishedWriter.getDataFile(), table.specs())); + } + if (writer == null) { + String fileName = + destination + .getFileFormat() + .addExtension(String.format("%s-%s", filePrefix, UUID.randomUUID())); + writer = + new RecordWriter( + table, destination.getFileFormat(), fileName, partitionData, writeProperties); + } Record record = IcebergUtils.beamRowToIcebergRecord(table.schema(), row); writer.write(record); } } finally { - writer.close(); + if (writer != null) { + writer.close(); + } + } + + if (writer != null) { + writtenFiles.add(SerializableDataFile.from(writer.getDataFile(), table.specs())); } - // 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()); + for (SerializableDataFile sdf : writtenFiles) { + out.output( + FileWriteResult.builder() + .setTableIdentifier(destination.getTableIdentifier()) + .setSerializableDataFile(sdf) + .build()); + } } private Map getPartitionFieldMap( 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 6d8f08a2ae56..95a926e71f86 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 @@ -62,6 +62,7 @@ class WriteToDestinations extends PTransform>, Icebe private final @Nullable Integer directWriteByteLimit; private final @Nullable Map writeProperties; private final @Nullable PCollectionView> metadataView; + private final boolean sortedWrites; WriteToDestinations( IcebergCatalogConfig catalogConfig, @@ -75,7 +76,8 @@ class WriteToDestinations extends PTransform>, Icebe triggeringFrequency, directWriteByteLimit, writeProperties, - null); + null, + false); } WriteToDestinations( @@ -85,12 +87,31 @@ class WriteToDestinations extends PTransform>, Icebe @Nullable Integer directWriteByteLimit, @Nullable Map writeProperties, @Nullable PCollectionView> metadataView) { + this( + catalogConfig, + dynamicDestinations, + triggeringFrequency, + directWriteByteLimit, + writeProperties, + metadataView, + false); + } + + WriteToDestinations( + IcebergCatalogConfig catalogConfig, + DynamicDestinations dynamicDestinations, + @Nullable Duration triggeringFrequency, + @Nullable Integer directWriteByteLimit, + @Nullable Map writeProperties, + @Nullable PCollectionView> metadataView, + boolean sortedWrites) { this.dynamicDestinations = dynamicDestinations; this.catalogConfig = catalogConfig; this.triggeringFrequency = triggeringFrequency; this.directWriteByteLimit = directWriteByteLimit; this.writeProperties = writeProperties; this.metadataView = metadataView; + this.sortedWrites = sortedWrites; // single unique prefix per write transform this.filePrefix = UUID.randomUUID().toString(); } @@ -101,7 +122,7 @@ public IcebergWriteResult expand(PCollection> input) { PCollection writtenFiles; if (IcebergUtils.isUnbounded(input)) { writtenFiles = - IcebergUtils.validDirectWriteLimit(directWriteByteLimit) + IcebergUtils.validDirectWriteLimit(directWriteByteLimit) && !sortedWrites ? writeTriggeredWithBundleLifting(input) : writeTriggered(input); } else { @@ -139,7 +160,8 @@ private PCollection groupAndWriteRecords(PCollection applyUserTriggering(PCollection input) { @@ -226,7 +248,8 @@ private PCollection writeUntriggered(PCollection writeGroupedResult = @@ -241,7 +264,8 @@ private PCollection writeUntriggered(PCollection>, 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; @@ -50,6 +52,7 @@ class WriteToPartitions extends PTransform>, IcebergWri private final boolean autoSharding; private final @Nullable Map writeProperties; private final @Nullable PCollectionView> metadataView; + private final boolean sortedWrites; WriteToPartitions( IcebergCatalogConfig catalogConfig, @@ -63,7 +66,8 @@ class WriteToPartitions extends PTransform>, IcebergWri triggeringFrequency, autoSharding, writeProperties, - null); + null, + false); } WriteToPartitions( @@ -73,6 +77,24 @@ class WriteToPartitions extends PTransform>, IcebergWri boolean autoSharding, @Nullable Map writeProperties, @Nullable PCollectionView> metadataView) { + this( + catalogConfig, + dynamicDestinations, + triggeringFrequency, + autoSharding, + writeProperties, + metadataView, + false); + } + + WriteToPartitions( + IcebergCatalogConfig catalogConfig, + DynamicDestinations dynamicDestinations, + @Nullable Duration triggeringFrequency, + boolean autoSharding, + @Nullable Map writeProperties, + @Nullable PCollectionView> metadataView, + boolean sortedWrites) { this.dynamicDestinations = dynamicDestinations; this.catalogConfig = catalogConfig; this.triggeringFrequency = triggeringFrequency; @@ -81,12 +103,36 @@ class WriteToPartitions extends PTransform>, IcebergWri this.autoSharding = autoSharding; this.writeProperties = writeProperties; this.metadataView = metadataView; + this.sortedWrites = sortedWrites; + } + + /** Recovers the grouping-key {@link RowCoder} that the upstream assignment step established. */ + private static RowCoder partitionKeyCoder(PCollection> input) { + Coder> inputCoder = input.getCoder(); + if (!(inputCoder instanceof KvCoder)) { + throw new IllegalArgumentException( + "Expected the partitioned write input to use a KvCoder, but got: " + inputCoder); + } + Coder keyCoder = ((KvCoder) inputCoder).getKeyCoder(); + if (!(keyCoder instanceof RowCoder)) { + throw new IllegalArgumentException( + "Expected the partition grouping key to use a RowCoder, but got: " + keyCoder); + } + return (RowCoder) keyCoder; } private PCollection>> groupByPartition(PCollection> input) { - RowCoder destinationCoder = RowCoder.of(AssignDestinationsAndPartitions.OUTPUT_SCHEMA); + // The grouping key shape depends on the upstream distribution mode (RANGE carries an extra + // shard field). Derive it from the input coder so the two can never drift apart. + RowCoder destinationCoder = partitionKeyCoder(input); RowCoder dataCoder = RowCoder.of(dynamicDestinations.getDataSchema()); + if (!IcebergUtils.isUnbounded(input) && !autoSharding) { + return input + .apply(GroupByKey.create()) + .setCoder(KvCoder.of(destinationCoder, IterableCoder.of(dataCoder))); + } + GroupIntoBatches groupIntoPartitions = GroupIntoBatches.ofByteSize(DEFAULT_BYTES_PER_FILE); if (IcebergUtils.isUnbounded(input) && triggeringFrequency != null) { @@ -119,7 +165,12 @@ public IcebergWriteResult expand(PCollection> input) { PCollection writtenFiles = groupedRows.apply( new WritePartitionedRowsToFiles( - catalogConfig, dynamicDestinations, filePrefix, writeProperties, metadataView)); + catalogConfig, + dynamicDestinations, + filePrefix, + writeProperties, + metadataView, + sortedWrites)); 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 8eb462f159b2..114c3ef12b95 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 @@ -48,6 +48,8 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps; +import org.apache.iceberg.Table; import org.checkerframework.checker.nullness.qual.Nullable; /** @@ -75,6 +77,7 @@ class WriteUngroupedRowsToFiles private final long maxBytesPerFile; private final @Nullable Map writeProperties; private final @Nullable PCollectionView> metadataView; + private final boolean sortedWrites; WriteUngroupedRowsToFiles( IcebergCatalogConfig catalogConfig, @@ -82,7 +85,14 @@ class WriteUngroupedRowsToFiles String filePrefix, long maxBytesPerFile, @Nullable Map writeProperties) { - this(catalogConfig, dynamicDestinations, filePrefix, maxBytesPerFile, writeProperties, null); + this( + catalogConfig, + dynamicDestinations, + filePrefix, + maxBytesPerFile, + writeProperties, + null, + false); } WriteUngroupedRowsToFiles( @@ -92,12 +102,31 @@ class WriteUngroupedRowsToFiles long maxBytesPerFile, @Nullable Map writeProperties, @Nullable PCollectionView> metadataView) { + this( + catalogConfig, + dynamicDestinations, + filePrefix, + maxBytesPerFile, + writeProperties, + metadataView, + false); + } + + WriteUngroupedRowsToFiles( + IcebergCatalogConfig catalogConfig, + DynamicDestinations dynamicDestinations, + String filePrefix, + long maxBytesPerFile, + @Nullable Map writeProperties, + @Nullable PCollectionView> metadataView, + boolean sortedWrites) { this.catalogConfig = catalogConfig; this.dynamicDestinations = dynamicDestinations; this.filePrefix = filePrefix; this.maxBytesPerFile = maxBytesPerFile; this.writeProperties = writeProperties; this.metadataView = metadataView; + this.sortedWrites = sortedWrites; } @Override @@ -112,7 +141,8 @@ public Result expand(PCollection> input) { DEFAULT_MAX_WRITERS_PER_BUNDLE, maxBytesPerFile, writeProperties, - metadataView)) + metadataView, + sortedWrites)) .withOutputTags( WRITTEN_FILES_TAG, TupleTagList.of(ImmutableList.of(WRITTEN_ROWS_TAG, SPILLED_ROWS_TAG))); @@ -216,7 +246,9 @@ private static class WriteUngroupedRowsToFilesDoFn private final IcebergCatalogConfig catalogConfig; private final @Nullable Map writeProperties; private final @Nullable PCollectionView> metadataView; + private final boolean sortedWrites; private transient @Nullable RecordWriterManager recordWriterManager; + private transient @Nullable Map destinationIsSorted; private int spilledShardNumber; public WriteUngroupedRowsToFilesDoFn( @@ -233,7 +265,8 @@ public WriteUngroupedRowsToFilesDoFn( maximumWritersPerBundle, maxFileSize, writeProperties, - null); + null, + false); } public WriteUngroupedRowsToFilesDoFn( @@ -244,6 +277,26 @@ public WriteUngroupedRowsToFilesDoFn( long maxFileSize, @Nullable Map writeProperties, @Nullable PCollectionView> metadataView) { + this( + catalogConfig, + dynamicDestinations, + filename, + maximumWritersPerBundle, + maxFileSize, + writeProperties, + metadataView, + false); + } + + public WriteUngroupedRowsToFilesDoFn( + IcebergCatalogConfig catalogConfig, + DynamicDestinations dynamicDestinations, + String filename, + int maximumWritersPerBundle, + long maxFileSize, + @Nullable Map writeProperties, + @Nullable PCollectionView> metadataView, + boolean sortedWrites) { this.catalogConfig = catalogConfig; this.dynamicDestinations = dynamicDestinations; this.filename = filename; @@ -251,6 +304,7 @@ public WriteUngroupedRowsToFilesDoFn( this.maxFileSize = maxFileSize; this.writeProperties = writeProperties; this.metadataView = metadataView; + this.sortedWrites = sortedWrites; } @StartBundle @@ -258,6 +312,7 @@ public void startBundle() { recordWriterManager = new RecordWriterManager( catalogConfig, filename, maxFileSize, maxWritersPerBundle, writeProperties); + destinationIsSorted = Maps.newHashMap(); this.spilledShardNumber = ThreadLocalRandom.current().nextInt(SPILLED_RECORD_SHARDING_FACTOR); } @@ -279,12 +334,25 @@ public void processElement( metadataView != null ? c.sideInput(metadataView) : null; // Attempt to write record. If the writer is saturated and cannot accept - // the record, spill it over to WriteGroupedRowsToFiles - boolean writeSuccess; + // the record, or if sorted writes are enabled and the target table is sorted, + // spill it over to WriteGroupedRowsToFiles + boolean writeSuccess = false; try { - writeSuccess = - Preconditions.checkNotNull(recordWriterManager) - .write(windowedDestination, data, sideInputs); + RecordWriterManager writerManager = Preconditions.checkNotNull(recordWriterManager); + boolean mustSpillForSorting = false; + if (sortedWrites) { + Map isSortedCache = Preconditions.checkNotNull(destinationIsSorted); + Boolean isSorted = isSortedCache.get(dest); + if (isSorted == null) { + Table table = writerManager.getOrCreateTable(destination, data.getSchema(), sideInputs); + isSorted = table.sortOrder().isSorted(); + isSortedCache.put(dest, isSorted); + } + mustSpillForSorting = isSorted; + } + if (!mustSpillForSorting) { + writeSuccess = writerManager.write(windowedDestination, data, sideInputs); + } } catch (Exception e) { try { Preconditions.checkNotNull(recordWriterManager).close(); diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergIOSideInputTableCacheTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergIOSideInputTableCacheTest.java index 957dafa78479..cb5c954aa777 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergIOSideInputTableCacheTest.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergIOSideInputTableCacheTest.java @@ -55,6 +55,7 @@ import org.apache.iceberg.PartitionSpec; import org.apache.iceberg.Snapshot; import org.apache.iceberg.SnapshotChanges; +import org.apache.iceberg.SortOrder; import org.apache.iceberg.Table; import org.apache.iceberg.catalog.TableIdentifier; import org.apache.iceberg.data.IcebergGenerics; @@ -77,10 +78,11 @@ public class IcebergIOSideInputTableCacheTest implements Serializable { private static final String NONE = "none"; private static final String HASH = "hash"; private static final String HASH_WITH_AUTOSHARDING = "hashWithAutoSharding"; + private static final String RANGE = "range"; @Parameterized.Parameters(name = "distributionMode={0}") public static Iterable data() { - return asList(new Object[][] {{NONE}, {HASH}, {HASH_WITH_AUTOSHARDING}}); + return asList(new Object[][] {{NONE}, {HASH}, {HASH_WITH_AUTOSHARDING}, {RANGE}}); } @Parameterized.Parameter(0) @@ -119,6 +121,9 @@ private IcebergIO.WriteRows applyDistribution(IcebergIO.WriteRows write) { if (distributionMode.equals(HASH_WITH_AUTOSHARDING)) { write = write.withAutosharding(); } + if (distributionMode.equals(RANGE)) { + write = write.withDistributionMode(DistributionMode.RANGE).withDistributionFunction(row -> 0); + } return write; } @@ -590,6 +595,79 @@ public void testDisplayData() { assertEquals("3", items.get("tableCachePollingBuckets")); } + @Test + public void testSortedWritesWithSideInputCache() throws Exception { + TableIdentifier tableId = + TableIdentifier.of( + "default_side_input", "sorted_side_" + Long.toString(UUID.randomUUID().hashCode(), 16)); + + Schema schema = Schema.builder().addInt64Field("id").addStringField("name").build(); + org.apache.iceberg.Schema icebergSchema = IcebergUtils.beamSchemaToIcebergSchema(schema); + + catalogConfig + .catalog() + .buildTable(tableId, icebergSchema) + .withPartitionSpec(PartitionSpec.unpartitioned()) + .withSortOrder(SortOrder.builderFor(icebergSchema).asc("name").desc("id").build()) + .create(); + + List inputRows = + asList( + Row.withSchema(schema).addValues(2L, "banana").build(), + Row.withSchema(schema).addValues(1L, "banana").build(), + Row.withSchema(schema).addValues(5L, "apple").build(), + Row.withSchema(schema).addValues(10L, "cherry").build()); + + IcebergIO.WriteRows write = + IcebergIO.writeRows(catalogConfig) + .to(tableId) + .withSideInputTableCache() + .withTableCachePollingBuckets(1) + .withSortedWrites(); + + testPipeline + .apply("ScrambledInput", Create.of(inputRows)) + .setRowSchema(schema) + .apply("WriteSortedWithSideInput", applyDistribution(write)); + + PipelineResult result = testPipeline.run(); + result.waitUntilFinish(); + + assertEquals(1L, getTablesPolledCount(result)); + + Table table = warehouse.loadTable(tableId); + List writtenRecords = ImmutableList.copyOf(IcebergGenerics.read(table).build()); + assertEquals(4, writtenRecords.size()); + + for (org.apache.iceberg.FileScanTask task : table.newScan().planFiles()) { + String path = task.file().path().toString(); + try (org.apache.iceberg.io.CloseableIterable reader = + org.apache.iceberg.parquet.Parquet.read(table.io().newInputFile(path)) + .project(table.schema()) + .createReaderFunc(org.apache.iceberg.data.parquet.GenericParquetReaders::buildReader) + .build()) { + List records = ImmutableList.copyOf(reader); + assertTrue("File must have at least one record", !records.isEmpty()); + for (int i = 1; i < records.size(); i++) { + Record prev = records.get(i - 1); + Record curr = records.get(i); + String prevName = (String) prev.getField("name"); + String currName = (String) curr.getField("name"); + int cmpName = prevName.compareTo(currName); + if (cmpName > 0) { + throw new AssertionError("File not sorted by name ASC: " + prevName + " > " + currName); + } else if (cmpName == 0) { + long prevId = (Long) prev.getField("id"); + long currId = (Long) curr.getField("id"); + if (prevId < currId) { + throw new AssertionError("File not sorted by id DESC: " + prevId + " < " + currId); + } + } + } + } + } + } + private static class EvolveSpecMidExecutionDoFn extends DoFn { private final IcebergCatalogConfig catalogConfig; private final String tableIdString; 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..d48f29428b32 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 @@ -829,4 +829,298 @@ public void testCreateTableWithPartitionSpecAndSortOrder() { List writtenRecords = ImmutableList.copyOf(IcebergGenerics.read(table).build()); assertThat(writtenRecords, Matchers.containsInAnyOrder(TestFixtures.FILE1SNAPSHOT1.toArray())); } + + @Test + public void testRangeDistribution() throws Exception { + assumeTrue(distributionMode.equals(HASH_WITH_AUTOSHARDING)); + + Schema schema = Schema.builder().addInt64Field("id").addStringField("name").build(); + + TableIdentifier tableId = + TableIdentifier.of("default", "range_" + Long.toString(UUID.randomUUID().hashCode(), 16)); + 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(); + + org.apache.iceberg.Schema icebergSchema = IcebergUtils.beamSchemaToIcebergSchema(schema); + catalog + .catalog() + .buildTable(tableId, icebergSchema) + .withPartitionSpec(PartitionSpec.unpartitioned()) + .withSortOrder(SortOrder.builderFor(icebergSchema).asc("name").desc("id").build()) + .create(); + + PCollection rows = + testPipeline + .apply(GenerateSequence.from(0).to(100)) + .apply( + "Make rows", + MapElements.into(TypeDescriptors.rows()) + .via(i -> Row.withSchema(schema).addValues(i, "name_" + (99 - i)).build())) + .setRowSchema(schema); + + rows.apply( + "range distribution write", + IcebergIO.writeRows(catalog) + .to(tableId) + .withDistributionMode(DistributionMode.RANGE) + .withDistributionFunction(row -> (int) (row.getInt64("id") % 5))); + + testPipeline.run().waitUntilFinish(); + + Table table = warehouse.loadTable(tableId); + List writtenRecords = ImmutableList.copyOf(IcebergGenerics.read(table).build()); + assertEquals(100, writtenRecords.size()); + assertFilesAreInternallySorted(table); + } + + @Test + public void testSortedWrite() { + TableIdentifier tableId = + TableIdentifier.of("default", "sorted_" + Long.toString(UUID.randomUUID().hashCode(), 16)); + + 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(); + + Schema schema = Schema.builder().addInt64Field("id").addStringField("name").build(); + org.apache.iceberg.Schema icebergSchema = IcebergUtils.beamSchemaToIcebergSchema(schema); + + catalog + .catalog() + .buildTable(tableId, icebergSchema) + .withPartitionSpec(PartitionSpec.unpartitioned()) + .withSortOrder(SortOrder.builderFor(icebergSchema).asc("name").desc("id").build()) + .create(); + + List inputRows = + Arrays.asList( + Row.withSchema(schema).addValues(2L, "banana").build(), + Row.withSchema(schema).addValues(1L, "banana").build(), + Row.withSchema(schema).addValues(5L, "apple").build(), + Row.withSchema(schema).addValues(10L, "cherry").build()); + + testPipeline + .apply("Scrambled Input", Create.of(inputRows)) + .setRowSchema(schema) + .apply("Append Sorted To Table", writeTransform(catalog, tableId).withSortedWrites()); + + testPipeline.run().waitUntilFinish(); + + Table table = warehouse.loadTable(tableId); + List writtenRecords = ImmutableList.copyOf(IcebergGenerics.read(table).build()); + + assertEquals(4, writtenRecords.size()); + + try { + assertFilesAreInternallySorted(table); + } catch (Exception e) { + throw new RuntimeException(e); + } + } + + @Test + public void testWritePartitionedRowsToFilesRollsNonOverlappingSortedFiles() throws Exception { + assumeTrue(distributionMode.equals(NONE)); + + TableIdentifier tableId = + TableIdentifier.of( + "default", "rolled_sorted_" + Long.toString(UUID.randomUUID().hashCode(), 16)); + + 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(); + + Schema schema = Schema.builder().addInt64Field("id").addStringField("name").build(); + org.apache.iceberg.Schema icebergSchema = IcebergUtils.beamSchemaToIcebergSchema(schema); + + catalog + .catalog() + .buildTable(tableId, icebergSchema) + .withPartitionSpec(PartitionSpec.unpartitioned()) + .withSortOrder(SortOrder.builderFor(icebergSchema).asc("name").desc("id").build()) + .create(); + + String tableIdString = IcebergUtils.tableIdentifierToString(tableId); + Row partitionKey = + Row.withSchema(AssignDestinationsAndPartitions.OUTPUT_SCHEMA) + .addValues(tableIdString, "") + .build(); + + List scrambledRows = + Arrays.asList( + Row.withSchema(schema).addValues(40L, "delta").build(), + Row.withSchema(schema).addValues(20L, "bravo").build(), + Row.withSchema(schema).addValues(10L, "alpha").build(), + Row.withSchema(schema).addValues(30L, "charlie").build()); + + DynamicDestinations dynamicDestinations = DynamicDestinations.singleTable(tableId, schema); + + // Set maxFileSize = 1 byte so WritePartitionedRowsToFiles rolls after each row into a new file + PCollection writtenFiles = + testPipeline + .apply( + "CreateGroupedPartition", + Create.of(KV.of(partitionKey, (Iterable) scrambledRows)) + .withCoder( + KvCoder.of( + org.apache.beam.sdk.coders.RowCoder.of( + AssignDestinationsAndPartitions.OUTPUT_SCHEMA), + org.apache.beam.sdk.coders.IterableCoder.of( + org.apache.beam.sdk.coders.RowCoder.of(schema))))) + .apply( + "WriteAndRollSortedFiles", + new WritePartitionedRowsToFiles( + catalog, dynamicDestinations, "roll-prefix", 1L, null, null, true)); + + writtenFiles.apply("CommitRolledFiles", new AppendFilesToTables(catalog, "roll-prefix")); + + testPipeline.run().waitUntilFinish(); + + Table table = warehouse.loadTable(tableId); + List tasks = ImmutableList.copyOf(table.newScan().planFiles()); + assertTrue("Expected multiple rolled files, got " + tasks.size(), tasks.size() > 1); + assertFilesAreInternallySorted(table); + + // Collect each file's [minName, maxName] range and verify ranges across files are disjoint + List fileNameRanges = new java.util.ArrayList<>(); + for (org.apache.iceberg.FileScanTask task : tasks) { + String path = task.file().path().toString(); + try (org.apache.iceberg.io.CloseableIterable reader = + Parquet.read(table.io().newInputFile(path)) + .project(table.schema()) + .createReaderFunc(org.apache.iceberg.data.parquet.GenericParquetReaders::buildReader) + .build()) { + List records = Lists.newArrayList(reader.iterator()); + String minName = (String) records.get(0).getField("name"); + String maxName = (String) records.get(records.size() - 1).getField("name"); + fileNameRanges.add(new String[] {minName, maxName}); + } + } + fileNameRanges.sort(java.util.Comparator.comparing(r -> r[0])); + for (int i = 1; i < fileNameRanges.size(); i++) { + String prevMax = fileNameRanges.get(i - 1)[1]; + String currMin = fileNameRanges.get(i)[0]; + assertTrue( + "Rolled files have overlapping sort key ranges: " + prevMax + " >= " + currMin, + prevMax.compareTo(currMin) < 0); + } + } + + @Test + public void testStreamingSortedWriteWithDirectWriteLimit() throws Exception { + assumeTrue(distributionMode.equals(NONE)); + + TableIdentifier tableId = + TableIdentifier.of( + "default", "stream_sorted_" + Long.toString(UUID.randomUUID().hashCode(), 16)); + + 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(); + + Schema schema = Schema.builder().addInt64Field("id").addStringField("name").build(); + org.apache.iceberg.Schema icebergSchema = IcebergUtils.beamSchemaToIcebergSchema(schema); + + catalog + .catalog() + .buildTable(tableId, icebergSchema) + .withPartitionSpec(PartitionSpec.unpartitioned()) + .withSortOrder(SortOrder.builderFor(icebergSchema).asc("name").desc("id").build()) + .create(); + + TestStream stream = + TestStream.create(schema) + .advanceWatermarkTo(new Instant(0)) + .addElements( + Row.withSchema(schema).addValues(2L, "banana").build(), + Row.withSchema(schema).addValues(1L, "banana").build(), + Row.withSchema(schema).addValues(5L, "apple").build(), + Row.withSchema(schema).addValues(10L, "cherry").build()) + .advanceProcessingTime(Duration.standardSeconds(10)) + .advanceWatermarkToInfinity(); + + // Even with directWriteByteLimit = 1 (which would normally lift large bundles to + // WriteDirectRowsToFiles), enabling withSortedWrites() skips bundle lifting so all streamed + // batches are sorted before writing. + testPipeline + .apply(stream) + .apply( + writeTransform(catalog, tableId) + .withTriggeringFrequency(Duration.standardSeconds(5)) + .withDirectWriteByteLimit(1) + .withSortedWrites()); + + testPipeline.run().waitUntilFinish(); + + Table table = warehouse.loadTable(tableId); + List writtenRecords = ImmutableList.copyOf(IcebergGenerics.read(table).build()); + assertEquals(4, writtenRecords.size()); + assertFilesAreInternallySorted(table); + } + + private void assertFilesAreInternallySorted(Table table) throws Exception { + for (org.apache.iceberg.FileScanTask task : table.newScan().planFiles()) { + String path = task.file().path().toString(); + try (org.apache.iceberg.io.CloseableIterable reader = + org.apache.iceberg.parquet.Parquet.read(table.io().newInputFile(path)) + .project(table.schema()) + .createReaderFunc(org.apache.iceberg.data.parquet.GenericParquetReaders::buildReader) + .build()) { + List records = + org.apache.commons.compress.utils.Lists.newArrayList(reader.iterator()); + assertTrue("File must have at least one record", records.size() > 0); + + for (int i = 1; i < records.size(); i++) { + Record prev = records.get(i - 1); + Record curr = records.get(i); + + String prevName = (String) prev.getField("name"); + String currName = (String) curr.getField("name"); + + int cmpName = prevName.compareTo(currName); + if (cmpName > 0) { + throw new AssertionError("File not sorted by name ASC: " + prevName + " > " + currName); + } else if (cmpName == 0) { + long prevId = (Long) prev.getField("id"); + long currId = (Long) curr.getField("id"); + if (prevId < currId) { + throw new AssertionError("File not sorted by id DESC: " + prevId + " < " + currId); + } + } + } + } + } + } } diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergRowSorterTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergRowSorterTest.java new file mode 100644 index 000000000000..1eae37f1e030 --- /dev/null +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergRowSorterTest.java @@ -0,0 +1,432 @@ +/* + * 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.beam.sdk.io.iceberg; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +import java.io.ByteArrayOutputStream; +import java.math.BigDecimal; +import java.time.LocalDate; +import java.time.LocalDateTime; +import java.time.LocalTime; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Comparator; +import java.util.List; +import java.util.Random; +import java.util.stream.Collectors; +import java.util.stream.StreamSupport; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.schemas.logicaltypes.SqlTypes; +import org.apache.beam.sdk.values.Row; +import org.apache.iceberg.NullOrder; +import org.apache.iceberg.SortOrder; +import org.apache.iceberg.SortOrderParser; +import org.apache.iceberg.expressions.Expressions; +import org.apache.iceberg.io.CloseableIterable; +import org.apache.iceberg.types.Type; +import org.apache.iceberg.types.Types; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class IcebergRowSorterTest { + + private static final Schema BEAM_SCHEMA = + Schema.builder() + .addInt32Field("id") + .addNullableField("name", Schema.FieldType.STRING) + .addNullableField("value", Schema.FieldType.DOUBLE) + .addNullableField("active", Schema.FieldType.BOOLEAN) + .build(); + + private static final org.apache.iceberg.Schema ICEBERG_SCHEMA = + new org.apache.iceberg.Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "name", Types.StringType.get()), + Types.NestedField.optional(3, "value", Types.DoubleType.get()), + Types.NestedField.optional(4, "active", Types.BooleanType.get())); + + private static final Comparator BYTE_ARR_COMPARATOR = + org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.primitives.UnsignedBytes + .lexicographicalComparator(); + + private static byte[] encodeSortKeyHelper(Row row, SortOrder sortOrder) throws Exception { + java.util.List fields = sortOrder.fields(); + String[] columnNames = new String[fields.size()]; + for (int i = 0; i < fields.size(); i++) { + columnNames[i] = ICEBERG_SCHEMA.findColumnName(fields.get(i).sourceId()); + } + ByteArrayOutputStream baos = new ByteArrayOutputStream(); + IcebergRowSorter.encodeSortKey(row, sortOrder, columnNames, baos, ICEBERG_SCHEMA, BEAM_SCHEMA); + return baos.toByteArray(); + } + + @Test + public void testStringKeyEncodingOrder() throws Exception { + SortOrder sortOrder = SortOrder.builderFor(ICEBERG_SCHEMA).asc("name").build(); + + Row r1 = Row.withSchema(BEAM_SCHEMA).addValues(1, "apple", 1.5, true).build(); + Row r2 = Row.withSchema(BEAM_SCHEMA).addValues(2, "banana", 2.0, true).build(); + Row r3 = Row.withSchema(BEAM_SCHEMA).addValues(3, "apricot", 3.0, false).build(); + + byte[] k1 = encodeSortKeyHelper(r1, sortOrder); + byte[] k2 = encodeSortKeyHelper(r2, sortOrder); + byte[] k3 = encodeSortKeyHelper(r3, sortOrder); + + assertTrue(BYTE_ARR_COMPARATOR.compare(k1, k2) < 0); // apple < banana + assertTrue(BYTE_ARR_COMPARATOR.compare(k1, k3) < 0); // apple < apricot + assertTrue(BYTE_ARR_COMPARATOR.compare(k3, k2) < 0); // apricot < banana + } + + @Test + public void testStringCollisionProofing() throws Exception { + SortOrder sortOrder = SortOrder.builderFor(ICEBERG_SCHEMA).asc("name").asc("value").build(); + + Row r1 = Row.withSchema(BEAM_SCHEMA).addValues(1, "abc", 1.0, true).build(); + Row r2 = Row.withSchema(BEAM_SCHEMA).addValues(2, "abcdef", null, true).build(); + + byte[] k1 = encodeSortKeyHelper(r1, sortOrder); + byte[] k2 = encodeSortKeyHelper(r2, sortOrder); + + assertTrue(BYTE_ARR_COMPARATOR.compare(k1, k2) < 0); + } + + @Test + public void testDescInversion() throws Exception { + SortOrder sortOrderAsc = SortOrder.builderFor(ICEBERG_SCHEMA).asc("id").build(); + SortOrder sortOrderDesc = SortOrder.builderFor(ICEBERG_SCHEMA).desc("id").build(); + + Row r1 = Row.withSchema(BEAM_SCHEMA).addValues(10, "test", 1.5, true).build(); + Row r2 = Row.withSchema(BEAM_SCHEMA).addValues(20, "test", 2.0, true).build(); + + byte[] k1Asc = encodeSortKeyHelper(r1, sortOrderAsc); + byte[] k2Asc = encodeSortKeyHelper(r2, sortOrderAsc); + + byte[] k1Desc = encodeSortKeyHelper(r1, sortOrderDesc); + byte[] k2Desc = encodeSortKeyHelper(r2, sortOrderDesc); + + assertTrue(BYTE_ARR_COMPARATOR.compare(k1Asc, k2Asc) < 0); + assertTrue(BYTE_ARR_COMPARATOR.compare(k1Desc, k2Desc) > 0); + } + + @Test + public void testNullOrderingMatrix() throws Exception { + Row rNonNull = Row.withSchema(BEAM_SCHEMA).addValues(1, "apple", 1.5, true).build(); + Row rNull = Row.withSchema(BEAM_SCHEMA).addValues(2, null, 2.0, true).build(); + + // 1. ASC, NULLS_FIRST + SortOrder ascFirst = + SortOrder.builderFor(ICEBERG_SCHEMA).asc("name", NullOrder.NULLS_FIRST).build(); + byte[] kNonNullAscFirst = encodeSortKeyHelper(rNonNull, ascFirst); + byte[] kNullAscFirst = encodeSortKeyHelper(rNull, ascFirst); + assertTrue( + "ASC NULLS_FIRST failed: null should sort before non-null", + BYTE_ARR_COMPARATOR.compare(kNullAscFirst, kNonNullAscFirst) < 0); + + // 2. ASC, NULLS_LAST + SortOrder ascLast = + SortOrder.builderFor(ICEBERG_SCHEMA).asc("name", NullOrder.NULLS_LAST).build(); + byte[] kNonNullAscLast = encodeSortKeyHelper(rNonNull, ascLast); + byte[] kNullAscLast = encodeSortKeyHelper(rNull, ascLast); + assertTrue( + "ASC NULLS_LAST failed: null should sort after non-null", + BYTE_ARR_COMPARATOR.compare(kNullAscLast, kNonNullAscLast) > 0); + + // 3. DESC, NULLS_FIRST + SortOrder descFirst = + SortOrder.builderFor(ICEBERG_SCHEMA).desc("name", NullOrder.NULLS_FIRST).build(); + byte[] kNonNullDescFirst = encodeSortKeyHelper(rNonNull, descFirst); + byte[] kNullDescFirst = encodeSortKeyHelper(rNull, descFirst); + assertTrue( + "DESC NULLS_FIRST failed: null should sort before non-null", + BYTE_ARR_COMPARATOR.compare(kNullDescFirst, kNonNullDescFirst) < 0); + + // 4. DESC, NULLS_LAST + SortOrder descLast = + SortOrder.builderFor(ICEBERG_SCHEMA).desc("name", NullOrder.NULLS_LAST).build(); + byte[] kNonNullDescLast = encodeSortKeyHelper(rNonNull, descLast); + byte[] kNullDescLast = encodeSortKeyHelper(rNull, descLast); + assertTrue( + "DESC NULLS_LAST failed: null should sort after non-null", + BYTE_ARR_COMPARATOR.compare(kNullDescLast, kNonNullDescLast) > 0); + } + + @Test + public void testEndToEndSorting() { + SortOrder sortOrder = SortOrder.builderFor(ICEBERG_SCHEMA).asc("name").desc("id").build(); + + List input = + Arrays.asList( + Row.withSchema(BEAM_SCHEMA).addValues(2, "banana", 2.0, true).build(), + Row.withSchema(BEAM_SCHEMA).addValues(1, "banana", 1.0, true).build(), + Row.withSchema(BEAM_SCHEMA).addValues(5, "apple", 1.5, true).build(), + Row.withSchema(BEAM_SCHEMA).addValues(10, "cherry", 3.0, false).build()); + + Iterable sorted = IcebergRowSorter.sortRows(input, sortOrder, ICEBERG_SCHEMA, BEAM_SCHEMA); + List sortedList = + StreamSupport.stream(sorted.spliterator(), false).collect(Collectors.toList()); + + assertEquals(4, sortedList.size()); + + assertEquals("apple", sortedList.get(0).getString("name")); + assertEquals(Integer.valueOf(5), sortedList.get(0).getInt32("id")); + + assertEquals("banana", sortedList.get(1).getString("name")); + assertEquals(Integer.valueOf(2), sortedList.get(1).getInt32("id")); + + assertEquals("banana", sortedList.get(2).getString("name")); + assertEquals(Integer.valueOf(1), sortedList.get(2).getInt32("id")); + + assertEquals("cherry", sortedList.get(3).getString("name")); + assertEquals(Integer.valueOf(10), sortedList.get(3).getInt32("id")); + } + + @Test + public void testScaleAndExternalDiskSpill() { + SortOrder sortOrder = SortOrder.builderFor(ICEBERG_SCHEMA).asc("id").build(); + + int count = 5000; + List input = new ArrayList<>(count); + Random rand = new Random(42); + + for (int i = 0; i < count; i++) { + int randomId = rand.nextInt(100_000); + input.add(Row.withSchema(BEAM_SCHEMA).addValues(randomId, "item" + i, 1.0, true).build()); + } + + Iterable sorted = IcebergRowSorter.sortRows(input, sortOrder, ICEBERG_SCHEMA, BEAM_SCHEMA); + List sortedList = + StreamSupport.stream(sorted.spliterator(), false).collect(Collectors.toList()); + + assertEquals(count, sortedList.size()); + + for (int i = 0; i < sortedList.size() - 1; i++) { + int idCurrent = sortedList.get(i).getInt32("id"); + int idNext = sortedList.get(i + 1).getInt32("id"); + assertTrue( + String.format("Sort violation at index %d: %d > %d", i, idCurrent, idNext), + idCurrent <= idNext); + } + } + + @Test + public void testUnsupportedComplexTypeSorting() { + org.apache.iceberg.Schema mapSchema = + new org.apache.iceberg.Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional( + 2, + "attributes", + Types.MapType.ofOptional(3, 4, Types.StringType.get(), Types.StringType.get()))); + + assertThrows( + org.apache.iceberg.exceptions.ValidationException.class, + () -> SortOrder.builderFor(mapSchema).asc("attributes").build()); + } + + /** Encodes a single-column row so each sortable primitive can be exercised in isolation. */ + private static byte[] encodeSingle( + Schema.FieldType beamType, Type icebergType, Object value, boolean desc) throws Exception { + Schema beamSchema = Schema.of(Schema.Field.nullable("v", beamType)); + org.apache.iceberg.Schema icebergSchema = + new org.apache.iceberg.Schema(Types.NestedField.optional(1, "v", icebergType)); + SortOrder sortOrder = + desc + ? SortOrder.builderFor(icebergSchema).desc("v").build() + : SortOrder.builderFor(icebergSchema).asc("v").build(); + + Row row = Row.withSchema(beamSchema).addValue(value).build(); + ByteArrayOutputStream baos = new ByteArrayOutputStream(); + IcebergRowSorter.encodeSortKey( + row, sortOrder, new String[] {"v"}, baos, icebergSchema, beamSchema); + return baos.toByteArray(); + } + + /** Asserts {@code smaller} sorts before {@code larger} ASC, and after it DESC. */ + private static void assertOrdered( + Schema.FieldType beamType, Type icebergType, Object smaller, Object larger) throws Exception { + assertTrue( + "ASC ordering failed for " + icebergType + ": " + smaller + " should precede " + larger, + BYTE_ARR_COMPARATOR.compare( + encodeSingle(beamType, icebergType, smaller, false), + encodeSingle(beamType, icebergType, larger, false)) + < 0); + assertTrue( + "DESC ordering failed for " + icebergType + ": " + smaller + " should follow " + larger, + BYTE_ARR_COMPARATOR.compare( + encodeSingle(beamType, icebergType, smaller, true), + encodeSingle(beamType, icebergType, larger, true)) + > 0); + } + + @Test + public void testDecimalKeyOrdering() throws Exception { + Type decimal = Types.DecimalType.of(12, 2); + assertOrdered( + Schema.FieldType.DECIMAL, decimal, new BigDecimal("-9999999.99"), new BigDecimal("-10.50")); + assertOrdered( + Schema.FieldType.DECIMAL, decimal, new BigDecimal("-10.50"), new BigDecimal("-0.01")); + assertOrdered( + Schema.FieldType.DECIMAL, decimal, new BigDecimal("-0.01"), new BigDecimal("0.00")); + assertOrdered( + Schema.FieldType.DECIMAL, decimal, new BigDecimal("0.00"), new BigDecimal("3.14")); + assertOrdered( + Schema.FieldType.DECIMAL, decimal, new BigDecimal("3.14"), new BigDecimal("9999999.99")); + } + + @Test + public void testDateAndTimeKeyOrdering() throws Exception { + assertOrdered( + Schema.FieldType.logicalType(SqlTypes.DATE), + Types.DateType.get(), + LocalDate.parse("1969-12-31"), + LocalDate.parse("2024-02-29")); + assertOrdered( + Schema.FieldType.logicalType(SqlTypes.TIME), + Types.TimeType.get(), + LocalTime.parse("00:00:01"), + LocalTime.parse("23:59:59.999999")); + } + + @Test + public void testTimestampKeyOrdering() throws Exception { + // timestamp without zone resolves to a LocalDateTime. + assertOrdered( + Schema.FieldType.logicalType(SqlTypes.DATETIME), + Types.TimestampType.withoutZone(), + LocalDateTime.parse("1969-07-20T20:17:40"), + LocalDateTime.parse("2031-06-05T12:00:00")); + // timestamptz resolves to an OffsetDateTime. + assertOrdered( + Schema.FieldType.DATETIME, + Types.TimestampType.withZone(), + org.joda.time.Instant.parse("1969-07-20T20:17:40Z"), + org.joda.time.Instant.parse("2031-06-05T12:00:00Z")); + } + + @Test + public void testNarrowIntegerKeyOrdering() throws Exception { + assertOrdered(Schema.FieldType.INT16, Types.IntegerType.get(), (short) -300, (short) 300); + assertOrdered(Schema.FieldType.BYTE, Types.IntegerType.get(), (byte) -5, (byte) 5); + } + + @Test + public void testUuidKeyEncodingIsDeterministic() throws Exception { + byte[] first = + encodeSingle(Schema.FieldType.BYTES, Types.UUIDType.get(), new byte[] {1, 2}, false); + byte[] repeat = + encodeSingle(Schema.FieldType.BYTES, Types.UUIDType.get(), new byte[] {1, 2}, false); + byte[] other = + encodeSingle(Schema.FieldType.BYTES, Types.UUIDType.get(), new byte[] {3, 4}, false); + + assertEquals(0, BYTE_ARR_COMPARATOR.compare(first, repeat)); + assertTrue(BYTE_ARR_COMPARATOR.compare(first, other) != 0); + } + + @Test + public void testDecimalRejectsValueExceedingDeclaredScale() { + assertThrows( + IllegalArgumentException.class, + () -> + encodeSingle( + Schema.FieldType.DECIMAL, + Types.DecimalType.of(12, 2), + new BigDecimal("1.234"), + false)); + } + + @Test + public void testNonIdentityTransformSortOrder() throws Exception { + Schema beamSchema = + Schema.builder() + .addInt32Field("id") + .addLogicalTypeField("d", SqlTypes.DATE) + .addLogicalTypeField("ts", SqlTypes.DATETIME) + .addDecimalField("dec") + .addStringField("s") + .build(); + org.apache.iceberg.Schema icebergSchema = + new org.apache.iceberg.Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.required(2, "d", Types.DateType.get()), + Types.NestedField.required(3, "ts", Types.TimestampType.withoutZone()), + Types.NestedField.required(4, "dec", Types.DecimalType.of(12, 2)), + Types.NestedField.required(5, "s", Types.StringType.get())); + + // Sort by month(d) ASC, hour(ts) ASC, truncate(dec, 100) ASC, truncate(s, 3) ASC, id ASC. + // Exercise both a directly built SortOrder and a JSON-deserialized SortOrder (which produces + // unbound transforms until .bind(sourceType) is invoked). + SortOrder builtOrder = + SortOrder.builderFor(icebergSchema) + .asc(Expressions.month("d")) + .asc(Expressions.hour("ts")) + .asc(Expressions.truncate("dec", 100)) + .asc(Expressions.truncate("s", 3)) + .asc("id") + .build(); + SortOrder jsonOrder = + SortOrderParser.fromJson(icebergSchema, SortOrderParser.toJson(builtOrder)); + + // r1 and r2 share the same month (2024-03), same hour (10:xx), same truncated decimal + // (1.99 and 1.05 both truncate to 1.00 with width 100 unscaled units), and same 3-char prefix + // ("abc..."), so they must tie on all 4 transform keys and order by id ASC (r2 before r1). + // r3 is in month 2024-01, so it must sort before both r1 and r2. + Row r1 = + Row.withSchema(beamSchema) + .addValues( + 20, + LocalDate.parse("2024-03-25"), + LocalDateTime.parse("2024-03-25T10:55:00"), + new BigDecimal("1.99"), + "abc_zzz") + .build(); + Row r2 = + Row.withSchema(beamSchema) + .addValues( + 10, + LocalDate.parse("2024-03-02"), + LocalDateTime.parse("2024-03-02T10:05:00"), + new BigDecimal("1.05"), + "abc_aaa") + .build(); + Row r3 = + Row.withSchema(beamSchema) + .addValues( + 30, + LocalDate.parse("2024-01-31"), + LocalDateTime.parse("2024-01-31T23:00:00"), + new BigDecimal("99.99"), + "zzz") + .build(); + + for (SortOrder order : Arrays.asList(builtOrder, jsonOrder)) { + try (CloseableIterable sorted = + IcebergRowSorter.sortRows(Arrays.asList(r1, r2, r3), order, icebergSchema, beamSchema)) { + List sortedList = + StreamSupport.stream(sorted.spliterator(), false).collect(Collectors.toList()); + assertEquals( + Arrays.asList(30, 10, 20), + sortedList.stream().map(r -> r.getInt32("id")).collect(Collectors.toList())); + } + } + } +} diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergUtilsTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergUtilsTest.java index aa8f6799a3eb..6bcdc3ff9049 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergUtilsTest.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergUtilsTest.java @@ -356,6 +356,19 @@ public void testMapOfRecords() { IcebergUtils.beamRowToIcebergRecord(RECORD_MAP_ICEBERG_SCHEMA, ROW_MAP_OF_ROWS); assertEquals(RECORD_MAP_OF_RECORDS, actual); } + + @Test + public void testNullStringConversion() { + Schema beamSchema = + Schema.of(Schema.Field.of("v", Schema.FieldType.STRING).withNullable(true)); + Row row = Row.withSchema(beamSchema).addValue(null).build(); + + org.apache.iceberg.Schema icebergSchema = + new org.apache.iceberg.Schema(optional(0, "v", Types.StringType.get())); + Record record = IcebergUtils.beamRowToIcebergRecord(icebergSchema, row); + + assertEquals(null, record.getField("v")); + } } @RunWith(JUnit4.class) diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergWriteSchemaTransformProviderTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergWriteSchemaTransformProviderTest.java index b408445053f2..750e680a40a7 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergWriteSchemaTransformProviderTest.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergWriteSchemaTransformProviderTest.java @@ -941,4 +941,58 @@ public void testDynamicWriteCreateTableWithTableProperties() { List writtenRecords = ImmutableList.copyOf(IcebergGenerics.read(table).build()); assertEquals(10, writtenRecords.size()); } + + @Test + public void testWriteWithSortOnWriteEnabled() throws Exception { + String identifier = "default.sorted_table_" + Long.toString(UUID.randomUUID().hashCode(), 16); + Schema schema = Schema.builder().addStringField("str").addInt32Field("int").build(); + + Map config = + ImmutableMap.builder() + .put("table", identifier) + .put( + "catalog_properties", + ImmutableMap.of("type", "hadoop", "warehouse", warehouse.location)) + .put("distribution_mode", distributionMode.name()) + .put("sort_fields", Collections.singletonList("str desc")) + .put("sort_on_write", true) + .build(); + + List rows = + Arrays.asList( + Row.withSchema(schema).addValues("apple", 1).build(), + Row.withSchema(schema).addValues("cherry", 3).build(), + Row.withSchema(schema).addValues("banana", 2).build()); + + PCollection result = + testPipeline + .apply("Records To Add", Create.of(rows)) + .setRowSchema(schema) + .apply(Managed.write(Managed.ICEBERG).withConfig(config)) + .get(SNAPSHOTS_TAG); + + PAssert.that(result) + .satisfies(new VerifyOutputs(Collections.singletonList(identifier), "append")); + testPipeline.run().waitUntilFinish(); + + Table table = warehouse.loadTable(TableIdentifier.parse(identifier)); + assertTrue(table.sortOrder().isSorted()); + for (org.apache.iceberg.FileScanTask task : table.newScan().planFiles()) { + String path = task.file().path().toString(); + try (org.apache.iceberg.io.CloseableIterable reader = + org.apache.iceberg.parquet.Parquet.read(table.io().newInputFile(path)) + .project(table.schema()) + .createReaderFunc(org.apache.iceberg.data.parquet.GenericParquetReaders::buildReader) + .build()) { + List records = + org.apache.commons.compress.utils.Lists.newArrayList(reader.iterator()); + for (int i = 1; i < records.size(); i++) { + String prev = (String) records.get(i - 1).getField("str"); + String curr = (String) records.get(i).getField("str"); + assertTrue( + "Expected descending sort order: " + prev + " >= " + curr, prev.compareTo(curr) >= 0); + } + } + } + } } From 7230247c84a99072b4be16523f72337c4d736e24 Mon Sep 17 00:00:00 2001 From: atognolag Date: Fri, 16 Oct 2026 14:10:33 +0000 Subject: [PATCH 2/3] Support zero-shuffle bundle sorting in NONE mode and dynamic quantile sharding in RANGE mode --- CHANGES.md | 3 +- .../AssignDestinationsAndPartitions.java | 814 +++++++++++++++--- .../beam/sdk/io/iceberg/BeamRowWrapper.java | 10 +- .../apache/beam/sdk/io/iceberg/IcebergIO.java | 153 ++-- .../beam/sdk/io/iceberg/IcebergRowSorter.java | 313 ++++--- .../beam/sdk/io/iceberg/IcebergUtils.java | 48 +- .../IcebergWriteSchemaTransformProvider.java | 8 +- .../sdk/io/iceberg/RecordWriterManager.java | 167 +++- .../io/iceberg/WriteDirectRowsToFiles.java | 93 +- .../io/iceberg/WriteGroupedRowsToFiles.java | 65 +- .../iceberg/WritePartitionedRowsToFiles.java | 128 +-- .../sdk/io/iceberg/WriteToDestinations.java | 47 +- .../io/iceberg/WriteUngroupedRowsToFiles.java | 101 +-- .../AssignDestinationsAndPartitionsTest.java | 47 + .../IcebergIOSideInputTableCacheTest.java | 2 +- .../sdk/io/iceberg/IcebergIOWriteTest.java | 162 +++- .../sdk/io/iceberg/IcebergRowSorterTest.java | 55 +- .../beam/sdk/io/iceberg/IcebergUtilsTest.java | 26 + ...ebergWriteSchemaTransformProviderTest.java | 5 +- .../en/documentation/io/built-in/iceberg.md | 49 ++ 20 files changed, 1639 insertions(+), 657 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index 601ae5f41047..d997b9fe1100 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -67,7 +67,7 @@ * Support for X source added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). * (Java) FirestoreIO: added `RpcQosOptions.Builder.withRampupThrottlingDisabled()`, mirroring the existing DatastoreIO option, to opt out of the 500/50/5 write ramp-up ([#27170](https://github.com/apache/beam/issues/27170)). * Added support for attributing BigQuery API quota and billing to a specific GCP project (quota project): `quota_project_id` parameter in `ReadFromBigQuery` or `--quota_project_id` pipeline option, honored by the Python BigQuery IO transforms (`WriteToBigQuery` with `FILE_LOADS` or `STREAMING_INSERTS`; not yet `STORAGE_WRITE_API`, which is Java-backed), `--bigQueryQuotaProjectId` pipeline option (Java), and `bigqueryio.WithQuotaProject` read/query option (Go) ([#37431](https://github.com/apache/beam/issues/37431)). -* [IcebergIO] Added opt-in sorted write support (`withSortedWrites()` / `sort_on_write`) to write data files in the table's sort order, spilling to local disk via Beam's external sorter for groups that exceed memory. Also added `DistributionMode.RANGE` with a user-supplied `withDistributionFunction(...)`, which shards a partition into non-overlapping key ranges so each shard becomes its own sorted file(s) (Java) ([#38406](https://github.com/apache/beam/issues/38406)). +* [IcebergIO] Added opt-in sorted write support (`withSortedWrites()` / `sort_on_write`) across `NONE`, `HASH`, and `RANGE` distribution modes to write data files in the table's sort order, spilling to local disk via Beam's external sorter when groups exceed memory. Also added `DistributionMode.RANGE`, which dynamically computes non-overlapping quantile key ranges per partition (or accepts a custom `withDistributionFunction(...)`) so each shard writes sorted, non-overlapping file(s) (Java) ([#38406](https://github.com/apache/beam/issues/38406)). ## New Features / Improvements @@ -136,6 +136,7 @@ * BigQueryIO now supports reading BigQuery Lakehouse runtime catalog (BigLake metastore) Iceberg tables with the Storage Read API, using 4-part `project.catalog.namespace.table` identifiers (or a `TableReference` with a composite `catalog.namespace` dataset id). Previously such references were silently mis-parsed (Java) ([#39597](https://github.com/apache/beam/issues/39597)) . * SolaceIO now supports reading and writing binary and text content data payload (Java) ([#39875](https://github.com/apache/beam/issues/39875)). * ClickHouseIO: support writing `Decimal(P, S)` / `Decimal32/64/128/256` columns (Java) ([#39840](https://github.com/apache/beam/issues/39840)). +* SolaceIO now supports reading and writing user properties (message metadata) (Java) ([#40099](https://github.com/apache/beam/issues/40099)). * [IcebergIO] AddFiles (`IcebergAddFiles` in YAML) can evolve the table schema before registering files, with `schema_evolution_options`, `required_columns`, `incompatible_schema_handling` and `unverifiable_file_handling` (Java/YAML, batch only) ([#40144](https://github.com/apache/beam/issues/40144)). * [IcebergIO] Added batch and streaming CDC writes that applies INSERT/UPDATE_BEFORE/UPDATE_AFTER/DELETE changes to Iceberg V2+ tables by primary key. Invoke with `IcebergIO.writeCdcRows` (Java) or by setting `mode: merge-on-read` on the Managed `ICEBERG` write (Java, Python, YAML) ([#39979](https://github.com/apache/beam/issues/39979)). * [IcebergIO] Added an optional side-input table cache for writes to significantly reduce catalog and table requests for large pipelines. A single worker polls the table and broadcasts it to other workers in the pipeline. Enable with `IcebergIO.writeRows(...).withSideInputTableCache()` (Java) or by setting `use_side_input_table_cache: true` on the Managed `ICEBERG` write (Java, Python) ([#39723](https://github.com/apache/beam/issues/39723)). diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/AssignDestinationsAndPartitions.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/AssignDestinationsAndPartitions.java index fa664e7c4aa1..1560de82f7f4 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/AssignDestinationsAndPartitions.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/AssignDestinationsAndPartitions.java @@ -19,25 +19,52 @@ import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.io.Serializable; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.Comparator; import java.util.HashMap; +import java.util.List; import java.util.Map; +import java.util.Objects; +import java.util.concurrent.ThreadLocalRandom; +import org.apache.beam.sdk.coders.ByteArrayCoder; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.CoderException; +import org.apache.beam.sdk.coders.CoderRegistry; import org.apache.beam.sdk.coders.KvCoder; +import org.apache.beam.sdk.coders.ListCoder; import org.apache.beam.sdk.coders.RowCoder; +import org.apache.beam.sdk.coders.StructuredCoder; +import org.apache.beam.sdk.coders.VarLongCoder; +import org.apache.beam.sdk.transforms.Combine; 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.SerializableFunction; +import org.apache.beam.sdk.transforms.View; 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.PCollectionTuple; import org.apache.beam.sdk.values.PCollectionView; import org.apache.beam.sdk.values.Row; +import org.apache.beam.sdk.values.TupleTag; +import org.apache.beam.sdk.values.TupleTagList; import org.apache.beam.sdk.values.ValueInSingleWindow; +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.primitives.UnsignedBytes; import org.apache.iceberg.DistributionMode; import org.apache.iceberg.PartitionKey; import org.apache.iceberg.PartitionSpec; import org.apache.iceberg.Schema; +import org.apache.iceberg.SortOrder; import org.apache.iceberg.Table; import org.apache.iceberg.exceptions.NoSuchTableException; import org.checkerframework.checker.nullness.qual.MonotonicNonNull; @@ -46,9 +73,11 @@ import org.joda.time.Instant; /** - * Assigns destination metadata for each input record. + * Assigns destination and partition metadata for each input record. * - *

The output will have the format { {destination, partition}, data } + *

The output has the format {@code { {destination, partition}, data }} for {@link + * DistributionMode#HASH} and {@code { {destination, partition, shard}, data }} for {@link + * DistributionMode#RANGE}. */ class AssignDestinationsAndPartitions extends PTransform, PCollection>> { @@ -58,11 +87,26 @@ class AssignDestinationsAndPartitions private final @Nullable PCollectionView> metadataView; private final DistributionMode distributionMode; private final @Nullable SerializableFunction distributionFunction; + private final long targetBytesPerRangeShard; static final String DESTINATION = "destination"; static final String PARTITION = "partition"; static final String SHARD = "shard"; + /** Maximum number of sort-key samples retained per partition within a single bundle. */ + static final int MAX_SAMPLES_PER_BUNDLE = 200; + + /** Maximum number of sort-key samples retained per partition across combined bundles. */ + static final int MAX_COMBINED_SAMPLES = 2048; + + /** Maximum number of dynamic range shards created for a single partition. */ + static final int MAX_RANGE_SHARDS = 1000; + + private static final TupleTag> PARTITIONED_ROWS_TAG = + new TupleTag>() {}; + private static final TupleTag> PARTITION_SAMPLES_TAG = + new TupleTag>() {}; + /** * Grouping key used for {@link DistributionMode#NONE} and {@link DistributionMode#HASH} writes. * @@ -108,26 +152,42 @@ public AssignDestinationsAndPartitions( public AssignDestinationsAndPartitions( DynamicDestinations dynamicDestinations, IcebergCatalogConfig catalogConfig, + @Nullable PCollectionView> metadataView, DistributionMode distributionMode, @Nullable SerializableFunction distributionFunction) { - this(dynamicDestinations, catalogConfig, null, distributionMode, distributionFunction); + this( + dynamicDestinations, + catalogConfig, + metadataView, + distributionMode, + distributionFunction, + WriteToPartitions.DEFAULT_BYTES_PER_FILE); } - public AssignDestinationsAndPartitions( + @VisibleForTesting + AssignDestinationsAndPartitions( DynamicDestinations dynamicDestinations, IcebergCatalogConfig catalogConfig, @Nullable PCollectionView> metadataView, DistributionMode distributionMode, - @Nullable SerializableFunction distributionFunction) { + @Nullable SerializableFunction distributionFunction, + long targetBytesPerRangeShard) { this.dynamicDestinations = dynamicDestinations; this.catalogConfig = catalogConfig; this.metadataView = metadataView; this.distributionMode = distributionMode; this.distributionFunction = distributionFunction; + this.targetBytesPerRangeShard = targetBytesPerRangeShard; } @Override public PCollection> expand(PCollection input) { + if (distributionMode == DistributionMode.RANGE + && distributionFunction == null + && !IcebergUtils.isUnbounded(input)) { + return expandWithDynamicRangeShards(input); + } + ParDo.SingleOutput> parDo = ParDo.of( new AssignDoFn( @@ -147,74 +207,144 @@ public PCollection> expand(PCollection input) { RowCoder.of(dynamicDestinations.getDataSchema()))); } - static class AssignDoFn extends DoFn> { + private PCollection> expandWithDynamicRangeShards(PCollection input) { + RowCoder partitionKeyCoder = RowCoder.of(OUTPUT_SCHEMA); + RowCoder rangeKeyCoder = RowCoder.of(OUTPUT_SCHEMA_WITH_SHARD); + RowCoder dataCoder = RowCoder.of(dynamicDestinations.getDataSchema()); - private static final Duration REFRESH_INTERVAL = Duration.standardMinutes(5); + ParDo.MultiOutput> extractParDo = + ParDo.of( + new ExtractPartitionAndSampleSortKeyDoFn( + dynamicDestinations, catalogConfig, metadataView)) + .withOutputTags(PARTITIONED_ROWS_TAG, TupleTagList.of(PARTITION_SAMPLES_TAG)); + if (metadataView != null) { + extractParDo = extractParDo.withSideInputs(metadataView); + } - private transient @MonotonicNonNull Map partitionKeys; - private transient @MonotonicNonNull Map wrappers; - private transient @MonotonicNonNull Map lastRefreshTimes; - private transient @MonotonicNonNull Map cachedSpecIds; + PCollectionTuple extracted = input.apply("ExtractPartitionAndSampleSortKey", extractParDo); - private final DynamicDestinations dynamicDestinations; - private final IcebergCatalogConfig catalogConfig; - private final @Nullable PCollectionView> metadataView; - private final DistributionMode distributionMode; - private final @Nullable SerializableFunction distributionFunction; + PCollection> partitionedRows = + extracted.get(PARTITIONED_ROWS_TAG).setCoder(KvCoder.of(partitionKeyCoder, dataCoder)); + + PCollection> partitionSamples = + extracted + .get(PARTITION_SAMPLES_TAG) + .setCoder(KvCoder.of(partitionKeyCoder, PartitionSampleCoder.of())); + + PCollectionView>> splitPointsView = + partitionSamples + .apply( + "ComputePartitionQuantileSplits", + Combine.perKey(new PartitionQuantilesCombineFn(targetBytesPerRangeShard))) + .setCoder(KvCoder.of(partitionKeyCoder, ListCoder.of(ByteArrayCoder.of()))) + .apply("SplitsAsMapView", View.asMap()); + + ParDo.SingleOutput, KV> assignShardParDo = + ParDo.of( + new AssignDynamicRangeShardDoFn( + dynamicDestinations, catalogConfig, metadataView, splitPointsView)) + .withSideInputs( + metadataView != null + ? Arrays.asList(splitPointsView, metadataView) + : Collections.singletonList(splitPointsView)); + + return partitionedRows + .apply("AssignDynamicRangeShard", assignShardParDo) + .setCoder(KvCoder.of(rangeKeyCoder, dataCoder)); + } - AssignDoFn(DynamicDestinations dynamicDestinations, IcebergCatalogConfig catalogConfig) { - this(dynamicDestinations, catalogConfig, null, DistributionMode.HASH, null); + /** + * Maps a {@code sortKey} to a shard index {@code 0 .. splits.size()} using half-open intervals + * {@code (-inf, splits[0]), [splits[0], splits[1]), ..., [splits[k-1], +inf)}. + */ + @VisibleForTesting + static int findRangeShard(@Nullable List splits, byte[] sortKey) { + if (splits == null || splits.isEmpty()) { + return 0; } + int pos = Collections.binarySearch(splits, sortKey, UnsignedBytes.lexicographicalComparator()); + return pos >= 0 ? (pos + 1) : -(pos + 1); + } - AssignDoFn( - DynamicDestinations dynamicDestinations, - IcebergCatalogConfig catalogConfig, - @Nullable PCollectionView> metadataView) { - this(dynamicDestinations, catalogConfig, metadataView, DistributionMode.HASH, null); + /** + * Resolved per-destination partition and sort-key metadata cached on the worker across bundles. + */ + private static class ResolvedDestination { + final String tableIdentifier; + final Row data; + final String partitionPath; + final IcebergRowSorter.@Nullable KeyEncoder keyEncoder; + + ResolvedDestination( + String tableIdentifier, + Row data, + String partitionPath, + IcebergRowSorter.@Nullable KeyEncoder keyEncoder) { + this.tableIdentifier = tableIdentifier; + this.data = data; + this.partitionPath = partitionPath; + this.keyEncoder = keyEncoder; } + } - AssignDoFn( - DynamicDestinations dynamicDestinations, - IcebergCatalogConfig catalogConfig, - DistributionMode distributionMode, - @Nullable SerializableFunction distributionFunction) { - this(dynamicDestinations, catalogConfig, null, distributionMode, distributionFunction); + private static class CachedTableMetadata { + final PartitionKey partitionKey; + final BeamRowWrapper wrapper; + final IcebergRowSorter.@Nullable KeyEncoder keyEncoder; + final Instant lastRefresh; + final @Nullable Integer specId; + final @Nullable Integer orderId; + + CachedTableMetadata( + PartitionKey partitionKey, + BeamRowWrapper wrapper, + IcebergRowSorter.@Nullable KeyEncoder keyEncoder, + Instant lastRefresh, + @Nullable Integer specId, + @Nullable Integer orderId) { + this.partitionKey = partitionKey; + this.wrapper = wrapper; + this.keyEncoder = keyEncoder; + this.lastRefresh = lastRefresh; + this.specId = specId; + this.orderId = orderId; } + } - AssignDoFn( + /** + * Shared helper that resolves a destination table's {@link PartitionSpec}, {@link SortOrder}, and + * {@link Schema} from the side-input cache, {@link TableCache}, or {@link + * IcebergTableCreateConfig}, and computes the record's partition path and optional {@link + * IcebergRowSorter.KeyEncoder}. + */ + private static class DestinationMetadataResolver implements Serializable { + private static final Duration REFRESH_INTERVAL = Duration.standardMinutes(5); + + private final DynamicDestinations dynamicDestinations; + private final IcebergCatalogConfig catalogConfig; + private final @Nullable PCollectionView> metadataView; + private final boolean includeSortKeyEncoder; + private transient @MonotonicNonNull Map cache; + + DestinationMetadataResolver( DynamicDestinations dynamicDestinations, IcebergCatalogConfig catalogConfig, @Nullable PCollectionView> metadataView, - DistributionMode distributionMode, - @Nullable SerializableFunction distributionFunction) { + boolean includeSortKeyEncoder) { this.dynamicDestinations = dynamicDestinations; this.catalogConfig = catalogConfig; this.metadataView = metadataView; - this.distributionMode = distributionMode; - this.distributionFunction = distributionFunction; + this.includeSortKeyEncoder = includeSortKeyEncoder; } - @Setup - public void setup() { - this.wrappers = new HashMap<>(); - this.partitionKeys = new HashMap<>(); - this.lastRefreshTimes = new HashMap<>(); - this.cachedSpecIds = new HashMap<>(); + void setup() { + this.cache = new HashMap<>(); } - @ProcessElement - public void processElement( - ProcessContext c, - @Element Row element, - BoundedWindow window, - PaneInfo paneInfo, - @Timestamp Instant timestamp, - OutputReceiver> out) { - - String tableIdentifier = - dynamicDestinations.getTableStringIdentifier( - ValueInSingleWindow.of(element, timestamp, window, paneInfo)); - + private CachedTableMetadata getOrRefreshMetadata( + DoFn.ProcessContext c, + String tableIdentifier, + org.apache.beam.sdk.schemas.Schema dataSchema) { SerializableTableSpec tableSpec = null; if (metadataView != null) { Map viewMap = c.sideInput(metadataView); @@ -223,92 +353,568 @@ public void processElement( } } - Row data = dynamicDestinations.getData(element); - - @Nullable PartitionKey partitionKey = checkStateNotNull(partitionKeys).get(tableIdentifier); - @Nullable BeamRowWrapper wrapper = checkStateNotNull(wrappers).get(tableIdentifier); - @Nullable Instant lastRefresh = checkStateNotNull(lastRefreshTimes).get(tableIdentifier); - @Nullable Integer cachedSpecId = checkStateNotNull(cachedSpecIds).get(tableIdentifier); + Map metadataCache = checkStateNotNull(cache); + @Nullable CachedTableMetadata cached = metadataCache.get(tableIdentifier); Instant now = Instant.now(); - - boolean specChanged = + boolean specOrOrderChanged = tableSpec != null - && (cachedSpecId == null || !cachedSpecId.equals(tableSpec.getSpecId())); + && (cached == null + || !Objects.equals(cached.specId, tableSpec.getSpecId()) + || !Objects.equals(cached.orderId, tableSpec.getOrderId())); boolean shouldRefresh = - partitionKey == null - || wrapper == null - || specChanged - || (tableSpec == null - && (lastRefresh == null || now.isAfter(lastRefresh.plus(REFRESH_INTERVAL)))); + cached == null + || specOrOrderChanged + || (tableSpec == null && now.isAfter(cached.lastRefresh.plus(REFRESH_INTERVAL))); if (shouldRefresh) { - PartitionSpec spec = PartitionSpec.unpartitioned(); - Schema schema = IcebergUtils.beamSchemaToIcebergSchema(data.getSchema()); - - @Nullable IcebergTableCreateConfig createConfig = - dynamicDestinations.instantiateDestination(tableIdentifier).getTableCreateConfig(); + Schema schema = IcebergUtils.beamSchemaToIcebergSchema(dataSchema); + Schema tableSchema = schema; + SortOrder sortOrder = SortOrder.unsorted(); + Integer specId = null; + Integer orderId = null; if (tableSpec != null) { spec = tableSpec.getPartitionSpec(); - if (data.getSchema().getFieldCount() == tableSpec.getSchema().columns().size()) { - schema = tableSpec.getSchema(); + sortOrder = tableSpec.getSortOrder(); + tableSchema = tableSpec.getSchema(); + if (dataSchema.getFieldCount() == tableSchema.columns().size()) { + schema = tableSchema; } - checkStateNotNull(cachedSpecIds).put(tableIdentifier, tableSpec.getSpecId()); - } else if (createConfig != null && createConfig.getPartitionFields() != null) { - spec = - PartitionUtils.toPartitionSpec(createConfig.getPartitionFields(), data.getSchema()); + specId = tableSpec.getSpecId(); + orderId = tableSpec.getOrderId(); } else { try { - // see if table already exists with a spec Table table = TableCache.getAndRefreshIfStale( catalogConfig, IcebergUtils.parseTableIdentifier(tableIdentifier)); spec = table.spec(); - if (data.getSchema().getFieldCount() == table.schema().columns().size()) { - schema = table.schema(); + sortOrder = table.sortOrder(); + tableSchema = table.schema(); + if (dataSchema.getFieldCount() == tableSchema.columns().size()) { + schema = tableSchema; } } catch (NoSuchTableException ignored) { - // no partition to apply + @Nullable IcebergTableCreateConfig createConfig = + dynamicDestinations.instantiateDestination(tableIdentifier).getTableCreateConfig(); + if (createConfig != null) { + if (createConfig.getPartitionFields() != null) { + spec = PartitionUtils.toPartitionSpec(createConfig.getPartitionFields(), schema); + } + if (createConfig.getSortOrder() != null) { + sortOrder = createConfig.getSortOrder(); + } + } } } - partitionKey = new PartitionKey(spec, schema); - wrapper = new BeamRowWrapper(data.getSchema(), schema.asStruct()); + PartitionKey partitionKey = new PartitionKey(spec, schema); + BeamRowWrapper wrapper = new BeamRowWrapper(dataSchema, schema.asStruct()); + IcebergRowSorter.@Nullable KeyEncoder encoder = + includeSortKeyEncoder && sortOrder != null && sortOrder.isSorted() + ? new IcebergRowSorter.KeyEncoder(sortOrder, tableSchema, dataSchema) + : null; - checkStateNotNull(partitionKeys).put(tableIdentifier, partitionKey); - checkStateNotNull(wrappers).put(tableIdentifier, wrapper); - checkStateNotNull(lastRefreshTimes).put(tableIdentifier, now); + cached = new CachedTableMetadata(partitionKey, wrapper, encoder, now, specId, orderId); + metadataCache.put(tableIdentifier, cached); } - partitionKey = checkStateNotNull(partitionKey); - wrapper = checkStateNotNull(wrapper); + return checkStateNotNull(cached); + } + + ResolvedDestination resolve( + DoFn.ProcessContext c, + Row element, + BoundedWindow window, + PaneInfo paneInfo, + Instant timestamp) { + String tableIdentifier = + dynamicDestinations.getTableStringIdentifier( + ValueInSingleWindow.of(element, timestamp, window, paneInfo)); + Row data = dynamicDestinations.getData(element); + CachedTableMetadata cached = getOrRefreshMetadata(c, tableIdentifier, data.getSchema()); + cached.partitionKey.partition(cached.wrapper.wrap(data)); + return new ResolvedDestination( + tableIdentifier, data, cached.partitionKey.toPath(), cached.keyEncoder); + } + + IcebergRowSorter.@Nullable KeyEncoder resolveKeyEncoder( + DoFn.ProcessContext c, + String tableIdentifier, + org.apache.beam.sdk.schemas.Schema dataSchema) { + return getOrRefreshMetadata(c, tableIdentifier, dataSchema).keyEncoder; + } + } + + static class AssignDoFn extends DoFn> { + + private final DestinationMetadataResolver metadataResolver; + private final DistributionMode distributionMode; + private final @Nullable SerializableFunction distributionFunction; + + AssignDoFn( + DynamicDestinations dynamicDestinations, + IcebergCatalogConfig catalogConfig, + @Nullable PCollectionView> metadataView, + DistributionMode distributionMode, + @Nullable SerializableFunction distributionFunction) { + this.metadataResolver = + new DestinationMetadataResolver(dynamicDestinations, catalogConfig, metadataView, false); + this.distributionMode = distributionMode; + this.distributionFunction = distributionFunction; + } - partitionKey.partition(wrapper.wrap(data)); + @Setup + public void setup() { + metadataResolver.setup(); + } + + @ProcessElement + public void processElement( + ProcessContext c, + @Element Row element, + BoundedWindow window, + PaneInfo paneInfo, + @Timestamp Instant timestamp, + OutputReceiver> out) { - String partitionPath = partitionKey.toPath(); + ResolvedDestination resolved = + metadataResolver.resolve(c, element, window, paneInfo, timestamp); Row.Builder keyBuilder = Row.withSchema(outputSchemaFor(distributionMode)) - .addValue(tableIdentifier) - .addValue(partitionPath); + .addValue(resolved.tableIdentifier) + .addValue(resolved.partitionPath); if (distributionMode == DistributionMode.RANGE) { - SerializableFunction shardFn = - checkStateNotNull( - distributionFunction, - "A distribution function is required when using RANGE distribution mode."); - keyBuilder = - keyBuilder.addValue( - checkStateNotNull( - shardFn.apply(data), - "The RANGE distribution function returned a null shard id. It must return a" - + " non-null Integer for every row.")); + int shard = 0; + if (distributionFunction != null) { + shard = + checkStateNotNull( + distributionFunction.apply(resolved.data), + "The RANGE distribution function returned a null shard id. It must return a" + + " non-null Integer for every row."); + } + keyBuilder = keyBuilder.addValue(shard); } - out.output(KV.of(keyBuilder.build(), data)); + out.output(KV.of(keyBuilder.build(), resolved.data)); + } + } + + /** + * Accumulator and wire representation of a partition's estimated byte count, total row count, and + * reservoir-sampled sort keys. + */ + @VisibleForTesting + static class PartitionSample { + long totalBytes; + long rowCount; + List samples; + + PartitionSample(long totalBytes, long rowCount, List samples) { + this.totalBytes = totalBytes; + this.rowCount = rowCount; + this.samples = samples; + } + } + + /** Deterministic {@link StructuredCoder} for {@link PartitionSample}. */ + static class PartitionSampleCoder extends StructuredCoder { + private static final PartitionSampleCoder INSTANCE = new PartitionSampleCoder(); + private static final VarLongCoder LONG_CODER = VarLongCoder.of(); + private static final ListCoder SAMPLES_CODER = ListCoder.of(ByteArrayCoder.of()); + + static PartitionSampleCoder of() { + return INSTANCE; + } + + @Override + public void encode(PartitionSample value, OutputStream outStream) + throws CoderException, IOException { + LONG_CODER.encode(value.totalBytes, outStream); + LONG_CODER.encode(value.rowCount, outStream); + SAMPLES_CODER.encode(value.samples, outStream); + } + + @Override + public PartitionSample decode(InputStream inStream) throws CoderException, IOException { + long totalBytes = LONG_CODER.decode(inStream); + long rowCount = LONG_CODER.decode(inStream); + List samples = new ArrayList<>(SAMPLES_CODER.decode(inStream)); + return new PartitionSample(totalBytes, rowCount, samples); + } + + @Override + public List> getCoderArguments() { + return Collections.emptyList(); + } + + @Override + public void verifyDeterministic() throws NonDeterministicException { + LONG_CODER.verifyDeterministic(); + SAMPLES_CODER.verifyDeterministic(); + } + } + + /** Per-bundle reservoir sample and byte counter for a single {@code (destination, partition)}. */ + private static class BundlePartitionSample { + long sampledBytes = 0L; + long sampledRows = 0L; + long rowCount = 0L; + final List reservoir = new ArrayList<>(MAX_SAMPLES_PER_BUNDLE); + final Row partitionKeyRow; + final BoundedWindow window; + final Instant timestamp; + + BundlePartitionSample(Row partitionKeyRow, BoundedWindow window, Instant timestamp) { + this.partitionKeyRow = partitionKeyRow; + this.window = window; + this.timestamp = timestamp; + } + + /** + * Returns the reservoir slot index {@code [0, MAX_SAMPLES_PER_BUNDLE)} if the next row should + * be sampled into the reservoir, or {@code -1} if it is skipped. + */ + int nextReservoirSlot() { + rowCount++; + if (reservoir.size() < MAX_SAMPLES_PER_BUNDLE) { + return reservoir.size(); + } + long replaceIdx = ThreadLocalRandom.current().nextLong(rowCount); + return replaceIdx < MAX_SAMPLES_PER_BUNDLE ? (int) replaceIdx : -1; + } + + void recordSample(int slotIndex, byte[] sortKeyBytes, long estimatedBytes) { + sampledBytes += estimatedBytes; + sampledRows++; + if (slotIndex == reservoir.size()) { + reservoir.add(sortKeyBytes); + } else { + reservoir.set(slotIndex, sortKeyBytes); + } + } + + long estimatedTotalBytes() { + if (sampledRows <= 0L) { + return 0L; + } + if (sampledRows == rowCount) { + return sampledBytes; + } + return Math.round((double) sampledBytes * rowCount / sampledRows); + } + } + + /** + * Computes {@code (destination, partition)} for each row while accumulating a lightweight + * per-bundle reservoir sample of sort keys and estimated partition bytes emitted once per bundle + * in {@link FinishBundle}. Sort keys and row byte sizes are only computed for rows selected into + * the reservoir. + */ + static class ExtractPartitionAndSampleSortKeyDoFn extends DoFn> { + + private static final byte[] EMPTY_SORT_KEY = new byte[0]; + + private final DestinationMetadataResolver metadataResolver; + private transient @MonotonicNonNull Map, BundlePartitionSample> + bundleSamples; + private transient @MonotonicNonNull ByteArrayOutputStream reusableKeyBaos; + + ExtractPartitionAndSampleSortKeyDoFn( + DynamicDestinations dynamicDestinations, + IcebergCatalogConfig catalogConfig, + @Nullable PCollectionView> metadataView) { + this.metadataResolver = + new DestinationMetadataResolver(dynamicDestinations, catalogConfig, metadataView, true); + } + + @Setup + public void setup() { + this.metadataResolver.setup(); + this.reusableKeyBaos = new ByteArrayOutputStream(); + } + + @StartBundle + public void startBundle() { + this.bundleSamples = new HashMap<>(); + } + + @ProcessElement + public void processElement( + ProcessContext c, + @Element Row element, + BoundedWindow window, + PaneInfo paneInfo, + @Timestamp Instant timestamp, + MultiOutputReceiver out) + throws IOException { + + ResolvedDestination resolved = + metadataResolver.resolve(c, element, window, paneInfo, timestamp); + Row data = resolved.data; + + Row partitionKeyRow = + Row.withSchema(OUTPUT_SCHEMA) + .addValue(resolved.tableIdentifier) + .addValue(resolved.partitionPath) + .build(); + + BundlePartitionSample sample = + checkStateNotNull(bundleSamples) + .computeIfAbsent( + KV.of(partitionKeyRow, window), + k -> new BundlePartitionSample(partitionKeyRow, window, timestamp)); + int slot = sample.nextReservoirSlot(); + if (slot >= 0) { + byte[] sortKeyBytes = + resolved.keyEncoder != null + ? resolved.keyEncoder.encodeToBytes(data, checkStateNotNull(reusableKeyBaos)) + : EMPTY_SORT_KEY; + long estimatedBytes = + Math.max(16L, IcebergUtils.estimateRowBytes(data) + sortKeyBytes.length); + sample.recordSample(slot, sortKeyBytes, estimatedBytes); + } + + out.get(PARTITIONED_ROWS_TAG).output(KV.of(partitionKeyRow, data)); + } + + @FinishBundle + public void finishBundle(FinishBundleContext c) { + if (bundleSamples == null || bundleSamples.isEmpty()) { + return; + } + for (BundlePartitionSample sample : bundleSamples.values()) { + c.output( + PARTITION_SAMPLES_TAG, + KV.of( + sample.partitionKeyRow, + new PartitionSample( + sample.estimatedTotalBytes(), sample.rowCount, sample.reservoir)), + sample.timestamp, + sample.window); + } + bundleSamples.clear(); + } + } + + /** + * Combines per-bundle partition byte counts, row counts, and sort-key reservoirs to compute + * right-sized, deduplicated quantile split points per {@code (destination, partition)}. + */ + static class PartitionQuantilesCombineFn + extends Combine.CombineFn> { + + private final long targetBytesPerShard; + + PartitionQuantilesCombineFn(long targetBytesPerShard) { + this.targetBytesPerShard = Math.max(1L, targetBytesPerShard); + } + + @Override + public PartitionSample createAccumulator() { + return new PartitionSample(0L, 0L, new ArrayList<>()); + } + + @Override + public PartitionSample addInput(PartitionSample accum, PartitionSample input) { + mergeIntoReservoir(accum, input); + return accum; + } + + @Override + public PartitionSample mergeAccumulators(Iterable accumulators) { + PartitionSample merged = createAccumulator(); + for (PartitionSample accum : accumulators) { + mergeIntoReservoir(merged, accum); + } + return merged; + } + + /** + * Merges {@code incoming} into {@code accum} in place. When the combined sample count exceeds + * {@link #MAX_COMBINED_SAMPLES} or the two reservoirs have different per-sample weights (e.g. a + * small bundle merged with a large bundle that exceeded {@link #MAX_SAMPLES_PER_BUNDLE}), uses + * Efraimidis-Spirakis weighted reservoir sampling without replacement so each retained sample + * represents an equal number of rows and no sample is ever duplicated. + */ + private static void mergeIntoReservoir(PartitionSample accum, PartitionSample incoming) { + accum.totalBytes += Math.max(0L, incoming.totalBytes); + List validIncoming = new ArrayList<>(incoming.samples.size()); + for (byte[] sample : incoming.samples) { + if (sample != null && sample.length > 0) { + validIncoming.add(sample); + } + } + if (validIncoming.isEmpty()) { + accum.rowCount += Math.max(0L, incoming.rowCount); + return; + } + + long effectiveIncomingRows = Math.max((long) validIncoming.size(), incoming.rowCount); + if (accum.samples.isEmpty()) { + if (validIncoming.size() <= MAX_COMBINED_SAMPLES) { + accum.samples.addAll(validIncoming); + } else { + accum.samples.addAll(validIncoming.subList(0, MAX_COMBINED_SAMPLES)); + } + accum.rowCount += effectiveIncomingRows; + return; + } + + long effectiveAccumRows = Math.max((long) accum.samples.size(), accum.rowCount); + long combinedRows = effectiveAccumRows + effectiveIncomingRows; + int nA = accum.samples.size(); + int nB = validIncoming.size(); + int totalCandidates = nA + nB; + double weightA = (double) effectiveAccumRows / nA; + double weightB = (double) effectiveIncomingRows / nB; + + if (totalCandidates <= MAX_COMBINED_SAMPLES && Math.abs(weightA - weightB) < 1e-9) { + accum.samples.addAll(validIncoming); + accum.rowCount = combinedRows; + return; + } + + double targetWeight = + Math.max(Math.max(weightA, weightB), (double) combinedRows / MAX_COMBINED_SAMPLES); + int targetSamples = + (int) + Math.min( + totalCandidates, + Math.min( + MAX_COMBINED_SAMPLES, Math.max(1L, Math.round(combinedRows / targetWeight)))); + + double[] keys = new double[totalCandidates]; + Integer[] indices = new Integer[totalCandidates]; + ThreadLocalRandom rng = ThreadLocalRandom.current(); + for (int i = 0; i < nA; i++) { + double u = Math.max(Double.MIN_NORMAL, rng.nextDouble()); + keys[i] = -Math.log(u) / weightA; + indices[i] = i; + } + for (int j = 0; j < nB; j++) { + double u = Math.max(Double.MIN_NORMAL, rng.nextDouble()); + int idx = nA + j; + keys[idx] = -Math.log(u) / weightB; + indices[idx] = idx; + } + Arrays.sort(indices, Comparator.comparingDouble(i -> keys[i])); + + List selected = new ArrayList<>(targetSamples); + for (int k = 0; k < targetSamples; k++) { + int chosen = indices[k]; + selected.add(chosen < nA ? accum.samples.get(chosen) : validIncoming.get(chosen - nA)); + } + accum.samples = selected; + accum.rowCount = combinedRows; + } + + @Override + public List extractOutput(PartitionSample accum) { + long totalBytes = accum.totalBytes; + List samples = new ArrayList<>(accum.samples); + int numShards = + (int) + Math.min( + MAX_RANGE_SHARDS, + Math.max(1L, (totalBytes + targetBytesPerShard - 1) / targetBytesPerShard)); + if (numShards <= 1 || samples.isEmpty()) { + return Collections.emptyList(); + } + + Comparator cmp = UnsignedBytes.lexicographicalComparator(); + samples.sort(cmp); + + List splits = new ArrayList<>(numShards - 1); + int n = samples.size(); + for (int i = 1; i < numShards; i++) { + int idx = (int) ((long) i * n / numShards); + if (idx >= n) { + idx = n - 1; + } + byte[] candidate = samples.get(idx); + if (splits.isEmpty() || cmp.compare(splits.get(splits.size() - 1), candidate) < 0) { + splits.add(candidate); + } + } + return splits; + } + + @Override + public Coder getAccumulatorCoder( + CoderRegistry registry, Coder inputCoder) { + return PartitionSampleCoder.of(); + } + + @Override + public Coder> getDefaultOutputCoder( + CoderRegistry registry, Coder inputCoder) { + return ListCoder.of(ByteArrayCoder.of()); + } + } + + /** + * Assigns the final {@code (destination, partition, shard)} key using the dynamically computed + * quantile split points for each partition. Sort keys are only encoded for partitions that span + * multiple range shards. + */ + private static class AssignDynamicRangeShardDoFn extends DoFn, KV> { + + private static final byte[] EMPTY_SORT_KEY = new byte[0]; + + private final DestinationMetadataResolver metadataResolver; + private final PCollectionView>> splitPointsView; + private transient @MonotonicNonNull ByteArrayOutputStream reusableKeyBaos; + + AssignDynamicRangeShardDoFn( + DynamicDestinations dynamicDestinations, + IcebergCatalogConfig catalogConfig, + @Nullable PCollectionView> metadataView, + PCollectionView>> splitPointsView) { + this.metadataResolver = + new DestinationMetadataResolver(dynamicDestinations, catalogConfig, metadataView, true); + this.splitPointsView = splitPointsView; + } + + @Setup + public void setup() { + this.metadataResolver.setup(); + this.reusableKeyBaos = new ByteArrayOutputStream(); + } + + @ProcessElement + public void processElement( + ProcessContext c, @Element KV element, OutputReceiver> out) + throws IOException { + Row partitionKey = element.getKey(); + Row data = element.getValue(); + String tableIdentifier = checkStateNotNull(partitionKey.getString(DESTINATION)); + String partitionPath = checkStateNotNull(partitionKey.getString(PARTITION)); + + Map> splitMap = c.sideInput(splitPointsView); + @Nullable List splits = splitMap != null ? splitMap.get(partitionKey) : null; + int shard = 0; + if (splits != null && !splits.isEmpty()) { + IcebergRowSorter.@Nullable KeyEncoder keyEncoder = + metadataResolver.resolveKeyEncoder(c, tableIdentifier, data.getSchema()); + byte[] sortKeyBytes = + keyEncoder != null + ? keyEncoder.encodeToBytes(data, checkStateNotNull(reusableKeyBaos)) + : EMPTY_SORT_KEY; + shard = findRangeShard(splits, sortKeyBytes); + } + + Row rangeKey = + Row.withSchema(OUTPUT_SCHEMA_WITH_SHARD) + .addValue(tableIdentifier) + .addValue(partitionPath) + .addValue(shard) + .build(); + out.output(KV.of(rangeKey, data)); } } } 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..96e51f832eda 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 @@ -32,9 +32,11 @@ import org.apache.beam.sdk.schemas.Schema.FieldType; import org.apache.beam.sdk.schemas.logicaltypes.Date; import org.apache.beam.sdk.schemas.logicaltypes.DateTime; +import org.apache.beam.sdk.schemas.logicaltypes.EnumerationType; 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 +160,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; @@ -177,6 +180,11 @@ private interface PositionalGetter { checkStateNotNull(row.getLogicalTypeValue(pos, LocalTime.class))); } else if (beamType.isLogicalType(FixedPrecisionNumeric.IDENTIFIER)) { return (row, pos) -> row.getLogicalTypeValue(pos, BigDecimal.class); + } else if (beamType.isLogicalType(EnumerationType.IDENTIFIER)) { + EnumerationType enumType = beamType.getLogicalType(EnumerationType.class); + return (row, pos) -> + enumType.toString( + checkStateNotNull(row.getLogicalTypeValue(pos, EnumerationType.Value.class))); } else { return null; } 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 8bf26962e231..17e19e21ec99 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 @@ -238,6 +238,43 @@ * * } * + *

Sorting and Distribution Modes

+ * + *

You can control how records are grouped, sharded, and sorted before writing to Iceberg data + * files using {@code distribution_mode}, {@code sort_fields}, {@code + * sort_on_write}, and {@code autosharding}: + * + *

    + *
  • {@code distribution_mode}: {@code "none"} (default; writes directly within worker + * bundles without a shuffle), {@code "hash"} (shuffles by Iceberg partition key), or {@code + * "range"} (shuffles by partition key and sort-order key). In batch mode, {@code "range"} + * samples sort keys per partition to compute balanced quantile boundaries sized to {@code + * write.target-file-size-bytes}. + *
  • {@code sort_fields}: List of sort expressions used when creating a new table (e.g., + * {@code ["ts DESC NULLS LAST", "id ASC"]}, {@code "bucket(id, 16) ASC"}). + *
  • {@code sort_on_write}: When {@code true}, sorts records within each output file + * using the table's {@link org.apache.iceberg.SortOrder} via external merge sort (spilling to + * local disk when buffered records exceed memory limits). + *
  • {@code autosharding}: Dynamically scales the number of parallel writers per + * partition when using {@code "hash"} or {@code "range"}. + *
+ * + * 
{@code
+ * Map config = Map.of(
+ *         "table", "my_db.events",
+ *         "catalog_properties", Map.of(...),
+ *         "partition_fields", List.of("day(ts)"),
+ *         "sort_fields", List.of("ts ASC NULLS LAST", "user_id ASC"),
+ *         "sort_on_write", true,
+ *         "distribution_mode", "range");
+ *
+ * pipeline
+ *     .apply(Managed.write(ICEBERG).withConfig(config));
+ * }
+ * + *

See {@link WriteRows#withDistributionMode(DistributionMode)} for a detailed comparison of + * distribution modes and sorting trade-offs. + * *

Output Snapshots

* *

When records are written and committed to a table, a snapshot is produced. A batch pipeline @@ -520,18 +557,22 @@ public WriteRows withDirectWriteByteLimit(Integer directWriteByteLimit) { } /** - * The default distribution mode is {@link DistributionMode#NONE}, which writes without a - * shuffle unless {@link #withSortedWrites()} is enabled on a table with a sort order. + * The default distribution mode is {@link DistributionMode#NONE}, which writes directly within + * each worker bundle without a partition shuffle. When {@link #withSortedWrites()} is enabled + * on a table with a {@link org.apache.iceberg.SortOrder}, records are buffered and sorted + * locally per bundle by {@code (partition, sortKey)} before writing (or per triggered batch in + * streaming). * *

Warning on HASH mode: Utilizing {@code HASH} distribution mode (with or without * auto-sharding) can suffer from large unpartitioned or skewed writes if key spaces are not * uniformly distributed. This can bottleneck workers and produce fragmented layout files. * - *

Note on RANGE mode: When utilizing {@code RANGE} distribution mode, it is - * recommended that the custom distribution function is designed to produce adequately sized and - * strictly non-overlapping ranges of the sorting column to optimize downstream read - * performance. {@code RANGE} mode automatically enables sorted writes when the target table has - * a {@link org.apache.iceberg.SortOrder}. + *

Note on RANGE mode: {@code RANGE} mode automatically enables sorted writes when the + * target table has a {@link org.apache.iceberg.SortOrder}. In bounded (batch) pipelines, if no + * {@link #withDistributionFunction} is supplied, {@code RANGE} mode automatically samples sort + * keys and partition byte volume to compute non-overlapping quantile range shards sized to the + * target file size. Users may also supply a custom {@link #withDistributionFunction} to control + * range shard assignment manually (or for streaming pipelines). * *

Comparison of Distribution Modes:

* @@ -545,21 +586,21 @@ public WriteRows withDirectWriteByteLimit(Integer directWriteByteLimit) { * * * {@link DistributionMode#NONE} - * No partition shuffle is performed by default. When {@link #withSortedWrites()} is enabled on a sorted table, records are grouped by destination shard and sorted prior to writing. - * Highly lightweight with zero shuffle overhead for unsorted writes. - * Writers on different workers can write to overlapping min/max key ranges across multiple files. Relies heavily on post-fact compaction or query time merges. + * No partition shuffle is performed. When {@link #withSortedWrites()} is enabled on a sorted table, records are sorted locally per bundle by partition and sort key prior to writing. + * Highly lightweight with zero network shuffle overhead. + * Writers on different workers can write to the same partition or overlapping min/max key ranges across multiple files. Relies on pre-aligned input bundles or downstream compaction. * * * {@link DistributionMode#HASH} - * Data is shuffled and consolidated by partition key. All records for a partition are routed to a single worker. - * Consolidates partition files, eliminating cross-worker file overlapping for partition keys. Excellent worker stability. - * Can suffer from severe data skew if a single partition contains significantly more data than others (hot partitions). + * Data is shuffled and consolidated by partition key. All records for a partition are routed to a single worker (or dynamically sub-sharded when {@link #withAutosharding()} is enabled). + * Consolidates partition files, eliminating cross-worker file fragmentation for partition keys. + * Without {@link #withAutosharding()}, can suffer from data skew if a single partition contains significantly more data than others (hot partitions). * * * {@link DistributionMode#RANGE} - * Data is shuffled based on a user-provided shard/bucket function (e.g., hashing/binning continuous keys) and sorted within each shard. - * Distributes writes for hot partitions across multiple workers. Eliminates skew while keeping file min/max key ranges tight and non-overlapping. - * Requires providing a custom {@link SerializableFunction} mapping rows to integer shard/bucket IDs. + * Data is range-partitioned along the table's sort order (dynamically via quantile sampling in batch by default, or via a custom {@link #withDistributionFunction}) and sorted within each shard. + * Distributes writes for hot partitions across multiple workers while keeping file min/max key ranges tight and non-overlapping. + * Incurs a range-sampling pass and shuffle in batch; requires a user-supplied {@link #withDistributionFunction} to fan out hot partitions in streaming. * * * @@ -589,15 +630,15 @@ public WriteRows withDirectWriteByteLimit(Integer directWriteByteLimit) { * Medium / Large * Low Write Latency * {@link DistributionMode#NONE} + {@link #withSortedWrites()} - * Shards by destination without partition shuffle. Results in overlapping key ranges across files, which requires downstream compaction. + * Sorts locally per bundle without a partition shuffle. Results in overlapping key ranges across files unless upstream bundles are pre-partitioned, requiring downstream compaction. * * * Partitioned * Sorted * Medium / Large * Low Read Latency - * {@link DistributionMode#HASH} with auto-sharding + {@link #withSortedWrites()} OR {@link DistributionMode#RANGE} - * HASH with auto-sharding scales writes for hot partitions but can result in overlapping file ranges requiring query-time sort merges. RANGE sharding distributes hot partitions into sequential, non-overlapping files to optimize reads. + * {@link DistributionMode#RANGE} OR {@link DistributionMode#HASH} with auto-sharding + {@link #withSortedWrites()} + * RANGE sharding distributes hot partitions into sequential, non-overlapping files to optimize reads. HASH with auto-sharding scales writes for hot partitions in streaming/batch but can result in overlapping file ranges requiring query-time sort merges. * * * Partitioned @@ -621,7 +662,7 @@ public WriteRows withDirectWriteByteLimit(Integer directWriteByteLimit) { * Small * Any * {@link DistributionMode#NONE} + {@link #withSortedWrites()} - * Groups by destination shard for fast, low-volume local sorting. + * Sorts locally per bundle by sort key with zero shuffle overhead. * * * Unpartitioned @@ -629,15 +670,15 @@ public WriteRows withDirectWriteByteLimit(Integer directWriteByteLimit) { * Medium / Large * Low Write Latency * {@link DistributionMode#NONE} + {@link #withSortedWrites()} - * Parallel worker writes across destination shards. Requires downstream compaction to resolve overlapping file ranges. + * Parallel worker writes sorted locally per bundle. Requires downstream compaction to resolve overlapping file ranges across bundles. * * * Unpartitioned * Sorted * Medium / Large * Low Read Latency - * {@link DistributionMode#RANGE} (with custom sharding function) - * Shards continuous keys into non-overlapping worker ranges. Eliminates single-worker bottlenecks and guarantees zero file overlap for fast queries. + * {@link DistributionMode#RANGE} + * Shards continuous keys into non-overlapping worker ranges (dynamically via quantile sampling in batch, or via a custom sharding function). Eliminates single-worker bottlenecks and guarantees zero file overlap for fast queries. * * * Unpartitioned @@ -666,17 +707,13 @@ public WriteRows withDirectWriteByteLimit(Integer directWriteByteLimit) { * .withDistributionMode(DistributionMode.HASH) * .withSortedWrites()); * - * // 3. Using RANGE distribution mode with a custom shard/bucket function to avoid data skew + * // 3. Using RANGE distribution mode (automatically computes quantile shards in batch, + * // or optionally accepts a custom shard function) * pipeline * .apply(Create.of(BEAM_ROWS)) * .apply(IcebergIO.writeRows(catalogConfig) * .to(tableId) - * .withDistributionMode(DistributionMode.RANGE) - * .withDistributionFunction(row -> { - * // Group continuous IDs into parallel, non-overlapping shards of 10,000 IDs each - * long id = row.getInt64("id"); - * return (int) (id / 10000); - * })); + * .withDistributionMode(DistributionMode.RANGE)); * } */ public WriteRows withDistributionMode(DistributionMode mode) { @@ -684,10 +721,12 @@ public WriteRows withDistributionMode(DistributionMode mode) { } /** - * Sets the custom range-distribution function. + * Sets an optional custom range-distribution function. * *

Only applicable when the distribution mode is set to {@link DistributionMode#RANGE}. The - * function maps a Beam {@link Row} to an Integer representing a shard/bucket ID. + * function maps a Beam {@link Row} to an Integer representing a shard/bucket ID. If omitted in + * {@link DistributionMode#RANGE}, Beam dynamically computes range shards from the table's sort + * order and partition byte volume. */ public WriteRows withDistributionFunction(SerializableFunction shardFn) { return toBuilder().setDistributionFunction(shardFn).build(); @@ -818,6 +857,16 @@ public WriteRows withTableCachePollingBuckets(int pollingBuckets) { @Override public void populateDisplayData(DisplayData.Builder builder) { super.populateDisplayData(builder); + builder.add( + DisplayData.item("distributionMode", getDistributionMode().modeName()) + .withLabel("Distribution Mode")); + builder.add( + DisplayData.item("autoSharding", getAutoSharding()).withLabel("Using Auto-Sharding")); + builder.add( + DisplayData.item( + "sortedWrites", + getSortedWrites() || getDistributionMode() == DistributionMode.RANGE) + .withLabel("Using Sorted Writes")); builder.add( DisplayData.item("useSideInputTableCache", getUseSideInputTableCache()) .withLabel("Using Side-Input Table Cache")); @@ -887,12 +936,15 @@ public IcebergWriteResult expand(PCollection input) { } boolean sortedWrites = getSortedWrites() || getDistributionMode() == DistributionMode.RANGE; + Preconditions.checkArgument( + !getAutoSharding() || getDistributionMode() == DistributionMode.HASH, + "Autosharding option is only available with 'hash' distribution mode."); + Preconditions.checkArgument( + getDistributionFunction() == null || getDistributionMode() == DistributionMode.RANGE, + "Distribution function is only available with 'range' distribution mode."); switch (getDistributionMode()) { case NONE: - Preconditions.checkArgument( - !getAutoSharding(), - "Autosharding option is only available with " + "'hash' distribution mode."); return input .apply("Assign Table Destinations", new AssignDestinations(destinations)) .apply( @@ -906,35 +958,14 @@ public IcebergWriteResult expand(PCollection input) { metadataView, sortedWrites)); case HASH: - return input - .apply( - "AssignDestinationAndPartition", - new AssignDestinationsAndPartitions( - destinations, - getCatalogConfig(), - metadataView, - getDistributionMode(), - getDistributionFunction())) - .apply( - "Write Rows to Partitions", - new WriteToPartitions( - getCatalogConfig(), - destinations, - getTriggeringFrequency(), - getAutoSharding(), - getWriteProperties(), - metadataView, - sortedWrites)); case RANGE: - Preconditions.checkArgument( - !getAutoSharding(), - "Autosharding option is only available with 'hash' distribution mode."); - Preconditions.checkArgument( - getDistributionFunction() != null, - "Must provide a distribution function when using RANGE distribution mode."); + String assignStepName = + getDistributionMode() == DistributionMode.RANGE + ? "AssignDestinationAndPartitionWithRange" + : "AssignDestinationAndPartition"; return input .apply( - "AssignDestinationAndPartitionWithRange", + assignStepName, new AssignDestinationsAndPartitions( destinations, getCatalogConfig(), diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergRowSorter.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergRowSorter.java index 64362ca9f217..e3b838552881 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergRowSorter.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergRowSorter.java @@ -41,8 +41,10 @@ import org.apache.beam.sdk.coders.RowCoder; import org.apache.beam.sdk.extensions.sorter.BufferedExternalSorter; import org.apache.beam.sdk.extensions.sorter.ExternalSorter; +import org.apache.beam.sdk.schemas.logicaltypes.EnumerationType; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.Row; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; import org.apache.iceberg.NullOrder; import org.apache.iceberg.Schema; import org.apache.iceberg.SortDirection; @@ -63,112 +65,155 @@ * A utility class to sort Beam {@link Row}s based on an Iceberg {@link SortOrder}. Leverages {@link * BufferedExternalSorter} to spill to local disk when elements exceed memory limit. */ -class IcebergRowSorter { +class IcebergRowSorter implements AutoCloseable { private static final Logger LOG = LoggerFactory.getLogger(IcebergRowSorter.class); /** Width of an encoded decimal key. 16 bytes covers Iceberg's maximum precision of 38 digits. */ private static final int DECIMAL_KEY_BYTES = 16; + /** Default in-memory buffer limit (in MB) before spilling sorted keys/rows to local disk. */ + static final int DEFAULT_SORTER_MEMORY_MB = 64; + + private final Path sortTempDir; + private final BufferedExternalSorter sorter; + private final RowCoder rowCoder; + private final KeyEncoder keyEncoder; + private final ByteArrayOutputStream keyBaos = new ByteArrayOutputStream(); + private final ByteArrayOutputStream valBaos = new ByteArrayOutputStream(); + /** - * Best-effort cleanup of a temporary spill directory. Failures are logged rather than propagated. + * Pre-bound encoder for an Iceberg {@link SortOrder} that avoids re-binding transforms and field + * name lookups per row. */ - private static void deleteTempDir(@Nullable Path tempDir) { - if (tempDir == null) { - return; + @SuppressWarnings({"unchecked", "rawtypes"}) + static class KeyEncoder { + private final SortOrder sortOrder; + private final String[] columnNames; + private final int[] beamFieldIndices; + private final @Nullable EnumerationType[] enumTypes; + private final Type[] sourceTypes; + private final Type[] resultTypes; + private final @Nullable SerializableFunction[] boundTransforms; + + KeyEncoder(SortOrder sortOrder, Schema icebergSchema) { + this(sortOrder, icebergSchema, null); } - try (DirectoryStream entries = Files.newDirectoryStream(tempDir)) { - for (Path entry : entries) { - if (!Files.deleteIfExists(entry)) { - LOG.debug("Temporary sort file {} was already removed", entry); + + KeyEncoder( + SortOrder sortOrder, + Schema icebergSchema, + org.apache.beam.sdk.schemas.@Nullable Schema beamSchema) { + this.sortOrder = sortOrder; + List fields = sortOrder.fields(); + this.columnNames = new String[fields.size()]; + this.beamFieldIndices = new int[fields.size()]; + this.enumTypes = new EnumerationType[fields.size()]; + this.sourceTypes = new Type[fields.size()]; + this.resultTypes = new Type[fields.size()]; + this.boundTransforms = new SerializableFunction[fields.size()]; + for (int i = 0; i < fields.size(); i++) { + SortField field = fields.get(i); + String colName = icebergSchema.findColumnName(field.sourceId()); + columnNames[i] = colName; + if (beamSchema != null && beamSchema.hasField(colName)) { + beamFieldIndices[i] = beamSchema.indexOf(colName); + org.apache.beam.sdk.schemas.Schema.FieldType fieldType = + beamSchema.getField(beamFieldIndices[i]).getType(); + if (fieldType.isLogicalType(EnumerationType.IDENTIFIER)) { + enumTypes[i] = fieldType.getLogicalType(EnumerationType.class); + } + } else { + beamFieldIndices[i] = -1; + } + Type sourceType = icebergSchema.findType(field.sourceId()); + sourceTypes[i] = sourceType; + if (field.transform().isIdentity()) { + resultTypes[i] = sourceType; + boundTransforms[i] = null; + } else { + Transform transform = (Transform) field.transform(); + resultTypes[i] = transform.getResultType(sourceType); + boundTransforms[i] = transform.bind(sourceType); } } - } catch (IOException | RuntimeException e) { - LOG.debug("Failed to list temporary sort directory {}", tempDir, e); - return; } - try { - if (!Files.deleteIfExists(tempDir)) { - LOG.debug("Temporary sort directory {} was already removed", tempDir); - } - } catch (IOException | RuntimeException e) { - LOG.debug("Failed to delete temporary sort directory {}", tempDir, e); + + void encode(Row row, ByteArrayOutputStream baos) throws IOException { + encodeSortKey( + row, + sortOrder, + columnNames, + beamFieldIndices, + enumTypes, + sourceTypes, + resultTypes, + boundTransforms, + baos); + } + + byte[] encodeToBytes(Row row, ByteArrayOutputStream reusableBaos) throws IOException { + reusableBaos.reset(); + encode(row, reusableBaos); + return reusableBaos.toByteArray(); } } - /** - * Sorts {@code rows} according to {@code sortOrder}, spilling to local disk beyond the in-memory - * buffer. - * - *

The caller must close the returned iterable (ideally via try-with-resources). Closing - * reclaims the temporary spill directory; skipping it leaks that directory on the worker's local - * disk, which matters on long-lived runner harnesses that process many bundles. - * - *

Note that spilled rows are written unencrypted to the worker's local disk, in a directory - * created with owner-only permissions. - */ - @SuppressWarnings({"unchecked", "rawtypes"}) - public static CloseableIterable sortRows( - Iterable rows, + IcebergRowSorter( + SortOrder sortOrder, Schema icebergSchema, org.apache.beam.sdk.schemas.Schema beamSchema) { + this(sortOrder, icebergSchema, beamSchema, DEFAULT_SORTER_MEMORY_MB); + } + + IcebergRowSorter( SortOrder sortOrder, Schema icebergSchema, - org.apache.beam.sdk.schemas.Schema beamSchema) { - - if (sortOrder == null || !sortOrder.isSorted()) { - return CloseableIterable.withNoopClose(rows); - } - - final Path sortTempDir; + org.apache.beam.sdk.schemas.Schema beamSchema, + int memoryMb) { try { - sortTempDir = Files.createTempDirectory("iceberg-row-sorter-"); + this.sortTempDir = Files.createTempDirectory("iceberg-row-sorter-"); + this.sortTempDir.toFile().deleteOnExit(); } catch (IOException e) { throw new RuntimeException("Failed to create temporary directory for IcebergRowSorter", e); } - BufferedExternalSorter.Options sorterOptions = BufferedExternalSorter.options() .withTempLocation(sortTempDir.toString()) + .withMemoryMB(memoryMb) .withExternalSorterType(ExternalSorter.Options.SorterType.NATIVE); - BufferedExternalSorter sorter = BufferedExternalSorter.create(sorterOptions); - RowCoder rowCoder = RowCoder.of(beamSchema); + this.sorter = BufferedExternalSorter.create(sorterOptions); + this.rowCoder = RowCoder.of(beamSchema); + this.keyEncoder = new KeyEncoder(sortOrder, icebergSchema, beamSchema); + } - List fields = sortOrder.fields(); - String[] columnNames = new String[fields.size()]; - Type[] sourceTypes = new Type[fields.size()]; - Type[] resultTypes = new Type[fields.size()]; - @Nullable SerializableFunction[] boundTransforms = - new SerializableFunction[fields.size()]; - for (int i = 0; i < fields.size(); i++) { - SortField field = fields.get(i); - columnNames[i] = icebergSchema.findColumnName(field.sourceId()); - Type sourceType = icebergSchema.findType(field.sourceId()); - sourceTypes[i] = sourceType; - if (field.transform().isIdentity()) { - resultTypes[i] = sourceType; - boundTransforms[i] = null; - } else { - Transform transform = (Transform) field.transform(); - resultTypes[i] = transform.getResultType(sourceType); - boundTransforms[i] = transform.bind(sourceType); + /** + * Adds a row to the sorter, optionally prefixing the sort key with {@code partitionPath} so that + * rows are grouped by partition first and sorted by the table's {@link SortOrder} within each + * partition. + */ + public void add(@Nullable String partitionPath, Row row) { + try { + keyBaos.reset(); + valBaos.reset(); + if (partitionPath != null) { + writeString(partitionPath, keyBaos, false); } - } + keyEncoder.encode(row, keyBaos); + byte[] keyBytes = keyBaos.toByteArray(); - // Create reusable ByteArrayOutputStreams for key and value encoding - ByteArrayOutputStream keyBaos = new ByteArrayOutputStream(); - ByteArrayOutputStream valBaos = new ByteArrayOutputStream(); + rowCoder.encode(row, valBaos); + byte[] valBytes = valBaos.toByteArray(); + sorter.add(KV.of(keyBytes, valBytes)); + } catch (IOException | RuntimeException e) { + deleteTempDir(sortTempDir); + throw new RuntimeException("Failed to add row to external sorter", e); + } + } + /** + * Returns a sorted {@link CloseableIterable} over all added rows. Closing the returned iterable + * cleans up the temporary spill directory. + */ + public CloseableIterable sort() { try { - for (Row row : rows) { - keyBaos.reset(); - valBaos.reset(); - encodeSortKey( - row, sortOrder, columnNames, sourceTypes, resultTypes, boundTransforms, keyBaos); - byte[] keyBytes = keyBaos.toByteArray(); - - rowCoder.encode(row, valBaos); - byte[] valBytes = valBaos.toByteArray(); - sorter.add(KV.of(keyBytes, valBytes)); - } - Iterable> sortedKVs = sorter.sort(); Iterable decodedRows = new Iterable() { @@ -196,42 +241,81 @@ public Row next() { }; return CloseableIterable.combine(decodedRows, () -> deleteTempDir(sortTempDir)); - } catch (IOException | RuntimeException e) { deleteTempDir(sortTempDir); throw new RuntimeException("Failed to sort rows with external sorter", e); } } - @SuppressWarnings({"nullness", "unchecked", "rawtypes"}) - public static void encodeSortKey( - Row row, + @Override + public void close() { + deleteTempDir(sortTempDir); + } + + /** + * Best-effort cleanup of a temporary spill directory. Failures are logged rather than propagated. + */ + private static void deleteTempDir(@Nullable Path tempDir) { + if (tempDir == null) { + return; + } + try (DirectoryStream entries = Files.newDirectoryStream(tempDir)) { + for (Path entry : entries) { + if (!Files.deleteIfExists(entry)) { + LOG.debug("Temporary sort file {} was already removed", entry); + } + } + } catch (IOException | RuntimeException e) { + LOG.debug("Failed to list temporary sort directory {}", tempDir, e); + return; + } + try { + if (!Files.deleteIfExists(tempDir)) { + LOG.debug("Temporary sort directory {} was already removed", tempDir); + } + } catch (IOException | RuntimeException e) { + LOG.debug("Failed to delete temporary sort directory {}", tempDir, e); + } + } + + /** + * Sorts {@code rows} according to {@code sortOrder}, spilling to local disk beyond the in-memory + * buffer. + * + *

The caller must close the returned iterable (ideally via try-with-resources). Closing + * reclaims the temporary spill directory; skipping it leaks that directory on the worker's local + * disk, which matters on long-lived runner harnesses that process many bundles. + * + *

Note that spilled rows are written unencrypted to the worker's local disk, in a directory + * created with owner-only permissions. + */ + public static CloseableIterable sortRows( + Iterable rows, SortOrder sortOrder, - String[] columnNames, - ByteArrayOutputStream baos, Schema icebergSchema, - org.apache.beam.sdk.schemas.Schema beamSchema) - throws IOException { + org.apache.beam.sdk.schemas.Schema beamSchema) { - List fields = sortOrder.fields(); - Type[] sourceTypes = new Type[fields.size()]; - Type[] resultTypes = new Type[fields.size()]; - @Nullable SerializableFunction[] boundTransforms = - new SerializableFunction[fields.size()]; - for (int i = 0; i < fields.size(); i++) { - SortField field = fields.get(i); - Type sourceType = icebergSchema.findType(field.sourceId()); - sourceTypes[i] = sourceType; - if (field.transform().isIdentity()) { - resultTypes[i] = sourceType; - boundTransforms[i] = null; - } else { - Transform transform = (Transform) field.transform(); - resultTypes[i] = transform.getResultType(sourceType); - boundTransforms[i] = transform.bind(sourceType); + if (sortOrder == null || !sortOrder.isSorted()) { + return CloseableIterable.withNoopClose(rows); + } + + IcebergRowSorter rowSorter = new IcebergRowSorter(sortOrder, icebergSchema, beamSchema); + try { + for (Row row : rows) { + rowSorter.add(null, row); } + return rowSorter.sort(); + } catch (RuntimeException e) { + rowSorter.close(); + throw e; } - encodeSortKey(row, sortOrder, columnNames, sourceTypes, resultTypes, boundTransforms, baos); + } + + @VisibleForTesting + static void encodeSortKey( + Row row, SortOrder sortOrder, ByteArrayOutputStream baos, Schema icebergSchema) + throws IOException { + new KeyEncoder(sortOrder, icebergSchema, row.getSchema()).encode(row, baos); } @SuppressWarnings("nullness") @@ -239,6 +323,8 @@ private static void encodeSortKey( Row row, SortOrder sortOrder, String[] columnNames, + int[] beamFieldIndices, + @Nullable EnumerationType[] enumTypes, Type[] sourceTypes, Type[] resultTypes, @Nullable SerializableFunction[] boundTransforms, @@ -249,13 +335,24 @@ private static void encodeSortKey( for (int i = 0; i < fields.size(); i++) { SortField field = fields.get(i); - String colName = columnNames[i]; // Normalize into the Iceberg value domain up front so that the encoder only ever deals with // Iceberg's representation, regardless of which Beam logical type the row carried. Type sourceType = sourceTypes[i]; Type valueType = resultTypes[i]; - Object val = IcebergUtils.beamValueToIcebergValue(sourceType, row.getValue(colName)); + int fieldIdx = beamFieldIndices[i]; + Object rawVal = fieldIdx >= 0 ? row.getValue(fieldIdx) : row.getValue(columnNames[i]); + if (rawVal instanceof EnumerationType.Value) { + EnumerationType enumType = + enumTypes[i] != null + ? enumTypes[i] + : row.getSchema() + .getField(columnNames[i]) + .getType() + .getLogicalType(EnumerationType.class); + rawVal = enumType.toString((EnumerationType.Value) rawVal); + } + Object val = IcebergUtils.beamValueToIcebergValue(sourceType, rawVal); SerializableFunction boundTransform = boundTransforms[i]; if (boundTransform != null && val != null) { @@ -392,7 +489,7 @@ private static void writeValue( } else if (val instanceof LocalDate) { writeLong(((LocalDate) val).toEpochDay() ^ Long.MIN_VALUE, baos, invert); } else if (val instanceof LocalTime) { - writeLong(((LocalTime) val).toNanoOfDay() ^ Long.MIN_VALUE, baos, invert); + writeLong(DateTimeUtil.microsFromTime((LocalTime) val) ^ Long.MIN_VALUE, baos, invert); } else if (val instanceof LocalDateTime) { writeLong( DateTimeUtil.microsFromTimestamp((LocalDateTime) val) ^ Long.MIN_VALUE, baos, invert); @@ -400,11 +497,11 @@ private static void writeValue( writeLong( DateTimeUtil.microsFromTimestamptz((OffsetDateTime) val) ^ Long.MIN_VALUE, baos, invert); } else if (val instanceof ReadableInstant) { - writeLong(((ReadableInstant) val).getMillis() ^ Long.MIN_VALUE, baos, invert); + writeLong((((ReadableInstant) val).getMillis() * 1000L) ^ Long.MIN_VALUE, baos, invert); } else if (val instanceof Instant) { - writeLong(((Instant) val).toEpochMilli() ^ Long.MIN_VALUE, baos, invert); + writeLong(DateTimeUtil.microsFromInstant((Instant) val) ^ Long.MIN_VALUE, baos, invert); } else if (val instanceof Date) { - writeLong(((Date) val).getTime() ^ Long.MIN_VALUE, baos, invert); + writeLong((((Date) val).getTime() * 1000L) ^ Long.MIN_VALUE, baos, invert); } else if (val instanceof UUID) { // Iceberg orders UUIDs by their 16-byte big-endian representation, so flip the sign bit of // each half to turn Java's signed longs into an unsigned lexicographic ordering. diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergUtils.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergUtils.java index 87aa73d716a9..509c95057350 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergUtils.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergUtils.java @@ -49,6 +49,7 @@ 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.collect.ImmutableList; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.primitives.Ints; import org.apache.iceberg.StructLike; import org.apache.iceberg.catalog.TableIdentifier; import org.apache.iceberg.catalog.TableIdentifierParser; @@ -354,6 +355,7 @@ public static org.apache.iceberg.Schema beamSchemaToIcebergSchema(final Schema s case TIME: case DECIMAL: case STRING: + case FIXED: return value; case TIMESTAMP: Types.TimestampType ts = (Types.TimestampType) type.asPrimitiveType(); @@ -368,8 +370,6 @@ public static org.apache.iceberg.Schema beamSchemaToIcebergSchema(final Schema s return ByteBuffer.wrap((byte[]) value); } return value; - case FIXED: - throw new UnsupportedOperationException("Fixed-precision fields are not yet supported."); default: return value; } @@ -814,4 +814,48 @@ static boolean isUnbounded(PCollection input) { static boolean validDirectWriteLimit(@Nullable Integer directWriteByteLimit) { return directWriteByteLimit != null && directWriteByteLimit >= 0; } + + /** + * Estimates the in-memory/encoded byte size of a Beam {@link Row} without the overhead of coder + * serialization. + */ + static int estimateRowBytes(@Nullable Row row) { + if (row == null) { + return 0; + } + long size = 0L; + int fieldCount = row.getFieldCount(); + for (int i = 0; i < fieldCount; i++) { + size += estimateObjectBytes(row.getValue(i)); + } + return Ints.saturatedCast(size); + } + + private static long estimateObjectBytes(@Nullable Object value) { + if (value == null) { + return 1L; + } + if (value instanceof CharSequence) { + return org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Utf8.encodedLength( + (CharSequence) value); + } else if (value instanceof byte[]) { + return ((byte[]) value).length; + } else if (value instanceof Row) { + return estimateRowBytes((Row) value); + } else if (value instanceof Iterable) { + long listSize = 0L; + for (Object item : (Iterable) value) { + listSize += estimateObjectBytes(item); + } + return listSize; + } else if (value instanceof Map) { + long mapSize = 0L; + for (Map.Entry entry : ((Map) value).entrySet()) { + mapSize += estimateObjectBytes(entry.getKey()) + estimateObjectBytes(entry.getValue()); + } + return mapSize; + } else { + return 8L; + } + } } 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 47c9469f8b1d..7bb06ad16a18 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 @@ -194,13 +194,16 @@ public static Builder builder() { @SchemaFieldDescription( "When true, physically sorts records within each written data file according to the target" - + " table's sort order. Defaults to false.") + + " table's sort order. Automatically enabled when distribution_mode is set to 'range'." + + " Defaults to false.") public abstract @Nullable Boolean getSortOnWrite(); @SchemaFieldDescription( "Defines distribution of write data. Supported distributions:" + "\n- none: don't shuffle rows (default)" - + "\n- hash: shuffle rows by partition key before writing data") + + "\n- hash: shuffle rows by partition key before writing data" + + "\n- range: shuffle rows into non-overlapping sort-key ranges per partition and sort" + + " within each file") public abstract @Nullable String getDistributionMode(); @SchemaFieldDescription( @@ -425,6 +428,7 @@ void validateModeOptions() { requireMode(Mode.APPEND, "direct_write_byte_limit", getDirectWriteByteLimit(), unsupported); requireMode(Mode.APPEND, "distribution_mode", getDistributionMode(), unsupported); requireMode(Mode.APPEND, "autosharding", getAutosharding(), unsupported); + requireMode(Mode.APPEND, "sort_on_write", getSortOnWrite(), unsupported); requireMode(Mode.APPEND, "write_properties", getWriteProperties(), unsupported); requireMode( Mode.APPEND, "using_side_input_table_cache", getUseSideInputTableCache(), unsupported); 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..5019a1fc5467 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 @@ -32,6 +32,7 @@ import java.util.UUID; import java.util.concurrent.TimeUnit; import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.util.Preconditions; import org.apache.beam.sdk.values.Row; import org.apache.beam.sdk.values.WindowedValue; @@ -40,10 +41,10 @@ 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; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; 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; @@ -57,6 +58,7 @@ import org.apache.iceberg.data.Record; import org.apache.iceberg.exceptions.AlreadyExistsException; import org.apache.iceberg.exceptions.NoSuchTableException; +import org.apache.iceberg.io.CloseableIterable; import org.apache.iceberg.transforms.Transforms; import org.checkerframework.checker.nullness.qual.Nullable; import org.slf4j.Logger; @@ -72,11 +74,14 @@ * *

A {@link DestinationState} maintains its writers in a {@link Cache}. If a {@link RecordWriter} * is inactive for 1 minute, the {@link DestinationState} will automatically close it to free up - * resources. When a data writer is closed, its resulting {@link DataFile} gets written. Calling - * {@link #close()} on this {@link RecordWriterManager} will do the following for each {@link - * DestinationState}: + * resources. When a data writer is closed, its resulting {@link DataFile} gets written. When {@code + * sortedWrites} is enabled on a sorted table, incoming rows are buffered in an {@link + * IcebergRowSorter} during {@link #write} without opening any {@link RecordWriter}s, and flushed + * one partition at a time during {@link #close()}. Calling {@link #close()} on this {@link + * RecordWriterManager} will do the following for each {@link DestinationState}: * *

    + *
  1. Flush any buffered sorted rows one partition at a time, rolling at {@code maxFileSize} *
  2. Close all underlying {@link RecordWriter}s *
  3. Collect all {@link DataFile}s as {@link SerializableDataFile}s (a more Beam-friendly type) *
@@ -87,6 +92,12 @@ class RecordWriterManager implements AutoCloseable { private static final Logger LOG = LoggerFactory.getLogger(RecordWriterManager.class); + /** + * Per-destination in-memory buffer limit (in MB) before {@link IcebergRowSorter} spills to local + * disk when multiple destinations are buffered concurrently within a bundle. + */ + private static final int DEFAULT_SORTER_MEMORY_MB = 16; + /** * Represents the state of one Iceberg table destination. Creates one {@link RecordWriter} per * partition and manages them in a {@link Cache}. @@ -107,6 +118,8 @@ final class DestinationState { @VisibleForTesting final Map writerCounts = Maps.newHashMap(); private final List exceptions = Lists.newArrayList(); private final InternalRecordWrapper wrapper; // wrapper that facilitates partitioning + private @Nullable BeamRowWrapper beamRowWrapper; + private @Nullable IcebergRowSorter rowSorter; DestinationState(IcebergDestination icebergDestination, Table table) { this.icebergDestination = icebergDestination; @@ -169,6 +182,70 @@ boolean write(Record record) { return true; } + boolean isSorted() { + return sortedWrites && table.sortOrder().isSorted(); + } + + boolean bufferForSortedWrite(Row row) { + IcebergRowSorter sorter = rowSorter; + if (sorter == null) { + if (openWriters >= maxNumWriters) { + writers.cleanUp(); + if (openWriters >= maxNumWriters) { + return false; + } + } + sorter = + new IcebergRowSorter( + table.sortOrder(), schema, row.getSchema(), DEFAULT_SORTER_MEMORY_MB); + rowSorter = sorter; + openWriters++; + } + BeamRowWrapper rowWrapper = beamRowWrapper; + if (rowWrapper == null) { + rowWrapper = new BeamRowWrapper(row.getSchema(), schema.asStruct()); + beamRowWrapper = rowWrapper; + } + routingPartitionKey.partition(rowWrapper.wrap(row)); + sorter.add(routingPartitionKey.toPath(), row); + return true; + } + + void flushSortedRows() throws IOException { + IcebergRowSorter sorter = rowSorter; + if (sorter == null) { + return; + } + rowSorter = null; + openWriters--; + try (CloseableIterable sortedRows = sorter.sort()) { + String activePartitionPath = null; + for (Row row : sortedRows) { + Record record = IcebergUtils.beamRowToIcebergRecord(schema, row); + routingPartitionKey.partition(wrapper.wrap(record)); + String partitionPath = routingPartitionKey.toPath(); + if (activePartitionPath != null && !partitionPath.equals(activePartitionPath)) { + // Rows are ordered by partitionPath first; once the partition path changes, the + // previous partition group is complete and its writer(s) can be closed immediately so + // only the active partition's writer stays open in memory at a time. + writers.invalidateAll(); + } + activePartitionPath = partitionPath; + @Nullable RecordWriter writer = writers.getIfPresent(routingPartitionKey); + writer = fetchWriterForPartition(routingPartitionKey, writer); + writer.write(record); + } + } + } + + void closeSorterQuietly() { + if (rowSorter != null) { + rowSorter.close(); + rowSorter = null; + openWriters--; + } + } + /** * Checks if a viable {@link RecordWriter} already exists for this partition and returns it. If * no {@link RecordWriter} exists or if it has reached the maximum limit of bytes written, a new @@ -251,6 +328,7 @@ static String getPartitionDataPath( private final long maxFileSize; private final int maxNumWriters; private final @Nullable Map writeProperties; + private final boolean sortedWrites; private volatile @Nullable Map sideInputTableSpecs; @VisibleForTesting int openWriters = 0; @@ -264,7 +342,7 @@ static String getPartitionDataPath( RecordWriterManager( IcebergCatalogConfig catalogConfig, String filePrefix, long maxFileSize, int maxNumWriters) { - this(catalogConfig, filePrefix, maxFileSize, maxNumWriters, null, null); + this(catalogConfig, filePrefix, maxFileSize, maxNumWriters, null, null, false); } RecordWriterManager( @@ -273,7 +351,7 @@ static String getPartitionDataPath( long maxFileSize, int maxNumWriters, @Nullable Map writeProperties) { - this(catalogConfig, filePrefix, maxFileSize, maxNumWriters, writeProperties, null); + this(catalogConfig, filePrefix, maxFileSize, maxNumWriters, writeProperties, null, false); } RecordWriterManager( @@ -283,12 +361,31 @@ static String getPartitionDataPath( int maxNumWriters, @Nullable Map writeProperties, @Nullable Map sideInputTableSpecs) { + this( + catalogConfig, + filePrefix, + maxFileSize, + maxNumWriters, + writeProperties, + sideInputTableSpecs, + false); + } + + RecordWriterManager( + IcebergCatalogConfig catalogConfig, + String filePrefix, + long maxFileSize, + int maxNumWriters, + @Nullable Map writeProperties, + @Nullable Map sideInputTableSpecs, + boolean sortedWrites) { this.catalogConfig = catalogConfig; this.filePrefix = filePrefix; this.maxFileSize = maxFileSize; this.maxNumWriters = maxNumWriters; this.writeProperties = writeProperties; this.sideInputTableSpecs = sideInputTableSpecs; + this.sortedWrites = sortedWrites; } @VisibleForTesting @@ -307,8 +404,8 @@ void setSideInputTableSpecs(@Nullable Map sideInp * implementation. Although it is expected, some implementations may not support creating a table * using the Iceberg API. */ - @VisibleForTesting - Table getOrCreateTable( + static Table getOrCreateTable( + IcebergCatalogConfig catalogConfig, IcebergDestination destination, Schema dataSchema, @Nullable Map sideInputTableSpecs) { @@ -321,15 +418,24 @@ Table getOrCreateTable( } } return TableCache.getAndRefreshIfStale( - catalogConfig, identifier, () -> loadOrCreateTable(destination, dataSchema)); + catalogConfig, identifier, () -> loadOrCreateTable(catalogConfig, destination, dataSchema)); + } + + @VisibleForTesting + Table getOrCreateTable( + IcebergDestination destination, + Schema dataSchema, + @Nullable Map sideInputTableSpecs) { + return getOrCreateTable(catalogConfig, destination, dataSchema, sideInputTableSpecs); } @VisibleForTesting Table getOrCreateTable(IcebergDestination destination, Schema dataSchema) { - return getOrCreateTable(destination, dataSchema, this.sideInputTableSpecs); + return getOrCreateTable(catalogConfig, destination, dataSchema, this.sideInputTableSpecs); } - private Table loadOrCreateTable(IcebergDestination destination, Schema dataSchema) { + private static Table loadOrCreateTable( + IcebergCatalogConfig catalogConfig, IcebergDestination destination, Schema dataSchema) { Catalog catalog = catalogConfig.catalog(); TableIdentifier identifier = destination.getTableIdentifier(); Namespace namespace = identifier.namespace(); @@ -420,13 +526,17 @@ public boolean write(WindowedValue icebergDestination, Row r return new DestinationState(dest, table); }); + if (destinationState.isSorted()) { + return destinationState.bufferForSortedWrite(row); + } + Record icebergRecord = IcebergUtils.beamRowToIcebergRecord(destinationState.schema, row); return destinationState.write(icebergRecord); } /** - * Closes all remaining writers and collects all their {@link DataFile}s. Writes one {@link - * ManifestFile} per windowed table destination. + * Flushes any buffered sorted rows, closes all remaining writers, and collects all produced + * {@link DataFile}s as {@link SerializableDataFile}s per windowed table destination. */ @Override public void close() throws IOException { @@ -435,6 +545,8 @@ public void close() throws IOException { windowedDestinationAndState : destinations.entrySet()) { DestinationState state = windowedDestinationAndState.getValue(); + state.flushSortedRows(); + // removing writers from the state's cache will trigger the logic to collect each writer's // data file. state.writers.invalidateAll(); @@ -458,6 +570,10 @@ public void close() throws IOException { state.dataFiles.clear(); } } finally { + for (DestinationState state : destinations.values()) { + state.closeSorterQuietly(); + state.writers.invalidateAll(); + } destinations.clear(); } checkArgument( @@ -480,4 +596,29 @@ public void close() throws IOException { getClass().getSimpleName()); return totalSerializableDataFiles; } + + /** + * Closes {@code manager} and emits all produced {@link SerializableDataFile}s to {@code context}. + */ + static void closeAndOutputFiles( + @Nullable RecordWriterManager manager, DoFn.FinishBundleContext context) + throws IOException { + if (manager == null) { + return; + } + manager.close(); + for (Map.Entry, List> + destinationAndFiles : manager.getSerializableDataFiles().entrySet()) { + WindowedValue windowedDestination = destinationAndFiles.getKey(); + for (SerializableDataFile dataFile : destinationAndFiles.getValue()) { + context.output( + FileWriteResult.builder() + .setSerializableDataFile(dataFile) + .setTableIdentifier(windowedDestination.getValue().getTableIdentifier()) + .build(), + windowedDestination.getTimestamp(), + Iterables.getFirst(windowedDestination.getWindows(), null)); + } + } + } } 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 8bafe9eeaef9..5b68bad2d329 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 @@ -17,7 +17,6 @@ */ 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; @@ -31,7 +30,6 @@ 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.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; import org.checkerframework.checker.nullness.qual.Nullable; class WriteDirectRowsToFiles @@ -43,6 +41,7 @@ class WriteDirectRowsToFiles private final long maxBytesPerFile; private final @Nullable Map writeProperties; private final @Nullable PCollectionView> metadataView; + private final boolean sortedWrites; WriteDirectRowsToFiles( IcebergCatalogConfig catalogConfig, @@ -50,7 +49,14 @@ class WriteDirectRowsToFiles String filePrefix, long maxBytesPerFile, @Nullable Map writeProperties) { - this(catalogConfig, dynamicDestinations, filePrefix, maxBytesPerFile, writeProperties, null); + this( + catalogConfig, + dynamicDestinations, + filePrefix, + maxBytesPerFile, + writeProperties, + null, + false); } WriteDirectRowsToFiles( @@ -60,12 +66,31 @@ class WriteDirectRowsToFiles long maxBytesPerFile, @Nullable Map writeProperties, @Nullable PCollectionView> metadataView) { + this( + catalogConfig, + dynamicDestinations, + filePrefix, + maxBytesPerFile, + writeProperties, + metadataView, + false); + } + + WriteDirectRowsToFiles( + IcebergCatalogConfig catalogConfig, + DynamicDestinations dynamicDestinations, + String filePrefix, + long maxBytesPerFile, + @Nullable Map writeProperties, + @Nullable PCollectionView> metadataView, + boolean sortedWrites) { this.catalogConfig = catalogConfig; this.dynamicDestinations = dynamicDestinations; this.filePrefix = filePrefix; this.maxBytesPerFile = maxBytesPerFile; this.writeProperties = writeProperties; this.metadataView = metadataView; + this.sortedWrites = sortedWrites; } @Override @@ -78,7 +103,8 @@ public PCollection expand(PCollection> input) { maxBytesPerFile, filePrefix, writeProperties, - metadataView)); + metadataView, + sortedWrites)); if (metadataView != null) { parDo = parDo.withSideInputs(metadataView); } @@ -93,30 +119,24 @@ private static class WriteDirectRowsToFilesDoFn extends DoFn, Fi private final long maxFileSize; private final @Nullable Map writeProperties; private final @Nullable PCollectionView> metadataView; + private final boolean sortedWrites; private transient @Nullable RecordWriterManager recordWriterManager; - WriteDirectRowsToFilesDoFn( - IcebergCatalogConfig catalogConfig, - DynamicDestinations dynamicDestinations, - long maxFileSize, - String filePrefix, - @Nullable Map writeProperties) { - this(catalogConfig, dynamicDestinations, maxFileSize, filePrefix, writeProperties, null); - } - WriteDirectRowsToFilesDoFn( IcebergCatalogConfig catalogConfig, DynamicDestinations dynamicDestinations, long maxFileSize, String filePrefix, @Nullable Map writeProperties, - @Nullable PCollectionView> metadataView) { + @Nullable PCollectionView> metadataView, + boolean sortedWrites) { this.catalogConfig = catalogConfig; this.dynamicDestinations = dynamicDestinations; this.filePrefix = filePrefix; this.maxFileSize = maxFileSize; this.writeProperties = writeProperties; this.metadataView = metadataView; + this.sortedWrites = sortedWrites; this.recordWriterManager = null; } @@ -124,7 +144,13 @@ private static class WriteDirectRowsToFilesDoFn extends DoFn, Fi public void startBundle() { recordWriterManager = new RecordWriterManager( - catalogConfig, filePrefix, maxFileSize, Integer.MAX_VALUE, writeProperties); + catalogConfig, + filePrefix, + maxFileSize, + Integer.MAX_VALUE, + writeProperties, + null, + sortedWrites); } @ProcessElement @@ -140,34 +166,23 @@ public void processElement( WindowedValues.of(destination, window.maxTimestamp(), window, paneInfo); Map sideInputs = metadataView != null ? context.sideInput(metadataView) : null; - Preconditions.checkNotNull(recordWriterManager) - .write(windowedDestination, element.getValue(), sideInputs); + try { + Preconditions.checkNotNull(recordWriterManager) + .write(windowedDestination, element.getValue(), sideInputs); + } catch (Exception e) { + try { + Preconditions.checkNotNull(recordWriterManager).close(); + } catch (Exception closeException) { + e.addSuppressed(closeException); + } + recordWriterManager = null; + throw e; + } } @FinishBundle public void finishBundle(FinishBundleContext context) throws Exception { - if (recordWriterManager == null) { - return; - } - recordWriterManager.close(); - - for (Map.Entry, List> - destinationAndFiles : - Preconditions.checkNotNull(recordWriterManager) - .getSerializableDataFiles() - .entrySet()) { - WindowedValue windowedDestination = destinationAndFiles.getKey(); - - for (SerializableDataFile dataFile : destinationAndFiles.getValue()) { - context.output( - FileWriteResult.builder() - .setSerializableDataFile(dataFile) - .setTableIdentifier(windowedDestination.getValue().getTableIdentifier()) - .build(), - windowedDestination.getTimestamp(), - Iterables.getFirst(windowedDestination.getWindows(), null)); - } - } + RecordWriterManager.closeAndOutputFiles(recordWriterManager, context); recordWriterManager = null; } } 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 7b850103b1cf..8c83d8daffbb 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 @@ -17,9 +17,9 @@ */ package org.apache.beam.sdk.io.iceberg; +import java.util.Collections; import java.util.List; import java.util.Map; -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; @@ -32,9 +32,6 @@ 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.apache.iceberg.Table; -import org.apache.iceberg.io.CloseableIterable; import org.checkerframework.checker.nullness.qual.Nullable; class WriteGroupedRowsToFiles @@ -102,7 +99,6 @@ class WriteGroupedRowsToFiles @Override public PCollection expand( PCollection, Iterable>> input) { - Schema dataSchema = dynamicDestinations.getDataSchema(); ParDo.SingleOutput, Iterable>, FileWriteResult> parDo = ParDo.of( new WriteGroupedRowsToFilesDoFn( @@ -110,7 +106,6 @@ public PCollection expand( dynamicDestinations, maxBytesPerFile, filePrefix, - dataSchema, writeProperties, metadataView, sortedWrites)); @@ -127,7 +122,6 @@ private static class WriteGroupedRowsToFilesDoFn private final IcebergCatalogConfig catalogConfig; private final String filePrefix; private final long maxFileSize; - private final Schema dataSchema; private final @Nullable Map writeProperties; private final @Nullable PCollectionView> metadataView; private final boolean sortedWrites; @@ -137,44 +131,6 @@ private static class WriteGroupedRowsToFilesDoFn DynamicDestinations dynamicDestinations, long maxFileSize, String filePrefix, - Schema dataSchema, - @Nullable Map writeProperties) { - this( - catalogConfig, - dynamicDestinations, - maxFileSize, - filePrefix, - dataSchema, - writeProperties, - null, - false); - } - - WriteGroupedRowsToFilesDoFn( - IcebergCatalogConfig catalogConfig, - DynamicDestinations dynamicDestinations, - long maxFileSize, - String filePrefix, - Schema dataSchema, - @Nullable Map writeProperties, - @Nullable PCollectionView> metadataView) { - this( - catalogConfig, - dynamicDestinations, - maxFileSize, - filePrefix, - dataSchema, - writeProperties, - metadataView, - false); - } - - WriteGroupedRowsToFilesDoFn( - IcebergCatalogConfig catalogConfig, - DynamicDestinations dynamicDestinations, - long maxFileSize, - String filePrefix, - Schema dataSchema, @Nullable Map writeProperties, @Nullable PCollectionView> metadataView, boolean sortedWrites) { @@ -182,7 +138,6 @@ private static class WriteGroupedRowsToFilesDoFn this.dynamicDestinations = dynamicDestinations; this.filePrefix = filePrefix; this.maxFileSize = maxFileSize; - this.dataSchema = dataSchema; this.writeProperties = writeProperties; this.metadataView = metadataView; this.sortedWrites = sortedWrites; @@ -210,22 +165,18 @@ public void processElement( maxFileSize, Integer.MAX_VALUE, writeProperties, - sideInputs)) { + sideInputs, + sortedWrites)) { writer = openWriter; - Table table = sortedWrites ? openWriter.getOrCreateTable(destination, dataSchema) : null; - try (CloseableIterable sortedOrUnsortedRows = - table != null && table.sortOrder().isSorted() - ? IcebergRowSorter.sortRows( - element.getValue(), table.sortOrder(), table.schema(), dataSchema) - : CloseableIterable.withNoopClose(element.getValue())) { - for (Row e : sortedOrUnsortedRows) { - writer.write(windowedDestination, e); - } + for (Row e : element.getValue()) { + writer.write(windowedDestination, e); } } List serializableDataFiles = - Preconditions.checkNotNull(writer.getSerializableDataFiles().get(windowedDestination)); + writer + .getSerializableDataFiles() + .getOrDefault(windowedDestination, Collections.emptyList()); for (SerializableDataFile dataFile : serializableDataFiles) { c.output( FileWriteResult.builder() 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 fef0ab034d7d..f4dee7c9a8cc 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 @@ -38,26 +38,16 @@ 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.apache.iceberg.io.CloseableIterable; import org.checkerframework.checker.nullness.qual.MonotonicNonNull; 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; @@ -165,41 +155,6 @@ private static class WriteDoFn extends DoFn>, FileWriteRes private transient @MonotonicNonNull Map> partitionFieldMaps; - WriteDoFn( - IcebergCatalogConfig catalogConfig, - DynamicDestinations dynamicDestinations, - String filePrefix, - Schema dataSchema, - @Nullable Map writeProperties) { - this( - catalogConfig, - dynamicDestinations, - filePrefix, - WriteToPartitions.DEFAULT_BYTES_PER_FILE, - dataSchema, - writeProperties, - null, - false); - } - - WriteDoFn( - IcebergCatalogConfig catalogConfig, - DynamicDestinations dynamicDestinations, - String filePrefix, - Schema dataSchema, - @Nullable Map writeProperties, - @Nullable PCollectionView> metadataView) { - this( - catalogConfig, - dynamicDestinations, - filePrefix, - WriteToPartitions.DEFAULT_BYTES_PER_FILE, - dataSchema, - writeProperties, - metadataView, - false); - } - WriteDoFn( IcebergCatalogConfig catalogConfig, DynamicDestinations dynamicDestinations, @@ -274,16 +229,18 @@ public void processElement( Record record = IcebergUtils.beamRowToIcebergRecord(table.schema(), row); writer.write(record); } + if (writer != null) { + RecordWriter finishedWriter = writer; + writer = null; + finishedWriter.close(); + writtenFiles.add(SerializableDataFile.from(finishedWriter.getDataFile(), table.specs())); + } } finally { if (writer != null) { writer.close(); } } - if (writer != null) { - writtenFiles.add(SerializableDataFile.from(writer.getDataFile(), table.specs())); - } - for (SerializableDataFile sdf : writtenFiles) { out.output( FileWriteResult.builder() @@ -312,81 +269,12 @@ 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)); + return RecordWriterManager.getOrCreateTable( + catalogConfig, destination, dataSchema, sideInputTableSpecs); } 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); - } - } - } } } 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 95a926e71f86..17e4e8f4a85b 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 @@ -17,10 +17,8 @@ */ package org.apache.beam.sdk.io.iceberg; -import static java.nio.charset.StandardCharsets.UTF_8; import static org.apache.beam.sdk.util.Preconditions.checkArgumentNotNull; -import java.util.List; import java.util.Map; import java.util.UUID; import org.apache.beam.sdk.coders.IterableCoder; @@ -122,7 +120,7 @@ public IcebergWriteResult expand(PCollection> input) { PCollection writtenFiles; if (IcebergUtils.isUnbounded(input)) { writtenFiles = - IcebergUtils.validDirectWriteLimit(directWriteByteLimit) && !sortedWrites + IcebergUtils.validDirectWriteLimit(directWriteByteLimit) ? writeTriggeredWithBundleLifting(input) : writeTriggered(input); } else { @@ -211,7 +209,8 @@ private PCollection writeTriggeredWithBundleLifting( filePrefix, DEFAULT_MAX_BYTES_PER_FILE, writeProperties, - metadataView)); + metadataView, + sortedWrites)); PCollection groupedFileWrites = groupAndWriteRecords(smallBatches); @@ -279,45 +278,7 @@ private PCollection writeUntriggered(PCollection, Integer> { @Override public Integer apply(KV element) { - return estimateRowSize(element.getValue()); - } - - private int estimateRowSize(Row row) { - if (row == null) { - return 0; - } - int size = 0; - for (Object value : row.getValues()) { - size += estimateObjectSize(value); - } - return size; - } - - private int estimateObjectSize(@Nullable Object value) { - if (value == null) { - return 0; - } - if (value instanceof String) { - return ((String) value).getBytes(UTF_8).length; - } else if (value instanceof byte[]) { - return ((byte[]) value).length; - } else if (value instanceof Row) { - return estimateRowSize((Row) value); - } else if (value instanceof List) { - int listSize = 0; - for (Object item : (List) value) { - listSize += estimateObjectSize(item); - } - return listSize; - } else if (value instanceof Map) { - int mapSize = 0; - for (Map.Entry entry : ((Map) value).entrySet()) { - mapSize += estimateObjectSize(entry.getKey()) + estimateObjectSize(entry.getValue()); - } - return mapSize; - } else { - return 8; // Approximation for other fields - } + return IcebergUtils.estimateRowBytes(element.getValue()); } } } 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 114c3ef12b95..2e60c4f10f47 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 @@ -18,7 +18,6 @@ package org.apache.beam.sdk.io.iceberg; import java.nio.ByteBuffer; -import java.util.List; import java.util.Map; import java.util.concurrent.ThreadLocalRandom; import org.apache.beam.sdk.Pipeline; @@ -47,9 +46,6 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps; -import org.apache.iceberg.Table; import org.checkerframework.checker.nullness.qual.Nullable; /** @@ -248,47 +244,9 @@ private static class WriteUngroupedRowsToFilesDoFn private final @Nullable PCollectionView> metadataView; private final boolean sortedWrites; private transient @Nullable RecordWriterManager recordWriterManager; - private transient @Nullable Map destinationIsSorted; private int spilledShardNumber; - public WriteUngroupedRowsToFilesDoFn( - IcebergCatalogConfig catalogConfig, - DynamicDestinations dynamicDestinations, - String filename, - int maximumWritersPerBundle, - long maxFileSize, - @Nullable Map writeProperties) { - this( - catalogConfig, - dynamicDestinations, - filename, - maximumWritersPerBundle, - maxFileSize, - writeProperties, - null, - false); - } - - public WriteUngroupedRowsToFilesDoFn( - IcebergCatalogConfig catalogConfig, - DynamicDestinations dynamicDestinations, - String filename, - int maximumWritersPerBundle, - long maxFileSize, - @Nullable Map writeProperties, - @Nullable PCollectionView> metadataView) { - this( - catalogConfig, - dynamicDestinations, - filename, - maximumWritersPerBundle, - maxFileSize, - writeProperties, - metadataView, - false); - } - - public WriteUngroupedRowsToFilesDoFn( + WriteUngroupedRowsToFilesDoFn( IcebergCatalogConfig catalogConfig, DynamicDestinations dynamicDestinations, String filename, @@ -311,8 +269,13 @@ public WriteUngroupedRowsToFilesDoFn( public void startBundle() { recordWriterManager = new RecordWriterManager( - catalogConfig, filename, maxFileSize, maxWritersPerBundle, writeProperties); - destinationIsSorted = Maps.newHashMap(); + catalogConfig, + filename, + maxFileSize, + maxWritersPerBundle, + writeProperties, + null, + sortedWrites); this.spilledShardNumber = ThreadLocalRandom.current().nextInt(SPILLED_RECORD_SHARDING_FACTOR); } @@ -334,25 +297,14 @@ public void processElement( metadataView != null ? c.sideInput(metadataView) : null; // Attempt to write record. If the writer is saturated and cannot accept - // the record, or if sorted writes are enabled and the target table is sorted, - // spill it over to WriteGroupedRowsToFiles - boolean writeSuccess = false; + // the record, spill it over to WriteGroupedRowsToFiles. When sortedWrites + // is enabled, RecordWriterManager buffers into IcebergRowSorter during + // processElement and flushes one partition at a time in finishBundle. + boolean writeSuccess; try { - RecordWriterManager writerManager = Preconditions.checkNotNull(recordWriterManager); - boolean mustSpillForSorting = false; - if (sortedWrites) { - Map isSortedCache = Preconditions.checkNotNull(destinationIsSorted); - Boolean isSorted = isSortedCache.get(dest); - if (isSorted == null) { - Table table = writerManager.getOrCreateTable(destination, data.getSchema(), sideInputs); - isSorted = table.sortOrder().isSorted(); - isSortedCache.put(dest, isSorted); - } - mustSpillForSorting = isSorted; - } - if (!mustSpillForSorting) { - writeSuccess = writerManager.write(windowedDestination, data, sideInputs); - } + writeSuccess = + Preconditions.checkNotNull(recordWriterManager) + .write(windowedDestination, data, sideInputs); } catch (Exception e) { try { Preconditions.checkNotNull(recordWriterManager).close(); @@ -373,28 +325,7 @@ public void processElement( @FinishBundle public void finishBundle(FinishBundleContext c) throws Exception { - if (recordWriterManager == null) { - return; - } - recordWriterManager.close(); - - for (Map.Entry, List> - destinationAndFiles : - Preconditions.checkNotNull(recordWriterManager) - .getSerializableDataFiles() - .entrySet()) { - WindowedValue windowedDestination = destinationAndFiles.getKey(); - - for (SerializableDataFile dataFile : destinationAndFiles.getValue()) { - c.output( - FileWriteResult.builder() - .setSerializableDataFile(dataFile) - .setTableIdentifier(windowedDestination.getValue().getTableIdentifier()) - .build(), - windowedDestination.getTimestamp(), - Iterables.getFirst(windowedDestination.getWindows(), null)); - } - } + RecordWriterManager.closeAndOutputFiles(recordWriterManager, c); recordWriterManager = null; } } diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/AssignDestinationsAndPartitionsTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/AssignDestinationsAndPartitionsTest.java index a753c7b947a9..8883f95a6163 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/AssignDestinationsAndPartitionsTest.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/AssignDestinationsAndPartitionsTest.java @@ -17,7 +17,13 @@ */ package org.apache.beam.sdk.io.iceberg; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + import java.io.Serializable; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.List; import java.util.Map; import org.apache.beam.sdk.coders.KvCoder; import org.apache.beam.sdk.coders.StringUtf8Coder; @@ -197,4 +203,45 @@ public void testAssignDestinationsWithMetadataViewMissFallsBack() { pipeline.run(); } + + @Test + public void testWeightedReservoirSamplingMergesSkewedBundlesCorrectly() { + AssignDestinationsAndPartitions.PartitionQuantilesCombineFn combineFn = + new AssignDestinationsAndPartitions.PartitionQuantilesCombineFn(1000L); + AssignDestinationsAndPartitions.PartitionSample accum = combineFn.createAccumulator(); + + // First, fill the 2048-slot reservoir using 11 small bundles of 200 rows each (2,200 rows total + // with keys "small_0000" .. "small_0199"). + for (int b = 0; b < 11; b++) { + List smallSamples = new ArrayList<>(); + for (int i = 0; i < 200; i++) { + smallSamples.add(String.format("small_%04d", i).getBytes(StandardCharsets.UTF_8)); + } + accum = + combineFn.addInput( + accum, new AssignDestinationsAndPartitions.PartitionSample(200L, 200L, smallSamples)); + } + + // Next, merge a single huge bundle representing 1,000,000 rows (with 200 reservoir samples + // "z_huge_0000" .. "z_huge_0199"). Because the huge bundle represents >99.7% of all rows, + // weighted reservoir sampling must ensure that the resulting quantile splits come from the + // huge bundle's key domain rather than being drowned out by the earlier small bundles. + List hugeSamples = new ArrayList<>(); + for (int i = 0; i < 200; i++) { + hugeSamples.add(String.format("z_huge_%04d", i).getBytes(StandardCharsets.UTF_8)); + } + accum = + combineFn.addInput( + accum, + new AssignDestinationsAndPartitions.PartitionSample(10_000L, 1_000_000L, hugeSamples)); + + List splits = combineFn.extractOutput(accum); + assertFalse("Expected non-empty split points", splits.isEmpty()); + for (byte[] split : splits) { + String splitKey = new String(split, StandardCharsets.UTF_8); + assertTrue( + "Expected split point to reflect the 1M-row bundle, but got: " + splitKey, + splitKey.startsWith("z_huge_")); + } + } } diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergIOSideInputTableCacheTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergIOSideInputTableCacheTest.java index cb5c954aa777..4b7e44e5e935 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergIOSideInputTableCacheTest.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergIOSideInputTableCacheTest.java @@ -122,7 +122,7 @@ private IcebergIO.WriteRows applyDistribution(IcebergIO.WriteRows write) { write = write.withAutosharding(); } if (distributionMode.equals(RANGE)) { - write = write.withDistributionMode(DistributionMode.RANGE).withDistributionFunction(row -> 0); + write = write.withDistributionMode(DistributionMode.RANGE); } return write; } 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 d48f29428b32..fd5495857d81 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 @@ -1071,9 +1071,8 @@ public void testStreamingSortedWriteWithDirectWriteLimit() throws Exception { .advanceProcessingTime(Duration.standardSeconds(10)) .advanceWatermarkToInfinity(); - // Even with directWriteByteLimit = 1 (which would normally lift large bundles to - // WriteDirectRowsToFiles), enabling withSortedWrites() skips bundle lifting so all streamed - // batches are sorted before writing. + // With directWriteByteLimit = 1, BundleLifter lifts the bundle to WriteDirectRowsToFiles, + // which sorts the bundle locally in finishBundle before writing. testPipeline .apply(stream) .apply( @@ -1090,6 +1089,163 @@ public void testStreamingSortedWriteWithDirectWriteLimit() throws Exception { assertFilesAreInternallySorted(table); } + @Test + public void testRangeDistributionDynamicQuantileShards() throws Exception { + assumeTrue(distributionMode.equals(NONE)); + + Schema schema = Schema.builder().addInt64Field("id").addStringField("name").build(); + TableIdentifier tableId = + TableIdentifier.of( + "default", "range_dynamic_" + Long.toString(UUID.randomUUID().hashCode(), 16)); + 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(); + + org.apache.iceberg.Schema icebergSchema = IcebergUtils.beamSchemaToIcebergSchema(schema); + catalog + .catalog() + .buildTable(tableId, icebergSchema) + .withPartitionSpec(PartitionSpec.unpartitioned()) + .withSortOrder(SortOrder.builderFor(icebergSchema).asc("name").desc("id").build()) + .create(); + + List scrambledRows = new java.util.ArrayList<>(); + for (long i = 0; i < 40; i++) { + scrambledRows.add( + Row.withSchema(schema).addValues(i, String.format("key_%03d", 39 - i)).build()); + } + + // Use PortableIcebergDestinations (whose createConfig is non-null with null partitionFields + // and sortFields) to verify that DestinationMetadataResolver loads the pre-existing table's + // SortOrder from TableCache. + DynamicDestinations dynamicDestinations = + new PortableIcebergDestinations( + tableId.toString(), + FileFormat.PARQUET.toString(), + schema, + null, + null, + null, + null, + null, + null); + + // Target 150 bytes per range shard so the 40 rows (~800+ bytes) dynamically split into + // multiple non-overlapping quantile shards without any user-provided distributionFunction. + PCollection> assigned = + testPipeline + .apply("CreateScrambled", Create.of(scrambledRows).withRowSchema(schema)) + .apply( + "DynamicRangeAssign", + new AssignDestinationsAndPartitions( + dynamicDestinations, catalog, null, DistributionMode.RANGE, null, 150L)); + + assigned.apply( + "WriteRangeShards", + new WriteToPartitions(catalog, dynamicDestinations, null, false, null, null, true)); + + testPipeline.run().waitUntilFinish(); + + Table table = warehouse.loadTable(tableId); + List tasks = ImmutableList.copyOf(table.newScan().planFiles()); + assertTrue( + "Expected multiple dynamically range-sharded files, got " + tasks.size(), tasks.size() > 1); + assertFilesAreInternallySorted(table); + + List fileNameRanges = new java.util.ArrayList<>(); + for (org.apache.iceberg.FileScanTask task : tasks) { + String path = task.file().path().toString(); + try (org.apache.iceberg.io.CloseableIterable reader = + Parquet.read(table.io().newInputFile(path)) + .project(table.schema()) + .createReaderFunc(org.apache.iceberg.data.parquet.GenericParquetReaders::buildReader) + .build()) { + List records = Lists.newArrayList(reader.iterator()); + String minName = (String) records.get(0).getField("name"); + String maxName = (String) records.get(records.size() - 1).getField("name"); + fileNameRanges.add(new String[] {minName, maxName}); + } + } + fileNameRanges.sort(java.util.Comparator.comparing(r -> r[0])); + for (int i = 1; i < fileNameRanges.size(); i++) { + String prevMax = fileNameRanges.get(i - 1)[1]; + String currMin = fileNameRanges.get(i)[0]; + assertTrue( + "Dynamic range-sharded files have overlapping sort key ranges: " + + prevMax + + " >= " + + currMin, + prevMax.compareTo(currMin) < 0); + } + } + + @Test + public void testWriteUngroupedRowsToFilesSortsWithoutSpilling() throws Exception { + assumeTrue(distributionMode.equals(NONE)); + + Schema schema = Schema.builder().addInt64Field("id").addStringField("name").build(); + TableIdentifier tableId = + TableIdentifier.of( + "default", "ungrouped_sorted_" + Long.toString(UUID.randomUUID().hashCode(), 16)); + 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(); + + org.apache.iceberg.Schema icebergSchema = IcebergUtils.beamSchemaToIcebergSchema(schema); + catalog + .catalog() + .buildTable(tableId, icebergSchema) + .withPartitionSpec(PartitionSpec.unpartitioned()) + .withSortOrder(SortOrder.builderFor(icebergSchema).asc("name").desc("id").build()) + .create(); + + String tableIdString = IcebergUtils.tableIdentifierToString(tableId); + List> inputKVs = + Arrays.asList( + KV.of(tableIdString, Row.withSchema(schema).addValues(40L, "delta").build()), + KV.of(tableIdString, Row.withSchema(schema).addValues(20L, "bravo").build()), + KV.of(tableIdString, Row.withSchema(schema).addValues(10L, "alpha").build()), + KV.of(tableIdString, Row.withSchema(schema).addValues(30L, "charlie").build())); + + DynamicDestinations dynamicDestinations = DynamicDestinations.singleTable(tableId, schema); + + WriteUngroupedRowsToFiles.Result result = + testPipeline + .apply( + "CreateUngrouped", + Create.of(inputKVs) + .withCoder( + KvCoder.of( + org.apache.beam.sdk.coders.StringUtf8Coder.of(), + org.apache.beam.sdk.coders.RowCoder.of(schema)))) + .apply( + "FastPathSortedWrite", + new WriteUngroupedRowsToFiles( + catalog, dynamicDestinations, "fast-sorted", 1L, null, null, true)); + + org.apache.beam.sdk.testing.PAssert.that(result.getSpilledRows()).empty(); + result.getWrittenFiles().apply("CommitFiles", new AppendFilesToTables(catalog, "fast-sorted")); + + testPipeline.run().waitUntilFinish(); + + Table table = warehouse.loadTable(tableId); + assertFilesAreInternallySorted(table); + } + private void assertFilesAreInternallySorted(Table table) throws Exception { for (org.apache.iceberg.FileScanTask task : table.newScan().planFiles()) { String path = task.file().path().toString(); diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergRowSorterTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergRowSorterTest.java index 1eae37f1e030..e08c1a3aa0b3 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergRowSorterTest.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergRowSorterTest.java @@ -70,13 +70,8 @@ public class IcebergRowSorterTest { .lexicographicalComparator(); private static byte[] encodeSortKeyHelper(Row row, SortOrder sortOrder) throws Exception { - java.util.List fields = sortOrder.fields(); - String[] columnNames = new String[fields.size()]; - for (int i = 0; i < fields.size(); i++) { - columnNames[i] = ICEBERG_SCHEMA.findColumnName(fields.get(i).sourceId()); - } ByteArrayOutputStream baos = new ByteArrayOutputStream(); - IcebergRowSorter.encodeSortKey(row, sortOrder, columnNames, baos, ICEBERG_SCHEMA, BEAM_SCHEMA); + IcebergRowSorter.encodeSortKey(row, sortOrder, baos, ICEBERG_SCHEMA); return baos.toByteArray(); } @@ -171,7 +166,7 @@ public void testNullOrderingMatrix() throws Exception { } @Test - public void testEndToEndSorting() { + public void testEndToEndSorting() throws Exception { SortOrder sortOrder = SortOrder.builderFor(ICEBERG_SCHEMA).asc("name").desc("id").build(); List input = @@ -181,9 +176,11 @@ public void testEndToEndSorting() { Row.withSchema(BEAM_SCHEMA).addValues(5, "apple", 1.5, true).build(), Row.withSchema(BEAM_SCHEMA).addValues(10, "cherry", 3.0, false).build()); - Iterable sorted = IcebergRowSorter.sortRows(input, sortOrder, ICEBERG_SCHEMA, BEAM_SCHEMA); - List sortedList = - StreamSupport.stream(sorted.spliterator(), false).collect(Collectors.toList()); + List sortedList; + try (CloseableIterable sorted = + IcebergRowSorter.sortRows(input, sortOrder, ICEBERG_SCHEMA, BEAM_SCHEMA)) { + sortedList = StreamSupport.stream(sorted.spliterator(), false).collect(Collectors.toList()); + } assertEquals(4, sortedList.size()); @@ -201,7 +198,7 @@ public void testEndToEndSorting() { } @Test - public void testScaleAndExternalDiskSpill() { + public void testScaleAndExternalDiskSpill() throws Exception { SortOrder sortOrder = SortOrder.builderFor(ICEBERG_SCHEMA).asc("id").build(); int count = 5000; @@ -213,9 +210,11 @@ public void testScaleAndExternalDiskSpill() { input.add(Row.withSchema(BEAM_SCHEMA).addValues(randomId, "item" + i, 1.0, true).build()); } - Iterable sorted = IcebergRowSorter.sortRows(input, sortOrder, ICEBERG_SCHEMA, BEAM_SCHEMA); - List sortedList = - StreamSupport.stream(sorted.spliterator(), false).collect(Collectors.toList()); + List sortedList; + try (CloseableIterable sorted = + IcebergRowSorter.sortRows(input, sortOrder, ICEBERG_SCHEMA, BEAM_SCHEMA)) { + sortedList = StreamSupport.stream(sorted.spliterator(), false).collect(Collectors.toList()); + } assertEquals(count, sortedList.size()); @@ -256,8 +255,7 @@ private static byte[] encodeSingle( Row row = Row.withSchema(beamSchema).addValue(value).build(); ByteArrayOutputStream baos = new ByteArrayOutputStream(); - IcebergRowSorter.encodeSortKey( - row, sortOrder, new String[] {"v"}, baos, icebergSchema, beamSchema); + IcebergRowSorter.encodeSortKey(row, sortOrder, baos, icebergSchema); return baos.toByteArray(); } @@ -329,6 +327,17 @@ public void testNarrowIntegerKeyOrdering() throws Exception { assertOrdered(Schema.FieldType.BYTE, Types.IntegerType.get(), (byte) -5, (byte) 5); } + @Test + public void testEnumerationTypeKeyOrdering() throws Exception { + org.apache.beam.sdk.schemas.logicaltypes.EnumerationType enumType = + org.apache.beam.sdk.schemas.logicaltypes.EnumerationType.create("ALPHA", "BETA", "GAMMA"); + assertOrdered( + Schema.FieldType.logicalType(enumType), + Types.StringType.get(), + enumType.valueOf("ALPHA"), + enumType.valueOf("GAMMA")); + } + @Test public void testUuidKeyEncodingIsDeterministic() throws Exception { byte[] first = @@ -342,6 +351,20 @@ public void testUuidKeyEncodingIsDeterministic() throws Exception { assertTrue(BYTE_ARR_COMPARATOR.compare(first, other) != 0); } + @Test + public void testFixedAndBinaryKeyOrdering() throws Exception { + assertOrdered( + Schema.FieldType.BYTES, + Types.FixedType.ofLength(4), + new byte[] {0x00, 0x01, 0x02, 0x03}, + new byte[] {0x00, 0x01, 0x02, 0x04}); + assertOrdered( + Schema.FieldType.BYTES, + Types.BinaryType.get(), + new byte[] {0x00, 0x01}, + new byte[] {0x00, 0x01, 0x00}); + } + @Test public void testDecimalRejectsValueExceedingDeclaredScale() { assertThrows( diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergUtilsTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergUtilsTest.java index 6bcdc3ff9049..5975c936f229 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergUtilsTest.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergUtilsTest.java @@ -369,6 +369,32 @@ public void testNullStringConversion() { assertEquals(null, record.getField("v")); } + + @Test + public void testEstimateRowBytes() { + Schema nestedSchema = Schema.builder().addInt64Field("x").addStringField("s").build(); + Schema outerSchema = + Schema.builder() + .addNullableField("nullable_str", Schema.FieldType.STRING) + .addByteArrayField("bytes") + .addRowField("nested", nestedSchema) + .addArrayField("arr", Schema.FieldType.INT32) + .addMapField("map", Schema.FieldType.STRING, Schema.FieldType.INT32) + .build(); + Row nested = Row.withSchema(nestedSchema).addValues(42L, "hello").build(); + Row outer = + Row.withSchema(outerSchema) + .addValues( + null, + new byte[] {1, 2, 3}, + nested, + Arrays.asList(10, 20), + ImmutableMap.of("k", 99)) + .build(); + + // null (1) + byte[3] (3) + nested (8 + 5) + array (8 + 8) + map ("k"=1 + 8) = 42 + assertEquals(42L, IcebergUtils.estimateRowBytes(outer)); + } } @RunWith(JUnit4.class) diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergWriteSchemaTransformProviderTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergWriteSchemaTransformProviderTest.java index 750e680a40a7..c1a2ec7ed06b 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergWriteSchemaTransformProviderTest.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergWriteSchemaTransformProviderTest.java @@ -90,7 +90,10 @@ public class IcebergWriteSchemaTransformProviderTest { @Parameterized.Parameters public static Iterable data() { - return asList(new Object[][] {{DistributionMode.NONE}, {DistributionMode.HASH}}); + return asList( + new Object[][] { + {DistributionMode.NONE}, {DistributionMode.HASH}, {DistributionMode.RANGE} + }); } @Parameterized.Parameter(0) diff --git a/website/www/site/content/en/documentation/io/built-in/iceberg.md b/website/www/site/content/en/documentation/io/built-in/iceberg.md index 3f3b15b57867..6ad5a0e1ac2c 100644 --- a/website/www/site/content/en/documentation/io/built-in/iceberg.md +++ b/website/www/site/content/en/documentation/io/built-in/iceberg.md @@ -307,6 +307,55 @@ pipeline: catalog_properties: *catalog_props {{< /highlight >}} +### Sorting and Distribution Modes + +You can control how records are partitioned, sharded, and sorted before writing to Iceberg data files using `partition_fields`, `sort_fields`, `sort_on_write`, `distribution_mode`, and `autosharding`: + +- **`distribution_mode`**: + - `none` (default): Writes directly from worker bundles without a shuffle. When combined with `sort_on_write: true`, each worker locally sorts its bundle per partition (producing sorted files whose min/max key ranges may overlap across workers). + - `hash`: Shuffles records by Iceberg partition key before writing, consolidating files per partition. + - `range`: Shuffles records by partition key and sort-order key. In batch mode, `range` automatically samples sort keys per partition to compute balanced quantile boundaries sized to `write.target-file-size-bytes`. When combined with `sort_on_write: true`, files within each partition have non-overlapping sort-key ranges. +- **`sort_fields`**: List of sort expressions applied when creating a new table (for example, `["ts DESC NULLS LAST", "id ASC"]` or `"bucket(id, 16) ASC"`). +- **`sort_on_write`**: When `true`, sorts records within each output file using the table's Iceberg `SortOrder` via an external merge sort that spills to local disk when needed. +- **`autosharding`**: Dynamically scales the number of parallel writers per partition when using `hash` or `range`. + +{{< highlight java >}} +Map config = Map.of( + "table", "my_db.my_table", + "catalog_properties", catalogProps, + "partition_fields", List.of("truncate(name, 1)"), + "sort_fields", List.of("age ASC NULLS LAST", "id ASC"), + "sort_on_write", true, + "distribution_mode", "range"); + +pipeline.apply(Managed.write(Managed.ICEBERG).withConfig(config)); +{{< /highlight >}} +{{< highlight py >}} +pipeline | beam.managed.Write( + beam.managed.ICEBERG, + config={ + "table": "my_db.my_table", + "catalog_properties": catalog_props, + "partition_fields": ["truncate(name, 1)"], + "sort_fields": ["age ASC NULLS LAST", "id ASC"], + "sort_on_write": True, + "distribution_mode": "range", + }) +{{< /highlight >}} +{{< highlight yaml >}} +- type: WriteToIceberg + config: + table: "my_db.my_table" + catalog_properties: *catalog_props + partition_fields: + - "truncate(name, 1)" + sort_fields: + - "age ASC NULLS LAST" + - "id ASC" + sort_on_write: true + distribution_mode: "range" +{{< /highlight >}} + ### View Namespaces and Tables If you're on Beam SQL, you can view the newly created resources: From b19c382372d4a616ec232a309e3ec66e25f900f5 Mon Sep 17 00:00:00 2001 From: atognolag Date: Mon, 19 Oct 2026 11:44:22 +0000 Subject: [PATCH 3/3] Fix SpotBugs CT_CONSTRUCTOR_THROW by making IcebergRowSorter final --- .../java/org/apache/beam/sdk/io/iceberg/IcebergRowSorter.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergRowSorter.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergRowSorter.java index e3b838552881..7dab3648206a 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergRowSorter.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergRowSorter.java @@ -65,7 +65,7 @@ * A utility class to sort Beam {@link Row}s based on an Iceberg {@link SortOrder}. Leverages {@link * BufferedExternalSorter} to spill to local disk when elements exceed memory limit. */ -class IcebergRowSorter implements AutoCloseable { +final class IcebergRowSorter implements AutoCloseable { private static final Logger LOG = LoggerFactory.getLogger(IcebergRowSorter.class); /** Width of an encoded decimal key. 16 bytes covers Iceberg's maximum precision of 38 digits. */ @@ -86,7 +86,7 @@ class IcebergRowSorter implements AutoCloseable { * name lookups per row. */ @SuppressWarnings({"unchecked", "rawtypes"}) - static class KeyEncoder { + static final class KeyEncoder { private final SortOrder sortOrder; private final String[] columnNames; private final int[] beamFieldIndices;