diff --git a/CHANGES.md b/CHANGES.md
index c10b7262209c..d997b9fe1100 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`) 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
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..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,23 +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;
@@ -44,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>> {
@@ -54,34 +85,117 @@ class AssignDestinationsAndPartitions
private final DynamicDestinations dynamicDestinations;
private final IcebergCatalogConfig catalogConfig;
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.
+ *
+ * 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,
+ @Nullable PCollectionView> metadataView,
+ DistributionMode distributionMode,
+ @Nullable SerializableFunction distributionFunction) {
+ this(
+ dynamicDestinations,
+ catalogConfig,
+ metadataView,
+ distributionMode,
+ distributionFunction,
+ WriteToPartitions.DEFAULT_BYTES_PER_FILE);
+ }
+
+ @VisibleForTesting
+ AssignDestinationsAndPartitions(
+ DynamicDestinations dynamicDestinations,
+ IcebergCatalogConfig catalogConfig,
+ @Nullable PCollectionView> metadataView,
+ DistributionMode distributionMode,
+ @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(dynamicDestinations, catalogConfig, metadataView));
+ ParDo.of(
+ new AssignDoFn(
+ dynamicDestinations,
+ catalogConfig,
+ metadataView,
+ distributionMode,
+ distributionFunction));
if (metadataView != null) {
parDo = parDo.withSideInputs(metadataView);
}
@@ -89,56 +203,148 @@ 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> {
+ 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);
+ }
+
+ PCollectionTuple extracted = input.apply("ExtractPartitionAndSampleSortKey", extractParDo);
+
+ 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));
+ }
- private transient @MonotonicNonNull Map partitionKeys;
- private transient @MonotonicNonNull Map wrappers;
- private transient @MonotonicNonNull Map lastRefreshTimes;
- private transient @MonotonicNonNull Map cachedSpecIds;
+ /**
+ * 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);
+ }
+
+ /**
+ * 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;
+ }
+ }
+
+ 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;
+ }
+ }
+
+ /**
+ * 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;
- AssignDoFn(DynamicDestinations dynamicDestinations, IcebergCatalogConfig catalogConfig) {
- this(dynamicDestinations, catalogConfig, null);
- }
-
- AssignDoFn(
+ DestinationMetadataResolver(
DynamicDestinations dynamicDestinations,
IcebergCatalogConfig catalogConfig,
- @Nullable PCollectionView> metadataView) {
+ @Nullable PCollectionView> metadataView,
+ boolean includeSortKeyEncoder) {
this.dynamicDestinations = dynamicDestinations;
this.catalogConfig = catalogConfig;
this.metadataView = metadataView;
+ 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);
@@ -147,77 +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;
+
+ cached = new CachedTableMetadata(partitionKey, wrapper, encoder, now, specId, orderId);
+ metadataCache.put(tableIdentifier, cached);
+ }
+
+ 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;
+ }
+
+ @Setup
+ public void setup() {
+ metadataResolver.setup();
+ }
+
+ @ProcessElement
+ public void processElement(
+ ProcessContext c,
+ @Element Row element,
+ BoundedWindow window,
+ PaneInfo paneInfo,
+ @Timestamp Instant timestamp,
+ OutputReceiver> out) {
- checkStateNotNull(partitionKeys).put(tableIdentifier, partitionKey);
- checkStateNotNull(wrappers).put(tableIdentifier, wrapper);
- checkStateNotNull(lastRefreshTimes).put(tableIdentifier, now);
+ ResolvedDestination resolved =
+ metadataResolver.resolve(c, element, window, paneInfo, timestamp);
+
+ Row.Builder keyBuilder =
+ Row.withSchema(outputSchemaFor(distributionMode))
+ .addValue(resolved.tableIdentifier)
+ .addValue(resolved.partitionPath);
+
+ if (distributionMode == DistributionMode.RANGE) {
+ 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);
}
- partitionKey = checkStateNotNull(partitionKey);
- wrapper = checkStateNotNull(wrapper);
+ out.output(KV.of(keyBuilder.build(), resolved.data));
+ }
+ }
- partitionKey.partition(wrapper.wrap(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());
- String partitionPath = partitionKey.toPath();
+ 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);
+ }
- Row destAndPartition =
- Row.withSchema(OUTPUT_SCHEMA).addValues(tableIdentifier, partitionPath).build();
+ @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);
+ }
- out.output(KV.of(destAndPartition, data));
+ 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 d98d9944bf3c..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
@@ -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;
@@ -237,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
@@ -394,6 +432,7 @@ public static WriteRows writeRows(IcebergCatalogConfig catalog) {
.setDistributionMode(DistributionMode.NONE)
.setAutoSharding(false)
.setUseSideInputTableCache(false)
+ .setSortedWrites(false)
.build();
}
@@ -429,8 +468,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 +504,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 +557,222 @@ public WriteRows withDirectWriteByteLimit(Integer directWriteByteLimit) {
}
/**
- * Defines distribution of write data. Supported distributions:
+ * 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: {@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:
*
- *
- * - {@link DistributionMode#NONE}: don't shuffle rows (default)
- *
- {@link DistributionMode#HASH}: shuffle rows by partition key before writing data
- *
+ *
+ * Comparison of Distribution Modes
+ *
+ * | Mode |
+ * Description |
+ * Pros |
+ * Cons |
+ *
+ *
+ * | {@link DistributionMode#NONE} |
+ * 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 (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 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. |
+ *
+ *
*
- * {@link DistributionMode#RANGE} is not supported yet
+ * Recommendation Matrix (Sorting & Partitioning vs. Scale):
+ *
+ *
+ * Recommendation Matrix
+ *
+ * | Partitioning |
+ * Sorting |
+ * Scale / Volume |
+ * Latency Priority |
+ * Recommended Mode |
+ * Operational Impact |
+ *
+ *
+ * | Partitioned |
+ * Sorted |
+ * Small |
+ * Any |
+ * {@link DistributionMode#HASH} + {@link #withSortedWrites()} |
+ * Consolidates partition files and sorts them locally. Avoids file overlaps for small volumes. |
+ *
+ *
+ * | Partitioned |
+ * Sorted |
+ * Medium / Large |
+ * Low Write Latency |
+ * {@link DistributionMode#NONE} + {@link #withSortedWrites()} |
+ * 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#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 |
+ * Unsorted |
+ * Small |
+ * Any |
+ * {@link DistributionMode#HASH} |
+ * Consolidates data files into single partition directories to prevent file fragmentation. |
+ *
+ *
+ * | Partitioned |
+ * Unsorted |
+ * Medium / Large |
+ * Any |
+ * {@link DistributionMode#HASH} with auto-sharding |
+ * Consolidates partition files while dynamically balancing hot partition writes across parallel workers. |
+ *
+ *
+ * | Unpartitioned |
+ * Sorted |
+ * Small |
+ * Any |
+ * {@link DistributionMode#NONE} + {@link #withSortedWrites()} |
+ * Sorts locally per bundle by sort key with zero shuffle overhead. |
+ *
+ *
+ * | Unpartitioned |
+ * Sorted |
+ * Medium / Large |
+ * Low Write Latency |
+ * {@link DistributionMode#NONE} + {@link #withSortedWrites()} |
+ * 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} |
+ * 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 |
+ * Unsorted |
+ * Any |
+ * Any |
+ * {@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 (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));
+ * }
*/
public WriteRows withDistributionMode(DistributionMode mode) {
return toBuilder().setDistributionMode(mode).build();
}
+ /**
+ * 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. 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();
+ }
+
+ /**
+ * 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
@@ -611,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"));
@@ -679,11 +935,16 @@ public IcebergWriteResult expand(PCollection input) {
metadataView = input.apply("GenerateTableMetadataView", driverBuilder.build().asView());
}
+ 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(
@@ -694,13 +955,23 @@ public IcebergWriteResult expand(PCollection input) {
getTriggeringFrequency(),
getDirectWriteByteLimit(),
getWriteProperties(),
- metadataView));
+ metadataView,
+ sortedWrites));
case HASH:
+ case RANGE:
+ String assignStepName =
+ getDistributionMode() == DistributionMode.RANGE
+ ? "AssignDestinationAndPartitionWithRange"
+ : "AssignDestinationAndPartition";
return input
.apply(
- "AssignDestinationAndPartition",
+ assignStepName,
new AssignDestinationsAndPartitions(
- destinations, getCatalogConfig(), metadataView))
+ destinations,
+ getCatalogConfig(),
+ metadataView,
+ getDistributionMode(),
+ getDistributionFunction()))
.apply(
"Write Rows to Partitions",
new WriteToPartitions(
@@ -709,7 +980,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..7dab3648206a
--- /dev/null
+++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergRowSorter.java
@@ -0,0 +1,589 @@
+/*
+ * 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.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;
+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.
+ */
+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. */
+ 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();
+
+ /**
+ * Pre-bound encoder for an Iceberg {@link SortOrder} that avoids re-binding transforms and field
+ * name lookups per row.
+ */
+ @SuppressWarnings({"unchecked", "rawtypes"})
+ static final 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);
+ }
+
+ 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);
+ }
+ }
+ }
+
+ 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();
+ }
+ }
+
+ 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,
+ int memoryMb) {
+ try {
+ 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);
+ this.sorter = BufferedExternalSorter.create(sorterOptions);
+ this.rowCoder = RowCoder.of(beamSchema);
+ this.keyEncoder = new KeyEncoder(sortOrder, icebergSchema, beamSchema);
+ }
+
+ /**
+ * 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();
+
+ 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 {
+ 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);
+ }
+ }
+
+ @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,
+ Schema icebergSchema,
+ org.apache.beam.sdk.schemas.Schema beamSchema) {
+
+ 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;
+ }
+ }
+
+ @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")
+ private static void encodeSortKey(
+ Row row,
+ SortOrder sortOrder,
+ String[] columnNames,
+ int[] beamFieldIndices,
+ @Nullable EnumerationType[] enumTypes,
+ 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);
+
+ // 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];
+ 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) {
+ 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(DateTimeUtil.microsFromTime((LocalTime) val) ^ 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() * 1000L) ^ Long.MIN_VALUE, baos, invert);
+ } else if (val instanceof Instant) {
+ writeLong(DateTimeUtil.microsFromInstant((Instant) val) ^ Long.MIN_VALUE, baos, invert);
+ } else if (val instanceof Date) {
+ 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.
+ 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..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;
@@ -336,6 +337,44 @@ 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:
+ case FIXED:
+ 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;
+ 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()) {
@@ -775,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 91d0ed73835e..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
@@ -188,14 +188,22 @@ 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. 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(
@@ -310,6 +318,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);
@@ -418,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);
@@ -530,6 +541,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/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}:
*
*
+ * - Flush any buffered sorted rows one partition at a time, rolling at {@code maxFileSize}
*
- Close all underlying {@link RecordWriter}s
*
- 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 a2ed40c87f99..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,6 +17,7 @@
*/
package org.apache.beam.sdk.io.iceberg;
+import java.util.Collections;
import java.util.List;
import java.util.Map;
import org.apache.beam.sdk.transforms.DoFn;
@@ -31,7 +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.checkerframework.checker.nullness.qual.Nullable;
class WriteGroupedRowsToFiles
@@ -44,6 +44,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 +52,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,12 +69,31 @@ 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
@@ -80,7 +107,8 @@ public PCollection expand(
maxBytesPerFile,
filePrefix,
writeProperties,
- metadataView));
+ metadataView,
+ sortedWrites));
if (metadataView != null) {
parDo = parDo.withSideInputs(metadataView);
}
@@ -96,15 +124,7 @@ private static class WriteGroupedRowsToFilesDoFn
private final long maxFileSize;
private final @Nullable Map writeProperties;
private final @Nullable PCollectionView> metadataView;
-
- WriteGroupedRowsToFilesDoFn(
- IcebergCatalogConfig catalogConfig,
- DynamicDestinations dynamicDestinations,
- long maxFileSize,
- String filePrefix,
- @Nullable Map writeProperties) {
- this(catalogConfig, dynamicDestinations, maxFileSize, filePrefix, writeProperties, null);
- }
+ private final boolean sortedWrites;
WriteGroupedRowsToFilesDoFn(
IcebergCatalogConfig catalogConfig,
@@ -112,13 +132,15 @@ private static class WriteGroupedRowsToFilesDoFn
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;
}
@ProcessElement
@@ -143,7 +165,8 @@ public void processElement(
maxFileSize,
Integer.MAX_VALUE,
writeProperties,
- sideInputs)) {
+ sideInputs,
+ sortedWrites)) {
writer = openWriter;
for (Row e : element.getValue()) {
writer.write(windowedDestination, e);
@@ -151,7 +174,9 @@ public void processElement(
}
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 fbb6cb08aca4..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
@@ -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;
@@ -39,37 +38,37 @@
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;
+ 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 +77,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 +146,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;
@@ -124,24 +159,19 @@ private static class WriteDoFn extends DoFn>, FileWriteRes
IcebergCatalogConfig catalogConfig,
DynamicDestinations dynamicDestinations,
String filePrefix,
- Schema dataSchema,
- @Nullable Map writeProperties) {
- this(catalogConfig, dynamicDestinations, filePrefix, dataSchema, writeProperties, null);
- }
-
- WriteDoFn(
- IcebergCatalogConfig catalogConfig,
- DynamicDestinations dynamicDestinations,
- String filePrefix,
+ long maxFileSize,
Schema dataSchema,
@Nullable Map writeProperties,
- @Nullable PCollectionView> metadataView) {
+ @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 +202,52 @@ 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);
}
+ if (writer != null) {
+ RecordWriter finishedWriter = writer;
+ writer = null;
+ finishedWriter.close();
+ writtenFiles.add(SerializableDataFile.from(finishedWriter.getDataFile(), table.specs()));
+ }
} finally {
- writer.close();
+ if (writer != null) {
+ writer.close();
+ }
}
- // Serialize against the file's own spec
- SerializableDataFile sdf = SerializableDataFile.from(writer.getDataFile(), table.specs());
- out.output(
- FileWriteResult.builder()
- .setTableIdentifier(destination.getTableIdentifier())
- .setSerializableDataFile(sdf)
- .build());
+ for (SerializableDataFile sdf : writtenFiles) {
+ out.output(
+ FileWriteResult.builder()
+ .setTableIdentifier(destination.getTableIdentifier())
+ .setSerializableDataFile(sdf)
+ .build());
+ }
}
private Map getPartitionFieldMap(
@@ -217,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 6d8f08a2ae56..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;
@@ -62,6 +60,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 +74,8 @@ class WriteToDestinations extends PTransform>, Icebe
triggeringFrequency,
directWriteByteLimit,
writeProperties,
- null);
+ null,
+ false);
}
WriteToDestinations(
@@ -85,12 +85,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();
}
@@ -139,7 +158,8 @@ private PCollection groupAndWriteRecords(PCollection applyUserTriggering(PCollection input) {
@@ -189,7 +209,8 @@ private PCollection writeTriggeredWithBundleLifting(
filePrefix,
DEFAULT_MAX_BYTES_PER_FILE,
writeProperties,
- metadataView));
+ metadataView,
+ sortedWrites));
PCollection groupedFileWrites = groupAndWriteRecords(smallBatches);
@@ -226,7 +247,8 @@ private PCollection writeUntriggered(PCollection writeGroupedResult =
@@ -241,7 +263,8 @@ private PCollection writeUntriggered(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/WriteToPartitions.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WriteToPartitions.java
index 7e9d56df98a9..4054ca731cb9 100644
--- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WriteToPartitions.java
+++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/WriteToPartitions.java
@@ -24,9 +24,11 @@
import java.util.Map;
import java.util.UUID;
+import org.apache.beam.sdk.coders.Coder;
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.transforms.GroupByKey;
import org.apache.beam.sdk.transforms.GroupIntoBatches;
import org.apache.beam.sdk.transforms.MapElements;
import org.apache.beam.sdk.transforms.PTransform;
@@ -42,7 +44,7 @@
import org.joda.time.Duration;
class WriteToPartitions extends PTransform>, IcebergWriteResult> {
- private static final long DEFAULT_BYTES_PER_FILE = (1L << 29); // 512mb
+ static final long DEFAULT_BYTES_PER_FILE = (1L << 29); // 512mb
private final IcebergCatalogConfig catalogConfig;
private final DynamicDestinations dynamicDestinations;
private final @Nullable Duration triggeringFrequency;
@@ -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..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,7 +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.checkerframework.checker.nullness.qual.Nullable;
/**
@@ -75,6 +73,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 +81,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 +98,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 +137,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,34 +242,19 @@ 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 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);
- }
-
- public WriteUngroupedRowsToFilesDoFn(
+ WriteUngroupedRowsToFilesDoFn(
IcebergCatalogConfig catalogConfig,
DynamicDestinations dynamicDestinations,
String filename,
int maximumWritersPerBundle,
long maxFileSize,
@Nullable Map writeProperties,
- @Nullable PCollectionView