You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@iceberg.apache.org by pv...@apache.org on 2025/06/04 12:07:21 UTC
(iceberg) branch main updated: Flink: Backport IcerbegSink RANGE distribution to Flink 1.19 and 2.0 (#13228)
This is an automated email from the ASF dual-hosted git repository.
pvary pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/iceberg.git
The following commit(s) were added to refs/heads/main by this push:
new 73802f63c5 Flink: Backport IcerbegSink RANGE distribution to Flink 1.19 and 2.0 (#13228)
73802f63c5 is described below
commit 73802f63c5d3c57349c92ff8112692ec4b67e897
Author: Rodrigo <rm...@apple.com>
AuthorDate: Wed Jun 4 05:07:16 2025 -0700
Flink: Backport IcerbegSink RANGE distribution to Flink 1.19 and 2.0 (#13228)
backports #12071
---
.../org/apache/iceberg/flink/sink/IcebergSink.java | 251 +++++++---
.../TestFlinkIcebergSinkV2DistributionMode.java | 539 +++++++++++++++++++++
.../apache/iceberg/flink/sink/TestIcebergSink.java | 70 ---
.../org/apache/iceberg/flink/sink/IcebergSink.java | 251 +++++++---
.../TestFlinkIcebergSinkV2DistributionMode.java | 539 +++++++++++++++++++++
.../apache/iceberg/flink/sink/TestIcebergSink.java | 70 ---
6 files changed, 1458 insertions(+), 262 deletions(-)
diff --git a/flink/v1.19/flink/src/main/java/org/apache/iceberg/flink/sink/IcebergSink.java b/flink/v1.19/flink/src/main/java/org/apache/iceberg/flink/sink/IcebergSink.java
index 8956b03e0a..83ee85d773 100644
--- a/flink/v1.19/flink/src/main/java/org/apache/iceberg/flink/sink/IcebergSink.java
+++ b/flink/v1.19/flink/src/main/java/org/apache/iceberg/flink/sink/IcebergSink.java
@@ -34,6 +34,7 @@ import java.util.Map;
import java.util.UUID;
import java.util.function.Function;
import org.apache.flink.annotation.Experimental;
+import org.apache.flink.api.common.functions.FlatMapFunction;
import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.connector.sink2.Committer;
@@ -63,13 +64,20 @@ import org.apache.iceberg.DistributionMode;
import org.apache.iceberg.FileFormat;
import org.apache.iceberg.PartitionField;
import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Partitioning;
import org.apache.iceberg.Schema;
import org.apache.iceberg.SerializableTable;
+import org.apache.iceberg.SortOrder;
import org.apache.iceberg.Table;
import org.apache.iceberg.flink.FlinkSchemaUtil;
import org.apache.iceberg.flink.FlinkWriteConf;
import org.apache.iceberg.flink.FlinkWriteOptions;
import org.apache.iceberg.flink.TableLoader;
+import org.apache.iceberg.flink.sink.shuffle.DataStatisticsOperatorFactory;
+import org.apache.iceberg.flink.sink.shuffle.RangePartitioner;
+import org.apache.iceberg.flink.sink.shuffle.StatisticsOrRecord;
+import org.apache.iceberg.flink.sink.shuffle.StatisticsOrRecordTypeInformation;
+import org.apache.iceberg.flink.sink.shuffle.StatisticsType;
import org.apache.iceberg.flink.util.FlinkCompatibilityUtil;
import org.apache.iceberg.io.WriteResult;
import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
@@ -369,21 +377,75 @@ public class IcebergSink
/**
* Configure the write {@link DistributionMode} that the IcebergSink will use. Currently, flink
- * support {@link DistributionMode#NONE} and {@link DistributionMode#HASH}.
+ * support {@link DistributionMode#NONE} and {@link DistributionMode#HASH} and {@link
+ * DistributionMode#RANGE}
*
* @param mode to specify the write distribution mode.
* @return {@link IcebergSink.Builder} to connect the iceberg table.
*/
public Builder distributionMode(DistributionMode mode) {
- Preconditions.checkArgument(
- !DistributionMode.RANGE.equals(mode),
- "Flink does not support 'range' write distribution mode now.");
if (mode != null) {
writeOptions.put(FlinkWriteOptions.DISTRIBUTION_MODE.key(), mode.modeName());
}
return this;
}
+ /**
+ * Range distribution needs to collect statistics about data distribution to properly shuffle
+ * the records in relatively balanced way. In general, low cardinality should use {@link
+ * StatisticsType#Map} and high cardinality should use {@link StatisticsType#Sketch} Refer to
+ * {@link StatisticsType} Javadoc for more details.
+ *
+ * <p>Default is {@link StatisticsType#Auto} where initially Map statistics is used. But if
+ * cardinality is higher than the threshold (currently 10K) as defined in {@code
+ * SketchUtil#OPERATOR_SKETCH_SWITCH_THRESHOLD}, statistics collection automatically switches to
+ * the sketch reservoir sampling.
+ *
+ * <p>Explicit set the statistics type if the default behavior doesn't work.
+ *
+ * @param type to specify the statistics type for range distribution.
+ * @return {@link IcebergSink.Builder} to connect the iceberg table.
+ */
+ public IcebergSink.Builder rangeDistributionStatisticsType(StatisticsType type) {
+ if (type != null) {
+ writeOptions.put(FlinkWriteOptions.RANGE_DISTRIBUTION_STATISTICS_TYPE.key(), type.name());
+ }
+ return this;
+ }
+
+ /**
+ * If sort order contains partition columns, each sort key would map to one partition and data
+ * file. This relative weight can avoid placing too many small files for sort keys with low
+ * traffic. It is a double value that defines the minimal weight for each sort key. `0.02` means
+ * each key has a base weight of `2%` of the targeted traffic weight per writer task.
+ *
+ * <p>E.g. the sink Iceberg table is partitioned daily by event time. Assume the data stream
+ * contains events from now up to 180 days ago. With event time, traffic weight distribution
+ * across different days typically has a long tail pattern. Current day contains the most
+ * traffic. The older days (long tail) contain less and less traffic. Assume writer parallelism
+ * is `10`. The total weight across all 180 days is `10,000`. Target traffic weight per writer
+ * task would be `1,000`. Assume the weight sum for the oldest 150 days is `1,000`. Normally,
+ * the range partitioner would put all the oldest 150 days in one writer task. That writer task
+ * would write to 150 small files (one per day). Keeping 150 open files can potentially consume
+ * large amount of memory. Flushing and uploading 150 files (however small) at checkpoint time
+ * can also be potentially slow. If this config is set to `0.02`. It means every sort key has a
+ * base weight of `2%` of targeted weight of `1,000` for every write task. It would essentially
+ * avoid placing more than `50` data files (one per day) on one writer task no matter how small
+ * they are.
+ *
+ * <p>This is only applicable to {@link StatisticsType#Map} for low-cardinality scenario. For
+ * {@link StatisticsType#Sketch} high-cardinality sort columns, they are usually not used as
+ * partition columns. Otherwise, too many partitions and small files may be generated during
+ * write. Sketch range partitioner simply splits high-cardinality keys into ordered ranges.
+ *
+ * <p>Default is {@code 0.0%}.
+ */
+ public Builder rangeDistributionSortKeyBaseWeight(double weight) {
+ writeOptions.put(
+ FlinkWriteOptions.RANGE_DISTRIBUTION_SORT_KEY_BASE_WEIGHT.key(), Double.toString(weight));
+ return this;
+ }
+
/**
* Configuring the write parallel number for iceberg stream writer.
*
@@ -549,6 +611,10 @@ public class IcebergSink
}
}
+ private String operatorName(String suffix) {
+ return uidSuffix != null ? suffix + "-" + uidSuffix : suffix;
+ }
+
private static String defaultSuffix(String uidSuffix, String defaultSuffix) {
if (uidSuffix == null || uidSuffix.isEmpty()) {
return defaultSuffix;
@@ -635,73 +701,136 @@ public class IcebergSink
DistributionMode mode = flinkWriteConf.distributionMode();
Schema schema = table.schema();
PartitionSpec spec = table.spec();
+ SortOrder sortOrder = table.sortOrder();
+
LOG.info("Write distribution mode is '{}'", mode.modeName());
switch (mode) {
case NONE:
- if (equalityFieldIds.isEmpty()) {
- return input;
- } else {
- LOG.info("Distribute rows by equality fields, because there are equality fields set");
- return input.keyBy(new EqualityFieldKeySelector(schema, flinkRowType, equalityFieldIds));
- }
-
+ return distributeDataStreamByNoneDistributionMode(input, schema);
case HASH:
- if (equalityFieldIds.isEmpty()) {
- if (table.spec().isUnpartitioned()) {
- LOG.warn(
- "Fallback to use 'none' distribution mode, because there are no equality fields set "
- + "and table is unpartitioned");
- return input;
- } else {
- if (BucketPartitionerUtil.hasOneBucketField(spec)) {
- return input.partitionCustom(
- new BucketPartitioner(spec),
- new BucketPartitionKeySelector(spec, schema, flinkRowType));
- } else {
- return input.keyBy(new PartitionKeySelector(spec, schema, flinkRowType));
- }
- }
- } else {
- if (spec.isUnpartitioned()) {
- LOG.info(
- "Distribute rows by equality fields, because there are equality fields set "
- + "and table is unpartitioned");
- return input.keyBy(
- new EqualityFieldKeySelector(schema, flinkRowType, equalityFieldIds));
- } else {
- for (PartitionField partitionField : spec.fields()) {
- Preconditions.checkState(
- equalityFieldIds.contains(partitionField.sourceId()),
- "In 'hash' distribution mode with equality fields set, source column '%s' of partition field '%s' "
- + "should be included in equality fields: '%s'",
- table.schema().findColumnName(partitionField.sourceId()),
- partitionField,
- equalityFieldColumns);
- }
- return input.keyBy(new PartitionKeySelector(spec, schema, flinkRowType));
- }
- }
-
+ return distributeDataStreamByHashDistributionMode(input, schema, spec);
case RANGE:
- if (equalityFieldIds.isEmpty()) {
- LOG.warn(
- "Fallback to use 'none' distribution mode, because there are no equality fields set "
- + "and {}=range is not supported yet in flink",
- WRITE_DISTRIBUTION_MODE);
- return input;
- } else {
- LOG.info(
- "Distribute rows by equality fields, because there are equality fields set "
- + "and{}=range is not supported yet in flink",
- WRITE_DISTRIBUTION_MODE);
- return input.keyBy(new EqualityFieldKeySelector(schema, flinkRowType, equalityFieldIds));
- }
-
+ return distributeDataStreamByRangeDistributionMode(input, schema, spec, sortOrder);
default:
throw new RuntimeException("Unrecognized " + WRITE_DISTRIBUTION_MODE + ": " + mode);
}
}
+ private DataStream<RowData> distributeDataStreamByNoneDistributionMode(
+ DataStream<RowData> input, Schema iSchema) {
+ if (equalityFieldIds.isEmpty()) {
+ return input;
+ } else {
+ LOG.info("Distribute rows by equality fields, because there are equality fields set");
+ return input.keyBy(new EqualityFieldKeySelector(iSchema, flinkRowType, equalityFieldIds));
+ }
+ }
+
+ private DataStream<RowData> distributeDataStreamByHashDistributionMode(
+ DataStream<RowData> input, Schema iSchema, PartitionSpec partitionSpec) {
+ if (equalityFieldIds.isEmpty()) {
+ if (partitionSpec.isUnpartitioned()) {
+ LOG.warn(
+ "Fallback to use 'none' distribution mode, because there are no equality fields set "
+ + "and table is unpartitioned");
+ return input;
+ } else {
+ return input.keyBy(new PartitionKeySelector(partitionSpec, iSchema, flinkRowType));
+ }
+ } else {
+ if (partitionSpec.isUnpartitioned()) {
+ LOG.info(
+ "Distribute rows by equality fields, because there are equality fields set "
+ + "and table is unpartitioned");
+ return input.keyBy(new EqualityFieldKeySelector(iSchema, flinkRowType, equalityFieldIds));
+ } else {
+ for (PartitionField partitionField : partitionSpec.fields()) {
+ Preconditions.checkState(
+ equalityFieldIds.contains(partitionField.sourceId()),
+ "In 'hash' distribution mode with equality fields set, source column '%s' of partition field '%s' "
+ + "should be included in equality fields: '%s'",
+ table.schema().findColumnName(partitionField.sourceId()),
+ partitionField,
+ equalityFieldColumns);
+ }
+ return input.keyBy(new PartitionKeySelector(partitionSpec, iSchema, flinkRowType));
+ }
+ }
+ }
+
+ private DataStream<RowData> distributeDataStreamByRangeDistributionMode(
+ DataStream<RowData> input,
+ Schema iSchema,
+ PartitionSpec partitionSpec,
+ SortOrder sortOrderParam) {
+
+ int writerParallelism =
+ flinkWriteConf.writeParallelism() == null
+ ? input.getParallelism()
+ : flinkWriteConf.writeParallelism();
+
+ // needed because of checkStyle not allowing us to change the value of an argument
+ SortOrder sortOrder = sortOrderParam;
+
+ // Ideally, exception should be thrown in the combination of range distribution and
+ // equality fields. Primary key case should use hash distribution mode.
+ // Keep the current behavior of falling back to keyBy for backward compatibility.
+ if (!equalityFieldIds.isEmpty()) {
+ LOG.warn(
+ "Hash distribute rows by equality fields, even though {}=range is set. "
+ + "Range distribution for primary keys are not always safe in "
+ + "Flink streaming writer.",
+ WRITE_DISTRIBUTION_MODE);
+ return input.keyBy(new EqualityFieldKeySelector(iSchema, flinkRowType, equalityFieldIds));
+ }
+
+ // range distribute by partition key or sort key if table has an SortOrder
+ Preconditions.checkState(
+ sortOrder.isSorted() || partitionSpec.isPartitioned(),
+ "Invalid write distribution mode: range. Need to define sort order or partition spec.");
+ if (sortOrder.isUnsorted()) {
+ sortOrder = Partitioning.sortOrderFor(partitionSpec);
+ LOG.info("Construct sort order from partition spec");
+ }
+
+ LOG.info("Range distribute rows by sort order: {}", sortOrder);
+ StatisticsOrRecordTypeInformation statisticsOrRecordTypeInformation =
+ new StatisticsOrRecordTypeInformation(flinkRowType, iSchema, sortOrder);
+ StatisticsType statisticsType = flinkWriteConf.rangeDistributionStatisticsType();
+ SingleOutputStreamOperator<StatisticsOrRecord> shuffleStream =
+ input
+ .transform(
+ operatorName("range-shuffle"),
+ statisticsOrRecordTypeInformation,
+ new DataStatisticsOperatorFactory(
+ iSchema,
+ sortOrder,
+ writerParallelism,
+ statisticsType,
+ flinkWriteConf.rangeDistributionSortKeyBaseWeight()))
+ // Set the parallelism same as input operator to encourage chaining
+ .setParallelism(input.getParallelism());
+
+ if (uidSuffix != null) {
+ shuffleStream = shuffleStream.uid("shuffle-" + uidSuffix);
+ }
+
+ return shuffleStream
+ .partitionCustom(new RangePartitioner(iSchema, sortOrder), r -> r)
+ .flatMap(
+ (FlatMapFunction<StatisticsOrRecord, RowData>)
+ (statisticsOrRecord, out) -> {
+ if (statisticsOrRecord.hasRecord()) {
+ out.collect(statisticsOrRecord.record());
+ }
+ })
+ // Set slot sharing group and the parallelism same as writerParallelism to
+ // promote operator chaining with the downstream writer operator
+ .slotSharingGroup("shuffle-partition-custom-group")
+ .setParallelism(writerParallelism)
+ .returns(RowData.class);
+ }
+
/**
* Initialize a {@link Builder} to export the data from generic input data stream into iceberg
* table. We use {@link RowData} inside the sink connector, so users need to provide a mapper
diff --git a/flink/v1.19/flink/src/test/java/org/apache/iceberg/flink/sink/TestFlinkIcebergSinkV2DistributionMode.java b/flink/v1.19/flink/src/test/java/org/apache/iceberg/flink/sink/TestFlinkIcebergSinkV2DistributionMode.java
new file mode 100644
index 0000000000..47217db206
--- /dev/null
+++ b/flink/v1.19/flink/src/test/java/org/apache/iceberg/flink/sink/TestFlinkIcebergSinkV2DistributionMode.java
@@ -0,0 +1,539 @@
+/*
+ * 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.iceberg.flink.sink;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.assertj.core.api.Assumptions.assumeThat;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.test.junit5.MiniClusterExtension;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.DataFile;
+import org.apache.iceberg.DistributionMode;
+import org.apache.iceberg.FileFormat;
+import org.apache.iceberg.Parameter;
+import org.apache.iceberg.ParameterizedTestExtension;
+import org.apache.iceberg.Parameters;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Snapshot;
+import org.apache.iceberg.TableProperties;
+import org.apache.iceberg.flink.FlinkWriteOptions;
+import org.apache.iceberg.flink.HadoopCatalogExtension;
+import org.apache.iceberg.flink.MiniFlinkClusterExtension;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.flink.TestFixtures;
+import org.apache.iceberg.flink.sink.shuffle.StatisticsType;
+import org.apache.iceberg.flink.source.BoundedTestSource;
+import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList;
+import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+import org.apache.iceberg.types.Conversions;
+import org.apache.iceberg.types.Types;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.TestTemplate;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+/**
+ * This tests the distribution mode of the IcebergSink. Extract them separately since it is
+ * unnecessary to test different file formats (Avro, Orc, Parquet) like in {@link TestIcebergSink}.
+ * Removing the file format dimension reduces the number of combinations from 12 to 4, which helps
+ * reduce test run time.
+ */
+@ExtendWith(ParameterizedTestExtension.class)
+public class TestFlinkIcebergSinkV2DistributionMode extends TestFlinkIcebergSinkBase {
+
+ @RegisterExtension
+ public static MiniClusterExtension miniClusterResource =
+ MiniFlinkClusterExtension.createWithClassloaderCheckDisabled();
+
+ @RegisterExtension
+ private static final HadoopCatalogExtension CATALOG_EXTENSION =
+ new HadoopCatalogExtension(TestFixtures.DATABASE, TestFixtures.TABLE);
+
+ private final FileFormat format = FileFormat.PARQUET;
+
+ @Parameter(index = 0)
+ private int parallelism;
+
+ @Parameter(index = 1)
+ private boolean partitioned;
+
+ @Parameter(index = 2)
+ private int writeParallelism;
+
+ @Parameters(name = "parallelism = {0}, partitioned = {1}, writeParallelism = {2}")
+ public static Object[][] parameters() {
+ return new Object[][] {
+ {1, true, 1},
+ {1, false, 1},
+ {2, true, 2},
+ {2, false, 2},
+ {1, true, 2},
+ {1, false, 2},
+ };
+ }
+
+ @BeforeEach
+ public void before() throws IOException {
+ this.table =
+ CATALOG_EXTENSION
+ .catalog()
+ .createTable(
+ TestFixtures.TABLE_IDENTIFIER,
+ SimpleDataUtil.SCHEMA,
+ partitioned
+ ? PartitionSpec.builderFor(SimpleDataUtil.SCHEMA).identity("data").build()
+ : PartitionSpec.unpartitioned(),
+ ImmutableMap.of(TableProperties.DEFAULT_FILE_FORMAT, format.name()));
+
+ this.env =
+ StreamExecutionEnvironment.getExecutionEnvironment(
+ MiniFlinkClusterExtension.DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .enableCheckpointing(100)
+ .setParallelism(parallelism)
+ .setMaxParallelism(Math.max(parallelism, writeParallelism));
+
+ this.tableLoader = CATALOG_EXTENSION.tableLoader();
+ }
+
+ @TestTemplate
+ public void testShuffleByPartitionWithSchema() throws Exception {
+ testWriteRow(parallelism, SimpleDataUtil.FLINK_SCHEMA, DistributionMode.HASH);
+ if (partitioned) {
+ assertThat(partitionFiles("aaa")).isEqualTo(1);
+ assertThat(partitionFiles("bbb")).isEqualTo(1);
+ assertThat(partitionFiles("ccc")).isEqualTo(1);
+ }
+ }
+
+ @TestTemplate
+ public void testJobNoneDistributeMode() throws Exception {
+ table
+ .updateProperties()
+ .set(TableProperties.WRITE_DISTRIBUTION_MODE, DistributionMode.HASH.modeName())
+ .commit();
+
+ testWriteRow(parallelism, null, DistributionMode.NONE);
+
+ if (parallelism > 1) {
+ if (partitioned) {
+ int files = partitionFiles("aaa") + partitionFiles("bbb") + partitionFiles("ccc");
+ assertThat(files).isGreaterThan(3);
+ }
+ }
+ }
+
+ @TestTemplate
+ public void testJobNullDistributionMode() throws Exception {
+ table
+ .updateProperties()
+ .set(TableProperties.WRITE_DISTRIBUTION_MODE, DistributionMode.HASH.modeName())
+ .commit();
+
+ testWriteRow(parallelism, null, null);
+
+ if (partitioned) {
+ assertThat(partitionFiles("aaa")).isEqualTo(1);
+ assertThat(partitionFiles("bbb")).isEqualTo(1);
+ assertThat(partitionFiles("ccc")).isEqualTo(1);
+ }
+ }
+
+ @TestTemplate
+ public void testPartitionWriteMode() throws Exception {
+ testWriteRow(parallelism, null, DistributionMode.HASH);
+ if (partitioned) {
+ assertThat(partitionFiles("aaa")).isEqualTo(1);
+ assertThat(partitionFiles("bbb")).isEqualTo(1);
+ assertThat(partitionFiles("ccc")).isEqualTo(1);
+ }
+ }
+
+ @TestTemplate
+ public void testOverrideWriteConfigWithUnknownDistributionMode() {
+ Map<String, String> newProps = Maps.newHashMap();
+ newProps.put(FlinkWriteOptions.DISTRIBUTION_MODE.key(), "UNRECOGNIZED");
+
+ List<Row> rows = createRows("");
+ DataStream<Row> dataStream = env.addSource(createBoundedSource(rows), ROW_TYPE_INFO);
+
+ IcebergSink.forRow(dataStream, SimpleDataUtil.FLINK_SCHEMA)
+ .table(table)
+ .tableLoader(tableLoader)
+ .writeParallelism(writeParallelism)
+ .setAll(newProps)
+ .append();
+
+ assertThatThrownBy(env::execute)
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessage("Invalid distribution mode: UNRECOGNIZED");
+ }
+
+ @TestTemplate
+ public void testRangeDistributionWithoutSortOrderUnpartitioned() throws Exception {
+ assumeThat(partitioned).isFalse();
+
+ table
+ .updateProperties()
+ .set(TableProperties.WRITE_DISTRIBUTION_MODE, DistributionMode.RANGE.modeName())
+ .commit();
+
+ int numOfCheckpoints = 6;
+ DataStream<Row> dataStream =
+ env.addSource(
+ createRangeDistributionBoundedSource(createCharRows(numOfCheckpoints, 10)),
+ ROW_TYPE_INFO);
+
+ IcebergSink.forRow(dataStream, SimpleDataUtil.FLINK_SCHEMA)
+ .table(table)
+ .tableLoader(tableLoader)
+ .writeParallelism(writeParallelism)
+ .append();
+
+ // Range distribution requires either sort order or partition spec defined
+ assertThatThrownBy(env::execute)
+ .isInstanceOf(IllegalStateException.class)
+ .hasMessage(
+ "Invalid write distribution mode: range. Need to define sort order or partition spec.");
+ }
+
+ @TestTemplate
+ public void testRangeDistributionWithoutSortOrderPartitioned() throws Exception {
+ assumeThat(partitioned).isTrue();
+
+ table
+ .updateProperties()
+ .set(TableProperties.WRITE_DISTRIBUTION_MODE, DistributionMode.RANGE.modeName())
+ .commit();
+
+ int numOfCheckpoints = 6;
+ DataStream<Row> dataStream =
+ env.addSource(
+ createRangeDistributionBoundedSource(createCharRows(numOfCheckpoints, 10)),
+ ROW_TYPE_INFO);
+
+ IcebergSink.forRow(dataStream, SimpleDataUtil.FLINK_SCHEMA)
+ .table(table)
+ .tableLoader(tableLoader)
+ .writeParallelism(writeParallelism)
+ .append();
+
+ // sort based on partition columns
+ env.execute(getClass().getSimpleName());
+
+ table.refresh();
+ // ordered in reverse timeline from the newest snapshot to the oldest snapshot
+ List<Snapshot> snapshots = Lists.newArrayList(table.snapshots().iterator());
+ // only keep the snapshots with added data files
+ snapshots =
+ snapshots.stream()
+ .filter(snapshot -> snapshot.addedDataFiles(table.io()).iterator().hasNext())
+ .collect(Collectors.toList());
+
+ // Sometimes we will have more checkpoints than the bounded source if we pass the
+ // auto checkpoint interval. Thus producing multiple snapshots.
+ assertThat(snapshots).hasSizeGreaterThanOrEqualTo(numOfCheckpoints);
+ }
+
+ @TestTemplate
+ public void testRangeDistributionWithNullValue() throws Exception {
+ assumeThat(partitioned).isTrue();
+
+ table
+ .updateProperties()
+ .set(TableProperties.WRITE_DISTRIBUTION_MODE, DistributionMode.RANGE.modeName())
+ .commit();
+
+ int numOfCheckpoints = 6;
+ List<List<Row>> charRows = createCharRows(numOfCheckpoints, 10);
+ charRows.add(ImmutableList.of(Row.of(1, null)));
+ DataStream<Row> dataStream =
+ env.addSource(createRangeDistributionBoundedSource(charRows), ROW_TYPE_INFO);
+
+ IcebergSink.forRow(dataStream, SimpleDataUtil.FLINK_SCHEMA)
+ .table(table)
+ .tableLoader(tableLoader)
+ .writeParallelism(parallelism)
+ .append();
+
+ // sort based on partition columns
+ env.execute(getClass().getSimpleName());
+
+ table.refresh();
+ // ordered in reverse timeline from the newest snapshot to the oldest snapshot
+ List<Snapshot> snapshots = Lists.newArrayList(table.snapshots().iterator());
+ // only keep the snapshots with added data files
+ snapshots =
+ snapshots.stream()
+ .filter(snapshot -> snapshot.addedDataFiles(table.io()).iterator().hasNext())
+ .collect(Collectors.toList());
+
+ // Sometimes we will have more checkpoints than the bounded source if we pass the
+ // auto checkpoint interval. Thus producing multiple snapshots.
+ assertThat(snapshots).hasSizeGreaterThanOrEqualTo(numOfCheckpoints);
+ }
+
+ @TestTemplate
+ public void testRangeDistributionWithSortOrder() throws Exception {
+ table
+ .updateProperties()
+ .set(TableProperties.WRITE_DISTRIBUTION_MODE, DistributionMode.RANGE.modeName())
+ .commit();
+ table.replaceSortOrder().asc("data").commit();
+
+ int numOfCheckpoints = 6;
+ DataStream<Row> dataStream =
+ env.addSource(
+ createRangeDistributionBoundedSource(createCharRows(numOfCheckpoints, 10)),
+ ROW_TYPE_INFO);
+ IcebergSink.forRow(dataStream, SimpleDataUtil.FLINK_SCHEMA)
+ .table(table)
+ .tableLoader(tableLoader)
+ .writeParallelism(writeParallelism)
+ .rangeDistributionStatisticsType(StatisticsType.Map)
+ .append();
+ env.execute(getClass().getSimpleName());
+
+ table.refresh();
+ // ordered in reverse timeline from the newest snapshot to the oldest snapshot
+ List<Snapshot> snapshots = Lists.newArrayList(table.snapshots().iterator());
+ // only keep the snapshots with added data files
+ snapshots =
+ snapshots.stream()
+ .filter(snapshot -> snapshot.addedDataFiles(table.io()).iterator().hasNext())
+ .collect(Collectors.toList());
+
+ // Sometimes we will have more checkpoints than the bounded source if we pass the
+ // auto checkpoint interval. Thus producing multiple snapshots.
+ assertThat(snapshots).hasSizeGreaterThanOrEqualTo(numOfCheckpoints);
+
+ // It takes 2 checkpoint cycle for statistics collection and application
+ // of the globally aggregated statistics in the range partitioner.
+ // The last two checkpoints should have range shuffle applied
+ List<Snapshot> rangePartitionedCycles =
+ snapshots.subList(snapshots.size() - 2, snapshots.size());
+
+ if (partitioned) {
+ for (Snapshot snapshot : rangePartitionedCycles) {
+ List<DataFile> addedDataFiles =
+ Lists.newArrayList(snapshot.addedDataFiles(table.io()).iterator());
+ // up to 26 partitions
+ assertThat(addedDataFiles).hasSizeLessThanOrEqualTo(26);
+ }
+ } else {
+ for (Snapshot snapshot : rangePartitionedCycles) {
+ List<DataFile> addedDataFiles =
+ Lists.newArrayList(snapshot.addedDataFiles(table.io()).iterator());
+ // each writer task should only write one file for non-partition sort column
+ assertThat(addedDataFiles).hasSize(writeParallelism);
+ // verify there is no overlap in min-max stats range
+ if (parallelism > 1) {
+ assertIdColumnStatsNoRangeOverlap(addedDataFiles.get(0), addedDataFiles.get(1));
+ }
+ }
+ }
+ }
+
+ @TestTemplate
+ public void testRangeDistributionSketchWithSortOrder() throws Exception {
+ table
+ .updateProperties()
+ .set(TableProperties.WRITE_DISTRIBUTION_MODE, DistributionMode.RANGE.modeName())
+ .commit();
+ table.replaceSortOrder().asc("id").commit();
+
+ int numOfCheckpoints = 6;
+ DataStream<Row> dataStream =
+ env.addSource(
+ createRangeDistributionBoundedSource(createIntRows(numOfCheckpoints, 1_000)),
+ ROW_TYPE_INFO);
+ IcebergSink.forRow(dataStream, SimpleDataUtil.FLINK_SCHEMA)
+ .table(table)
+ .tableLoader(tableLoader)
+ .writeParallelism(writeParallelism)
+ .rangeDistributionStatisticsType(StatisticsType.Sketch)
+ .append();
+ env.execute(getClass().getSimpleName());
+
+ table.refresh();
+ // ordered in reverse timeline from the newest snapshot to the oldest snapshot
+ List<Snapshot> snapshots = Lists.newArrayList(table.snapshots().iterator());
+ // only keep the snapshots with added data files
+ snapshots =
+ snapshots.stream()
+ .filter(snapshot -> snapshot.addedDataFiles(table.io()).iterator().hasNext())
+ .collect(Collectors.toList());
+
+ // Sometimes we will have more checkpoints than the bounded source if we pass the
+ // auto checkpoint interval. Thus producing multiple snapshots.
+ assertThat(snapshots).hasSizeGreaterThanOrEqualTo(numOfCheckpoints);
+
+ // It takes 2 checkpoint cycle for statistics collection and application
+ // of the globally aggregated statistics in the range partitioner.
+ // The last two checkpoints should have range shuffle applied
+ List<Snapshot> rangePartitionedCycles =
+ snapshots.subList(snapshots.size() - 2, snapshots.size());
+
+ // since the input has a single value for the data column,
+ // it is always the same partition. Hence there is no difference
+ // for partitioned or not
+ for (Snapshot snapshot : rangePartitionedCycles) {
+ List<DataFile> addedDataFiles =
+ Lists.newArrayList(snapshot.addedDataFiles(table.io()).iterator());
+ // each writer task should only write one file for non-partition sort column
+ assertThat(addedDataFiles).hasSize(writeParallelism);
+ // verify there is no overlap in min-max stats range
+ if (writeParallelism > 2) {
+ assertIdColumnStatsNoRangeOverlap(addedDataFiles.get(0), addedDataFiles.get(1));
+ }
+ }
+ }
+
+ /** Test migration from Map stats to Sketch stats */
+ @TestTemplate
+ public void testRangeDistributionStatisticsMigration() throws Exception {
+ table
+ .updateProperties()
+ .set(TableProperties.WRITE_DISTRIBUTION_MODE, DistributionMode.RANGE.modeName())
+ .commit();
+ table.replaceSortOrder().asc("id").commit();
+
+ int numOfCheckpoints = 4;
+ List<List<Row>> rowsPerCheckpoint = Lists.newArrayListWithCapacity(numOfCheckpoints);
+ for (int checkpointId = 0; checkpointId < numOfCheckpoints; ++checkpointId) {
+ // checkpointId 2 would emit 11_000 records which is larger than
+ // the OPERATOR_SKETCH_SWITCH_THRESHOLD of 10_000.
+ // This should trigger the stats migration.
+ int maxId = checkpointId < 1 ? 1_000 : 11_000;
+ List<Row> rows = Lists.newArrayListWithCapacity(maxId);
+ for (int j = 0; j < maxId; ++j) {
+ // fixed value "a" for the data (possible partition column)
+ rows.add(Row.of(j, "a"));
+ }
+
+ rowsPerCheckpoint.add(rows);
+ }
+
+ DataStream<Row> dataStream =
+ env.addSource(createRangeDistributionBoundedSource(rowsPerCheckpoint), ROW_TYPE_INFO);
+ IcebergSink.forRow(dataStream, SimpleDataUtil.FLINK_SCHEMA)
+ .table(table)
+ .tableLoader(tableLoader)
+ .writeParallelism(writeParallelism)
+ .rangeDistributionStatisticsType(StatisticsType.Auto)
+ .append();
+ env.execute(getClass().getSimpleName());
+
+ table.refresh();
+ // ordered in reverse timeline from the newest snapshot to the oldest snapshot
+ List<Snapshot> snapshots = Lists.newArrayList(table.snapshots().iterator());
+ // only keep the snapshots with added data files
+ snapshots =
+ snapshots.stream()
+ .filter(snapshot -> snapshot.addedDataFiles(table.io()).iterator().hasNext())
+ .collect(Collectors.toList());
+
+ // Sometimes we will have more checkpoints than the bounded source if we pass the
+ // auto checkpoint interval. Thus producing multiple snapshots.
+ assertThat(snapshots).hasSizeGreaterThanOrEqualTo(numOfCheckpoints);
+
+ // It takes 2 checkpoint cycle for statistics collection and application
+ // of the globally aggregated statistics in the range partitioner.
+ // The last two checkpoints should have range shuffle applied
+ List<Snapshot> rangePartitionedCycles =
+ snapshots.subList(snapshots.size() - 2, snapshots.size());
+
+ // since the input has a single value for the data column,
+ // it is always the same partition. Hence there is no difference
+ // for partitioned or not
+ for (Snapshot snapshot : rangePartitionedCycles) {
+ List<DataFile> addedDataFiles =
+ Lists.newArrayList(snapshot.addedDataFiles(table.io()).iterator());
+ // each writer task should only write one file for non-partition sort column
+ // sometimes
+ assertThat(addedDataFiles).hasSize(writeParallelism);
+ // verify there is no overlap in min-max stats range
+ if (writeParallelism > 1) {
+ assertIdColumnStatsNoRangeOverlap(addedDataFiles.get(0), addedDataFiles.get(1));
+ }
+ }
+ }
+
+ private BoundedTestSource<Row> createRangeDistributionBoundedSource(
+ List<List<Row>> rowsPerCheckpoint) {
+ return new BoundedTestSource<>(rowsPerCheckpoint);
+ }
+
+ private List<List<Row>> createCharRows(int numOfCheckpoints, int countPerChar) {
+ List<List<Row>> rowsPerCheckpoint = Lists.newArrayListWithCapacity(numOfCheckpoints);
+ for (int checkpointId = 0; checkpointId < numOfCheckpoints; ++checkpointId) {
+ List<Row> rows = Lists.newArrayListWithCapacity(26 * countPerChar);
+ for (int j = 0; j < countPerChar; ++j) {
+ for (char c = 'a'; c <= 'z'; ++c) {
+ rows.add(Row.of(1, String.valueOf(c)));
+ }
+ }
+
+ rowsPerCheckpoint.add(rows);
+ }
+
+ return rowsPerCheckpoint;
+ }
+
+ private List<List<Row>> createIntRows(int numOfCheckpoints, int maxId) {
+ List<List<Row>> rowsPerCheckpoint = Lists.newArrayListWithCapacity(numOfCheckpoints);
+ for (int checkpointId = 0; checkpointId < numOfCheckpoints; ++checkpointId) {
+ List<Row> rows = Lists.newArrayListWithCapacity(maxId);
+ for (int j = 0; j < maxId; ++j) {
+ // fixed value "a" for the data (possible partition column)
+ rows.add(Row.of(j, "a"));
+ }
+
+ rowsPerCheckpoint.add(rows);
+ }
+
+ return rowsPerCheckpoint;
+ }
+
+ private void assertIdColumnStatsNoRangeOverlap(DataFile file1, DataFile file2) {
+ // id column has fieldId 1
+ int file1LowerBound =
+ Conversions.fromByteBuffer(Types.IntegerType.get(), file1.lowerBounds().get(1));
+ int file1UpperBound =
+ Conversions.fromByteBuffer(Types.IntegerType.get(), file1.upperBounds().get(1));
+ int file2LowerBound =
+ Conversions.fromByteBuffer(Types.IntegerType.get(), file2.lowerBounds().get(1));
+ int file2UpperBound =
+ Conversions.fromByteBuffer(Types.IntegerType.get(), file2.upperBounds().get(1));
+
+ if (file1LowerBound < file2LowerBound) {
+ assertThat(file1UpperBound).isLessThanOrEqualTo(file2LowerBound);
+ } else {
+ assertThat(file2UpperBound).isLessThanOrEqualTo(file1LowerBound);
+ }
+ }
+}
diff --git a/flink/v1.19/flink/src/test/java/org/apache/iceberg/flink/sink/TestIcebergSink.java b/flink/v1.19/flink/src/test/java/org/apache/iceberg/flink/sink/TestIcebergSink.java
index 828ee5fb87..c7ad594412 100644
--- a/flink/v1.19/flink/src/test/java/org/apache/iceberg/flink/sink/TestIcebergSink.java
+++ b/flink/v1.19/flink/src/test/java/org/apache/iceberg/flink/sink/TestIcebergSink.java
@@ -145,57 +145,6 @@ public class TestIcebergSink extends TestFlinkIcebergSinkBase {
testWriteRow(SimpleDataUtil.FLINK_SCHEMA, DistributionMode.NONE);
}
- @TestTemplate
- void testJobNoneDistributeMode() throws Exception {
- table
- .updateProperties()
- .set(TableProperties.WRITE_DISTRIBUTION_MODE, DistributionMode.HASH.modeName())
- .commit();
-
- testWriteRow(null, DistributionMode.NONE);
-
- if (parallelism > 1) {
- if (partitioned) {
- int files = partitionFiles("aaa") + partitionFiles("bbb") + partitionFiles("ccc");
- assertThat(files).as("Should have more than 3 files in iceberg table.").isGreaterThan(3);
- }
- }
- }
-
- @TestTemplate
- void testJobHashDistributionMode() {
- table
- .updateProperties()
- .set(TableProperties.WRITE_DISTRIBUTION_MODE, DistributionMode.HASH.modeName())
- .commit();
-
- assertThatThrownBy(() -> testWriteRow(null, DistributionMode.RANGE))
- .isInstanceOf(IllegalArgumentException.class)
- .hasMessage("Flink does not support 'range' write distribution mode now.");
- }
-
- @TestTemplate
- void testJobNullDistributionMode() throws Exception {
- table
- .updateProperties()
- .set(TableProperties.WRITE_DISTRIBUTION_MODE, DistributionMode.HASH.modeName())
- .commit();
-
- testWriteRow(null, null);
-
- if (partitioned) {
- assertThat(partitionFiles("aaa"))
- .as("There should be only 1 data file in partition 'aaa'")
- .isEqualTo(1);
- assertThat(partitionFiles("bbb"))
- .as("There should be only 1 data file in partition 'bbb'")
- .isEqualTo(1);
- assertThat(partitionFiles("ccc"))
- .as("There should be only 1 data file in partition 'ccc'")
- .isEqualTo(1);
- }
- }
-
@TestTemplate
void testPartitionWriteMode() throws Exception {
testWriteRow(null, DistributionMode.HASH);
@@ -311,25 +260,6 @@ public class TestIcebergSink extends TestFlinkIcebergSinkBase {
assertThat(rightTable.currentSnapshot().summary().get("direction")).isEqualTo("rightTable");
}
- @TestTemplate
- void testOverrideWriteConfigWithUnknownDistributionMode() {
- Map<String, String> newProps = Maps.newHashMap();
- newProps.put(FlinkWriteOptions.DISTRIBUTION_MODE.key(), "UNRECOGNIZED");
-
- List<Row> rows = createRows("");
- DataStream<Row> dataStream = env.addSource(createBoundedSource(rows), ROW_TYPE_INFO);
- IcebergSink.forRow(dataStream, SimpleDataUtil.FLINK_SCHEMA)
- .table(table)
- .tableLoader(tableLoader)
- .writeParallelism(parallelism)
- .setAll(newProps)
- .append();
-
- assertThatThrownBy(() -> env.execute("Test Iceberg DataStream"))
- .isInstanceOf(IllegalArgumentException.class)
- .hasMessage("Invalid distribution mode: UNRECOGNIZED");
- }
-
@TestTemplate
void testOverrideWriteConfigWithUnknownFileFormat() {
Map<String, String> newProps = Maps.newHashMap();
diff --git a/flink/v2.0/flink/src/main/java/org/apache/iceberg/flink/sink/IcebergSink.java b/flink/v2.0/flink/src/main/java/org/apache/iceberg/flink/sink/IcebergSink.java
index 1c1f11c387..373568c1b3 100644
--- a/flink/v2.0/flink/src/main/java/org/apache/iceberg/flink/sink/IcebergSink.java
+++ b/flink/v2.0/flink/src/main/java/org/apache/iceberg/flink/sink/IcebergSink.java
@@ -34,6 +34,7 @@ import java.util.Map;
import java.util.UUID;
import java.util.function.Function;
import org.apache.flink.annotation.Experimental;
+import org.apache.flink.api.common.functions.FlatMapFunction;
import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.connector.sink2.Committer;
@@ -64,13 +65,20 @@ import org.apache.iceberg.DistributionMode;
import org.apache.iceberg.FileFormat;
import org.apache.iceberg.PartitionField;
import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Partitioning;
import org.apache.iceberg.Schema;
import org.apache.iceberg.SerializableTable;
+import org.apache.iceberg.SortOrder;
import org.apache.iceberg.Table;
import org.apache.iceberg.flink.FlinkSchemaUtil;
import org.apache.iceberg.flink.FlinkWriteConf;
import org.apache.iceberg.flink.FlinkWriteOptions;
import org.apache.iceberg.flink.TableLoader;
+import org.apache.iceberg.flink.sink.shuffle.DataStatisticsOperatorFactory;
+import org.apache.iceberg.flink.sink.shuffle.RangePartitioner;
+import org.apache.iceberg.flink.sink.shuffle.StatisticsOrRecord;
+import org.apache.iceberg.flink.sink.shuffle.StatisticsOrRecordTypeInformation;
+import org.apache.iceberg.flink.sink.shuffle.StatisticsType;
import org.apache.iceberg.flink.util.FlinkCompatibilityUtil;
import org.apache.iceberg.io.WriteResult;
import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
@@ -376,22 +384,76 @@ public class IcebergSink
/**
* Configure the write {@link DistributionMode} that the IcebergSink will use. Currently, flink
- * support {@link DistributionMode#NONE} and {@link DistributionMode#HASH}.
+ * support {@link DistributionMode#NONE} and {@link DistributionMode#HASH} and {@link
+ * DistributionMode#RANGE}
*
* @param mode to specify the write distribution mode.
* @return {@link IcebergSink.Builder} to connect the iceberg table.
*/
@Override
public Builder distributionMode(DistributionMode mode) {
- Preconditions.checkArgument(
- !DistributionMode.RANGE.equals(mode),
- "Flink does not support 'range' write distribution mode now.");
if (mode != null) {
writeOptions.put(FlinkWriteOptions.DISTRIBUTION_MODE.key(), mode.modeName());
}
return this;
}
+ /**
+ * Range distribution needs to collect statistics about data distribution to properly shuffle
+ * the records in relatively balanced way. In general, low cardinality should use {@link
+ * StatisticsType#Map} and high cardinality should use {@link StatisticsType#Sketch} Refer to
+ * {@link StatisticsType} Javadoc for more details.
+ *
+ * <p>Default is {@link StatisticsType#Auto} where initially Map statistics is used. But if
+ * cardinality is higher than the threshold (currently 10K) as defined in {@code
+ * SketchUtil#OPERATOR_SKETCH_SWITCH_THRESHOLD}, statistics collection automatically switches to
+ * the sketch reservoir sampling.
+ *
+ * <p>Explicit set the statistics type if the default behavior doesn't work.
+ *
+ * @param type to specify the statistics type for range distribution.
+ * @return {@link IcebergSink.Builder} to connect the iceberg table.
+ */
+ public IcebergSink.Builder rangeDistributionStatisticsType(StatisticsType type) {
+ if (type != null) {
+ writeOptions.put(FlinkWriteOptions.RANGE_DISTRIBUTION_STATISTICS_TYPE.key(), type.name());
+ }
+ return this;
+ }
+
+ /**
+ * If sort order contains partition columns, each sort key would map to one partition and data
+ * file. This relative weight can avoid placing too many small files for sort keys with low
+ * traffic. It is a double value that defines the minimal weight for each sort key. `0.02` means
+ * each key has a base weight of `2%` of the targeted traffic weight per writer task.
+ *
+ * <p>E.g. the sink Iceberg table is partitioned daily by event time. Assume the data stream
+ * contains events from now up to 180 days ago. With event time, traffic weight distribution
+ * across different days typically has a long tail pattern. Current day contains the most
+ * traffic. The older days (long tail) contain less and less traffic. Assume writer parallelism
+ * is `10`. The total weight across all 180 days is `10,000`. Target traffic weight per writer
+ * task would be `1,000`. Assume the weight sum for the oldest 150 days is `1,000`. Normally,
+ * the range partitioner would put all the oldest 150 days in one writer task. That writer task
+ * would write to 150 small files (one per day). Keeping 150 open files can potentially consume
+ * large amount of memory. Flushing and uploading 150 files (however small) at checkpoint time
+ * can also be potentially slow. If this config is set to `0.02`. It means every sort key has a
+ * base weight of `2%` of targeted weight of `1,000` for every write task. It would essentially
+ * avoid placing more than `50` data files (one per day) on one writer task no matter how small
+ * they are.
+ *
+ * <p>This is only applicable to {@link StatisticsType#Map} for low-cardinality scenario. For
+ * {@link StatisticsType#Sketch} high-cardinality sort columns, they are usually not used as
+ * partition columns. Otherwise, too many partitions and small files may be generated during
+ * write. Sketch range partitioner simply splits high-cardinality keys into ordered ranges.
+ *
+ * <p>Default is {@code 0.0%}.
+ */
+ public Builder rangeDistributionSortKeyBaseWeight(double weight) {
+ writeOptions.put(
+ FlinkWriteOptions.RANGE_DISTRIBUTION_SORT_KEY_BASE_WEIGHT.key(), Double.toString(weight));
+ return this;
+ }
+
/**
* Configuring the write parallel number for iceberg stream writer.
*
@@ -562,6 +624,10 @@ public class IcebergSink
}
}
+ private String operatorName(String suffix) {
+ return uidSuffix != null ? suffix + "-" + uidSuffix : suffix;
+ }
+
private static String defaultSuffix(String uidSuffix, String defaultSuffix) {
if (uidSuffix == null || uidSuffix.isEmpty()) {
return defaultSuffix;
@@ -648,73 +714,136 @@ public class IcebergSink
DistributionMode mode = flinkWriteConf.distributionMode();
Schema schema = table.schema();
PartitionSpec spec = table.spec();
+ SortOrder sortOrder = table.sortOrder();
+
LOG.info("Write distribution mode is '{}'", mode.modeName());
switch (mode) {
case NONE:
- if (equalityFieldIds.isEmpty()) {
- return input;
- } else {
- LOG.info("Distribute rows by equality fields, because there are equality fields set");
- return input.keyBy(new EqualityFieldKeySelector(schema, flinkRowType, equalityFieldIds));
- }
-
+ return distributeDataStreamByNoneDistributionMode(input, schema);
case HASH:
- if (equalityFieldIds.isEmpty()) {
- if (table.spec().isUnpartitioned()) {
- LOG.warn(
- "Fallback to use 'none' distribution mode, because there are no equality fields set "
- + "and table is unpartitioned");
- return input;
- } else {
- if (BucketPartitionerUtil.hasOneBucketField(spec)) {
- return input.partitionCustom(
- new BucketPartitioner(spec),
- new BucketPartitionKeySelector(spec, schema, flinkRowType));
- } else {
- return input.keyBy(new PartitionKeySelector(spec, schema, flinkRowType));
- }
- }
- } else {
- if (spec.isUnpartitioned()) {
- LOG.info(
- "Distribute rows by equality fields, because there are equality fields set "
- + "and table is unpartitioned");
- return input.keyBy(
- new EqualityFieldKeySelector(schema, flinkRowType, equalityFieldIds));
- } else {
- for (PartitionField partitionField : spec.fields()) {
- Preconditions.checkState(
- equalityFieldIds.contains(partitionField.sourceId()),
- "In 'hash' distribution mode with equality fields set, source column '%s' of partition field '%s' "
- + "should be included in equality fields: '%s'",
- table.schema().findColumnName(partitionField.sourceId()),
- partitionField,
- equalityFieldColumns);
- }
- return input.keyBy(new PartitionKeySelector(spec, schema, flinkRowType));
- }
- }
-
+ return distributeDataStreamByHashDistributionMode(input, schema, spec);
case RANGE:
- if (equalityFieldIds.isEmpty()) {
- LOG.warn(
- "Fallback to use 'none' distribution mode, because there are no equality fields set "
- + "and {}=range is not supported yet in flink",
- WRITE_DISTRIBUTION_MODE);
- return input;
- } else {
- LOG.info(
- "Distribute rows by equality fields, because there are equality fields set "
- + "and{}=range is not supported yet in flink",
- WRITE_DISTRIBUTION_MODE);
- return input.keyBy(new EqualityFieldKeySelector(schema, flinkRowType, equalityFieldIds));
- }
-
+ return distributeDataStreamByRangeDistributionMode(input, schema, spec, sortOrder);
default:
throw new RuntimeException("Unrecognized " + WRITE_DISTRIBUTION_MODE + ": " + mode);
}
}
+ private DataStream<RowData> distributeDataStreamByNoneDistributionMode(
+ DataStream<RowData> input, Schema iSchema) {
+ if (equalityFieldIds.isEmpty()) {
+ return input;
+ } else {
+ LOG.info("Distribute rows by equality fields, because there are equality fields set");
+ return input.keyBy(new EqualityFieldKeySelector(iSchema, flinkRowType, equalityFieldIds));
+ }
+ }
+
+ private DataStream<RowData> distributeDataStreamByHashDistributionMode(
+ DataStream<RowData> input, Schema iSchema, PartitionSpec partitionSpec) {
+ if (equalityFieldIds.isEmpty()) {
+ if (partitionSpec.isUnpartitioned()) {
+ LOG.warn(
+ "Fallback to use 'none' distribution mode, because there are no equality fields set "
+ + "and table is unpartitioned");
+ return input;
+ } else {
+ return input.keyBy(new PartitionKeySelector(partitionSpec, iSchema, flinkRowType));
+ }
+ } else {
+ if (partitionSpec.isUnpartitioned()) {
+ LOG.info(
+ "Distribute rows by equality fields, because there are equality fields set "
+ + "and table is unpartitioned");
+ return input.keyBy(new EqualityFieldKeySelector(iSchema, flinkRowType, equalityFieldIds));
+ } else {
+ for (PartitionField partitionField : partitionSpec.fields()) {
+ Preconditions.checkState(
+ equalityFieldIds.contains(partitionField.sourceId()),
+ "In 'hash' distribution mode with equality fields set, source column '%s' of partition field '%s' "
+ + "should be included in equality fields: '%s'",
+ table.schema().findColumnName(partitionField.sourceId()),
+ partitionField,
+ equalityFieldColumns);
+ }
+ return input.keyBy(new PartitionKeySelector(partitionSpec, iSchema, flinkRowType));
+ }
+ }
+ }
+
+ private DataStream<RowData> distributeDataStreamByRangeDistributionMode(
+ DataStream<RowData> input,
+ Schema iSchema,
+ PartitionSpec partitionSpec,
+ SortOrder sortOrderParam) {
+
+ int writerParallelism =
+ flinkWriteConf.writeParallelism() == null
+ ? input.getParallelism()
+ : flinkWriteConf.writeParallelism();
+
+ // needed because of checkStyle not allowing us to change the value of an argument
+ SortOrder sortOrder = sortOrderParam;
+
+ // Ideally, exception should be thrown in the combination of range distribution and
+ // equality fields. Primary key case should use hash distribution mode.
+ // Keep the current behavior of falling back to keyBy for backward compatibility.
+ if (!equalityFieldIds.isEmpty()) {
+ LOG.warn(
+ "Hash distribute rows by equality fields, even though {}=range is set. "
+ + "Range distribution for primary keys are not always safe in "
+ + "Flink streaming writer.",
+ WRITE_DISTRIBUTION_MODE);
+ return input.keyBy(new EqualityFieldKeySelector(iSchema, flinkRowType, equalityFieldIds));
+ }
+
+ // range distribute by partition key or sort key if table has an SortOrder
+ Preconditions.checkState(
+ sortOrder.isSorted() || partitionSpec.isPartitioned(),
+ "Invalid write distribution mode: range. Need to define sort order or partition spec.");
+ if (sortOrder.isUnsorted()) {
+ sortOrder = Partitioning.sortOrderFor(partitionSpec);
+ LOG.info("Construct sort order from partition spec");
+ }
+
+ LOG.info("Range distribute rows by sort order: {}", sortOrder);
+ StatisticsOrRecordTypeInformation statisticsOrRecordTypeInformation =
+ new StatisticsOrRecordTypeInformation(flinkRowType, iSchema, sortOrder);
+ StatisticsType statisticsType = flinkWriteConf.rangeDistributionStatisticsType();
+ SingleOutputStreamOperator<StatisticsOrRecord> shuffleStream =
+ input
+ .transform(
+ operatorName("range-shuffle"),
+ statisticsOrRecordTypeInformation,
+ new DataStatisticsOperatorFactory(
+ iSchema,
+ sortOrder,
+ writerParallelism,
+ statisticsType,
+ flinkWriteConf.rangeDistributionSortKeyBaseWeight()))
+ // Set the parallelism same as input operator to encourage chaining
+ .setParallelism(input.getParallelism());
+
+ if (uidSuffix != null) {
+ shuffleStream = shuffleStream.uid("shuffle-" + uidSuffix);
+ }
+
+ return shuffleStream
+ .partitionCustom(new RangePartitioner(iSchema, sortOrder), r -> r)
+ .flatMap(
+ (FlatMapFunction<StatisticsOrRecord, RowData>)
+ (statisticsOrRecord, out) -> {
+ if (statisticsOrRecord.hasRecord()) {
+ out.collect(statisticsOrRecord.record());
+ }
+ })
+ // Set slot sharing group and the parallelism same as writerParallelism to
+ // promote operator chaining with the downstream writer operator
+ .slotSharingGroup("shuffle-partition-custom-group")
+ .setParallelism(writerParallelism)
+ .returns(RowData.class);
+ }
+
/**
* Initialize a {@link Builder} to export the data from generic input data stream into iceberg
* table. We use {@link RowData} inside the sink connector, so users need to provide a mapper
diff --git a/flink/v2.0/flink/src/test/java/org/apache/iceberg/flink/sink/TestFlinkIcebergSinkV2DistributionMode.java b/flink/v2.0/flink/src/test/java/org/apache/iceberg/flink/sink/TestFlinkIcebergSinkV2DistributionMode.java
new file mode 100644
index 0000000000..47217db206
--- /dev/null
+++ b/flink/v2.0/flink/src/test/java/org/apache/iceberg/flink/sink/TestFlinkIcebergSinkV2DistributionMode.java
@@ -0,0 +1,539 @@
+/*
+ * 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.iceberg.flink.sink;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.assertj.core.api.Assumptions.assumeThat;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.test.junit5.MiniClusterExtension;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.DataFile;
+import org.apache.iceberg.DistributionMode;
+import org.apache.iceberg.FileFormat;
+import org.apache.iceberg.Parameter;
+import org.apache.iceberg.ParameterizedTestExtension;
+import org.apache.iceberg.Parameters;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Snapshot;
+import org.apache.iceberg.TableProperties;
+import org.apache.iceberg.flink.FlinkWriteOptions;
+import org.apache.iceberg.flink.HadoopCatalogExtension;
+import org.apache.iceberg.flink.MiniFlinkClusterExtension;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.flink.TestFixtures;
+import org.apache.iceberg.flink.sink.shuffle.StatisticsType;
+import org.apache.iceberg.flink.source.BoundedTestSource;
+import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList;
+import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+import org.apache.iceberg.types.Conversions;
+import org.apache.iceberg.types.Types;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.TestTemplate;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+/**
+ * This tests the distribution mode of the IcebergSink. Extract them separately since it is
+ * unnecessary to test different file formats (Avro, Orc, Parquet) like in {@link TestIcebergSink}.
+ * Removing the file format dimension reduces the number of combinations from 12 to 4, which helps
+ * reduce test run time.
+ */
+@ExtendWith(ParameterizedTestExtension.class)
+public class TestFlinkIcebergSinkV2DistributionMode extends TestFlinkIcebergSinkBase {
+
+ @RegisterExtension
+ public static MiniClusterExtension miniClusterResource =
+ MiniFlinkClusterExtension.createWithClassloaderCheckDisabled();
+
+ @RegisterExtension
+ private static final HadoopCatalogExtension CATALOG_EXTENSION =
+ new HadoopCatalogExtension(TestFixtures.DATABASE, TestFixtures.TABLE);
+
+ private final FileFormat format = FileFormat.PARQUET;
+
+ @Parameter(index = 0)
+ private int parallelism;
+
+ @Parameter(index = 1)
+ private boolean partitioned;
+
+ @Parameter(index = 2)
+ private int writeParallelism;
+
+ @Parameters(name = "parallelism = {0}, partitioned = {1}, writeParallelism = {2}")
+ public static Object[][] parameters() {
+ return new Object[][] {
+ {1, true, 1},
+ {1, false, 1},
+ {2, true, 2},
+ {2, false, 2},
+ {1, true, 2},
+ {1, false, 2},
+ };
+ }
+
+ @BeforeEach
+ public void before() throws IOException {
+ this.table =
+ CATALOG_EXTENSION
+ .catalog()
+ .createTable(
+ TestFixtures.TABLE_IDENTIFIER,
+ SimpleDataUtil.SCHEMA,
+ partitioned
+ ? PartitionSpec.builderFor(SimpleDataUtil.SCHEMA).identity("data").build()
+ : PartitionSpec.unpartitioned(),
+ ImmutableMap.of(TableProperties.DEFAULT_FILE_FORMAT, format.name()));
+
+ this.env =
+ StreamExecutionEnvironment.getExecutionEnvironment(
+ MiniFlinkClusterExtension.DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .enableCheckpointing(100)
+ .setParallelism(parallelism)
+ .setMaxParallelism(Math.max(parallelism, writeParallelism));
+
+ this.tableLoader = CATALOG_EXTENSION.tableLoader();
+ }
+
+ @TestTemplate
+ public void testShuffleByPartitionWithSchema() throws Exception {
+ testWriteRow(parallelism, SimpleDataUtil.FLINK_SCHEMA, DistributionMode.HASH);
+ if (partitioned) {
+ assertThat(partitionFiles("aaa")).isEqualTo(1);
+ assertThat(partitionFiles("bbb")).isEqualTo(1);
+ assertThat(partitionFiles("ccc")).isEqualTo(1);
+ }
+ }
+
+ @TestTemplate
+ public void testJobNoneDistributeMode() throws Exception {
+ table
+ .updateProperties()
+ .set(TableProperties.WRITE_DISTRIBUTION_MODE, DistributionMode.HASH.modeName())
+ .commit();
+
+ testWriteRow(parallelism, null, DistributionMode.NONE);
+
+ if (parallelism > 1) {
+ if (partitioned) {
+ int files = partitionFiles("aaa") + partitionFiles("bbb") + partitionFiles("ccc");
+ assertThat(files).isGreaterThan(3);
+ }
+ }
+ }
+
+ @TestTemplate
+ public void testJobNullDistributionMode() throws Exception {
+ table
+ .updateProperties()
+ .set(TableProperties.WRITE_DISTRIBUTION_MODE, DistributionMode.HASH.modeName())
+ .commit();
+
+ testWriteRow(parallelism, null, null);
+
+ if (partitioned) {
+ assertThat(partitionFiles("aaa")).isEqualTo(1);
+ assertThat(partitionFiles("bbb")).isEqualTo(1);
+ assertThat(partitionFiles("ccc")).isEqualTo(1);
+ }
+ }
+
+ @TestTemplate
+ public void testPartitionWriteMode() throws Exception {
+ testWriteRow(parallelism, null, DistributionMode.HASH);
+ if (partitioned) {
+ assertThat(partitionFiles("aaa")).isEqualTo(1);
+ assertThat(partitionFiles("bbb")).isEqualTo(1);
+ assertThat(partitionFiles("ccc")).isEqualTo(1);
+ }
+ }
+
+ @TestTemplate
+ public void testOverrideWriteConfigWithUnknownDistributionMode() {
+ Map<String, String> newProps = Maps.newHashMap();
+ newProps.put(FlinkWriteOptions.DISTRIBUTION_MODE.key(), "UNRECOGNIZED");
+
+ List<Row> rows = createRows("");
+ DataStream<Row> dataStream = env.addSource(createBoundedSource(rows), ROW_TYPE_INFO);
+
+ IcebergSink.forRow(dataStream, SimpleDataUtil.FLINK_SCHEMA)
+ .table(table)
+ .tableLoader(tableLoader)
+ .writeParallelism(writeParallelism)
+ .setAll(newProps)
+ .append();
+
+ assertThatThrownBy(env::execute)
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessage("Invalid distribution mode: UNRECOGNIZED");
+ }
+
+ @TestTemplate
+ public void testRangeDistributionWithoutSortOrderUnpartitioned() throws Exception {
+ assumeThat(partitioned).isFalse();
+
+ table
+ .updateProperties()
+ .set(TableProperties.WRITE_DISTRIBUTION_MODE, DistributionMode.RANGE.modeName())
+ .commit();
+
+ int numOfCheckpoints = 6;
+ DataStream<Row> dataStream =
+ env.addSource(
+ createRangeDistributionBoundedSource(createCharRows(numOfCheckpoints, 10)),
+ ROW_TYPE_INFO);
+
+ IcebergSink.forRow(dataStream, SimpleDataUtil.FLINK_SCHEMA)
+ .table(table)
+ .tableLoader(tableLoader)
+ .writeParallelism(writeParallelism)
+ .append();
+
+ // Range distribution requires either sort order or partition spec defined
+ assertThatThrownBy(env::execute)
+ .isInstanceOf(IllegalStateException.class)
+ .hasMessage(
+ "Invalid write distribution mode: range. Need to define sort order or partition spec.");
+ }
+
+ @TestTemplate
+ public void testRangeDistributionWithoutSortOrderPartitioned() throws Exception {
+ assumeThat(partitioned).isTrue();
+
+ table
+ .updateProperties()
+ .set(TableProperties.WRITE_DISTRIBUTION_MODE, DistributionMode.RANGE.modeName())
+ .commit();
+
+ int numOfCheckpoints = 6;
+ DataStream<Row> dataStream =
+ env.addSource(
+ createRangeDistributionBoundedSource(createCharRows(numOfCheckpoints, 10)),
+ ROW_TYPE_INFO);
+
+ IcebergSink.forRow(dataStream, SimpleDataUtil.FLINK_SCHEMA)
+ .table(table)
+ .tableLoader(tableLoader)
+ .writeParallelism(writeParallelism)
+ .append();
+
+ // sort based on partition columns
+ env.execute(getClass().getSimpleName());
+
+ table.refresh();
+ // ordered in reverse timeline from the newest snapshot to the oldest snapshot
+ List<Snapshot> snapshots = Lists.newArrayList(table.snapshots().iterator());
+ // only keep the snapshots with added data files
+ snapshots =
+ snapshots.stream()
+ .filter(snapshot -> snapshot.addedDataFiles(table.io()).iterator().hasNext())
+ .collect(Collectors.toList());
+
+ // Sometimes we will have more checkpoints than the bounded source if we pass the
+ // auto checkpoint interval. Thus producing multiple snapshots.
+ assertThat(snapshots).hasSizeGreaterThanOrEqualTo(numOfCheckpoints);
+ }
+
+ @TestTemplate
+ public void testRangeDistributionWithNullValue() throws Exception {
+ assumeThat(partitioned).isTrue();
+
+ table
+ .updateProperties()
+ .set(TableProperties.WRITE_DISTRIBUTION_MODE, DistributionMode.RANGE.modeName())
+ .commit();
+
+ int numOfCheckpoints = 6;
+ List<List<Row>> charRows = createCharRows(numOfCheckpoints, 10);
+ charRows.add(ImmutableList.of(Row.of(1, null)));
+ DataStream<Row> dataStream =
+ env.addSource(createRangeDistributionBoundedSource(charRows), ROW_TYPE_INFO);
+
+ IcebergSink.forRow(dataStream, SimpleDataUtil.FLINK_SCHEMA)
+ .table(table)
+ .tableLoader(tableLoader)
+ .writeParallelism(parallelism)
+ .append();
+
+ // sort based on partition columns
+ env.execute(getClass().getSimpleName());
+
+ table.refresh();
+ // ordered in reverse timeline from the newest snapshot to the oldest snapshot
+ List<Snapshot> snapshots = Lists.newArrayList(table.snapshots().iterator());
+ // only keep the snapshots with added data files
+ snapshots =
+ snapshots.stream()
+ .filter(snapshot -> snapshot.addedDataFiles(table.io()).iterator().hasNext())
+ .collect(Collectors.toList());
+
+ // Sometimes we will have more checkpoints than the bounded source if we pass the
+ // auto checkpoint interval. Thus producing multiple snapshots.
+ assertThat(snapshots).hasSizeGreaterThanOrEqualTo(numOfCheckpoints);
+ }
+
+ @TestTemplate
+ public void testRangeDistributionWithSortOrder() throws Exception {
+ table
+ .updateProperties()
+ .set(TableProperties.WRITE_DISTRIBUTION_MODE, DistributionMode.RANGE.modeName())
+ .commit();
+ table.replaceSortOrder().asc("data").commit();
+
+ int numOfCheckpoints = 6;
+ DataStream<Row> dataStream =
+ env.addSource(
+ createRangeDistributionBoundedSource(createCharRows(numOfCheckpoints, 10)),
+ ROW_TYPE_INFO);
+ IcebergSink.forRow(dataStream, SimpleDataUtil.FLINK_SCHEMA)
+ .table(table)
+ .tableLoader(tableLoader)
+ .writeParallelism(writeParallelism)
+ .rangeDistributionStatisticsType(StatisticsType.Map)
+ .append();
+ env.execute(getClass().getSimpleName());
+
+ table.refresh();
+ // ordered in reverse timeline from the newest snapshot to the oldest snapshot
+ List<Snapshot> snapshots = Lists.newArrayList(table.snapshots().iterator());
+ // only keep the snapshots with added data files
+ snapshots =
+ snapshots.stream()
+ .filter(snapshot -> snapshot.addedDataFiles(table.io()).iterator().hasNext())
+ .collect(Collectors.toList());
+
+ // Sometimes we will have more checkpoints than the bounded source if we pass the
+ // auto checkpoint interval. Thus producing multiple snapshots.
+ assertThat(snapshots).hasSizeGreaterThanOrEqualTo(numOfCheckpoints);
+
+ // It takes 2 checkpoint cycle for statistics collection and application
+ // of the globally aggregated statistics in the range partitioner.
+ // The last two checkpoints should have range shuffle applied
+ List<Snapshot> rangePartitionedCycles =
+ snapshots.subList(snapshots.size() - 2, snapshots.size());
+
+ if (partitioned) {
+ for (Snapshot snapshot : rangePartitionedCycles) {
+ List<DataFile> addedDataFiles =
+ Lists.newArrayList(snapshot.addedDataFiles(table.io()).iterator());
+ // up to 26 partitions
+ assertThat(addedDataFiles).hasSizeLessThanOrEqualTo(26);
+ }
+ } else {
+ for (Snapshot snapshot : rangePartitionedCycles) {
+ List<DataFile> addedDataFiles =
+ Lists.newArrayList(snapshot.addedDataFiles(table.io()).iterator());
+ // each writer task should only write one file for non-partition sort column
+ assertThat(addedDataFiles).hasSize(writeParallelism);
+ // verify there is no overlap in min-max stats range
+ if (parallelism > 1) {
+ assertIdColumnStatsNoRangeOverlap(addedDataFiles.get(0), addedDataFiles.get(1));
+ }
+ }
+ }
+ }
+
+ @TestTemplate
+ public void testRangeDistributionSketchWithSortOrder() throws Exception {
+ table
+ .updateProperties()
+ .set(TableProperties.WRITE_DISTRIBUTION_MODE, DistributionMode.RANGE.modeName())
+ .commit();
+ table.replaceSortOrder().asc("id").commit();
+
+ int numOfCheckpoints = 6;
+ DataStream<Row> dataStream =
+ env.addSource(
+ createRangeDistributionBoundedSource(createIntRows(numOfCheckpoints, 1_000)),
+ ROW_TYPE_INFO);
+ IcebergSink.forRow(dataStream, SimpleDataUtil.FLINK_SCHEMA)
+ .table(table)
+ .tableLoader(tableLoader)
+ .writeParallelism(writeParallelism)
+ .rangeDistributionStatisticsType(StatisticsType.Sketch)
+ .append();
+ env.execute(getClass().getSimpleName());
+
+ table.refresh();
+ // ordered in reverse timeline from the newest snapshot to the oldest snapshot
+ List<Snapshot> snapshots = Lists.newArrayList(table.snapshots().iterator());
+ // only keep the snapshots with added data files
+ snapshots =
+ snapshots.stream()
+ .filter(snapshot -> snapshot.addedDataFiles(table.io()).iterator().hasNext())
+ .collect(Collectors.toList());
+
+ // Sometimes we will have more checkpoints than the bounded source if we pass the
+ // auto checkpoint interval. Thus producing multiple snapshots.
+ assertThat(snapshots).hasSizeGreaterThanOrEqualTo(numOfCheckpoints);
+
+ // It takes 2 checkpoint cycle for statistics collection and application
+ // of the globally aggregated statistics in the range partitioner.
+ // The last two checkpoints should have range shuffle applied
+ List<Snapshot> rangePartitionedCycles =
+ snapshots.subList(snapshots.size() - 2, snapshots.size());
+
+ // since the input has a single value for the data column,
+ // it is always the same partition. Hence there is no difference
+ // for partitioned or not
+ for (Snapshot snapshot : rangePartitionedCycles) {
+ List<DataFile> addedDataFiles =
+ Lists.newArrayList(snapshot.addedDataFiles(table.io()).iterator());
+ // each writer task should only write one file for non-partition sort column
+ assertThat(addedDataFiles).hasSize(writeParallelism);
+ // verify there is no overlap in min-max stats range
+ if (writeParallelism > 2) {
+ assertIdColumnStatsNoRangeOverlap(addedDataFiles.get(0), addedDataFiles.get(1));
+ }
+ }
+ }
+
+ /** Test migration from Map stats to Sketch stats */
+ @TestTemplate
+ public void testRangeDistributionStatisticsMigration() throws Exception {
+ table
+ .updateProperties()
+ .set(TableProperties.WRITE_DISTRIBUTION_MODE, DistributionMode.RANGE.modeName())
+ .commit();
+ table.replaceSortOrder().asc("id").commit();
+
+ int numOfCheckpoints = 4;
+ List<List<Row>> rowsPerCheckpoint = Lists.newArrayListWithCapacity(numOfCheckpoints);
+ for (int checkpointId = 0; checkpointId < numOfCheckpoints; ++checkpointId) {
+ // checkpointId 2 would emit 11_000 records which is larger than
+ // the OPERATOR_SKETCH_SWITCH_THRESHOLD of 10_000.
+ // This should trigger the stats migration.
+ int maxId = checkpointId < 1 ? 1_000 : 11_000;
+ List<Row> rows = Lists.newArrayListWithCapacity(maxId);
+ for (int j = 0; j < maxId; ++j) {
+ // fixed value "a" for the data (possible partition column)
+ rows.add(Row.of(j, "a"));
+ }
+
+ rowsPerCheckpoint.add(rows);
+ }
+
+ DataStream<Row> dataStream =
+ env.addSource(createRangeDistributionBoundedSource(rowsPerCheckpoint), ROW_TYPE_INFO);
+ IcebergSink.forRow(dataStream, SimpleDataUtil.FLINK_SCHEMA)
+ .table(table)
+ .tableLoader(tableLoader)
+ .writeParallelism(writeParallelism)
+ .rangeDistributionStatisticsType(StatisticsType.Auto)
+ .append();
+ env.execute(getClass().getSimpleName());
+
+ table.refresh();
+ // ordered in reverse timeline from the newest snapshot to the oldest snapshot
+ List<Snapshot> snapshots = Lists.newArrayList(table.snapshots().iterator());
+ // only keep the snapshots with added data files
+ snapshots =
+ snapshots.stream()
+ .filter(snapshot -> snapshot.addedDataFiles(table.io()).iterator().hasNext())
+ .collect(Collectors.toList());
+
+ // Sometimes we will have more checkpoints than the bounded source if we pass the
+ // auto checkpoint interval. Thus producing multiple snapshots.
+ assertThat(snapshots).hasSizeGreaterThanOrEqualTo(numOfCheckpoints);
+
+ // It takes 2 checkpoint cycle for statistics collection and application
+ // of the globally aggregated statistics in the range partitioner.
+ // The last two checkpoints should have range shuffle applied
+ List<Snapshot> rangePartitionedCycles =
+ snapshots.subList(snapshots.size() - 2, snapshots.size());
+
+ // since the input has a single value for the data column,
+ // it is always the same partition. Hence there is no difference
+ // for partitioned or not
+ for (Snapshot snapshot : rangePartitionedCycles) {
+ List<DataFile> addedDataFiles =
+ Lists.newArrayList(snapshot.addedDataFiles(table.io()).iterator());
+ // each writer task should only write one file for non-partition sort column
+ // sometimes
+ assertThat(addedDataFiles).hasSize(writeParallelism);
+ // verify there is no overlap in min-max stats range
+ if (writeParallelism > 1) {
+ assertIdColumnStatsNoRangeOverlap(addedDataFiles.get(0), addedDataFiles.get(1));
+ }
+ }
+ }
+
+ private BoundedTestSource<Row> createRangeDistributionBoundedSource(
+ List<List<Row>> rowsPerCheckpoint) {
+ return new BoundedTestSource<>(rowsPerCheckpoint);
+ }
+
+ private List<List<Row>> createCharRows(int numOfCheckpoints, int countPerChar) {
+ List<List<Row>> rowsPerCheckpoint = Lists.newArrayListWithCapacity(numOfCheckpoints);
+ for (int checkpointId = 0; checkpointId < numOfCheckpoints; ++checkpointId) {
+ List<Row> rows = Lists.newArrayListWithCapacity(26 * countPerChar);
+ for (int j = 0; j < countPerChar; ++j) {
+ for (char c = 'a'; c <= 'z'; ++c) {
+ rows.add(Row.of(1, String.valueOf(c)));
+ }
+ }
+
+ rowsPerCheckpoint.add(rows);
+ }
+
+ return rowsPerCheckpoint;
+ }
+
+ private List<List<Row>> createIntRows(int numOfCheckpoints, int maxId) {
+ List<List<Row>> rowsPerCheckpoint = Lists.newArrayListWithCapacity(numOfCheckpoints);
+ for (int checkpointId = 0; checkpointId < numOfCheckpoints; ++checkpointId) {
+ List<Row> rows = Lists.newArrayListWithCapacity(maxId);
+ for (int j = 0; j < maxId; ++j) {
+ // fixed value "a" for the data (possible partition column)
+ rows.add(Row.of(j, "a"));
+ }
+
+ rowsPerCheckpoint.add(rows);
+ }
+
+ return rowsPerCheckpoint;
+ }
+
+ private void assertIdColumnStatsNoRangeOverlap(DataFile file1, DataFile file2) {
+ // id column has fieldId 1
+ int file1LowerBound =
+ Conversions.fromByteBuffer(Types.IntegerType.get(), file1.lowerBounds().get(1));
+ int file1UpperBound =
+ Conversions.fromByteBuffer(Types.IntegerType.get(), file1.upperBounds().get(1));
+ int file2LowerBound =
+ Conversions.fromByteBuffer(Types.IntegerType.get(), file2.lowerBounds().get(1));
+ int file2UpperBound =
+ Conversions.fromByteBuffer(Types.IntegerType.get(), file2.upperBounds().get(1));
+
+ if (file1LowerBound < file2LowerBound) {
+ assertThat(file1UpperBound).isLessThanOrEqualTo(file2LowerBound);
+ } else {
+ assertThat(file2UpperBound).isLessThanOrEqualTo(file1LowerBound);
+ }
+ }
+}
diff --git a/flink/v2.0/flink/src/test/java/org/apache/iceberg/flink/sink/TestIcebergSink.java b/flink/v2.0/flink/src/test/java/org/apache/iceberg/flink/sink/TestIcebergSink.java
index 9990146a77..50a220de2d 100644
--- a/flink/v2.0/flink/src/test/java/org/apache/iceberg/flink/sink/TestIcebergSink.java
+++ b/flink/v2.0/flink/src/test/java/org/apache/iceberg/flink/sink/TestIcebergSink.java
@@ -145,57 +145,6 @@ public class TestIcebergSink extends TestFlinkIcebergSinkBase {
testWriteRow(SimpleDataUtil.FLINK_SCHEMA, DistributionMode.NONE);
}
- @TestTemplate
- void testJobNoneDistributeMode() throws Exception {
- table
- .updateProperties()
- .set(TableProperties.WRITE_DISTRIBUTION_MODE, DistributionMode.HASH.modeName())
- .commit();
-
- testWriteRow(null, DistributionMode.NONE);
-
- if (parallelism > 1) {
- if (partitioned) {
- int files = partitionFiles("aaa") + partitionFiles("bbb") + partitionFiles("ccc");
- assertThat(files).as("Should have more than 3 files in iceberg table.").isGreaterThan(3);
- }
- }
- }
-
- @TestTemplate
- void testJobHashDistributionMode() {
- table
- .updateProperties()
- .set(TableProperties.WRITE_DISTRIBUTION_MODE, DistributionMode.HASH.modeName())
- .commit();
-
- assertThatThrownBy(() -> testWriteRow(null, DistributionMode.RANGE))
- .isInstanceOf(IllegalArgumentException.class)
- .hasMessage("Flink does not support 'range' write distribution mode now.");
- }
-
- @TestTemplate
- void testJobNullDistributionMode() throws Exception {
- table
- .updateProperties()
- .set(TableProperties.WRITE_DISTRIBUTION_MODE, DistributionMode.HASH.modeName())
- .commit();
-
- testWriteRow(null, null);
-
- if (partitioned) {
- assertThat(partitionFiles("aaa"))
- .as("There should be only 1 data file in partition 'aaa'")
- .isEqualTo(1);
- assertThat(partitionFiles("bbb"))
- .as("There should be only 1 data file in partition 'bbb'")
- .isEqualTo(1);
- assertThat(partitionFiles("ccc"))
- .as("There should be only 1 data file in partition 'ccc'")
- .isEqualTo(1);
- }
- }
-
@TestTemplate
void testPartitionWriteMode() throws Exception {
testWriteRow(null, DistributionMode.HASH);
@@ -311,25 +260,6 @@ public class TestIcebergSink extends TestFlinkIcebergSinkBase {
assertThat(rightTable.currentSnapshot().summary().get("direction")).isEqualTo("rightTable");
}
- @TestTemplate
- void testOverrideWriteConfigWithUnknownDistributionMode() {
- Map<String, String> newProps = Maps.newHashMap();
- newProps.put(FlinkWriteOptions.DISTRIBUTION_MODE.key(), "UNRECOGNIZED");
-
- List<Row> rows = createRows("");
- DataStream<Row> dataStream = env.addSource(createBoundedSource(rows), ROW_TYPE_INFO);
- IcebergSink.forRow(dataStream, SimpleDataUtil.FLINK_SCHEMA)
- .table(table)
- .tableLoader(tableLoader)
- .writeParallelism(parallelism)
- .setAll(newProps)
- .append();
-
- assertThatThrownBy(() -> env.execute("Test Iceberg DataStream"))
- .isInstanceOf(IllegalArgumentException.class)
- .hasMessage("Invalid distribution mode: UNRECOGNIZED");
- }
-
@TestTemplate
void testOverrideWriteConfigWithUnknownFileFormat() {
Map<String, String> newProps = Maps.newHashMap();