You are viewing a plain text version of this content. The canonical link for it is here.
Posted to issues@iceberg.apache.org by "kengtin (via GitHub)" <gi...@apache.org> on 2023/03/21 17:15:26 UTC
[GitHub] [iceberg] kengtin opened a new pull request, #7161: Custom Bucket Partitioning
kengtin opened a new pull request, #7161:
URL: https://github.com/apache/iceberg/pull/7161
Description pending...
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] kengtin commented on a diff in pull request #7161: Flink: Custom partitioner for bucket partitions
Posted by "kengtin (via GitHub)" <gi...@apache.org>.
kengtin commented on code in PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#discussion_r1209429494
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerFlinkIcebergSink.java:
##########
@@ -0,0 +1,253 @@
+/*
+ * 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.apache.iceberg.flink.MiniClusterResource.DISABLE_CLASSLOADER_CHECK_CONFIG;
+import static org.apache.iceberg.flink.TestFixtures.DATABASE;
+import static org.apache.iceberg.flink.TestFixtures.TABLE_IDENTIFIER;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.TreeMap;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.typeutils.RowTypeInfo;
+import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.data.util.DataFormatConverters;
+import org.apache.flink.test.junit5.MiniClusterExtension;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.DistributionMode;
+import org.apache.iceberg.FileFormat;
+import org.apache.iceberg.FileScanTask;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Table;
+import org.apache.iceberg.TableProperties;
+import org.apache.iceberg.flink.HadoopCatalogExtension;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.flink.TableLoader;
+import org.apache.iceberg.flink.TestFixtures;
+import org.apache.iceberg.flink.sink.TestBucketPartitionerUtils.TableSchemaType;
+import org.apache.iceberg.flink.source.BoundedTestSource;
+import org.apache.iceberg.flink.util.FlinkCompatibilityUtil;
+import org.apache.iceberg.io.CloseableIterable;
+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.assertj.core.api.Assertions;
+import org.jetbrains.annotations.NotNull;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+public class TestBucketPartitionerFlinkIcebergSink {
+
+ private static final int NUMBER_TASK_MANAGERS = 1;
+ private static final int SLOTS_PER_TASK_MANAGER = 8;
+
+ @RegisterExtension
+ private static final MiniClusterExtension MINI_CLUSTER_RESOURCE =
+ new MiniClusterExtension(
+ new MiniClusterResourceConfiguration.Builder()
+ .setNumberTaskManagers(NUMBER_TASK_MANAGERS)
+ .setNumberSlotsPerTaskManager(SLOTS_PER_TASK_MANAGER)
+ .setConfiguration(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .build());
+
+ @RegisterExtension
+ private static final HadoopCatalogExtension catalogExtension =
+ new HadoopCatalogExtension(DATABASE, TestFixtures.TABLE);
+
+ private static final TypeInformation<Row> ROW_TYPE_INFO =
+ new RowTypeInfo(SimpleDataUtil.FLINK_SCHEMA.getFieldTypes());
+
+ // Parallelism = 8 (parallelism > numBuckets) throughout the test suite
+ private final int parallelism = NUMBER_TASK_MANAGERS * SLOTS_PER_TASK_MANAGER;
+ private final FileFormat format = FileFormat.PARQUET;
+ private final int numBuckets = 4;
+
+ private Table table;
+ private StreamExecutionEnvironment env;
+ private TableLoader tableLoader;
+
+ private void setupEnvironment(TableSchemaType tableSchemaType) {
+ table = getTable(tableSchemaType);
+ env =
+ StreamExecutionEnvironment.getExecutionEnvironment(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .enableCheckpointing(100)
+ .setParallelism(parallelism)
+ .setMaxParallelism(parallelism * 2);
+ tableLoader = catalogExtension.tableLoader();
+ }
+
+ private Table getTable(TableSchemaType tableSchemaType) {
+ PartitionSpec partitionSpec = TableSchemaType.getPartitionSpec(tableSchemaType, numBuckets);
+
+ return catalogExtension
+ .catalog()
+ .createTable(
+ TABLE_IDENTIFIER,
+ SimpleDataUtil.SCHEMA,
+ partitionSpec,
+ ImmutableMap.of(TableProperties.DEFAULT_FILE_FORMAT, format.name()));
+ }
+
+ private void testWriteRowData(List<RowData> allRows) throws Exception {
+ DataFormatConverters.RowConverter converter =
+ new DataFormatConverters.RowConverter(SimpleDataUtil.FLINK_SCHEMA.getFieldDataTypes());
+
+ DataStream<RowData> dataStream =
+ env.addSource(
+ new BoundedTestSource<>(
+ allRows.stream().map(converter::toExternal).toArray(Row[]::new)),
+ ROW_TYPE_INFO)
+ .map(converter::toInternal, FlinkCompatibilityUtil.toTypeInfo(SimpleDataUtil.ROW_TYPE));
+
+ FlinkSink.forRowData(dataStream)
+ .table(table)
+ .tableLoader(tableLoader)
+ .writeParallelism(parallelism)
+ .distributionMode(DistributionMode.HASH)
+ .append();
+
+ // Execute the program.
+ env.execute("Test Iceberg DataStream");
+
+ // Assert the iceberg table's records.
+ SimpleDataUtil.assertTableRows(table, allRows);
+ }
+
+ @ParameterizedTest
+ @EnumSource(
+ value = TableSchemaType.class,
+ names = {"ONE_BUCKET", "IDENTITY_AND_BUCKET"})
+ public void testSendRecordsToAllBucketsEvenly(TableSchemaType tableSchemaType) throws Exception {
+ setupEnvironment(tableSchemaType);
+ List<RowData> rows = generateTestDataRows();
+
+ testWriteRowData(rows);
+ TableTestStats stats = extractPartitionResults(tableSchemaType);
+
+ Assertions.assertThat(stats.totalRowCount).isEqualTo(rows.size());
+ // All 4 buckets should've been written to
+ Assertions.assertThat(stats.writersPerBucket.size()).isEqualTo(numBuckets);
+ Assertions.assertThat(stats.numFilesPerBucket.size()).isEqualTo(numBuckets);
+ // Writer expectation (2 writers per bucket):
+ // - Bucket0 -> Writers [0, 4]
+ // - Bucket1 -> Writers [1, 5]
+ // - Bucket2 -> Writers [2, 6]
+ // - Bucket3 -> Writers [3, 7]
+ for (int i = 0, j = numBuckets; i < numBuckets; i++, j++) {
+ Assertions.assertThat(stats.writersPerBucket.get(i)).isEqualTo(Arrays.asList(i, j));
+ // 2 files per bucket (one file is created by each writer)
+ Assertions.assertThat((int) stats.numFilesPerBucket.get(i)).isEqualTo(2);
+ // 2 rows per file (total of 16 rows across 8 files)
+ Assertions.assertThat((long) stats.rowsPerFile.get(i)).isEqualTo(2);
+ }
+ }
+
+ /**
+ * Verifies the BucketPartitioner is not used when the PartitionSpec has more than 1 bucket, and
+ * that it should fallback to input.keyBy
+ */
+ @ParameterizedTest
+ @EnumSource(value = TableSchemaType.class, names = "TWO_BUCKETS")
+ public void testMultipleBucketsFallback(TableSchemaType tableSchemaType) throws Exception {
+ setupEnvironment(tableSchemaType);
+ List<RowData> rows = generateTestDataRows();
+
+ testWriteRowData(rows);
+ TableTestStats stats = extractPartitionResults(tableSchemaType);
+
+ Assertions.assertThat(stats.totalRowCount).isEqualTo(rows.size());
+ for (int i = 0, j = numBuckets; i < numBuckets; i++, j++) {
+ // Only 1 file per bucket will be created when falling back to input.keyBy(...)
+ Assertions.assertThat((int) stats.numFilesPerBucket.get(i)).isEqualTo(1);
+ }
+ }
+
+ /**
+ * Generating 16 rows to be sent uniformly to all writers (round-robin across 8 writers -> 4
+ * buckets)
+ */
+ @NotNull
+ private List<RowData> generateTestDataRows() {
+ int totalNumRows = parallelism * 2;
+ int numRowsPerBucket = totalNumRows / numBuckets;
+ return TestBucketPartitionerUtils.generateRowsForBucketIdRange(numRowsPerBucket, numBuckets);
+ }
+
+ private TableTestStats extractPartitionResults(TableSchemaType tableSchemaType)
+ throws IOException {
+ int totalRecordCount = 0;
+ Map<Integer, List<Integer>> writersPerBucket = Maps.newHashMap(); // <BucketId, List<WriterId>>
Review Comment:
I agree, and even though the names of the variables accurately describe their purpose, those comments helped my brain quite a bit every time I went back to the code (either to debug or to remember details). I'd like to leave these 3 if that's ok.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] kengtin commented on a diff in pull request #7161: Flink: Custom partitioner for bucket partitions
Posted by "kengtin (via GitHub)" <gi...@apache.org>.
kengtin commented on code in PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#discussion_r1209417568
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/HadoopCatalogExtension.java:
##########
@@ -0,0 +1,104 @@
+/*
+ * 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;
+
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.UUID;
+import org.apache.commons.io.FileUtils;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.iceberg.CatalogProperties;
+import org.apache.iceberg.catalog.Catalog;
+import org.apache.iceberg.catalog.TableIdentifier;
+import org.apache.iceberg.hadoop.HadoopCatalog;
+import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
+import org.assertj.core.api.Assertions;
+import org.junit.jupiter.api.extension.AfterAllCallback;
+import org.junit.jupiter.api.extension.AfterEachCallback;
+import org.junit.jupiter.api.extension.BeforeAllCallback;
+import org.junit.jupiter.api.extension.BeforeEachCallback;
+import org.junit.jupiter.api.extension.ExtensionContext;
+
+public class HadoopCatalogExtension
+ implements BeforeAllCallback, BeforeEachCallback, AfterAllCallback, AfterEachCallback {
+ protected final String database;
+ protected final String tableName;
+
+ protected Path temporaryFolder;
Review Comment:
It could but then the management and awareness of this TemporaryFolder is again outside of this class and left to the user to initialize. I prefer to have all this fully encapsulated over removing a few lines of code.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] kengtin commented on a diff in pull request #7161: Flink: Custom partitioner for bucket partitions
Posted by "kengtin (via GitHub)" <gi...@apache.org>.
kengtin commented on code in PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#discussion_r1210563737
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerFlinkIcebergSink.java:
##########
@@ -0,0 +1,253 @@
+/*
+ * 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.apache.iceberg.flink.MiniClusterResource.DISABLE_CLASSLOADER_CHECK_CONFIG;
+import static org.apache.iceberg.flink.TestFixtures.DATABASE;
+import static org.apache.iceberg.flink.TestFixtures.TABLE_IDENTIFIER;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.TreeMap;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.typeutils.RowTypeInfo;
+import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.data.util.DataFormatConverters;
+import org.apache.flink.test.junit5.MiniClusterExtension;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.DistributionMode;
+import org.apache.iceberg.FileFormat;
+import org.apache.iceberg.FileScanTask;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Table;
+import org.apache.iceberg.TableProperties;
+import org.apache.iceberg.flink.HadoopCatalogExtension;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.flink.TableLoader;
+import org.apache.iceberg.flink.TestFixtures;
+import org.apache.iceberg.flink.sink.TestBucketPartitionerUtils.TableSchemaType;
+import org.apache.iceberg.flink.source.BoundedTestSource;
+import org.apache.iceberg.flink.util.FlinkCompatibilityUtil;
+import org.apache.iceberg.io.CloseableIterable;
+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.assertj.core.api.Assertions;
+import org.jetbrains.annotations.NotNull;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+public class TestBucketPartitionerFlinkIcebergSink {
+
+ private static final int NUMBER_TASK_MANAGERS = 1;
+ private static final int SLOTS_PER_TASK_MANAGER = 8;
+
+ @RegisterExtension
+ private static final MiniClusterExtension MINI_CLUSTER_RESOURCE =
+ new MiniClusterExtension(
+ new MiniClusterResourceConfiguration.Builder()
+ .setNumberTaskManagers(NUMBER_TASK_MANAGERS)
+ .setNumberSlotsPerTaskManager(SLOTS_PER_TASK_MANAGER)
+ .setConfiguration(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .build());
+
+ @RegisterExtension
+ private static final HadoopCatalogExtension catalogExtension =
+ new HadoopCatalogExtension(DATABASE, TestFixtures.TABLE);
+
+ private static final TypeInformation<Row> ROW_TYPE_INFO =
+ new RowTypeInfo(SimpleDataUtil.FLINK_SCHEMA.getFieldTypes());
+
+ // Parallelism = 8 (parallelism > numBuckets) throughout the test suite
+ private final int parallelism = NUMBER_TASK_MANAGERS * SLOTS_PER_TASK_MANAGER;
+ private final FileFormat format = FileFormat.PARQUET;
+ private final int numBuckets = 4;
+
+ private Table table;
+ private StreamExecutionEnvironment env;
+ private TableLoader tableLoader;
+
+ private void setupEnvironment(TableSchemaType tableSchemaType) {
+ table = getTable(tableSchemaType);
+ env =
+ StreamExecutionEnvironment.getExecutionEnvironment(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .enableCheckpointing(100)
+ .setParallelism(parallelism)
+ .setMaxParallelism(parallelism * 2);
+ tableLoader = catalogExtension.tableLoader();
+ }
+
+ private Table getTable(TableSchemaType tableSchemaType) {
+ PartitionSpec partitionSpec = TableSchemaType.getPartitionSpec(tableSchemaType, numBuckets);
+
+ return catalogExtension
+ .catalog()
+ .createTable(
+ TABLE_IDENTIFIER,
+ SimpleDataUtil.SCHEMA,
+ partitionSpec,
+ ImmutableMap.of(TableProperties.DEFAULT_FILE_FORMAT, format.name()));
+ }
+
+ private void testWriteRowData(List<RowData> allRows) throws Exception {
+ DataFormatConverters.RowConverter converter =
+ new DataFormatConverters.RowConverter(SimpleDataUtil.FLINK_SCHEMA.getFieldDataTypes());
+
+ DataStream<RowData> dataStream =
+ env.addSource(
+ new BoundedTestSource<>(
+ allRows.stream().map(converter::toExternal).toArray(Row[]::new)),
Review Comment:
I don't like this back/forth conversion either but couldn't find a better way. Trying to have the source return `RowData` directly results in: `Caused by: java.io.NotSerializableException: org.apache.flink.table.data.GenericRowData`. I didn't see a "serializable" version of `RowData`, thoughts?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] kengtin commented on pull request #7161: Flink: Custom partitioner for bucket partitions
Posted by "kengtin (via GitHub)" <gi...@apache.org>.
kengtin commented on PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#issuecomment-1689273814
Hi @advancedxy, thanks for bringing it up. If I recall correctly @stevenzwu and I discussed that point during the early stages and at least back then we determined it would be much simpler. However I see your point about not restricting the user's design to 1 bucket only. Perhaps we could allow multi-bucket/key definitions but I'd still restrict the routing/distribution to be based only the first bucket (or something like that) to not overcomplicate it. Thoughts?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] kengtin commented on a diff in pull request #7161: Flink: Custom partitioner for bucket partitions
Posted by "kengtin (via GitHub)" <gi...@apache.org>.
kengtin commented on code in PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#discussion_r1199631556
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerUtils.java:
##########
@@ -0,0 +1,121 @@
+/*
+ * 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 java.util.List;
+import java.util.UUID;
+import org.apache.flink.table.data.util.DataFormatConverters;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.apache.iceberg.util.BucketUtil;
+
+final class TestBucketPartitionerUtils {
+
+ enum TableSchemaType {
+ ONE_BUCKET,
+ IDENTITY_AND_BUCKET,
+ TWO_BUCKETS;
+ }
+
+ private TestBucketPartitionerUtils() {}
+
+ static final DataFormatConverters.RowConverter CONVERTER =
+ new DataFormatConverters.RowConverter(SimpleDataUtil.FLINK_SCHEMA.getFieldDataTypes());
+
+ static PartitionSpec getPartitionSpec(TableSchemaType tableSchemaType, int numBuckets) {
+ PartitionSpec partitionSpec = null;
+
+ switch (tableSchemaType) {
+ case ONE_BUCKET:
+ partitionSpec =
+ PartitionSpec.builderFor(SimpleDataUtil.SCHEMA).bucket("data", numBuckets).build();
+ break;
+ case IDENTITY_AND_BUCKET:
+ partitionSpec =
+ PartitionSpec.builderFor(SimpleDataUtil.SCHEMA)
+ .identity("id")
+ .bucket("data", numBuckets)
+ .build();
+ break;
+ case TWO_BUCKETS:
+ partitionSpec =
+ PartitionSpec.builderFor(SimpleDataUtil.SCHEMA)
+ .bucket("id", numBuckets)
+ .bucket("data", numBuckets)
+ .build();
+ break;
+ }
+
+ Preconditions.checkNotNull(
+ partitionSpec, "Invalid tableSchemaType provided: " + tableSchemaType);
+ return partitionSpec;
+ }
+
+ /**
+ * Utility method to generate rows whose values will "hash" to a range of bucketIds (from 0 to
+ * numBuckets - 1)
+ *
+ * @param numRowsPerBucket how many different rows should be generated per bucket
+ * @param numBuckets max number of buckets to consider
+ * @return the list of rows whose data "hashes" to the desired bucketId
+ */
+ static List<Row> generateRowsForBucketIdRange(int numRowsPerBucket, int numBuckets) {
Review Comment:
I made this change, however while creating the `BoundedTestSource` I got an exception that `RowData` was not Serializable. At least I was able to use `RowData` across all tests and remove the `CONVERTER` from the general utils, but I had to keep it inside that function to provide `Row` during the source creation (`TestBucketPartitionerFlinkIcebergSink` line 116). Thoughts?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] kengtin commented on a diff in pull request #7161: Flink: Custom partitioner for bucket partitions
Posted by "kengtin (via GitHub)" <gi...@apache.org>.
kengtin commented on code in PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#discussion_r1209425526
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerFlinkIcebergSink.java:
##########
@@ -0,0 +1,253 @@
+/*
+ * 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.apache.iceberg.flink.MiniClusterResource.DISABLE_CLASSLOADER_CHECK_CONFIG;
+import static org.apache.iceberg.flink.TestFixtures.DATABASE;
+import static org.apache.iceberg.flink.TestFixtures.TABLE_IDENTIFIER;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.TreeMap;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.typeutils.RowTypeInfo;
+import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.data.util.DataFormatConverters;
+import org.apache.flink.test.junit5.MiniClusterExtension;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.DistributionMode;
+import org.apache.iceberg.FileFormat;
+import org.apache.iceberg.FileScanTask;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Table;
+import org.apache.iceberg.TableProperties;
+import org.apache.iceberg.flink.HadoopCatalogExtension;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.flink.TableLoader;
+import org.apache.iceberg.flink.TestFixtures;
+import org.apache.iceberg.flink.sink.TestBucketPartitionerUtils.TableSchemaType;
+import org.apache.iceberg.flink.source.BoundedTestSource;
+import org.apache.iceberg.flink.util.FlinkCompatibilityUtil;
+import org.apache.iceberg.io.CloseableIterable;
+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.assertj.core.api.Assertions;
+import org.jetbrains.annotations.NotNull;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+public class TestBucketPartitionerFlinkIcebergSink {
+
+ private static final int NUMBER_TASK_MANAGERS = 1;
+ private static final int SLOTS_PER_TASK_MANAGER = 8;
+
+ @RegisterExtension
+ private static final MiniClusterExtension MINI_CLUSTER_RESOURCE =
+ new MiniClusterExtension(
+ new MiniClusterResourceConfiguration.Builder()
+ .setNumberTaskManagers(NUMBER_TASK_MANAGERS)
+ .setNumberSlotsPerTaskManager(SLOTS_PER_TASK_MANAGER)
+ .setConfiguration(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .build());
+
+ @RegisterExtension
+ private static final HadoopCatalogExtension catalogExtension =
+ new HadoopCatalogExtension(DATABASE, TestFixtures.TABLE);
+
+ private static final TypeInformation<Row> ROW_TYPE_INFO =
+ new RowTypeInfo(SimpleDataUtil.FLINK_SCHEMA.getFieldTypes());
+
+ // Parallelism = 8 (parallelism > numBuckets) throughout the test suite
+ private final int parallelism = NUMBER_TASK_MANAGERS * SLOTS_PER_TASK_MANAGER;
+ private final FileFormat format = FileFormat.PARQUET;
+ private final int numBuckets = 4;
+
+ private Table table;
+ private StreamExecutionEnvironment env;
+ private TableLoader tableLoader;
+
+ private void setupEnvironment(TableSchemaType tableSchemaType) {
+ table = getTable(tableSchemaType);
+ env =
+ StreamExecutionEnvironment.getExecutionEnvironment(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .enableCheckpointing(100)
+ .setParallelism(parallelism)
+ .setMaxParallelism(parallelism * 2);
+ tableLoader = catalogExtension.tableLoader();
+ }
+
+ private Table getTable(TableSchemaType tableSchemaType) {
+ PartitionSpec partitionSpec = TableSchemaType.getPartitionSpec(tableSchemaType, numBuckets);
+
+ return catalogExtension
+ .catalog()
+ .createTable(
+ TABLE_IDENTIFIER,
+ SimpleDataUtil.SCHEMA,
+ partitionSpec,
+ ImmutableMap.of(TableProperties.DEFAULT_FILE_FORMAT, format.name()));
+ }
+
+ private void testWriteRowData(List<RowData> allRows) throws Exception {
Review Comment:
Done. I had borrowed this directly from `TestFlinkIcebergSink`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] kengtin commented on a diff in pull request #7161: Flink: Custom partitioner for bucket partitions
Posted by "kengtin (via GitHub)" <gi...@apache.org>.
kengtin commented on code in PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#discussion_r1209423448
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerFlinkIcebergSink.java:
##########
@@ -0,0 +1,253 @@
+/*
+ * 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.apache.iceberg.flink.MiniClusterResource.DISABLE_CLASSLOADER_CHECK_CONFIG;
+import static org.apache.iceberg.flink.TestFixtures.DATABASE;
+import static org.apache.iceberg.flink.TestFixtures.TABLE_IDENTIFIER;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.TreeMap;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.typeutils.RowTypeInfo;
+import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.data.util.DataFormatConverters;
+import org.apache.flink.test.junit5.MiniClusterExtension;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.DistributionMode;
+import org.apache.iceberg.FileFormat;
+import org.apache.iceberg.FileScanTask;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Table;
+import org.apache.iceberg.TableProperties;
+import org.apache.iceberg.flink.HadoopCatalogExtension;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.flink.TableLoader;
+import org.apache.iceberg.flink.TestFixtures;
+import org.apache.iceberg.flink.sink.TestBucketPartitionerUtils.TableSchemaType;
+import org.apache.iceberg.flink.source.BoundedTestSource;
+import org.apache.iceberg.flink.util.FlinkCompatibilityUtil;
+import org.apache.iceberg.io.CloseableIterable;
+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.assertj.core.api.Assertions;
+import org.jetbrains.annotations.NotNull;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+public class TestBucketPartitionerFlinkIcebergSink {
+
+ private static final int NUMBER_TASK_MANAGERS = 1;
+ private static final int SLOTS_PER_TASK_MANAGER = 8;
+
+ @RegisterExtension
+ private static final MiniClusterExtension MINI_CLUSTER_RESOURCE =
+ new MiniClusterExtension(
+ new MiniClusterResourceConfiguration.Builder()
+ .setNumberTaskManagers(NUMBER_TASK_MANAGERS)
+ .setNumberSlotsPerTaskManager(SLOTS_PER_TASK_MANAGER)
+ .setConfiguration(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .build());
+
+ @RegisterExtension
+ private static final HadoopCatalogExtension catalogExtension =
+ new HadoopCatalogExtension(DATABASE, TestFixtures.TABLE);
+
+ private static final TypeInformation<Row> ROW_TYPE_INFO =
+ new RowTypeInfo(SimpleDataUtil.FLINK_SCHEMA.getFieldTypes());
+
+ // Parallelism = 8 (parallelism > numBuckets) throughout the test suite
+ private final int parallelism = NUMBER_TASK_MANAGERS * SLOTS_PER_TASK_MANAGER;
+ private final FileFormat format = FileFormat.PARQUET;
+ private final int numBuckets = 4;
+
+ private Table table;
+ private StreamExecutionEnvironment env;
+ private TableLoader tableLoader;
+
+ private void setupEnvironment(TableSchemaType tableSchemaType) {
Review Comment:
It does run before each test, but its parameter is different for each.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] kengtin commented on a diff in pull request #7161: Flink: Custom partitioner for a more even data distribution when bucket partitions are present
Posted by "kengtin (via GitHub)" <gi...@apache.org>.
kengtin commented on code in PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#discussion_r1169275167
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitionKeySelector.java:
##########
@@ -0,0 +1,80 @@
+/*
+ * 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 java.util.List;
+import java.util.stream.IntStream;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.types.logical.RowType;
+import org.apache.iceberg.PartitionKey;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Schema;
+import org.apache.iceberg.flink.RowDataWrapper;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/**
+ * Very similar to the {@link PartitionKeySelector}, but optimized to extract and return an Integer
+ * bucketId as the key.
+ */
+class BucketPartitionKeySelector implements KeySelector<RowData, Integer> {
+
+ private final Schema schema;
+ private final PartitionKey partitionKey;
+ private final RowType flinkSchema;
+ private final int bucketFieldPosition;
+
+ private transient RowDataWrapper rowDataWrapper;
+
+ BucketPartitionKeySelector(PartitionSpec spec, Schema schema, RowType flinkSchema) {
+ List<Tuple2<Integer, Integer>> bucketFields = BucketPartitionerUtils.getBucketFields(spec);
+
+ // The current implementation only supports ONE bucket
+ Preconditions.checkArgument(
+ bucketFields.size() == 1,
+ BucketPartitionerUtils.BAD_NUMBER_OF_BUCKETS_ERROR_MESSAGE + bucketFields.size());
Review Comment:
The main point here is that I'm using the function in 3 places, two of which need to fail if is not 1 bucket, but the other (in `FlinkSink`) only falls back to the default partitioning instead of failing. However all 3 places need to "check" we're only dealing with 1 bucket.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] kengtin commented on a diff in pull request #7161: Flink: Custom partitioner for bucket partitions
Posted by "kengtin (via GitHub)" <gi...@apache.org>.
kengtin commented on code in PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#discussion_r1199491147
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitioner.java:
##########
@@ -0,0 +1,99 @@
+/*
+ * 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.apache.iceberg.flink.sink.BucketPartitioner.BUCKET_OUT_OF_RANGE_MESSAGE_PREFIX;
+
+import org.apache.iceberg.PartitionSpec;
+import org.assertj.core.api.Assertions;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+import org.junit.jupiter.params.provider.ValueSource;
+
+public class TestBucketPartitioner {
+
+ static final int NUM_BUCKETS = 60;
+
+ @ParameterizedTest
+ @EnumSource(
+ value = TestBucketPartitionerUtils.TableSchemaType.class,
+ names = {"ONE_BUCKET", "IDENTITY_AND_BUCKET"})
+ public void testPartitioningParallelismGreaterThanBuckets(
+ TestBucketPartitionerUtils.TableSchemaType tableSchemaType) {
+ final int numPartitions = 500;
+
+ PartitionSpec partitionSpec =
+ TestBucketPartitionerUtils.getPartitionSpec(tableSchemaType, NUM_BUCKETS);
+
+ BucketPartitioner bucketPartitioner = new BucketPartitioner(partitionSpec);
+
+ for (int expectedIdx = 0, bucketId = 0; expectedIdx < numPartitions; expectedIdx++) {
+ int actualIdx = bucketPartitioner.partition(bucketId, numPartitions);
+ Assertions.assertThat(actualIdx).isEqualTo(expectedIdx);
+ if (++bucketId == NUM_BUCKETS) {
Review Comment:
Not sure I fully understand this comment but I separated the increment from the conditional evaluation for clarity.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] stevenzwu commented on pull request #7161: Flink: Custom partitioner for bucket partitions
Posted by "stevenzwu (via GitHub)" <gi...@apache.org>.
stevenzwu commented on PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#issuecomment-1672557847
thanks @kengtin for the contribution
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] kengtin commented on a diff in pull request #7161: Flink: Custom partitioner for bucket partitions
Posted by "kengtin (via GitHub)" <gi...@apache.org>.
kengtin commented on code in PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#discussion_r1204800602
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerFlinkIcebergSink.java:
##########
@@ -0,0 +1,256 @@
+/*
+ * 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.apache.iceberg.flink.MiniClusterResource.DISABLE_CLASSLOADER_CHECK_CONFIG;
+import static org.apache.iceberg.flink.TestFixtures.DATABASE;
+import static org.apache.iceberg.flink.TestFixtures.TABLE_IDENTIFIER;
+import static org.apache.iceberg.flink.sink.TestBucketPartitionerUtils.CONVERTER;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.TreeMap;
+import java.util.stream.Collectors;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.typeutils.RowTypeInfo;
+import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.test.junit5.MiniClusterExtension;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.DistributionMode;
+import org.apache.iceberg.FileFormat;
+import org.apache.iceberg.FileScanTask;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Table;
+import org.apache.iceberg.TableProperties;
+import org.apache.iceberg.flink.HadoopCatalogExtension;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.flink.TableLoader;
+import org.apache.iceberg.flink.TestFixtures;
+import org.apache.iceberg.flink.sink.TestBucketPartitionerUtils.TableSchemaType;
+import org.apache.iceberg.flink.source.BoundedTestSource;
+import org.apache.iceberg.flink.util.FlinkCompatibilityUtil;
+import org.apache.iceberg.io.CloseableIterable;
+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.assertj.core.api.Assertions;
+import org.jetbrains.annotations.NotNull;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+public class TestBucketPartitionerFlinkIcebergSink {
+
+ private static final int NUMBER_TASK_MANAGERS = 1;
+ private static final int SLOTS_PER_TASK_MANAGER = 8;
+
+ @RegisterExtension
+ private static final MiniClusterExtension MINI_CLUSTER_RESOURCE =
+ new MiniClusterExtension(
+ new MiniClusterResourceConfiguration.Builder()
+ .setNumberTaskManagers(NUMBER_TASK_MANAGERS)
+ .setNumberSlotsPerTaskManager(SLOTS_PER_TASK_MANAGER)
+ .setConfiguration(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .build());
+
+ @RegisterExtension
+ private static final HadoopCatalogExtension catalogExtension =
+ new HadoopCatalogExtension(DATABASE, TestFixtures.TABLE);
+
+ private static final TypeInformation<Row> ROW_TYPE_INFO =
+ new RowTypeInfo(SimpleDataUtil.FLINK_SCHEMA.getFieldTypes());
+
+ // Parallelism = 8 (parallelism > numBuckets) throughout the test suite
+ private final int parallelism = NUMBER_TASK_MANAGERS * SLOTS_PER_TASK_MANAGER;
+ private final FileFormat format = FileFormat.PARQUET;
+ private final int numBuckets = 4;
+
+ private Table table;
+ private StreamExecutionEnvironment env;
+ private TableLoader tableLoader;
+
+ private void setupEnvironment(TableSchemaType tableSchemaType) {
+ table = getTable(tableSchemaType);
+ env =
+ StreamExecutionEnvironment.getExecutionEnvironment(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .enableCheckpointing(100)
+ .setParallelism(parallelism)
+ .setMaxParallelism(parallelism * 2);
+ tableLoader = catalogExtension.tableLoader();
+ }
+
+ private Table getTable(TableSchemaType tableSchemaType) {
+ PartitionSpec partitionSpec =
+ TestBucketPartitionerUtils.getPartitionSpec(tableSchemaType, numBuckets);
+
+ return catalogExtension
+ .catalog()
+ .createTable(
+ TABLE_IDENTIFIER,
+ SimpleDataUtil.SCHEMA,
+ partitionSpec,
+ ImmutableMap.of(TableProperties.DEFAULT_FILE_FORMAT, format.name()));
+ }
+
+ private List<RowData> convertToRowData(List<Row> rows) {
+ return rows.stream().map(CONVERTER::toInternal).collect(Collectors.toList());
+ }
+
+ private BoundedTestSource<Row> createBoundedSource(List<Row> rows) {
+ return new BoundedTestSource<>(rows.toArray(new Row[0]));
+ }
+
+ private TableTestStats extractTableTestStats(TableSchemaType tableSchemaType) throws IOException {
+ int totalRecordCount = 0;
+ Map<Integer, List<Integer>> writersPerBucket = Maps.newHashMap(); // <BucketId, List<WriterId>>
+ Map<Integer, Integer> filesPerBucket = Maps.newHashMap(); // <BucketId, NumFiles>
+ Map<Integer, Long> recordsPerFile = new TreeMap<>(); // <WriterId, NumRecords>
+
+ try (CloseableIterable<FileScanTask> fileScanTasks = table.newScan().planFiles()) {
+ for (FileScanTask scanTask : fileScanTasks) {
+ long recordCountInFile = scanTask.file().recordCount();
+
+ String[] splitFilePath = scanTask.file().path().toString().split("/");
+ String filename = splitFilePath[splitFilePath.length - 1];
+ int writerId = Integer.parseInt(filename.split("-")[0]);
Review Comment:
This is interesting, I'm already covering the "fallback to keyBy" in one of the tests so we should be good there. Are you suggesting as an alternative?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] stevenzwu commented on a diff in pull request #7161: Flink: Custom partitioner for a more even data distribution when bucket partitions are present
Posted by "stevenzwu (via GitHub)" <gi...@apache.org>.
stevenzwu commented on code in PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#discussion_r1161160019
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitioner.java:
##########
@@ -0,0 +1,86 @@
+/*
+ * 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 java.util.Optional;
+import org.apache.flink.api.common.functions.Partitioner;
+import org.apache.iceberg.PartitionField;
+import org.apache.iceberg.PartitionSpec;
+
+/**
+ * This partitioner will redirect elements to writers deterministically so that each writer only
+ * targets 1 bucket. If the number of writers > number of buckets each partitioner will keep a state
+ * of multiple writers per bucket as evenly as possible, and will round-robin the requests across
+ * them.
+ */
+public class BucketPartitioner implements Partitioner<Integer> {
+
+ private final Integer maxBuckets;
+
+ private final int[] currentWriterOffset;
+
+ BucketPartitioner(PartitionSpec partitionSpec) {
+ // The current implementation redirects to writers based on the _FIRST_ bucket found
+ Optional<PartitionField> bucket =
+ partitionSpec.fields().stream()
+ .filter(f -> f.transform().dedupName().contains("bucket"))
Review Comment:
maybe we can try `PartitionSpecVisitor`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] kengtin commented on a diff in pull request #7161: Flink: Custom partitioner for bucket partitions
Posted by "kengtin (via GitHub)" <gi...@apache.org>.
kengtin commented on code in PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#discussion_r1199631556
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerUtils.java:
##########
@@ -0,0 +1,121 @@
+/*
+ * 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 java.util.List;
+import java.util.UUID;
+import org.apache.flink.table.data.util.DataFormatConverters;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.apache.iceberg.util.BucketUtil;
+
+final class TestBucketPartitionerUtils {
+
+ enum TableSchemaType {
+ ONE_BUCKET,
+ IDENTITY_AND_BUCKET,
+ TWO_BUCKETS;
+ }
+
+ private TestBucketPartitionerUtils() {}
+
+ static final DataFormatConverters.RowConverter CONVERTER =
+ new DataFormatConverters.RowConverter(SimpleDataUtil.FLINK_SCHEMA.getFieldDataTypes());
+
+ static PartitionSpec getPartitionSpec(TableSchemaType tableSchemaType, int numBuckets) {
+ PartitionSpec partitionSpec = null;
+
+ switch (tableSchemaType) {
+ case ONE_BUCKET:
+ partitionSpec =
+ PartitionSpec.builderFor(SimpleDataUtil.SCHEMA).bucket("data", numBuckets).build();
+ break;
+ case IDENTITY_AND_BUCKET:
+ partitionSpec =
+ PartitionSpec.builderFor(SimpleDataUtil.SCHEMA)
+ .identity("id")
+ .bucket("data", numBuckets)
+ .build();
+ break;
+ case TWO_BUCKETS:
+ partitionSpec =
+ PartitionSpec.builderFor(SimpleDataUtil.SCHEMA)
+ .bucket("id", numBuckets)
+ .bucket("data", numBuckets)
+ .build();
+ break;
+ }
+
+ Preconditions.checkNotNull(
+ partitionSpec, "Invalid tableSchemaType provided: " + tableSchemaType);
+ return partitionSpec;
+ }
+
+ /**
+ * Utility method to generate rows whose values will "hash" to a range of bucketIds (from 0 to
+ * numBuckets - 1)
+ *
+ * @param numRowsPerBucket how many different rows should be generated per bucket
+ * @param numBuckets max number of buckets to consider
+ * @return the list of rows whose data "hashes" to the desired bucketId
+ */
+ static List<Row> generateRowsForBucketIdRange(int numRowsPerBucket, int numBuckets) {
Review Comment:
I made this change, however the while creating the bounded test source I got an exception that `RowData` is not Serializable. At least I was able to use `RowData` across all tests and remove the `CONVERTER` from the general utils, but I had to keep it inside that function to provide `Row` during the source creation (`TestBucketPartitionerFlinkIcebergSink` line 116).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] kengtin commented on a diff in pull request #7161: Flink: Custom partitioner for bucket partitions
Posted by "kengtin (via GitHub)" <gi...@apache.org>.
kengtin commented on code in PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#discussion_r1209424711
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerFlinkIcebergSink.java:
##########
@@ -0,0 +1,253 @@
+/*
+ * 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.apache.iceberg.flink.MiniClusterResource.DISABLE_CLASSLOADER_CHECK_CONFIG;
+import static org.apache.iceberg.flink.TestFixtures.DATABASE;
+import static org.apache.iceberg.flink.TestFixtures.TABLE_IDENTIFIER;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.TreeMap;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.typeutils.RowTypeInfo;
+import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.data.util.DataFormatConverters;
+import org.apache.flink.test.junit5.MiniClusterExtension;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.DistributionMode;
+import org.apache.iceberg.FileFormat;
+import org.apache.iceberg.FileScanTask;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Table;
+import org.apache.iceberg.TableProperties;
+import org.apache.iceberg.flink.HadoopCatalogExtension;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.flink.TableLoader;
+import org.apache.iceberg.flink.TestFixtures;
+import org.apache.iceberg.flink.sink.TestBucketPartitionerUtils.TableSchemaType;
+import org.apache.iceberg.flink.source.BoundedTestSource;
+import org.apache.iceberg.flink.util.FlinkCompatibilityUtil;
+import org.apache.iceberg.io.CloseableIterable;
+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.assertj.core.api.Assertions;
+import org.jetbrains.annotations.NotNull;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+public class TestBucketPartitionerFlinkIcebergSink {
+
+ private static final int NUMBER_TASK_MANAGERS = 1;
+ private static final int SLOTS_PER_TASK_MANAGER = 8;
+
+ @RegisterExtension
+ private static final MiniClusterExtension MINI_CLUSTER_RESOURCE =
+ new MiniClusterExtension(
+ new MiniClusterResourceConfiguration.Builder()
+ .setNumberTaskManagers(NUMBER_TASK_MANAGERS)
+ .setNumberSlotsPerTaskManager(SLOTS_PER_TASK_MANAGER)
+ .setConfiguration(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .build());
+
+ @RegisterExtension
+ private static final HadoopCatalogExtension catalogExtension =
+ new HadoopCatalogExtension(DATABASE, TestFixtures.TABLE);
+
+ private static final TypeInformation<Row> ROW_TYPE_INFO =
+ new RowTypeInfo(SimpleDataUtil.FLINK_SCHEMA.getFieldTypes());
+
+ // Parallelism = 8 (parallelism > numBuckets) throughout the test suite
+ private final int parallelism = NUMBER_TASK_MANAGERS * SLOTS_PER_TASK_MANAGER;
+ private final FileFormat format = FileFormat.PARQUET;
+ private final int numBuckets = 4;
+
+ private Table table;
+ private StreamExecutionEnvironment env;
+ private TableLoader tableLoader;
+
+ private void setupEnvironment(TableSchemaType tableSchemaType) {
+ table = getTable(tableSchemaType);
+ env =
+ StreamExecutionEnvironment.getExecutionEnvironment(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .enableCheckpointing(100)
+ .setParallelism(parallelism)
+ .setMaxParallelism(parallelism * 2);
+ tableLoader = catalogExtension.tableLoader();
+ }
+
+ private Table getTable(TableSchemaType tableSchemaType) {
Review Comment:
Done. I think it was like that at some point but got split after multiple refactorings.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] kengtin commented on a diff in pull request #7161: Flink: Custom partitioner for bucket partitions
Posted by "kengtin (via GitHub)" <gi...@apache.org>.
kengtin commented on code in PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#discussion_r1210536611
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerFlinkIcebergSink.java:
##########
@@ -0,0 +1,253 @@
+/*
+ * 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.apache.iceberg.flink.MiniClusterResource.DISABLE_CLASSLOADER_CHECK_CONFIG;
+import static org.apache.iceberg.flink.TestFixtures.DATABASE;
+import static org.apache.iceberg.flink.TestFixtures.TABLE_IDENTIFIER;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.TreeMap;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.typeutils.RowTypeInfo;
+import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.data.util.DataFormatConverters;
+import org.apache.flink.test.junit5.MiniClusterExtension;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.DistributionMode;
+import org.apache.iceberg.FileFormat;
+import org.apache.iceberg.FileScanTask;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Table;
+import org.apache.iceberg.TableProperties;
+import org.apache.iceberg.flink.HadoopCatalogExtension;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.flink.TableLoader;
+import org.apache.iceberg.flink.TestFixtures;
+import org.apache.iceberg.flink.sink.TestBucketPartitionerUtils.TableSchemaType;
+import org.apache.iceberg.flink.source.BoundedTestSource;
+import org.apache.iceberg.flink.util.FlinkCompatibilityUtil;
+import org.apache.iceberg.io.CloseableIterable;
+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.assertj.core.api.Assertions;
+import org.jetbrains.annotations.NotNull;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+public class TestBucketPartitionerFlinkIcebergSink {
+
+ private static final int NUMBER_TASK_MANAGERS = 1;
+ private static final int SLOTS_PER_TASK_MANAGER = 8;
+
+ @RegisterExtension
+ private static final MiniClusterExtension MINI_CLUSTER_RESOURCE =
+ new MiniClusterExtension(
+ new MiniClusterResourceConfiguration.Builder()
+ .setNumberTaskManagers(NUMBER_TASK_MANAGERS)
+ .setNumberSlotsPerTaskManager(SLOTS_PER_TASK_MANAGER)
+ .setConfiguration(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .build());
+
+ @RegisterExtension
+ private static final HadoopCatalogExtension catalogExtension =
+ new HadoopCatalogExtension(DATABASE, TestFixtures.TABLE);
+
+ private static final TypeInformation<Row> ROW_TYPE_INFO =
+ new RowTypeInfo(SimpleDataUtil.FLINK_SCHEMA.getFieldTypes());
+
+ // Parallelism = 8 (parallelism > numBuckets) throughout the test suite
+ private final int parallelism = NUMBER_TASK_MANAGERS * SLOTS_PER_TASK_MANAGER;
+ private final FileFormat format = FileFormat.PARQUET;
+ private final int numBuckets = 4;
+
+ private Table table;
+ private StreamExecutionEnvironment env;
+ private TableLoader tableLoader;
+
+ private void setupEnvironment(TableSchemaType tableSchemaType) {
+ table = getTable(tableSchemaType);
+ env =
+ StreamExecutionEnvironment.getExecutionEnvironment(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .enableCheckpointing(100)
+ .setParallelism(parallelism)
+ .setMaxParallelism(parallelism * 2);
+ tableLoader = catalogExtension.tableLoader();
+ }
+
+ private Table getTable(TableSchemaType tableSchemaType) {
+ PartitionSpec partitionSpec = TableSchemaType.getPartitionSpec(tableSchemaType, numBuckets);
+
+ return catalogExtension
+ .catalog()
+ .createTable(
+ TABLE_IDENTIFIER,
+ SimpleDataUtil.SCHEMA,
+ partitionSpec,
+ ImmutableMap.of(TableProperties.DEFAULT_FILE_FORMAT, format.name()));
+ }
+
+ private void testWriteRowData(List<RowData> allRows) throws Exception {
+ DataFormatConverters.RowConverter converter =
+ new DataFormatConverters.RowConverter(SimpleDataUtil.FLINK_SCHEMA.getFieldDataTypes());
+
+ DataStream<RowData> dataStream =
+ env.addSource(
+ new BoundedTestSource<>(
+ allRows.stream().map(converter::toExternal).toArray(Row[]::new)),
+ ROW_TYPE_INFO)
+ .map(converter::toInternal, FlinkCompatibilityUtil.toTypeInfo(SimpleDataUtil.ROW_TYPE));
+
+ FlinkSink.forRowData(dataStream)
+ .table(table)
+ .tableLoader(tableLoader)
+ .writeParallelism(parallelism)
+ .distributionMode(DistributionMode.HASH)
+ .append();
+
+ // Execute the program.
+ env.execute("Test Iceberg DataStream");
+
+ // Assert the iceberg table's records.
+ SimpleDataUtil.assertTableRows(table, allRows);
+ }
+
+ @ParameterizedTest
+ @EnumSource(
+ value = TableSchemaType.class,
+ names = {"ONE_BUCKET", "IDENTITY_AND_BUCKET"})
+ public void testSendRecordsToAllBucketsEvenly(TableSchemaType tableSchemaType) throws Exception {
+ setupEnvironment(tableSchemaType);
+ List<RowData> rows = generateTestDataRows();
+
+ testWriteRowData(rows);
+ TableTestStats stats = extractPartitionResults(tableSchemaType);
+
+ Assertions.assertThat(stats.totalRowCount).isEqualTo(rows.size());
+ // All 4 buckets should've been written to
+ Assertions.assertThat(stats.writersPerBucket.size()).isEqualTo(numBuckets);
+ Assertions.assertThat(stats.numFilesPerBucket.size()).isEqualTo(numBuckets);
+ // Writer expectation (2 writers per bucket):
+ // - Bucket0 -> Writers [0, 4]
+ // - Bucket1 -> Writers [1, 5]
+ // - Bucket2 -> Writers [2, 6]
+ // - Bucket3 -> Writers [3, 7]
+ for (int i = 0, j = numBuckets; i < numBuckets; i++, j++) {
+ Assertions.assertThat(stats.writersPerBucket.get(i)).isEqualTo(Arrays.asList(i, j));
+ // 2 files per bucket (one file is created by each writer)
+ Assertions.assertThat((int) stats.numFilesPerBucket.get(i)).isEqualTo(2);
+ // 2 rows per file (total of 16 rows across 8 files)
+ Assertions.assertThat((long) stats.rowsPerFile.get(i)).isEqualTo(2);
+ }
+ }
+
+ /**
+ * Verifies the BucketPartitioner is not used when the PartitionSpec has more than 1 bucket, and
+ * that it should fallback to input.keyBy
+ */
+ @ParameterizedTest
+ @EnumSource(value = TableSchemaType.class, names = "TWO_BUCKETS")
+ public void testMultipleBucketsFallback(TableSchemaType tableSchemaType) throws Exception {
+ setupEnvironment(tableSchemaType);
+ List<RowData> rows = generateTestDataRows();
+
+ testWriteRowData(rows);
+ TableTestStats stats = extractPartitionResults(tableSchemaType);
+
+ Assertions.assertThat(stats.totalRowCount).isEqualTo(rows.size());
+ for (int i = 0, j = numBuckets; i < numBuckets; i++, j++) {
+ // Only 1 file per bucket will be created when falling back to input.keyBy(...)
+ Assertions.assertThat((int) stats.numFilesPerBucket.get(i)).isEqualTo(1);
+ }
+ }
+
+ /**
+ * Generating 16 rows to be sent uniformly to all writers (round-robin across 8 writers -> 4
+ * buckets)
+ */
+ @NotNull
+ private List<RowData> generateTestDataRows() {
+ int totalNumRows = parallelism * 2;
+ int numRowsPerBucket = totalNumRows / numBuckets;
+ return TestBucketPartitionerUtils.generateRowsForBucketIdRange(numRowsPerBucket, numBuckets);
+ }
+
+ private TableTestStats extractPartitionResults(TableSchemaType tableSchemaType)
+ throws IOException {
+ int totalRecordCount = 0;
+ Map<Integer, List<Integer>> writersPerBucket = Maps.newHashMap(); // <BucketId, List<WriterId>>
+ Map<Integer, Integer> filesPerBucket = Maps.newHashMap(); // <BucketId, NumFiles>
+ Map<Integer, Long> recordsPerFile = new TreeMap<>(); // <WriterId, NumRecords>
Review Comment:
Probably a leftover from a previous iteration, no sorting needed here.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] kengtin commented on a diff in pull request #7161: Flink: Custom partitioner for a more even data distribution when bucket partitions are present
Posted by "kengtin (via GitHub)" <gi...@apache.org>.
kengtin commented on code in PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#discussion_r1159010720
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitioner.java:
##########
@@ -0,0 +1,86 @@
+/*
+ * 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 java.util.Optional;
+import org.apache.flink.api.common.functions.Partitioner;
+import org.apache.iceberg.PartitionField;
+import org.apache.iceberg.PartitionSpec;
+
+/**
+ * This partitioner will redirect elements to writers deterministically so that each writer only
+ * targets 1 bucket. If the number of writers > number of buckets each partitioner will keep a state
+ * of multiple writers per bucket as evenly as possible, and will round-robin the requests across
+ * them.
+ */
+public class BucketPartitioner implements Partitioner<Integer> {
+
+ private final Integer maxBuckets;
+
+ private final int[] currentWriterOffset;
+
+ BucketPartitioner(PartitionSpec partitionSpec) {
+ // The current implementation redirects to writers based on the _FIRST_ bucket found
+ Optional<PartitionField> bucket =
+ partitionSpec.fields().stream()
+ .filter(f -> f.transform().dedupName().contains("bucket"))
+ .findFirst();
+
+ if (!bucket.isPresent()) {
Review Comment:
OK, we can fail here in case this partitioner is accessed directly. And have the `FlinkSink.distributeDataStream` only use this strategy if only 1 bucket is defined in the partition spec.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Flink: Custom partitioner for bucket partitions [iceberg]
Posted by "bendevera (via GitHub)" <gi...@apache.org>.
bendevera commented on PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#issuecomment-1861920740
@stevenzwu thank you for the quick response!
Okay, will run some `BucketPartitioner` tests for our use case by copying code manually. Smart shuffling sounds interesting and would certainly test out. A lot of the use cases we deal with fit well with `BucketPartitioner` conceptually and so can test to verify. Current `DistributionMode.HASH` implementation is too slow computing the data file path for each record and we've noticed processing rates take a huge hit when enabling. Glad to see features being extended! Will read up more regarding the smart shuffling design and see if we can get involved
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Flink: Custom partitioner for bucket partitions [iceberg]
Posted by "stevenzwu (via GitHub)" <gi...@apache.org>.
stevenzwu commented on PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#issuecomment-1861850713
It is reverted because there are users depending on the previous behavior of keyBy all partition columns. https://github.com/apache/iceberg/pull/7161#issuecomment-1761169778
We were assuming that if there is a bucket column, users only want to shuffle by the bucketing column. that is not the case from the user report linked in the above comment. so we decided to roll back for backward compatibility.
@bendevera you are right that `BucketPartitioner` isn't public and can't be used at the moment. Now we need to discuss what's the best way moving forward? we are working on a more comprehensive smart shuffling (range partition) feature: https://github.com/apache/iceberg/projects/27. I am thinking maybe we can expose this in `range` distribution mode.
before that, you may have to copy the code and manually apply the bucketing shuffling.
```
input.partitionCustom(
new BucketPartitioner(partitionSpec),
new BucketPartitionKeySelector(partitionSpec, iSchema, flinkRowType));
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Flink: Custom partitioner for bucket partitions [iceberg]
Posted by "stevenzwu (via GitHub)" <gi...@apache.org>.
stevenzwu commented on PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#issuecomment-1763706777
is the partition time an event time or ingestion/processing time? or asking in a different way, how many active minutes do the Flink writer job process for every commit cycle?
I feel this ongoing work might work better for you
https://docs.google.com/document/d/13N8cMqPi-ZPSKbkXGOBMPOzbv2Fua59j8bIjjtxLWqo
https://www.youtube.com/watch?v=GJplmOO7ULA
On the other hand, I also recognize this change imposed a behavior change that doesn't work for your use case. We can revert the distribution mode change in the `FlinkSink` class. To enable the custom bucketing partitioner, users will have to manually apply the custom `keyBy`.
```
input.keyBy(new PartitionKeySelector(partitionSpec, iSchema, flinkRowType))
```
@kengtin can you create a PR for reverting the FlinkSink change?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] kengtin commented on a diff in pull request #7161: Flink: Custom partitioner for bucket partitions
Posted by "kengtin (via GitHub)" <gi...@apache.org>.
kengtin commented on code in PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#discussion_r1261854165
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitioner.java:
##########
@@ -0,0 +1,107 @@
+/*
+ * 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 org.apache.flink.api.common.functions.Partitioner;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/**
+ * This partitioner will redirect records to writers deterministically based on the Bucket partition
+ * spec. It'll attempt to optimize the file size written depending on whether numPartitions is
+ * greater, less or equal than the maxNumBuckets. Note: The current implementation only supports ONE
+ * bucket in the partition spec.
+ */
+class BucketPartitioner implements Partitioner<Integer> {
+
+ static final String BUCKET_NULL_MESSAGE = "bucketId cannot be null";
+ static final String BUCKET_LESS_THAN_LOWER_BOUND_MESSAGE =
+ "Invalid bucket ID %s: must be non-negative.";
+ static final String BUCKET_GREATER_THAN_UPPER_BOUND_MESSAGE =
+ "Invalid bucket ID %s: must be less than bucket limit: %s.";
+
+ private final int maxNumBuckets;
+
+ // To hold the OFFSET of the next writer to use for any bucket, only used when writers > the
+ // number of buckets
+ private final int[] currentBucketWriterOffset;
+
+ BucketPartitioner(PartitionSpec partitionSpec) {
+ Tuple2<Integer, Integer> bucketFieldInfo =
+ BucketPartitionerUtil.getBucketFieldInfo(partitionSpec);
Review Comment:
Since the tuple represents <fieldId, maxNumBuckets>, what I can do is to expose 2 separate methods: `getBucketFieldId` and `getMaxNumBuckets` respectively, instead of the current `getBucketFieldInfo`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] kengtin commented on a diff in pull request #7161: Flink: Custom partitioner for bucket partitions
Posted by "kengtin (via GitHub)" <gi...@apache.org>.
kengtin commented on code in PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#discussion_r1261854165
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitioner.java:
##########
@@ -0,0 +1,107 @@
+/*
+ * 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 org.apache.flink.api.common.functions.Partitioner;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/**
+ * This partitioner will redirect records to writers deterministically based on the Bucket partition
+ * spec. It'll attempt to optimize the file size written depending on whether numPartitions is
+ * greater, less or equal than the maxNumBuckets. Note: The current implementation only supports ONE
+ * bucket in the partition spec.
+ */
+class BucketPartitioner implements Partitioner<Integer> {
+
+ static final String BUCKET_NULL_MESSAGE = "bucketId cannot be null";
+ static final String BUCKET_LESS_THAN_LOWER_BOUND_MESSAGE =
+ "Invalid bucket ID %s: must be non-negative.";
+ static final String BUCKET_GREATER_THAN_UPPER_BOUND_MESSAGE =
+ "Invalid bucket ID %s: must be less than bucket limit: %s.";
+
+ private final int maxNumBuckets;
+
+ // To hold the OFFSET of the next writer to use for any bucket, only used when writers > the
+ // number of buckets
+ private final int[] currentBucketWriterOffset;
+
+ BucketPartitioner(PartitionSpec partitionSpec) {
+ Tuple2<Integer, Integer> bucketFieldInfo =
+ BucketPartitionerUtil.getBucketFieldInfo(partitionSpec);
Review Comment:
Since the tuple represents <fieldId, maxNumBuckets>, what I can do is to expose 2 separate methods: `getFieldId` and `getMaxNumBuckets` respectively, instead of the current `getBucketFieldInfo`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] kengtin commented on a diff in pull request #7161: Flink: Custom partitioner for bucket partitions
Posted by "kengtin (via GitHub)" <gi...@apache.org>.
kengtin commented on code in PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#discussion_r1197218923
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitionKeySelector.java:
##########
@@ -0,0 +1,74 @@
+/*
+ * 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 java.util.stream.IntStream;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.types.logical.RowType;
+import org.apache.iceberg.PartitionKey;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Schema;
+import org.apache.iceberg.flink.RowDataWrapper;
+
+/**
+ * A {@link KeySelector} that extracts the bucketId from a data row's bucket partition as the key.
+ * To be used with the {@link BucketPartitioner}.
+ */
+class BucketPartitionKeySelector implements KeySelector<RowData, Integer> {
+
+ private final Schema schema;
+ private final PartitionKey partitionKey;
+ private final RowType flinkSchema;
+ private final int bucketFieldPosition;
+
+ private transient RowDataWrapper rowDataWrapper;
+
+ BucketPartitionKeySelector(PartitionSpec partitionSpec, Schema schema, RowType flinkSchema) {
+ Tuple2<Integer, Integer> bucketFieldInfo =
+ BucketPartitionerUtils.getBucketFieldInfo(partitionSpec);
+
+ int bucketFieldId = bucketFieldInfo.f0;
+ this.schema = schema;
+ this.partitionKey = new PartitionKey(partitionSpec, schema);
+ this.flinkSchema = flinkSchema;
+ this.bucketFieldPosition =
+ IntStream.range(0, partitionSpec.fields().size())
+ .filter(i -> partitionSpec.fields().get(i).fieldId() == bucketFieldId)
+ .toArray()[0];
+ }
+
+ /**
Review Comment:
Good point, this comment got copied from the original `PartitionKeySelector`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] kengtin commented on a diff in pull request #7161: Flink: Custom partitioner for bucket partitions
Posted by "kengtin (via GitHub)" <gi...@apache.org>.
kengtin commented on code in PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#discussion_r1209406397
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitioner.java:
##########
@@ -0,0 +1,108 @@
+/*
+ * 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 org.apache.flink.api.common.functions.Partitioner;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/**
+ * This partitioner will redirect records to writers deterministically based on the Bucket partition
+ * spec. It'll attempt to optimize the file size written depending on whether numPartitions is
+ * greater, less or equal than the maxNumBuckets. Note: The current implementation only supports ONE
+ * bucket in the partition spec.
+ */
+class BucketPartitioner implements Partitioner<Integer> {
+
+ static final String BUCKET_NULL_MESSAGE = "bucketId cannot be null";
+ static final String BUCKET_LESS_THAN_LOWER_BOUND_MESSAGE =
+ "bucketId out of range: %s, must be non-negative.";
Review Comment:
Overall I agree with this one, but I have my reservations on the double `:` in this case.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] kengtin commented on a diff in pull request #7161: Flink: Custom partitioner for bucket partitions
Posted by "kengtin (via GitHub)" <gi...@apache.org>.
kengtin commented on code in PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#discussion_r1209425526
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerFlinkIcebergSink.java:
##########
@@ -0,0 +1,253 @@
+/*
+ * 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.apache.iceberg.flink.MiniClusterResource.DISABLE_CLASSLOADER_CHECK_CONFIG;
+import static org.apache.iceberg.flink.TestFixtures.DATABASE;
+import static org.apache.iceberg.flink.TestFixtures.TABLE_IDENTIFIER;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.TreeMap;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.typeutils.RowTypeInfo;
+import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.data.util.DataFormatConverters;
+import org.apache.flink.test.junit5.MiniClusterExtension;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.DistributionMode;
+import org.apache.iceberg.FileFormat;
+import org.apache.iceberg.FileScanTask;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Table;
+import org.apache.iceberg.TableProperties;
+import org.apache.iceberg.flink.HadoopCatalogExtension;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.flink.TableLoader;
+import org.apache.iceberg.flink.TestFixtures;
+import org.apache.iceberg.flink.sink.TestBucketPartitionerUtils.TableSchemaType;
+import org.apache.iceberg.flink.source.BoundedTestSource;
+import org.apache.iceberg.flink.util.FlinkCompatibilityUtil;
+import org.apache.iceberg.io.CloseableIterable;
+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.assertj.core.api.Assertions;
+import org.jetbrains.annotations.NotNull;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+public class TestBucketPartitionerFlinkIcebergSink {
+
+ private static final int NUMBER_TASK_MANAGERS = 1;
+ private static final int SLOTS_PER_TASK_MANAGER = 8;
+
+ @RegisterExtension
+ private static final MiniClusterExtension MINI_CLUSTER_RESOURCE =
+ new MiniClusterExtension(
+ new MiniClusterResourceConfiguration.Builder()
+ .setNumberTaskManagers(NUMBER_TASK_MANAGERS)
+ .setNumberSlotsPerTaskManager(SLOTS_PER_TASK_MANAGER)
+ .setConfiguration(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .build());
+
+ @RegisterExtension
+ private static final HadoopCatalogExtension catalogExtension =
+ new HadoopCatalogExtension(DATABASE, TestFixtures.TABLE);
+
+ private static final TypeInformation<Row> ROW_TYPE_INFO =
+ new RowTypeInfo(SimpleDataUtil.FLINK_SCHEMA.getFieldTypes());
+
+ // Parallelism = 8 (parallelism > numBuckets) throughout the test suite
+ private final int parallelism = NUMBER_TASK_MANAGERS * SLOTS_PER_TASK_MANAGER;
+ private final FileFormat format = FileFormat.PARQUET;
+ private final int numBuckets = 4;
+
+ private Table table;
+ private StreamExecutionEnvironment env;
+ private TableLoader tableLoader;
+
+ private void setupEnvironment(TableSchemaType tableSchemaType) {
+ table = getTable(tableSchemaType);
+ env =
+ StreamExecutionEnvironment.getExecutionEnvironment(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .enableCheckpointing(100)
+ .setParallelism(parallelism)
+ .setMaxParallelism(parallelism * 2);
+ tableLoader = catalogExtension.tableLoader();
+ }
+
+ private Table getTable(TableSchemaType tableSchemaType) {
+ PartitionSpec partitionSpec = TableSchemaType.getPartitionSpec(tableSchemaType, numBuckets);
+
+ return catalogExtension
+ .catalog()
+ .createTable(
+ TABLE_IDENTIFIER,
+ SimpleDataUtil.SCHEMA,
+ partitionSpec,
+ ImmutableMap.of(TableProperties.DEFAULT_FILE_FORMAT, format.name()));
+ }
+
+ private void testWriteRowData(List<RowData> allRows) throws Exception {
Review Comment:
Done. I had borrowed this directly from `TestFlinkIcebergSink`, also that's where the "obvious" comments below came from.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] kengtin commented on a diff in pull request #7161: Flink: Custom partitioner for bucket partitions
Posted by "kengtin (via GitHub)" <gi...@apache.org>.
kengtin commented on code in PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#discussion_r1209406397
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitioner.java:
##########
@@ -0,0 +1,108 @@
+/*
+ * 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 org.apache.flink.api.common.functions.Partitioner;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/**
+ * This partitioner will redirect records to writers deterministically based on the Bucket partition
+ * spec. It'll attempt to optimize the file size written depending on whether numPartitions is
+ * greater, less or equal than the maxNumBuckets. Note: The current implementation only supports ONE
+ * bucket in the partition spec.
+ */
+class BucketPartitioner implements Partitioner<Integer> {
+
+ static final String BUCKET_NULL_MESSAGE = "bucketId cannot be null";
+ static final String BUCKET_LESS_THAN_LOWER_BOUND_MESSAGE =
+ "bucketId out of range: %s, must be non-negative.";
Review Comment:
Overall I agree with this one, but I have my reservations on the double `:` for each message, in these cases.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] stevenzwu commented on a diff in pull request #7161: Flink: Custom partitioner for a more even data distribution when bucket partitions are present
Posted by "stevenzwu (via GitHub)" <gi...@apache.org>.
stevenzwu commented on code in PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#discussion_r1169091096
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitionKeySelector.java:
##########
@@ -0,0 +1,80 @@
+/*
+ * 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 java.util.List;
+import java.util.stream.IntStream;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.types.logical.RowType;
+import org.apache.iceberg.PartitionKey;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Schema;
+import org.apache.iceberg.flink.RowDataWrapper;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/**
+ * Very similar to the {@link PartitionKeySelector}, but optimized to extract and return an Integer
Review Comment:
nit: remove `Very`. also is `optimized` a good fit here?
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitioner.java:
##########
@@ -0,0 +1,73 @@
+/*
+ * 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 java.util.List;
+import org.apache.flink.api.common.functions.Partitioner;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/**
+ * This partitioner will redirect elements to writers deterministically so that each writer only
+ * targets 1 bucket. If the number of writers is greater than the number of buckets each partitioner
+ * will keep a state of multiple writers per bucket as evenly as possible, and will round-robin the
+ * requests across them.
+ */
+class BucketPartitioner implements Partitioner<Integer> {
+
+ private final int maxBuckets;
+
+ private final int[] currentWriterOffset;
+
+ BucketPartitioner(PartitionSpec partitionSpec) {
+ List<Tuple2<Integer, Integer>> bucketFields =
+ BucketPartitionerUtils.getBucketFields(partitionSpec);
+
+ // The current implementation only supports ONE bucket
+ Preconditions.checkArgument(
+ bucketFields.size() == 1,
+ BucketPartitionerUtils.BAD_NUMBER_OF_BUCKETS_ERROR_MESSAGE + bucketFields.size());
+
+ this.maxBuckets = bucketFields.get(0).f1;
+ this.currentWriterOffset = new int[this.maxBuckets];
+ }
+
+ @Override
+ public int partition(Integer bucketId, int numPartitions) {
Review Comment:
nit: add a Preconditions check on `bucketId < maxBuckets`
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerFlinkIcebergSink.java:
##########
@@ -0,0 +1,420 @@
+/*
+ * 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.apache.iceberg.flink.MiniClusterResource.DISABLE_CLASSLOADER_CHECK_CONFIG;
+import static org.apache.iceberg.flink.TestFixtures.DATABASE;
+import static org.apache.iceberg.flink.TestFixtures.TABLE_IDENTIFIER;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.TreeMap;
+import java.util.UUID;
+import java.util.stream.Collectors;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.typeutils.RowTypeInfo;
+import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.data.util.DataFormatConverters;
+import org.apache.flink.test.util.MiniClusterWithClientResource;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.DistributionMode;
+import org.apache.iceberg.FileFormat;
+import org.apache.iceberg.FileScanTask;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Table;
+import org.apache.iceberg.TableProperties;
+import org.apache.iceberg.flink.HadoopCatalogResource;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.flink.TableLoader;
+import org.apache.iceberg.flink.TestFixtures;
+import org.apache.iceberg.flink.source.BoundedTestSource;
+import org.apache.iceberg.flink.util.FlinkCompatibilityUtil;
+import org.apache.iceberg.io.CloseableIterable;
+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.util.BucketUtil;
+import org.jetbrains.annotations.NotNull;
+import org.junit.Assert;
+import org.junit.Assume;
+import org.junit.Before;
+import org.junit.ClassRule;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+
+@RunWith(Parameterized.class)
+public class TestBucketPartitionerFlinkIcebergSink {
+
+ private static final int NUMBER_TASK_MANAGERS = 1;
+ private static final int SLOTS_PER_TASK_MANAGER = 8;
+
+ @ClassRule
+ public static final MiniClusterWithClientResource MINI_CLUSTER_RESOURCE =
+ new MiniClusterWithClientResource(
+ new MiniClusterResourceConfiguration.Builder()
+ .setNumberTaskManagers(NUMBER_TASK_MANAGERS)
+ .setNumberSlotsPerTaskManager(SLOTS_PER_TASK_MANAGER)
+ .setConfiguration(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .build());
+
+ @ClassRule public static final TemporaryFolder TEMPORARY_FOLDER = new TemporaryFolder();
+
+ @Rule
+ public final HadoopCatalogResource catalogResource =
+ new HadoopCatalogResource(TEMPORARY_FOLDER, DATABASE, TestFixtures.TABLE);
+
+ private static final TypeInformation<Row> ROW_TYPE_INFO =
+ new RowTypeInfo(SimpleDataUtil.FLINK_SCHEMA.getFieldTypes());
+ private static final DataFormatConverters.RowConverter CONVERTER =
+ new DataFormatConverters.RowConverter(SimpleDataUtil.FLINK_SCHEMA.getFieldDataTypes());
+
+ private Table table;
+ private StreamExecutionEnvironment env;
+ private TableLoader tableLoader;
+
+ // Simple parallelism = 8 throughout the test suite
+ private final int parallelism = NUMBER_TASK_MANAGERS * SLOTS_PER_TASK_MANAGER;
+ private final FileFormat format = FileFormat.fromString("parquet");
+ private final int numBuckets = 4;
+ private final TableSchemaType tableSchemaType;
+
+ @Parameterized.Parameters(name = "TableSchemaType = {0}")
+ public static Object[][] parameters() {
+ return new Object[][] {
+ new Object[] {TableSchemaType.ONE_BUCKET},
+ new Object[] {TableSchemaType.IDENTITY_AND_BUCKET},
+ new Object[] {TableSchemaType.TWO_BUCKETS},
+ };
+ }
+
+ public TestBucketPartitionerFlinkIcebergSink(TableSchemaType tableSchemaType) {
+ this.tableSchemaType = tableSchemaType;
+ }
+
+ @Before
+ public void before() throws IOException {
+ table = getTable();
+
Review Comment:
nit: empty line not needed
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerFlinkIcebergSink.java:
##########
@@ -0,0 +1,420 @@
+/*
+ * 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.apache.iceberg.flink.MiniClusterResource.DISABLE_CLASSLOADER_CHECK_CONFIG;
+import static org.apache.iceberg.flink.TestFixtures.DATABASE;
+import static org.apache.iceberg.flink.TestFixtures.TABLE_IDENTIFIER;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.TreeMap;
+import java.util.UUID;
+import java.util.stream.Collectors;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.typeutils.RowTypeInfo;
+import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.data.util.DataFormatConverters;
+import org.apache.flink.test.util.MiniClusterWithClientResource;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.DistributionMode;
+import org.apache.iceberg.FileFormat;
+import org.apache.iceberg.FileScanTask;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Table;
+import org.apache.iceberg.TableProperties;
+import org.apache.iceberg.flink.HadoopCatalogResource;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.flink.TableLoader;
+import org.apache.iceberg.flink.TestFixtures;
+import org.apache.iceberg.flink.source.BoundedTestSource;
+import org.apache.iceberg.flink.util.FlinkCompatibilityUtil;
+import org.apache.iceberg.io.CloseableIterable;
+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.util.BucketUtil;
+import org.jetbrains.annotations.NotNull;
+import org.junit.Assert;
+import org.junit.Assume;
+import org.junit.Before;
+import org.junit.ClassRule;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+
+@RunWith(Parameterized.class)
+public class TestBucketPartitionerFlinkIcebergSink {
+
+ private static final int NUMBER_TASK_MANAGERS = 1;
+ private static final int SLOTS_PER_TASK_MANAGER = 8;
+
+ @ClassRule
+ public static final MiniClusterWithClientResource MINI_CLUSTER_RESOURCE =
+ new MiniClusterWithClientResource(
+ new MiniClusterResourceConfiguration.Builder()
+ .setNumberTaskManagers(NUMBER_TASK_MANAGERS)
+ .setNumberSlotsPerTaskManager(SLOTS_PER_TASK_MANAGER)
+ .setConfiguration(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .build());
+
+ @ClassRule public static final TemporaryFolder TEMPORARY_FOLDER = new TemporaryFolder();
+
+ @Rule
+ public final HadoopCatalogResource catalogResource =
+ new HadoopCatalogResource(TEMPORARY_FOLDER, DATABASE, TestFixtures.TABLE);
+
+ private static final TypeInformation<Row> ROW_TYPE_INFO =
+ new RowTypeInfo(SimpleDataUtil.FLINK_SCHEMA.getFieldTypes());
+ private static final DataFormatConverters.RowConverter CONVERTER =
+ new DataFormatConverters.RowConverter(SimpleDataUtil.FLINK_SCHEMA.getFieldDataTypes());
+
+ private Table table;
+ private StreamExecutionEnvironment env;
+ private TableLoader tableLoader;
+
+ // Simple parallelism = 8 throughout the test suite
+ private final int parallelism = NUMBER_TASK_MANAGERS * SLOTS_PER_TASK_MANAGER;
+ private final FileFormat format = FileFormat.fromString("parquet");
Review Comment:
directly use the enum (not via `fromString`)
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitioner.java:
##########
@@ -0,0 +1,73 @@
+/*
+ * 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 java.util.List;
+import org.apache.flink.api.common.functions.Partitioner;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/**
+ * This partitioner will redirect elements to writers deterministically so that each writer only
+ * targets 1 bucket. If the number of writers is greater than the number of buckets each partitioner
+ * will keep a state of multiple writers per bucket as evenly as possible, and will round-robin the
+ * requests across them.
+ */
+class BucketPartitioner implements Partitioner<Integer> {
+
+ private final int maxBuckets;
+
+ private final int[] currentWriterOffset;
+
+ BucketPartitioner(PartitionSpec partitionSpec) {
+ List<Tuple2<Integer, Integer>> bucketFields =
+ BucketPartitionerUtils.getBucketFields(partitionSpec);
+
+ // The current implementation only supports ONE bucket
+ Preconditions.checkArgument(
Review Comment:
similar to the other comment, maybe this check can be put inside `BucketPartitionerUtils.getBucketFields`, which can also return a single bucket field as `getBucketField`
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitionKeySelector.java:
##########
@@ -0,0 +1,80 @@
+/*
+ * 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 java.util.List;
+import java.util.stream.IntStream;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.types.logical.RowType;
+import org.apache.iceberg.PartitionKey;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Schema;
+import org.apache.iceberg.flink.RowDataWrapper;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/**
+ * Very similar to the {@link PartitionKeySelector}, but optimized to extract and return an Integer
+ * bucketId as the key.
+ */
+class BucketPartitionKeySelector implements KeySelector<RowData, Integer> {
+
+ private final Schema schema;
+ private final PartitionKey partitionKey;
+ private final RowType flinkSchema;
+ private final int bucketFieldPosition;
+
+ private transient RowDataWrapper rowDataWrapper;
+
+ BucketPartitionKeySelector(PartitionSpec spec, Schema schema, RowType flinkSchema) {
+ List<Tuple2<Integer, Integer>> bucketFields = BucketPartitionerUtils.getBucketFields(spec);
+
+ // The current implementation only supports ONE bucket
+ Preconditions.checkArgument(
+ bucketFields.size() == 1,
+ BucketPartitionerUtils.BAD_NUMBER_OF_BUCKETS_ERROR_MESSAGE + bucketFields.size());
Review Comment:
it is weird that the assertion is repeated in two places and hence share the error msg. should the check be done inside the `getBucketFields`?
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitionerUtils.java:
##########
@@ -0,0 +1,97 @@
+/*
+ * 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 java.util.List;
+import java.util.Objects;
+import java.util.stream.Collectors;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.transforms.PartitionSpecVisitor;
+
+final class BucketPartitionerUtils {
+ static final String BAD_NUMBER_OF_BUCKETS_ERROR_MESSAGE =
+ "Expected 1 Bucket transform in the provided PartitionSpec, found: ";
+
+ private BucketPartitionerUtils() {}
+
+ public static boolean hasOneBucketField(PartitionSpec partitionSpec) {
+ List<Tuple2<Integer, Integer>> bucketFields = getBucketFields(partitionSpec);
+ return bucketFields != null && bucketFields.size() == 1;
+ }
+
+ public static List<Tuple2<Integer, Integer>> getBucketFields(PartitionSpec spec) {
+ return PartitionSpecVisitor.visit(
+ spec,
+ new PartitionSpecVisitor<Tuple2<Integer, Integer>>() {
Review Comment:
I found it a little harder to read when the anonymous class is large. probably better to defined it as a nested class
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitioner.java:
##########
@@ -0,0 +1,73 @@
+/*
+ * 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 java.util.List;
+import org.apache.flink.api.common.functions.Partitioner;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/**
+ * This partitioner will redirect elements to writers deterministically so that each writer only
+ * targets 1 bucket. If the number of writers is greater than the number of buckets each partitioner
+ * will keep a state of multiple writers per bucket as evenly as possible, and will round-robin the
+ * requests across them.
+ */
+class BucketPartitioner implements Partitioner<Integer> {
+
+ private final int maxBuckets;
+
+ private final int[] currentWriterOffset;
+
+ BucketPartitioner(PartitionSpec partitionSpec) {
+ List<Tuple2<Integer, Integer>> bucketFields =
+ BucketPartitionerUtils.getBucketFields(partitionSpec);
+
+ // The current implementation only supports ONE bucket
+ Preconditions.checkArgument(
+ bucketFields.size() == 1,
+ BucketPartitionerUtils.BAD_NUMBER_OF_BUCKETS_ERROR_MESSAGE + bucketFields.size());
+
+ this.maxBuckets = bucketFields.get(0).f1;
+ this.currentWriterOffset = new int[this.maxBuckets];
+ }
+
+ @Override
+ public int partition(Integer bucketId, int numPartitions) {
+ if (numPartitions > maxBuckets) {
+ return getPartitionIndex(bucketId, numPartitions);
+ } else {
+ return bucketId % numPartitions;
+ }
+ }
+
+ private int getPartitionIndex(int bucketId, int numPartitions) {
Review Comment:
method name should reflect it is for the case where parallelism (e.g. 4) is less than the number of buckets (e.g. 7). or just merge it into the method above. it is not very long anyway.
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerFlinkIcebergSink.java:
##########
@@ -0,0 +1,420 @@
+/*
+ * 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.apache.iceberg.flink.MiniClusterResource.DISABLE_CLASSLOADER_CHECK_CONFIG;
+import static org.apache.iceberg.flink.TestFixtures.DATABASE;
+import static org.apache.iceberg.flink.TestFixtures.TABLE_IDENTIFIER;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.TreeMap;
+import java.util.UUID;
+import java.util.stream.Collectors;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.typeutils.RowTypeInfo;
+import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.data.util.DataFormatConverters;
+import org.apache.flink.test.util.MiniClusterWithClientResource;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.DistributionMode;
+import org.apache.iceberg.FileFormat;
+import org.apache.iceberg.FileScanTask;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Table;
+import org.apache.iceberg.TableProperties;
+import org.apache.iceberg.flink.HadoopCatalogResource;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.flink.TableLoader;
+import org.apache.iceberg.flink.TestFixtures;
+import org.apache.iceberg.flink.source.BoundedTestSource;
+import org.apache.iceberg.flink.util.FlinkCompatibilityUtil;
+import org.apache.iceberg.io.CloseableIterable;
+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.util.BucketUtil;
+import org.jetbrains.annotations.NotNull;
+import org.junit.Assert;
+import org.junit.Assume;
+import org.junit.Before;
+import org.junit.ClassRule;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+
+@RunWith(Parameterized.class)
Review Comment:
please switch to junit 5. all new tests in Iceberg need to be written in junit 5
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitioner.java:
##########
@@ -0,0 +1,73 @@
+/*
+ * 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 java.util.List;
+import org.apache.flink.api.common.functions.Partitioner;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/**
+ * This partitioner will redirect elements to writers deterministically so that each writer only
+ * targets 1 bucket. If the number of writers is greater than the number of buckets each partitioner
+ * will keep a state of multiple writers per bucket as evenly as possible, and will round-robin the
+ * requests across them.
+ */
+class BucketPartitioner implements Partitioner<Integer> {
+
+ private final int maxBuckets;
+
+ private final int[] currentWriterOffset;
+
+ BucketPartitioner(PartitionSpec partitionSpec) {
+ List<Tuple2<Integer, Integer>> bucketFields =
+ BucketPartitionerUtils.getBucketFields(partitionSpec);
+
+ // The current implementation only supports ONE bucket
+ Preconditions.checkArgument(
+ bucketFields.size() == 1,
+ BucketPartitionerUtils.BAD_NUMBER_OF_BUCKETS_ERROR_MESSAGE + bucketFields.size());
+
+ this.maxBuckets = bucketFields.get(0).f1;
+ this.currentWriterOffset = new int[this.maxBuckets];
+ }
+
+ @Override
+ public int partition(Integer bucketId, int numPartitions) {
+ if (numPartitions > maxBuckets) {
+ return getPartitionIndex(bucketId, numPartitions);
+ } else {
+ return bucketId % numPartitions;
+ }
+ }
+
+ private int getPartitionIndex(int bucketId, int numPartitions) {
+ int currentOffset = currentWriterOffset[bucketId];
+ // When numPartitions is not evenly divisible by maxBuckets
Review Comment:
I don't quite understand the padding logic here. can you add some Javadoc to explain the logic?
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerFlinkIcebergSink.java:
##########
@@ -0,0 +1,420 @@
+/*
+ * 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.apache.iceberg.flink.MiniClusterResource.DISABLE_CLASSLOADER_CHECK_CONFIG;
+import static org.apache.iceberg.flink.TestFixtures.DATABASE;
+import static org.apache.iceberg.flink.TestFixtures.TABLE_IDENTIFIER;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.TreeMap;
+import java.util.UUID;
+import java.util.stream.Collectors;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.typeutils.RowTypeInfo;
+import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.data.util.DataFormatConverters;
+import org.apache.flink.test.util.MiniClusterWithClientResource;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.DistributionMode;
+import org.apache.iceberg.FileFormat;
+import org.apache.iceberg.FileScanTask;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Table;
+import org.apache.iceberg.TableProperties;
+import org.apache.iceberg.flink.HadoopCatalogResource;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.flink.TableLoader;
+import org.apache.iceberg.flink.TestFixtures;
+import org.apache.iceberg.flink.source.BoundedTestSource;
+import org.apache.iceberg.flink.util.FlinkCompatibilityUtil;
+import org.apache.iceberg.io.CloseableIterable;
+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.util.BucketUtil;
+import org.jetbrains.annotations.NotNull;
+import org.junit.Assert;
+import org.junit.Assume;
+import org.junit.Before;
+import org.junit.ClassRule;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+
+@RunWith(Parameterized.class)
+public class TestBucketPartitionerFlinkIcebergSink {
+
+ private static final int NUMBER_TASK_MANAGERS = 1;
+ private static final int SLOTS_PER_TASK_MANAGER = 8;
+
+ @ClassRule
+ public static final MiniClusterWithClientResource MINI_CLUSTER_RESOURCE =
+ new MiniClusterWithClientResource(
+ new MiniClusterResourceConfiguration.Builder()
+ .setNumberTaskManagers(NUMBER_TASK_MANAGERS)
+ .setNumberSlotsPerTaskManager(SLOTS_PER_TASK_MANAGER)
+ .setConfiguration(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .build());
+
+ @ClassRule public static final TemporaryFolder TEMPORARY_FOLDER = new TemporaryFolder();
+
+ @Rule
+ public final HadoopCatalogResource catalogResource =
+ new HadoopCatalogResource(TEMPORARY_FOLDER, DATABASE, TestFixtures.TABLE);
+
+ private static final TypeInformation<Row> ROW_TYPE_INFO =
+ new RowTypeInfo(SimpleDataUtil.FLINK_SCHEMA.getFieldTypes());
+ private static final DataFormatConverters.RowConverter CONVERTER =
+ new DataFormatConverters.RowConverter(SimpleDataUtil.FLINK_SCHEMA.getFieldDataTypes());
+
+ private Table table;
+ private StreamExecutionEnvironment env;
+ private TableLoader tableLoader;
+
+ // Simple parallelism = 8 throughout the test suite
+ private final int parallelism = NUMBER_TASK_MANAGERS * SLOTS_PER_TASK_MANAGER;
Review Comment:
nit: typically define final variables before non-final ones
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitionerUtils.java:
##########
@@ -0,0 +1,97 @@
+/*
+ * 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 java.util.List;
+import java.util.Objects;
+import java.util.stream.Collectors;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.transforms.PartitionSpecVisitor;
+
+final class BucketPartitionerUtils {
+ static final String BAD_NUMBER_OF_BUCKETS_ERROR_MESSAGE =
+ "Expected 1 Bucket transform in the provided PartitionSpec, found: ";
+
+ private BucketPartitionerUtils() {}
+
+ public static boolean hasOneBucketField(PartitionSpec partitionSpec) {
+ List<Tuple2<Integer, Integer>> bucketFields = getBucketFields(partitionSpec);
+ return bucketFields != null && bucketFields.size() == 1;
+ }
+
+ public static List<Tuple2<Integer, Integer>> getBucketFields(PartitionSpec spec) {
Review Comment:
as mentioned in other places, this method can be private. create another package default method `getBucketFieldInfo` that only return one bucket field. move the assertion inside this new method.
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerFlinkIcebergSink.java:
##########
@@ -0,0 +1,420 @@
+/*
+ * 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.apache.iceberg.flink.MiniClusterResource.DISABLE_CLASSLOADER_CHECK_CONFIG;
+import static org.apache.iceberg.flink.TestFixtures.DATABASE;
+import static org.apache.iceberg.flink.TestFixtures.TABLE_IDENTIFIER;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.TreeMap;
+import java.util.UUID;
+import java.util.stream.Collectors;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.typeutils.RowTypeInfo;
+import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.data.util.DataFormatConverters;
+import org.apache.flink.test.util.MiniClusterWithClientResource;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.DistributionMode;
+import org.apache.iceberg.FileFormat;
+import org.apache.iceberg.FileScanTask;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Table;
+import org.apache.iceberg.TableProperties;
+import org.apache.iceberg.flink.HadoopCatalogResource;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.flink.TableLoader;
+import org.apache.iceberg.flink.TestFixtures;
+import org.apache.iceberg.flink.source.BoundedTestSource;
+import org.apache.iceberg.flink.util.FlinkCompatibilityUtil;
+import org.apache.iceberg.io.CloseableIterable;
+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.util.BucketUtil;
+import org.jetbrains.annotations.NotNull;
+import org.junit.Assert;
+import org.junit.Assume;
+import org.junit.Before;
+import org.junit.ClassRule;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+
+@RunWith(Parameterized.class)
+public class TestBucketPartitionerFlinkIcebergSink {
+
+ private static final int NUMBER_TASK_MANAGERS = 1;
+ private static final int SLOTS_PER_TASK_MANAGER = 8;
+
+ @ClassRule
+ public static final MiniClusterWithClientResource MINI_CLUSTER_RESOURCE =
+ new MiniClusterWithClientResource(
+ new MiniClusterResourceConfiguration.Builder()
+ .setNumberTaskManagers(NUMBER_TASK_MANAGERS)
+ .setNumberSlotsPerTaskManager(SLOTS_PER_TASK_MANAGER)
+ .setConfiguration(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .build());
+
+ @ClassRule public static final TemporaryFolder TEMPORARY_FOLDER = new TemporaryFolder();
+
+ @Rule
+ public final HadoopCatalogResource catalogResource =
+ new HadoopCatalogResource(TEMPORARY_FOLDER, DATABASE, TestFixtures.TABLE);
+
+ private static final TypeInformation<Row> ROW_TYPE_INFO =
+ new RowTypeInfo(SimpleDataUtil.FLINK_SCHEMA.getFieldTypes());
+ private static final DataFormatConverters.RowConverter CONVERTER =
+ new DataFormatConverters.RowConverter(SimpleDataUtil.FLINK_SCHEMA.getFieldDataTypes());
+
+ private Table table;
+ private StreamExecutionEnvironment env;
+ private TableLoader tableLoader;
+
+ // Simple parallelism = 8 throughout the test suite
+ private final int parallelism = NUMBER_TASK_MANAGERS * SLOTS_PER_TASK_MANAGER;
+ private final FileFormat format = FileFormat.fromString("parquet");
+ private final int numBuckets = 4;
+ private final TableSchemaType tableSchemaType;
+
+ @Parameterized.Parameters(name = "TableSchemaType = {0}")
+ public static Object[][] parameters() {
+ return new Object[][] {
+ new Object[] {TableSchemaType.ONE_BUCKET},
+ new Object[] {TableSchemaType.IDENTITY_AND_BUCKET},
+ new Object[] {TableSchemaType.TWO_BUCKETS},
+ };
+ }
+
+ public TestBucketPartitionerFlinkIcebergSink(TableSchemaType tableSchemaType) {
+ this.tableSchemaType = tableSchemaType;
+ }
+
+ @Before
+ public void before() throws IOException {
+ table = getTable();
+
+ env =
+ StreamExecutionEnvironment.getExecutionEnvironment(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .enableCheckpointing(100)
+ .setParallelism(parallelism)
+ .setMaxParallelism(parallelism * 2);
+
+ tableLoader = catalogResource.tableLoader();
+ }
+
+ private Table getTable() {
+ PartitionSpec partitionSpec = null;
+
+ switch (tableSchemaType) {
+ case ONE_BUCKET:
+ partitionSpec =
+ PartitionSpec.builderFor(SimpleDataUtil.SCHEMA).bucket("data", numBuckets).build();
+ break;
+ case IDENTITY_AND_BUCKET:
+ partitionSpec =
+ PartitionSpec.builderFor(SimpleDataUtil.SCHEMA)
+ .identity("id")
+ .bucket("data", numBuckets)
+ .build();
+ break;
+ case TWO_BUCKETS:
+ partitionSpec =
+ PartitionSpec.builderFor(SimpleDataUtil.SCHEMA)
+ .bucket("id", numBuckets)
+ .bucket("data", numBuckets)
+ .build();
+ break;
+ }
+
+ return catalogResource
+ .catalog()
+ .createTable(
+ TABLE_IDENTIFIER,
+ SimpleDataUtil.SCHEMA,
+ partitionSpec,
+ ImmutableMap.of(TableProperties.DEFAULT_FILE_FORMAT, format.name()));
+ }
+
+ private List<RowData> convertToRowData(List<Row> rows) {
+ return rows.stream().map(CONVERTER::toInternal).collect(Collectors.toList());
+ }
+
+ private BoundedTestSource<Row> createBoundedSource(List<Row> rows) {
+ return new BoundedTestSource<>(rows.toArray(new Row[0]));
+ }
+
+ private TableTestStats extractTableTestStats() throws IOException {
+ // Assertions:
+ // - Total record count in table
+ // - All records belong to that bucket
+ // - Each bucket should've been written by valid writers
+ // - Each file should have the expected number of records
+
+ int totalRecordCount = 0;
+ Map<Integer, List<Integer>> writersPerBucket = Maps.newHashMap(); // <BucketId, Set<WriterId>>
+ Map<Integer, Integer> filesPerBucket = Maps.newHashMap(); // <BucketId, NumFiles>
+ Map<Integer, Long> recordsPerFile = new TreeMap<>(); // <WriterId, NumRecords>
+
+ try (CloseableIterable<FileScanTask> fileScanTasks = table.newScan().planFiles()) {
+ for (FileScanTask scanTask : fileScanTasks) {
+ long recordCountInFile = scanTask.file().recordCount();
+
+ String[] splitFilePath = scanTask.file().path().toString().split("/");
+ String filename = splitFilePath[splitFilePath.length - 1];
+ int writerId = Integer.parseInt(filename.split("-")[0]);
+
+ totalRecordCount += recordCountInFile;
+ int bucketId =
+ scanTask
+ .file()
+ .partition()
+ .get(tableSchemaType == TableSchemaType.ONE_BUCKET ? 0 : 1, Integer.class);
+ writersPerBucket.computeIfAbsent(bucketId, k -> Lists.newArrayList());
+ writersPerBucket.get(bucketId).add(writerId);
+ filesPerBucket.put(bucketId, filesPerBucket.getOrDefault(bucketId, 0) + 1);
+ recordsPerFile.put(writerId, recordsPerFile.getOrDefault(writerId, 0L) + recordCountInFile);
+ }
+ }
+
+ for (int k : writersPerBucket.keySet()) {
+ Collections.sort(writersPerBucket.get(k));
+ }
+
+ return new TableTestStats(totalRecordCount, writersPerBucket, filesPerBucket, recordsPerFile);
+ }
+
+ private void testWriteRowData(List<Row> allRows) throws Exception {
+ DataStream<RowData> dataStream =
+ env.addSource(createBoundedSource(allRows), ROW_TYPE_INFO)
+ .map(CONVERTER::toInternal, FlinkCompatibilityUtil.toTypeInfo(SimpleDataUtil.ROW_TYPE));
+
+ FlinkSink.forRowData(dataStream)
+ .table(table)
+ .tableLoader(tableLoader)
+ .writeParallelism(parallelism)
+ .distributionMode(DistributionMode.HASH)
+ .append();
+
+ // Execute the program.
+ env.execute("Test Iceberg DataStream");
+
+ // Assert the iceberg table's records.
+ SimpleDataUtil.assertTableRows(table, convertToRowData(allRows));
+ }
+
+ // Test 1A: N records sent to the same bucket across all partitions
+ // - 24 rows (3 batches w/overlap), 4 buckets
+ @Test
+ public void testSendToBucket0A() throws Exception {
+ Assume.assumeFalse(tableSchemaType == TableSchemaType.TWO_BUCKETS);
+
+ // Data generation
+ // We make the number of requests a factor of the parallelism for an easier verification
+ int totalNumRows = parallelism * 3; // 24 rows
+ TableTestStats stats = sendAndAssertBucket0(totalNumRows);
+ // Number of records per file
+ Assert.assertEquals(16, stats.recordsPerFile.get(0).intValue());
+ Assert.assertEquals(8, stats.recordsPerFile.get(4).intValue());
+ }
+
+ // Test 1B: N records sent to the same bucket across all partitions
+ // - 30 rows (4 batches w/o overlap, uneven first bucket), 4 buckets
+ @Test
+ public void testSendToBucket0B() throws Exception {
+ Assume.assumeTrue(tableSchemaType != TableSchemaType.TWO_BUCKETS);
+
+ // Data generation
+ // We make the number of requests a factor of the parallelism for an easier verification
+ int totalNumRows = 30;
+ TableTestStats stats = sendAndAssertBucket0(totalNumRows);
+ // Number of records per file
+ Assert.assertEquals(16, stats.recordsPerFile.get(0).intValue());
+ Assert.assertEquals(14, stats.recordsPerFile.get(4).intValue());
+ }
+
+ @NotNull
+ private TableTestStats sendAndAssertBucket0(int totalNumRows) throws Exception {
+ List<Row> rows = generateRowsForBucketId(totalNumRows, 0);
+ // Writing to the table
+ testWriteRowData(rows);
+ // Extracting stats
+ TableTestStats stats = extractTableTestStats();
+
+ // Assertions
+ // All the records were received
+ Assert.assertEquals(totalNumRows, stats.totalRecordCount);
+ // Only bucketId = 0 was written
+ Assert.assertEquals(1, stats.writersPerBucket.size());
+ Assert.assertEquals(1, stats.filesPerBucket.size());
+ Assert.assertTrue(stats.writersPerBucket.containsKey(0));
+ Assert.assertTrue(stats.filesPerBucket.containsKey(0));
+ // Valid writers
+ List<Integer> validWriters = Arrays.asList(0, 4);
+ Assert.assertEquals(validWriters, stats.writersPerBucket.get(0));
+ // Num writers = num created files
+ Assert.assertEquals(validWriters.size(), stats.filesPerBucket.get(0).intValue());
+ // TODO: verify all the records actually hash to bucketId = 0
+ return stats;
+ }
+
+ // Test 2: N records of each bucket value, to all buckets, through the same partitions
+ // - 16 rows (2 writers per bucket, uneven last bucket), 4 buckets
+ @Test
+ public void testSendRecordsToAllBucketsEvenly() throws Exception {
+ Assume.assumeTrue(tableSchemaType != TableSchemaType.TWO_BUCKETS);
+
+ // Data generation
+ // We make the number of requests a factor of the parallelism for an easier verification
+ int totalNumRows = parallelism * 2;
+ int numRowsPerBucket = totalNumRows / numBuckets;
+ List<Row> rows = Lists.newArrayListWithCapacity(totalNumRows);
+ List<Row> templateRows = Lists.newArrayListWithCapacity(numBuckets);
+ for (int i = 0; i < numBuckets; i++) {
+ templateRows.addAll(generateRowsForBucketId(1, i));
+ }
+ for (int i = 0; i < numRowsPerBucket; i++) {
+ rows.addAll(templateRows);
+ }
+
+ // Writing to the table
+ testWriteRowData(rows);
+ // Extracting stats
+ TableTestStats stats = extractTableTestStats();
+
+ // Assertions
+ // All the records were received
+ Assert.assertEquals(totalNumRows, stats.totalRecordCount);
+ // Only bucketId = 0 was written
+ Assert.assertEquals(numBuckets, stats.writersPerBucket.size());
+ Assert.assertEquals(numBuckets, stats.filesPerBucket.size());
+ for (int i = 0, j = 4; i < numBuckets; i++, j++) {
+ Assert.assertEquals(Arrays.asList(i, j), stats.writersPerBucket.get(i));
+ Assert.assertTrue(stats.filesPerBucket.get(i) == 2);
+ }
+ }
+
+ @Test
+ public void testBucketPartitionerFail() {
+ Assume.assumeTrue(tableSchemaType == TableSchemaType.TWO_BUCKETS);
+
+ Exception exception =
+ Assert.assertThrows(RuntimeException.class, () -> new BucketPartitioner(table.spec()));
+
+ Assert.assertTrue(
+ exception
+ .getMessage()
+ .contains(BucketPartitionerUtils.BAD_NUMBER_OF_BUCKETS_ERROR_MESSAGE));
+ }
+
+ @Test
+ public void testBucketPartitionKeySelectorFail() {
+ Assume.assumeTrue(tableSchemaType == TableSchemaType.TWO_BUCKETS);
+
+ Exception exception =
+ Assert.assertThrows(
+ RuntimeException.class,
+ () -> new BucketPartitionKeySelector(table.spec(), table.schema(), null));
+
+ Assert.assertTrue(
+ exception
+ .getMessage()
+ .contains(BucketPartitionerUtils.BAD_NUMBER_OF_BUCKETS_ERROR_MESSAGE));
+ }
+
+ /**
+ * Utility method to generate rows whose values will "hash" to a desired bucketId
+ *
+ * @param numRows the number of rows/requests to generate
+ * @param bucketId the desired bucketId
+ * @return the list of rows whose data "hashes" to the desired bucketId
+ */
+ private List<Row> generateRowsForBucketId(int numRows, int bucketId) {
+ List<Row> rows = Lists.newArrayListWithCapacity(numRows);
+ for (int i = 0; i < numRows; i++) {
+ String value = generateValueForBucketId(bucketId, numBuckets);
+ rows.add(Row.of(1, value));
+ }
+ return rows;
+ }
+
+ /**
+ * Utility method to generate a UUID string that will "hash" to a desired bucketId
+ *
+ * @param bucketId the desired bucketId
+ * @return the string data that "hashes" to the desired bucketId
+ */
+ public static String generateValueForBucketId(int bucketId, int numBuckets) {
+ String value = "";
+ while (true) {
+ String uuid = UUID.randomUUID().toString();
+ if (hash(numBuckets, uuid) == bucketId) {
+ value = uuid;
+ break;
+ }
+ }
+ return value;
+ }
+
+ /**
+ * Utility method that performs the same "hashing" mechanism used by Bucket.java
+ *
+ * @param numBuckets the maximum number of buckets
+ * @param uuid the UUID string value to hash
+ * @return the target BucketId
+ */
+ public static int hash(int numBuckets, String uuid) {
+ return (BucketUtil.hash(uuid) & Integer.MAX_VALUE) % numBuckets;
+ }
+
+ /** DTO to hold Test Stats */
+ private static class TableTestStats {
+ final int totalRecordCount;
+ final Map<Integer, List<Integer>> writersPerBucket; // <BucketId, Set<WriterId>>
+ final Map<Integer, Integer> filesPerBucket; // <BucketId, NumFiles>
+ final Map<Integer, Long> recordsPerFile; // <WriterId, NumRecords>
+
+ TableTestStats(
+ int totalRecordCount,
+ Map<Integer, List<Integer>> writersPerBucket,
+ Map<Integer, Integer> filesPerBucket,
+ Map<Integer, Long> recordsPerFile) {
+ this.totalRecordCount = totalRecordCount;
+ this.writersPerBucket = writersPerBucket;
+ this.filesPerBucket = filesPerBucket;
+ this.recordsPerFile = recordsPerFile;
+ }
+ }
+
+ private enum TableSchemaType {
Review Comment:
nit: why not define it near where it is used near the top of the class?
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitioner.java:
##########
@@ -0,0 +1,86 @@
+/*
+ * 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 java.util.Optional;
+import org.apache.flink.api.common.functions.Partitioner;
+import org.apache.iceberg.PartitionField;
+import org.apache.iceberg.PartitionSpec;
+
+/**
+ * This partitioner will redirect elements to writers deterministically so that each writer only
Review Comment:
one writer can handle multiple buckets if the number of buckets is larger than the writer parallelism.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] kengtin commented on a diff in pull request #7161: Flink: Custom partitioner for bucket partitions
Posted by "kengtin (via GitHub)" <gi...@apache.org>.
kengtin commented on code in PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#discussion_r1199626634
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitioner.java:
##########
@@ -0,0 +1,99 @@
+/*
+ * 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.apache.iceberg.flink.sink.BucketPartitioner.BUCKET_OUT_OF_RANGE_MESSAGE_PREFIX;
+
+import org.apache.iceberg.PartitionSpec;
+import org.assertj.core.api.Assertions;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+import org.junit.jupiter.params.provider.ValueSource;
+
+public class TestBucketPartitioner {
+
+ static final int NUM_BUCKETS = 60;
+
+ @ParameterizedTest
+ @EnumSource(
+ value = TestBucketPartitionerUtils.TableSchemaType.class,
+ names = {"ONE_BUCKET", "IDENTITY_AND_BUCKET"})
+ public void testPartitioningParallelismGreaterThanBuckets(
+ TestBucketPartitionerUtils.TableSchemaType tableSchemaType) {
+ final int numPartitions = 500;
+
+ PartitionSpec partitionSpec =
+ TestBucketPartitionerUtils.getPartitionSpec(tableSchemaType, NUM_BUCKETS);
+
+ BucketPartitioner bucketPartitioner = new BucketPartitioner(partitionSpec);
+
+ for (int expectedIdx = 0, bucketId = 0; expectedIdx < numPartitions; expectedIdx++) {
+ int actualIdx = bucketPartitioner.partition(bucketId, numPartitions);
+ Assertions.assertThat(actualIdx).isEqualTo(expectedIdx);
+ if (++bucketId == NUM_BUCKETS) {
+ bucketId = 0;
+ }
+ }
+ }
+
+ @ParameterizedTest
+ @EnumSource(
+ value = TestBucketPartitionerUtils.TableSchemaType.class,
+ names = {"ONE_BUCKET", "IDENTITY_AND_BUCKET"})
+ public void testPartitioningParallelismEqualLessThanBuckets(
+ TestBucketPartitionerUtils.TableSchemaType tableSchemaType) {
+ final int numPartitions = 30;
+
+ PartitionSpec partitionSpec =
+ TestBucketPartitionerUtils.getPartitionSpec(tableSchemaType, NUM_BUCKETS);
+
+ BucketPartitioner bucketPartitioner = new BucketPartitioner(partitionSpec);
+
+ for (int bucketId = 0; bucketId < NUM_BUCKETS; bucketId++) {
+ int actualIdx = bucketPartitioner.partition(bucketId, numPartitions);
+ Assertions.assertThat(actualIdx).isEqualTo(bucketId % numPartitions);
+ }
+ }
+
+ @ParameterizedTest
+ @EnumSource(value = TestBucketPartitionerUtils.TableSchemaType.class, names = "TWO_BUCKETS")
+ public void testPartitionerMultipleBucketsFail(
+ TestBucketPartitionerUtils.TableSchemaType tableSchemaType) {
+ PartitionSpec partitionSpec =
+ TestBucketPartitionerUtils.getPartitionSpec(tableSchemaType, NUM_BUCKETS);
+
+ Assertions.assertThatExceptionOfType(RuntimeException.class)
+ .isThrownBy(() -> new BucketPartitioner(partitionSpec))
+ .withMessageContaining(BucketPartitionerUtils.BAD_NUMBER_OF_BUCKETS_ERROR_MESSAGE);
+ }
+
+ @ParameterizedTest
+ @ValueSource(ints = {-1, NUM_BUCKETS})
+ public void testPartitionerBucketIdOutOfRangeFail(int bucketId) {
+ PartitionSpec partitionSpec =
+ TestBucketPartitionerUtils.getPartitionSpec(
+ TestBucketPartitionerUtils.TableSchemaType.ONE_BUCKET, NUM_BUCKETS);
+
+ BucketPartitioner bucketPartitioner = new BucketPartitioner(partitionSpec);
+
+ Assertions.assertThatExceptionOfType(IllegalArgumentException.class)
+ .isThrownBy(() -> bucketPartitioner.partition(bucketId, 1))
+ .withMessageContaining(BUCKET_OUT_OF_RANGE_MESSAGE_PREFIX);
Review Comment:
I actually changed all assertions to check for complete messages
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] kengtin commented on a diff in pull request #7161: Flink: Custom partitioner for bucket partitions
Posted by "kengtin (via GitHub)" <gi...@apache.org>.
kengtin commented on code in PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#discussion_r1261868277
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerUtil.java:
##########
@@ -0,0 +1,118 @@
+/*
+ * 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 java.util.List;
+import java.util.UUID;
+import org.apache.flink.table.data.GenericRowData;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.data.StringData;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.apache.iceberg.util.BucketUtil;
+
+final class TestBucketPartitionerUtil {
+
+ enum TableSchemaType {
+ ONE_BUCKET {
+ @Override
+ public int bucketPartitionColumnPosition() {
+ return 0;
+ }
+
+ @Override
+ public PartitionSpec getPartitionSpec(int numBuckets) {
+ return PartitionSpec.builderFor(SimpleDataUtil.SCHEMA).bucket("data", numBuckets).build();
+ }
+ },
+ IDENTITY_AND_BUCKET {
+ @Override
+ public PartitionSpec getPartitionSpec(int numBuckets) {
+ return PartitionSpec.builderFor(SimpleDataUtil.SCHEMA)
+ .identity("id")
+ .bucket("data", numBuckets)
+ .build();
+ }
+ },
+ TWO_BUCKETS {
+ @Override
+ public PartitionSpec getPartitionSpec(int numBuckets) {
+ return PartitionSpec.builderFor(SimpleDataUtil.SCHEMA)
+ .bucket("id", numBuckets)
+ .bucket("data", numBuckets)
+ .build();
+ }
+ };
+
+ public int bucketPartitionColumnPosition() {
Review Comment:
Sure, but if I make this abstract I'll have to explicitly override it in the 2 last enums, which will also add clutter I think.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] kengtin commented on a diff in pull request #7161: Flink: Custom partitioner for bucket partitions
Posted by "kengtin (via GitHub)" <gi...@apache.org>.
kengtin commented on code in PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#discussion_r1199621059
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerUtils.java:
##########
@@ -0,0 +1,121 @@
+/*
+ * 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 java.util.List;
+import java.util.UUID;
+import org.apache.flink.table.data.util.DataFormatConverters;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.apache.iceberg.util.BucketUtil;
+
+final class TestBucketPartitionerUtils {
+
+ enum TableSchemaType {
+ ONE_BUCKET,
+ IDENTITY_AND_BUCKET,
+ TWO_BUCKETS;
+ }
+
+ private TestBucketPartitionerUtils() {}
+
+ static final DataFormatConverters.RowConverter CONVERTER =
+ new DataFormatConverters.RowConverter(SimpleDataUtil.FLINK_SCHEMA.getFieldDataTypes());
+
+ static PartitionSpec getPartitionSpec(TableSchemaType tableSchemaType, int numBuckets) {
+ PartitionSpec partitionSpec = null;
+
+ switch (tableSchemaType) {
+ case ONE_BUCKET:
+ partitionSpec =
+ PartitionSpec.builderFor(SimpleDataUtil.SCHEMA).bucket("data", numBuckets).build();
+ break;
+ case IDENTITY_AND_BUCKET:
+ partitionSpec =
+ PartitionSpec.builderFor(SimpleDataUtil.SCHEMA)
+ .identity("id")
+ .bucket("data", numBuckets)
+ .build();
+ break;
+ case TWO_BUCKETS:
+ partitionSpec =
+ PartitionSpec.builderFor(SimpleDataUtil.SCHEMA)
+ .bucket("id", numBuckets)
+ .bucket("data", numBuckets)
+ .build();
+ break;
+ }
+
+ Preconditions.checkNotNull(
+ partitionSpec, "Invalid tableSchemaType provided: " + tableSchemaType);
+ return partitionSpec;
+ }
+
+ /**
+ * Utility method to generate rows whose values will "hash" to a range of bucketIds (from 0 to
+ * numBuckets - 1)
+ *
+ * @param numRowsPerBucket how many different rows should be generated per bucket
+ * @param numBuckets max number of buckets to consider
+ * @return the list of rows whose data "hashes" to the desired bucketId
+ */
+ static List<Row> generateRowsForBucketIdRange(int numRowsPerBucket, int numBuckets) {
+ List<Row> rows = Lists.newArrayListWithCapacity(numBuckets * numRowsPerBucket);
+ // For some of our tests, this order of the generated rows matters
+ for (int i = 0; i < numRowsPerBucket; i++) {
+ for (int bucketId = 0; bucketId < numBuckets; bucketId++) {
+ String value = generateValueForBucketId(bucketId, numBuckets);
+ rows.add(Row.of(1, value));
+ }
+ }
+ return rows;
+ }
+
+ /**
+ * Utility method to generate a UUID string that will "hash" to a desired bucketId
+ *
+ * @param bucketId the desired bucketId
+ * @return the string data that "hashes" to the desired bucketId
+ */
+ private static String generateValueForBucketId(int bucketId, int numBuckets) {
+ String value = "";
+ while (true) {
+ String uuid = UUID.randomUUID().toString();
+ if (computeBucketId(numBuckets, uuid) == bucketId) {
+ value = uuid;
+ break;
+ }
+ }
+ return value;
+ }
+
+ /**
+ * Utility that performs the same hashing/bucketing mechanism used by Bucket.java
Review Comment:
Done, I added the corresponding Precondition to the `BucketPartitioner` and wrote a test for it
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] kengtin commented on a diff in pull request #7161: Flink: Custom partitioner for bucket partitions
Posted by "kengtin (via GitHub)" <gi...@apache.org>.
kengtin commented on code in PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#discussion_r1209406397
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitioner.java:
##########
@@ -0,0 +1,108 @@
+/*
+ * 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 org.apache.flink.api.common.functions.Partitioner;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/**
+ * This partitioner will redirect records to writers deterministically based on the Bucket partition
+ * spec. It'll attempt to optimize the file size written depending on whether numPartitions is
+ * greater, less or equal than the maxNumBuckets. Note: The current implementation only supports ONE
+ * bucket in the partition spec.
+ */
+class BucketPartitioner implements Partitioner<Integer> {
+
+ static final String BUCKET_NULL_MESSAGE = "bucketId cannot be null";
+ static final String BUCKET_LESS_THAN_LOWER_BOUND_MESSAGE =
+ "bucketId out of range: %s, must be non-negative.";
Review Comment:
Overall I agree with this one, but I have my reservations on the double `:` for each message, in these cases. Also I find it way more readable to have the number of the invalid bucket Id right after its message (e.g. `Invalid bucket id %s:...`)
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitioner.java:
##########
@@ -0,0 +1,108 @@
+/*
+ * 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 org.apache.flink.api.common.functions.Partitioner;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/**
+ * This partitioner will redirect records to writers deterministically based on the Bucket partition
+ * spec. It'll attempt to optimize the file size written depending on whether numPartitions is
+ * greater, less or equal than the maxNumBuckets. Note: The current implementation only supports ONE
+ * bucket in the partition spec.
+ */
+class BucketPartitioner implements Partitioner<Integer> {
+
+ static final String BUCKET_NULL_MESSAGE = "bucketId cannot be null";
+ static final String BUCKET_LESS_THAN_LOWER_BOUND_MESSAGE =
+ "bucketId out of range: %s, must be non-negative.";
Review Comment:
Overall I agree with this one, but I have my reservations on the double `:` for each message, in these cases. Also I find it way more readable to have the number of the invalid bucket Id right after its message (e.g. `Invalid bucket id %s: ...`)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] kengtin commented on a diff in pull request #7161: Flink: Custom partitioner for a more even data distribution when bucket partitions are present
Posted by "kengtin (via GitHub)" <gi...@apache.org>.
kengtin commented on code in PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#discussion_r1159008230
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitioner.java:
##########
@@ -0,0 +1,86 @@
+/*
+ * 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 java.util.Optional;
+import org.apache.flink.api.common.functions.Partitioner;
+import org.apache.iceberg.PartitionField;
+import org.apache.iceberg.PartitionSpec;
+
+/**
+ * This partitioner will redirect elements to writers deterministically so that each writer only
Review Comment:
The idea here is Bucket affinity, in other words each writer should be "assigned" one bucket to write to at all times. Would a wording like this make it clearer?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Flink: Custom partitioner for bucket partitions [iceberg]
Posted by "chenwyi2 (via GitHub)" <gi...@apache.org>.
chenwyi2 commented on PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#issuecomment-1765530967
In normal conditition, only the data of current minute will be written. However, if the data is delayed, for example, at 11:50, the data has not been written until 11:55, then at 11:56 will commit data composed by 11:51,11:52,...11:56, in this situation some small fies can be created
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Flink: Custom partitioner for bucket partitions [iceberg]
Posted by "bendevera (via GitHub)" <gi...@apache.org>.
bendevera commented on PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#issuecomment-1861295520
@stevenzwu I see defaulting to `BucketPartitioner` was reverted here: https://github.com/apache/iceberg/pull/8848
We've found performance issues with `DistributionMode.HASH`, and wanted to test `BucketPartitioner` but it isn't `public`.
You mentioned above that to use the bucketing partitioner, users will need to manually apply `keyBy`. Believe the code snippet though is the default `HASH` logic. What is the recommended method for using the bucket partitioner and any issues reported by users that resulted in reverting?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] kengtin commented on a diff in pull request #7161: Flink: Custom partitioner for bucket partitions
Posted by "kengtin (via GitHub)" <gi...@apache.org>.
kengtin commented on code in PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#discussion_r1210582580
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerUtils.java:
##########
@@ -0,0 +1,120 @@
+/*
+ * 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 java.util.List;
+import java.util.UUID;
+import org.apache.flink.table.data.GenericRowData;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.data.StringData;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.apache.iceberg.util.BucketUtil;
+
+final class TestBucketPartitionerUtils {
+
+ enum TableSchemaType {
+ ONE_BUCKET,
+ IDENTITY_AND_BUCKET,
+ TWO_BUCKETS;
+
+ public static int bucketPartitionColumnPosition(TableSchemaType tableSchemaType) {
Review Comment:
This method should not throw an exception because the stats are still extracted for the `TWO_BUCKETS` case to validate the fallback to `.keyBy` (via the `testMultipleBucketsFallback` test).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] stevenzwu commented on a diff in pull request #7161: Flink: Custom partitioner for bucket partitions
Posted by "stevenzwu (via GitHub)" <gi...@apache.org>.
stevenzwu commented on code in PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#discussion_r1233384313
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerUtil.java:
##########
@@ -0,0 +1,118 @@
+/*
+ * 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 java.util.List;
+import java.util.UUID;
+import org.apache.flink.table.data.GenericRowData;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.data.StringData;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.apache.iceberg.util.BucketUtil;
+
+final class TestBucketPartitionerUtil {
+
+ enum TableSchemaType {
+ ONE_BUCKET {
+ @Override
+ public int bucketPartitionColumnPosition() {
+ return 0;
+ }
+
+ @Override
+ public PartitionSpec getPartitionSpec(int numBuckets) {
+ return PartitionSpec.builderFor(SimpleDataUtil.SCHEMA).bucket("data", numBuckets).build();
+ }
+ },
+ IDENTITY_AND_BUCKET {
+ @Override
+ public PartitionSpec getPartitionSpec(int numBuckets) {
+ return PartitionSpec.builderFor(SimpleDataUtil.SCHEMA)
+ .identity("id")
+ .bucket("data", numBuckets)
+ .build();
+ }
+ },
+ TWO_BUCKETS {
+ @Override
+ public PartitionSpec getPartitionSpec(int numBuckets) {
+ return PartitionSpec.builderFor(SimpleDataUtil.SCHEMA)
+ .bucket("id", numBuckets)
+ .bucket("data", numBuckets)
+ .build();
+ }
+ };
+
+ public int bucketPartitionColumnPosition() {
Review Comment:
nit: I find it easier to read to have this as abstract, because the position and schema are defined nearby/together. I know it will be the same for the last 2 enums.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Flink: Custom partitioner for bucket partitions [iceberg]
Posted by "chenwyi2 (via GitHub)" <gi...@apache.org>.
chenwyi2 commented on PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#issuecomment-1761169778
Hi @stevenzwu @kengtin this PR can be create too many small files when parition with dt,hout,minute and bucekt(id), suppose paralisim is 120 and bucke number is 8, then 15 writes can write into same one bucket, but there is problem, data from the previous few hours can be into one commit because of data latency, there can be 15000 and more data files if changed partition is up to 1000, can we use complete parition name instead of just bucket?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Flink: Custom partitioner for bucket partitions [iceberg]
Posted by "stevenzwu (via GitHub)" <gi...@apache.org>.
stevenzwu commented on PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#issuecomment-1762592732
@chenwyi2 Is your point that we shouldn't only consider bucketing column (like did in this PR). you just want a plain keyBy in this case? that would be a fair point. Do you get balanced traffic distribution among write tasks with simple keyBy?
I am also wondering if the partition spec of dt,hour,minute and bucekt(id) is the best option. especially the minute column as partition. do you really need minute level partition granularity. you are creating very fine grained partitions. even with the most optimal data distribution/shuffle. there are still a lot of partitions and data files.
you used 8 for bucket number. it seems quite small for bucketing. what's the reason of using 8 buckets?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] advancedxy commented on pull request #7161: Flink: Custom partitioner for bucket partitions
Posted by "advancedxy (via GitHub)" <gi...@apache.org>.
advancedxy commented on PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#issuecomment-1689268351
Hi @stevenzwu @kengtin this PR restricts that only one bucket transform in the partition spec. Do you think it's beneficial to support multiple bucket transforms or even the mutli-arg bucket transform(#8259) since it's common to have multiple keys to form a primary key?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] stevenzwu commented on a diff in pull request #7161: Flink: Custom partitioner for a more even data distribution when bucket partitions are present
Posted by "stevenzwu (via GitHub)" <gi...@apache.org>.
stevenzwu commented on code in PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#discussion_r1147806601
##########
api/src/main/java/org/apache/iceberg/PartitionSpec.java:
##########
@@ -289,6 +289,23 @@ public Set<Integer> identitySourceIds() {
return sourceIds;
}
+ /** Determine if this PartitionSpec contains at least 1 bucket definition */
+ public boolean hasBucketPartition() {
+ List<PartitionField> partFields = fields();
+ boolean hasBucketPartition = false;
+
+ for (int i = 0; i < partFields.size(); i++) {
+ PartitionField field = partFields.get(i);
+ // Is there a more elegant way to know? (all Bucket* constructs are package private)
+ if (field.transform().dedupName().toLowerCase().contains("bucket")) {
+ hasBucketPartition = true;
Review Comment:
nit: may not need this variable. just return true?
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitionKeySelector.java:
##########
@@ -0,0 +1,85 @@
+/*
+ * 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 java.util.Optional;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.types.logical.RowType;
+import org.apache.iceberg.PartitionKey;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Schema;
+import org.apache.iceberg.flink.RowDataWrapper;
+
+/**
+ * Very similar to the {@link PartitionKeySelector}, but optimized to extract and return an Integer
+ * bucketId as the key.
+ */
+public class BucketPartitionKeySelector implements KeySelector<RowData, Integer> {
+
+ private final Schema schema;
+ private final PartitionKey partitionKey;
+ private final RowType flinkSchema;
+ private static final Pattern intPattern = Pattern.compile("\\d+");
+
+ private transient RowDataWrapper rowDataWrapper;
+
+ BucketPartitionKeySelector(PartitionSpec spec, Schema schema, RowType flinkSchema) {
+ this.schema = schema;
+ this.partitionKey = new PartitionKey(spec, schema);
+ this.flinkSchema = flinkSchema;
+ }
+
+ /**
+ * Construct the {@link RowDataWrapper} lazily here because few members in it are not
+ * serializable. In this way, we don't have to serialize them with forcing.
+ */
+ private RowDataWrapper lazyRowDataWrapper() {
+ if (rowDataWrapper == null) {
+ rowDataWrapper = new RowDataWrapper(flinkSchema, schema.asStruct());
+ }
+ return rowDataWrapper;
+ }
+
+ @Override
+ public Integer getKey(RowData rowData) throws Exception {
+ partitionKey.partition(lazyRowDataWrapper().wrap(rowData));
+ Optional<Integer> bucketId = BucketPartitionKeySelector.extractInteger(partitionKey.toPath());
+
+ if (!bucketId.isPresent()) {
+ throw new IllegalStateException("Unable to extract bucket in partition key: " + partitionKey);
+ }
+
+ return bucketId.get();
+ }
+
+ /**
+ * Utility method to extract a bucketId from a string value. Input examples: bucket[5] or
+ * bucket_data[11], would return 5 and 11 bucket Ids respectively.
+ *
+ * @param value the string with the Bucket Id
+ * @return the Bucket Id as an Integer or Empty if not found
+ */
+ public static Optional<Integer> extractInteger(String value) {
Review Comment:
this is a very hack way to convert PartitionKey to a string and use regex match. we should figure out a more direct way. E.g., we can get the partition value from tuples. we can know which partition position is for bucketing partition
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitioner.java:
##########
@@ -0,0 +1,86 @@
+/*
+ * 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 java.util.Optional;
+import org.apache.flink.api.common.functions.Partitioner;
+import org.apache.iceberg.PartitionField;
+import org.apache.iceberg.PartitionSpec;
+
+/**
+ * This partitioner will redirect elements to writers deterministically so that each writer only
+ * targets 1 bucket. If the number of writers > number of buckets each partitioner will keep a state
+ * of multiple writers per bucket as evenly as possible, and will round-robin the requests across
+ * them.
+ */
+public class BucketPartitioner implements Partitioner<Integer> {
+
+ private final Integer maxBuckets;
+
+ private final int[] currentWriterOffset;
+
+ BucketPartitioner(PartitionSpec partitionSpec) {
+ // The current implementation redirects to writers based on the _FIRST_ bucket found
+ Optional<PartitionField> bucket =
+ partitionSpec.fields().stream()
+ .filter(f -> f.transform().dedupName().contains("bucket"))
+ .findFirst();
+
+ if (!bucket.isPresent()) {
Review Comment:
Iceberg coding style uses `Preconditions` for those check
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitioner.java:
##########
@@ -0,0 +1,86 @@
+/*
+ * 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 java.util.Optional;
+import org.apache.flink.api.common.functions.Partitioner;
+import org.apache.iceberg.PartitionField;
+import org.apache.iceberg.PartitionSpec;
+
+/**
+ * This partitioner will redirect elements to writers deterministically so that each writer only
+ * targets 1 bucket. If the number of writers > number of buckets each partitioner will keep a state
+ * of multiple writers per bucket as evenly as possible, and will round-robin the requests across
+ * them.
+ */
+public class BucketPartitioner implements Partitioner<Integer> {
+
+ private final Integer maxBuckets;
+
+ private final int[] currentWriterOffset;
+
+ BucketPartitioner(PartitionSpec partitionSpec) {
+ // The current implementation redirects to writers based on the _FIRST_ bucket found
+ Optional<PartitionField> bucket =
+ partitionSpec.fields().stream()
+ .filter(f -> f.transform().dedupName().contains("bucket"))
+ .findFirst();
+
+ if (!bucket.isPresent()) {
Review Comment:
since this impl doesn't support multiple bucket partitions, we should also validate and fail explicitly for that scenario. we should also update the document for this feature and explicitly point out this limitation.
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitionKeySelector.java:
##########
@@ -0,0 +1,85 @@
+/*
+ * 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 java.util.Optional;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.types.logical.RowType;
+import org.apache.iceberg.PartitionKey;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Schema;
+import org.apache.iceberg.flink.RowDataWrapper;
+
+/**
+ * Very similar to the {@link PartitionKeySelector}, but optimized to extract and return an Integer
+ * bucketId as the key.
+ */
+public class BucketPartitionKeySelector implements KeySelector<RowData, Integer> {
Review Comment:
make this class package private
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitioner.java:
##########
@@ -0,0 +1,86 @@
+/*
+ * 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 java.util.Optional;
+import org.apache.flink.api.common.functions.Partitioner;
+import org.apache.iceberg.PartitionField;
+import org.apache.iceberg.PartitionSpec;
+
+/**
+ * This partitioner will redirect elements to writers deterministically so that each writer only
Review Comment:
> each writer only targets 1 bucket
this doesn't seem accurate. at least, it can be interpreted incorrectly.
##########
api/src/main/java/org/apache/iceberg/PartitionSpec.java:
##########
@@ -289,6 +289,23 @@ public Set<Integer> identitySourceIds() {
return sourceIds;
}
+ /** Determine if this PartitionSpec contains at least 1 bucket definition */
+ public boolean hasBucketPartition() {
Review Comment:
Let's not change core for now. we can keep this util method inside the flink module for now
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitioner.java:
##########
@@ -0,0 +1,86 @@
+/*
+ * 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 java.util.Optional;
+import org.apache.flink.api.common.functions.Partitioner;
+import org.apache.iceberg.PartitionField;
+import org.apache.iceberg.PartitionSpec;
+
+/**
+ * This partitioner will redirect elements to writers deterministically so that each writer only
+ * targets 1 bucket. If the number of writers > number of buckets each partitioner will keep a state
+ * of multiple writers per bucket as evenly as possible, and will round-robin the requests across
+ * them.
+ */
+public class BucketPartitioner implements Partitioner<Integer> {
Review Comment:
please make the class package private.
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitioner.java:
##########
@@ -0,0 +1,86 @@
+/*
+ * 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 java.util.Optional;
+import org.apache.flink.api.common.functions.Partitioner;
+import org.apache.iceberg.PartitionField;
+import org.apache.iceberg.PartitionSpec;
+
+/**
+ * This partitioner will redirect elements to writers deterministically so that each writer only
+ * targets 1 bucket. If the number of writers > number of buckets each partitioner will keep a state
+ * of multiple writers per bucket as evenly as possible, and will round-robin the requests across
+ * them.
+ */
+public class BucketPartitioner implements Partitioner<Integer> {
+
+ private final Integer maxBuckets;
Review Comment:
nit: this can be primitive `int`
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitioner.java:
##########
@@ -0,0 +1,86 @@
+/*
+ * 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 java.util.Optional;
+import org.apache.flink.api.common.functions.Partitioner;
+import org.apache.iceberg.PartitionField;
+import org.apache.iceberg.PartitionSpec;
+
+/**
+ * This partitioner will redirect elements to writers deterministically so that each writer only
+ * targets 1 bucket. If the number of writers > number of buckets each partitioner will keep a state
+ * of multiple writers per bucket as evenly as possible, and will round-robin the requests across
+ * them.
+ */
+public class BucketPartitioner implements Partitioner<Integer> {
+
+ private final Integer maxBuckets;
+
+ private final int[] currentWriterOffset;
+
+ BucketPartitioner(PartitionSpec partitionSpec) {
+ // The current implementation redirects to writers based on the _FIRST_ bucket found
+ Optional<PartitionField> bucket =
+ partitionSpec.fields().stream()
+ .filter(f -> f.transform().dedupName().contains("bucket"))
Review Comment:
why don't we directly check the class type is an instance of `Bucket`?
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitioner.java:
##########
@@ -0,0 +1,86 @@
+/*
+ * 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 java.util.Optional;
+import org.apache.flink.api.common.functions.Partitioner;
+import org.apache.iceberg.PartitionField;
+import org.apache.iceberg.PartitionSpec;
+
+/**
+ * This partitioner will redirect elements to writers deterministically so that each writer only
+ * targets 1 bucket. If the number of writers > number of buckets each partitioner will keep a state
+ * of multiple writers per bucket as evenly as possible, and will round-robin the requests across
+ * them.
+ */
+public class BucketPartitioner implements Partitioner<Integer> {
+
+ private final Integer maxBuckets;
+
+ private final int[] currentWriterOffset;
+
+ BucketPartitioner(PartitionSpec partitionSpec) {
+ // The current implementation redirects to writers based on the _FIRST_ bucket found
+ Optional<PartitionField> bucket =
+ partitionSpec.fields().stream()
+ .filter(f -> f.transform().dedupName().contains("bucket"))
+ .findFirst();
+
+ if (!bucket.isPresent()) {
+ throw new IllegalStateException("No buckets found on the provided PartitionSpec");
+ }
+
+ // Extracting the max number of buckets defined in the partition spec
+ String transformName = bucket.get().transform().dedupName();
+ Optional<Integer> maxBucketsOpt = BucketPartitionKeySelector.extractInteger(transformName);
+
+ if (maxBucketsOpt.isPresent()) {
+ this.maxBuckets = maxBucketsOpt.get();
+ } else {
+ throw new IllegalStateException(
+ "Could not extract the max number of buckets from the transform name ("
+ + transformName
+ + ")");
+ }
+
+ this.currentWriterOffset = new int[this.maxBuckets];
+ }
+
+ @Override
+ public int partition(Integer bucketId, int numPartitions) {
+ if (numPartitions > maxBuckets) {
+ return getSubpartitionIdx(bucketId, numPartitions);
+ } else {
+ return bucketId % numPartitions;
+ }
+ }
+
+ private int getSubpartitionIdx(int bucketId, int numPartitions) {
Review Comment:
nit: getPartitionIndex? Iceberg coding style typically don't use abbreviation like `idx`.
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitioner.java:
##########
@@ -0,0 +1,86 @@
+/*
+ * 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 java.util.Optional;
+import org.apache.flink.api.common.functions.Partitioner;
+import org.apache.iceberg.PartitionField;
+import org.apache.iceberg.PartitionSpec;
+
+/**
+ * This partitioner will redirect elements to writers deterministically so that each writer only
+ * targets 1 bucket. If the number of writers > number of buckets each partitioner will keep a state
+ * of multiple writers per bucket as evenly as possible, and will round-robin the requests across
+ * them.
+ */
+public class BucketPartitioner implements Partitioner<Integer> {
+
+ private final Integer maxBuckets;
+
+ private final int[] currentWriterOffset;
+
+ BucketPartitioner(PartitionSpec partitionSpec) {
+ // The current implementation redirects to writers based on the _FIRST_ bucket found
+ Optional<PartitionField> bucket =
+ partitionSpec.fields().stream()
+ .filter(f -> f.transform().dedupName().contains("bucket"))
+ .findFirst();
+
+ if (!bucket.isPresent()) {
+ throw new IllegalStateException("No buckets found on the provided PartitionSpec");
+ }
+
+ // Extracting the max number of buckets defined in the partition spec
+ String transformName = bucket.get().transform().dedupName();
+ Optional<Integer> maxBucketsOpt = BucketPartitionKeySelector.extractInteger(transformName);
+
+ if (maxBucketsOpt.isPresent()) {
+ this.maxBuckets = maxBucketsOpt.get();
+ } else {
+ throw new IllegalStateException(
Review Comment:
Iceberg coding style uses `Preconditions` for this type of check
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitioner.java:
##########
@@ -0,0 +1,86 @@
+/*
+ * 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 java.util.Optional;
+import org.apache.flink.api.common.functions.Partitioner;
+import org.apache.iceberg.PartitionField;
+import org.apache.iceberg.PartitionSpec;
+
+/**
+ * This partitioner will redirect elements to writers deterministically so that each writer only
+ * targets 1 bucket. If the number of writers > number of buckets each partitioner will keep a state
+ * of multiple writers per bucket as evenly as possible, and will round-robin the requests across
+ * them.
+ */
+public class BucketPartitioner implements Partitioner<Integer> {
+
+ private final Integer maxBuckets;
+
+ private final int[] currentWriterOffset;
+
+ BucketPartitioner(PartitionSpec partitionSpec) {
+ // The current implementation redirects to writers based on the _FIRST_ bucket found
+ Optional<PartitionField> bucket =
+ partitionSpec.fields().stream()
+ .filter(f -> f.transform().dedupName().contains("bucket"))
+ .findFirst();
+
+ if (!bucket.isPresent()) {
+ throw new IllegalStateException("No buckets found on the provided PartitionSpec");
+ }
+
+ // Extracting the max number of buckets defined in the partition spec
+ String transformName = bucket.get().transform().dedupName();
+ Optional<Integer> maxBucketsOpt = BucketPartitionKeySelector.extractInteger(transformName);
+
+ if (maxBucketsOpt.isPresent()) {
+ this.maxBuckets = maxBucketsOpt.get();
+ } else {
+ throw new IllegalStateException(
+ "Could not extract the max number of buckets from the transform name ("
+ + transformName
+ + ")");
+ }
+
+ this.currentWriterOffset = new int[this.maxBuckets];
+ }
+
+ @Override
+ public int partition(Integer bucketId, int numPartitions) {
+ if (numPartitions > maxBuckets) {
+ return getSubpartitionIdx(bucketId, numPartitions);
+ } else {
+ return bucketId % numPartitions;
+ }
+ }
+
+ private int getSubpartitionIdx(int bucketId, int numPartitions) {
Review Comment:
also the logic in this method is not easy to understand. can you provide some illustration on how this works.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Flink: Custom partitioner for bucket partitions [iceberg]
Posted by "chenwyi2 (via GitHub)" <gi...@apache.org>.
chenwyi2 commented on PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#issuecomment-1763604926
yes, I am creating very fine grained partitions, because i want to query and comput some business metrics between minutes ss fast as possible. As for bucket number, i use a fomula QPS * 500B/条 / bucket number / 1024 /1024 = 10M/秒(The write-in traffic of one bucket), because i found that the write-in traffic of one bucket is large, the writer can be OOM or backpressure, i set 10M/秒 in each bucket.
In this pr hdfs can be influenced by many files, writting to files can be slow then data latency will be relatively large.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] stevenzwu commented on a diff in pull request #7161: Flink: Custom partitioner for bucket partitions
Posted by "stevenzwu (via GitHub)" <gi...@apache.org>.
stevenzwu commented on code in PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#discussion_r1233384927
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitioner.java:
##########
@@ -0,0 +1,108 @@
+/*
+ * 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 org.apache.flink.api.common.functions.Partitioner;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/**
+ * This partitioner will redirect records to writers deterministically based on the Bucket partition
+ * spec. It'll attempt to optimize the file size written depending on whether numPartitions is
+ * greater, less or equal than the maxNumBuckets. Note: The current implementation only supports ONE
+ * bucket in the partition spec.
+ */
+class BucketPartitioner implements Partitioner<Integer> {
+
+ static final String BUCKET_NULL_MESSAGE = "bucketId cannot be null";
+ static final String BUCKET_LESS_THAN_LOWER_BOUND_MESSAGE =
+ "bucketId out of range: %s, must be non-negative.";
Review Comment:
we probably should try to follow existing style where value is after the semicolon. maybe we can use the style as this code?
```
Preconditions.checkArgument(
numBuckets > 0, "Invalid number of buckets: %s (must be > 0)", numBuckets);
```
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerUtil.java:
##########
@@ -0,0 +1,118 @@
+/*
+ * 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 java.util.List;
+import java.util.UUID;
+import org.apache.flink.table.data.GenericRowData;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.data.StringData;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.apache.iceberg.util.BucketUtil;
+
+final class TestBucketPartitionerUtil {
+
+ enum TableSchemaType {
+ ONE_BUCKET {
+ @Override
+ public int bucketPartitionColumnPosition() {
+ return 0;
+ }
+
+ @Override
+ public PartitionSpec getPartitionSpec(int numBuckets) {
+ return PartitionSpec.builderFor(SimpleDataUtil.SCHEMA).bucket("data", numBuckets).build();
+ }
+ },
+ IDENTITY_AND_BUCKET {
+ @Override
+ public PartitionSpec getPartitionSpec(int numBuckets) {
+ return PartitionSpec.builderFor(SimpleDataUtil.SCHEMA)
+ .identity("id")
+ .bucket("data", numBuckets)
+ .build();
+ }
+ },
+ TWO_BUCKETS {
+ @Override
+ public PartitionSpec getPartitionSpec(int numBuckets) {
+ return PartitionSpec.builderFor(SimpleDataUtil.SCHEMA)
+ .bucket("id", numBuckets)
+ .bucket("data", numBuckets)
+ .build();
+ }
+ };
+
+ public int bucketPartitionColumnPosition() {
Review Comment:
nit: I find it easier to read to have this as abstract. I know it will be the same for the last 2 enums.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] kengtin commented on pull request #7161: Flink: Custom partitioner for bucket partitions
Posted by "kengtin (via GitHub)" <gi...@apache.org>.
kengtin commented on PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#issuecomment-1690744828
@advancedxy I promise to get up to speed with the multi-arg bucket transform (not yet familiar) but your suggestion makes sense. I'll also let @stevenzwu weigh in.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] kengtin commented on a diff in pull request #7161: Flink: Custom partitioner for bucket partitions
Posted by "kengtin (via GitHub)" <gi...@apache.org>.
kengtin commented on code in PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#discussion_r1209428672
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerFlinkIcebergSink.java:
##########
@@ -0,0 +1,253 @@
+/*
+ * 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.apache.iceberg.flink.MiniClusterResource.DISABLE_CLASSLOADER_CHECK_CONFIG;
+import static org.apache.iceberg.flink.TestFixtures.DATABASE;
+import static org.apache.iceberg.flink.TestFixtures.TABLE_IDENTIFIER;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.TreeMap;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.typeutils.RowTypeInfo;
+import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.data.util.DataFormatConverters;
+import org.apache.flink.test.junit5.MiniClusterExtension;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.DistributionMode;
+import org.apache.iceberg.FileFormat;
+import org.apache.iceberg.FileScanTask;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Table;
+import org.apache.iceberg.TableProperties;
+import org.apache.iceberg.flink.HadoopCatalogExtension;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.flink.TableLoader;
+import org.apache.iceberg.flink.TestFixtures;
+import org.apache.iceberg.flink.sink.TestBucketPartitionerUtils.TableSchemaType;
+import org.apache.iceberg.flink.source.BoundedTestSource;
+import org.apache.iceberg.flink.util.FlinkCompatibilityUtil;
+import org.apache.iceberg.io.CloseableIterable;
+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.assertj.core.api.Assertions;
+import org.jetbrains.annotations.NotNull;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+public class TestBucketPartitionerFlinkIcebergSink {
+
+ private static final int NUMBER_TASK_MANAGERS = 1;
+ private static final int SLOTS_PER_TASK_MANAGER = 8;
+
+ @RegisterExtension
+ private static final MiniClusterExtension MINI_CLUSTER_RESOURCE =
+ new MiniClusterExtension(
+ new MiniClusterResourceConfiguration.Builder()
+ .setNumberTaskManagers(NUMBER_TASK_MANAGERS)
+ .setNumberSlotsPerTaskManager(SLOTS_PER_TASK_MANAGER)
+ .setConfiguration(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .build());
+
+ @RegisterExtension
+ private static final HadoopCatalogExtension catalogExtension =
+ new HadoopCatalogExtension(DATABASE, TestFixtures.TABLE);
+
+ private static final TypeInformation<Row> ROW_TYPE_INFO =
+ new RowTypeInfo(SimpleDataUtil.FLINK_SCHEMA.getFieldTypes());
+
+ // Parallelism = 8 (parallelism > numBuckets) throughout the test suite
+ private final int parallelism = NUMBER_TASK_MANAGERS * SLOTS_PER_TASK_MANAGER;
+ private final FileFormat format = FileFormat.PARQUET;
+ private final int numBuckets = 4;
+
+ private Table table;
+ private StreamExecutionEnvironment env;
+ private TableLoader tableLoader;
+
+ private void setupEnvironment(TableSchemaType tableSchemaType) {
+ table = getTable(tableSchemaType);
+ env =
+ StreamExecutionEnvironment.getExecutionEnvironment(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .enableCheckpointing(100)
+ .setParallelism(parallelism)
+ .setMaxParallelism(parallelism * 2);
+ tableLoader = catalogExtension.tableLoader();
+ }
+
+ private Table getTable(TableSchemaType tableSchemaType) {
+ PartitionSpec partitionSpec = TableSchemaType.getPartitionSpec(tableSchemaType, numBuckets);
+
+ return catalogExtension
+ .catalog()
+ .createTable(
+ TABLE_IDENTIFIER,
+ SimpleDataUtil.SCHEMA,
+ partitionSpec,
+ ImmutableMap.of(TableProperties.DEFAULT_FILE_FORMAT, format.name()));
+ }
+
+ private void testWriteRowData(List<RowData> allRows) throws Exception {
+ DataFormatConverters.RowConverter converter =
+ new DataFormatConverters.RowConverter(SimpleDataUtil.FLINK_SCHEMA.getFieldDataTypes());
+
+ DataStream<RowData> dataStream =
+ env.addSource(
+ new BoundedTestSource<>(
+ allRows.stream().map(converter::toExternal).toArray(Row[]::new)),
+ ROW_TYPE_INFO)
+ .map(converter::toInternal, FlinkCompatibilityUtil.toTypeInfo(SimpleDataUtil.ROW_TYPE));
+
+ FlinkSink.forRowData(dataStream)
+ .table(table)
+ .tableLoader(tableLoader)
+ .writeParallelism(parallelism)
+ .distributionMode(DistributionMode.HASH)
+ .append();
+
+ // Execute the program.
+ env.execute("Test Iceberg DataStream");
+
+ // Assert the iceberg table's records.
+ SimpleDataUtil.assertTableRows(table, allRows);
+ }
+
+ @ParameterizedTest
+ @EnumSource(
+ value = TableSchemaType.class,
+ names = {"ONE_BUCKET", "IDENTITY_AND_BUCKET"})
+ public void testSendRecordsToAllBucketsEvenly(TableSchemaType tableSchemaType) throws Exception {
+ setupEnvironment(tableSchemaType);
+ List<RowData> rows = generateTestDataRows();
+
+ testWriteRowData(rows);
+ TableTestStats stats = extractPartitionResults(tableSchemaType);
+
+ Assertions.assertThat(stats.totalRowCount).isEqualTo(rows.size());
+ // All 4 buckets should've been written to
+ Assertions.assertThat(stats.writersPerBucket.size()).isEqualTo(numBuckets);
+ Assertions.assertThat(stats.numFilesPerBucket.size()).isEqualTo(numBuckets);
+ // Writer expectation (2 writers per bucket):
+ // - Bucket0 -> Writers [0, 4]
+ // - Bucket1 -> Writers [1, 5]
+ // - Bucket2 -> Writers [2, 6]
+ // - Bucket3 -> Writers [3, 7]
+ for (int i = 0, j = numBuckets; i < numBuckets; i++, j++) {
+ Assertions.assertThat(stats.writersPerBucket.get(i)).isEqualTo(Arrays.asList(i, j));
+ // 2 files per bucket (one file is created by each writer)
+ Assertions.assertThat((int) stats.numFilesPerBucket.get(i)).isEqualTo(2);
Review Comment:
Can't remember exactly, maybe a leftover from a previous assertion that returned Object.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] kengtin commented on a diff in pull request #7161: Flink: Custom partitioner for bucket partitions
Posted by "kengtin (via GitHub)" <gi...@apache.org>.
kengtin commented on code in PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#discussion_r1209406397
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitioner.java:
##########
@@ -0,0 +1,108 @@
+/*
+ * 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 org.apache.flink.api.common.functions.Partitioner;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/**
+ * This partitioner will redirect records to writers deterministically based on the Bucket partition
+ * spec. It'll attempt to optimize the file size written depending on whether numPartitions is
+ * greater, less or equal than the maxNumBuckets. Note: The current implementation only supports ONE
+ * bucket in the partition spec.
+ */
+class BucketPartitioner implements Partitioner<Integer> {
+
+ static final String BUCKET_NULL_MESSAGE = "bucketId cannot be null";
+ static final String BUCKET_LESS_THAN_LOWER_BOUND_MESSAGE =
+ "bucketId out of range: %s, must be non-negative.";
Review Comment:
Overall I agree with this one, but I have my reservations on the double `:` for each message, in these cases. Also I find it way more readable to have the number of the invalid bucket Id right after its message (e.g. `Invalid bucket id %s...`)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] kengtin commented on a diff in pull request #7161: Flink: Custom partitioner for bucket partitions
Posted by "kengtin (via GitHub)" <gi...@apache.org>.
kengtin commented on code in PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#discussion_r1197832480
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/HadoopCatalogExtension.java:
##########
@@ -0,0 +1,104 @@
+/*
+ * 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;
+
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.UUID;
+import org.apache.commons.io.FileUtils;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.iceberg.CatalogProperties;
+import org.apache.iceberg.catalog.Catalog;
+import org.apache.iceberg.catalog.TableIdentifier;
+import org.apache.iceberg.hadoop.HadoopCatalog;
+import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
+import org.assertj.core.api.Assertions;
+import org.junit.jupiter.api.extension.AfterAllCallback;
+import org.junit.jupiter.api.extension.AfterEachCallback;
+import org.junit.jupiter.api.extension.BeforeAllCallback;
+import org.junit.jupiter.api.extension.BeforeEachCallback;
+import org.junit.jupiter.api.extension.ExtensionContext;
+
+public class HadoopCatalogExtension
+ implements BeforeAllCallback, BeforeEachCallback, AfterAllCallback, AfterEachCallback {
+ protected final String database;
+ protected final String tableName;
+
+ protected Path temporaryFolder;
Review Comment:
I had originally tried that but it seems that @TempDirs are not supported in Extensions. Refs: `https://stackoverflow.com/questions/54647577/how-to-use-tempdir-in-extensions` and `https://github.com/junit-team/junit5/issues/1786`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] stevenzwu commented on a diff in pull request #7161: Flink: Custom partitioner for bucket partitions
Posted by "stevenzwu (via GitHub)" <gi...@apache.org>.
stevenzwu commented on code in PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#discussion_r1206006772
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitionKeySelector.java:
##########
@@ -0,0 +1,70 @@
+/*
+ * 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 java.util.stream.IntStream;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.types.logical.RowType;
+import org.apache.iceberg.PartitionKey;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Schema;
+import org.apache.iceberg.flink.RowDataWrapper;
+
+/**
+ * A {@link KeySelector} that extracts the bucketId from a data row's bucket partition as the key.
+ * To be used with the {@link BucketPartitioner}.
+ */
+class BucketPartitionKeySelector implements KeySelector<RowData, Integer> {
+
+ private final Schema schema;
+ private final PartitionKey partitionKey;
+ private final RowType flinkSchema;
+ private final int bucketFieldPosition;
+
+ private transient RowDataWrapper rowDataWrapper;
+
+ BucketPartitionKeySelector(PartitionSpec partitionSpec, Schema schema, RowType flinkSchema) {
Review Comment:
nit: usually do simple assignments like `this.schema = schema` first. Then start a new section with calculated values like `bucketFieldId`, `bucketFieldPosition`
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitioner.java:
##########
@@ -0,0 +1,108 @@
+/*
+ * 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 org.apache.flink.api.common.functions.Partitioner;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/**
+ * This partitioner will redirect records to writers deterministically based on the Bucket partition
+ * spec. It'll attempt to optimize the file size written depending on whether numPartitions is
+ * greater, less or equal than the maxNumBuckets. Note: The current implementation only supports ONE
+ * bucket in the partition spec.
+ */
+class BucketPartitioner implements Partitioner<Integer> {
+
+ static final String BUCKET_NULL_MESSAGE = "bucketId cannot be null";
+ static final String BUCKET_LESS_THAN_LOWER_BOUND_MESSAGE =
+ "bucketId out of range: %s, must be non-negative.";
+ static final String BUCKET_GREATER_THAN_UPPER_BOUND_MESSAGE =
+ "bucketId out of range: %s, must be less than bucket limit: %s.";
+
+ private final int maxNumBuckets;
+
+ // To hold the OFFSET of the next writer to use for any bucket, only used when writers > the
+ // number
Review Comment:
nit: fix formatting.
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitionerUtils.java:
##########
@@ -0,0 +1,117 @@
+/*
+ * 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 java.util.List;
+import java.util.Objects;
+import java.util.stream.Collectors;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+import org.apache.iceberg.transforms.PartitionSpecVisitor;
+
+final class BucketPartitionerUtils {
+ static final String BAD_NUMBER_OF_BUCKETS_ERROR_MESSAGE =
Review Comment:
nit on error msg `Unsupported partition spec: bucket partitioner expect 1 bucket partition: %s`
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitioner.java:
##########
@@ -0,0 +1,108 @@
+/*
+ * 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 org.apache.flink.api.common.functions.Partitioner;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/**
+ * This partitioner will redirect records to writers deterministically based on the Bucket partition
+ * spec. It'll attempt to optimize the file size written depending on whether numPartitions is
+ * greater, less or equal than the maxNumBuckets. Note: The current implementation only supports ONE
+ * bucket in the partition spec.
+ */
+class BucketPartitioner implements Partitioner<Integer> {
+
+ static final String BUCKET_NULL_MESSAGE = "bucketId cannot be null";
+ static final String BUCKET_LESS_THAN_LOWER_BOUND_MESSAGE =
+ "bucketId out of range: %s, must be non-negative.";
Review Comment:
Iceberg style on error msg. it should be readable to users (not Iceberg developers). Hence, variable or class name shouldn't be exposed. `bucketId` should be `Bucket ID`.
Iceberg error msg tends to follow the style of `Invalid ...`. Hence error msgs can be like `Invalid bucket ID: must be non-negative: %s` and `Invalid bucket ID: must be less than the bucket limit of %s: %s`
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerUtils.java:
##########
@@ -0,0 +1,120 @@
+/*
+ * 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 java.util.List;
+import java.util.UUID;
+import org.apache.flink.table.data.GenericRowData;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.data.StringData;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.apache.iceberg.util.BucketUtil;
+
+final class TestBucketPartitionerUtils {
+
+ enum TableSchemaType {
+ ONE_BUCKET,
+ IDENTITY_AND_BUCKET,
+ TWO_BUCKETS;
+
+ public static int bucketPartitionColumnPosition(TableSchemaType tableSchemaType) {
Review Comment:
it is weird that we are putting those static method in the enum class definition. they should be overriden in the enum definition so that we can avoid the switch-case like the method below. these two methods should be non static.
also `TWO_BUCKETS` should throw an exception for this method.
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitioner.java:
##########
@@ -0,0 +1,110 @@
+/*
+ * 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.apache.iceberg.flink.sink.BucketPartitioner.BUCKET_GREATER_THAN_UPPER_BOUND_MESSAGE;
+import static org.apache.iceberg.flink.sink.BucketPartitioner.BUCKET_LESS_THAN_LOWER_BOUND_MESSAGE;
+import static org.apache.iceberg.flink.sink.BucketPartitioner.BUCKET_NULL_MESSAGE;
+
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.flink.sink.TestBucketPartitionerUtils.TableSchemaType;
+import org.assertj.core.api.Assertions;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.CsvSource;
+
+public class TestBucketPartitioner {
+
+ static final int DEFAULT_NUM_BUCKETS = 60;
+
+ @ParameterizedTest
+ @CsvSource({"ONE_BUCKET,50", "IDENTITY_AND_BUCKET,50", "ONE_BUCKET,60", "IDENTITY_AND_BUCKET,60"})
+ public void testPartitioningParallelismGreaterThanBuckets(
+ String schemaTypeStr, String numBucketsStr) {
+ final int numPartitions = 500;
Review Comment:
nit: no `final` for local var
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerFlinkIcebergSink.java:
##########
@@ -0,0 +1,253 @@
+/*
+ * 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.apache.iceberg.flink.MiniClusterResource.DISABLE_CLASSLOADER_CHECK_CONFIG;
+import static org.apache.iceberg.flink.TestFixtures.DATABASE;
+import static org.apache.iceberg.flink.TestFixtures.TABLE_IDENTIFIER;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.TreeMap;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.typeutils.RowTypeInfo;
+import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.data.util.DataFormatConverters;
+import org.apache.flink.test.junit5.MiniClusterExtension;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.DistributionMode;
+import org.apache.iceberg.FileFormat;
+import org.apache.iceberg.FileScanTask;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Table;
+import org.apache.iceberg.TableProperties;
+import org.apache.iceberg.flink.HadoopCatalogExtension;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.flink.TableLoader;
+import org.apache.iceberg.flink.TestFixtures;
+import org.apache.iceberg.flink.sink.TestBucketPartitionerUtils.TableSchemaType;
+import org.apache.iceberg.flink.source.BoundedTestSource;
+import org.apache.iceberg.flink.util.FlinkCompatibilityUtil;
+import org.apache.iceberg.io.CloseableIterable;
+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.assertj.core.api.Assertions;
+import org.jetbrains.annotations.NotNull;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+public class TestBucketPartitionerFlinkIcebergSink {
+
+ private static final int NUMBER_TASK_MANAGERS = 1;
+ private static final int SLOTS_PER_TASK_MANAGER = 8;
+
+ @RegisterExtension
+ private static final MiniClusterExtension MINI_CLUSTER_RESOURCE =
+ new MiniClusterExtension(
+ new MiniClusterResourceConfiguration.Builder()
+ .setNumberTaskManagers(NUMBER_TASK_MANAGERS)
+ .setNumberSlotsPerTaskManager(SLOTS_PER_TASK_MANAGER)
+ .setConfiguration(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .build());
+
+ @RegisterExtension
+ private static final HadoopCatalogExtension catalogExtension =
+ new HadoopCatalogExtension(DATABASE, TestFixtures.TABLE);
+
+ private static final TypeInformation<Row> ROW_TYPE_INFO =
+ new RowTypeInfo(SimpleDataUtil.FLINK_SCHEMA.getFieldTypes());
+
+ // Parallelism = 8 (parallelism > numBuckets) throughout the test suite
+ private final int parallelism = NUMBER_TASK_MANAGERS * SLOTS_PER_TASK_MANAGER;
+ private final FileFormat format = FileFormat.PARQUET;
+ private final int numBuckets = 4;
+
+ private Table table;
+ private StreamExecutionEnvironment env;
+ private TableLoader tableLoader;
+
+ private void setupEnvironment(TableSchemaType tableSchemaType) {
+ table = getTable(tableSchemaType);
+ env =
+ StreamExecutionEnvironment.getExecutionEnvironment(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .enableCheckpointing(100)
+ .setParallelism(parallelism)
+ .setMaxParallelism(parallelism * 2);
+ tableLoader = catalogExtension.tableLoader();
+ }
+
+ private Table getTable(TableSchemaType tableSchemaType) {
+ PartitionSpec partitionSpec = TableSchemaType.getPartitionSpec(tableSchemaType, numBuckets);
+
Review Comment:
nit this empty line is not needed
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerFlinkIcebergSink.java:
##########
@@ -0,0 +1,253 @@
+/*
+ * 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.apache.iceberg.flink.MiniClusterResource.DISABLE_CLASSLOADER_CHECK_CONFIG;
+import static org.apache.iceberg.flink.TestFixtures.DATABASE;
+import static org.apache.iceberg.flink.TestFixtures.TABLE_IDENTIFIER;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.TreeMap;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.typeutils.RowTypeInfo;
+import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.data.util.DataFormatConverters;
+import org.apache.flink.test.junit5.MiniClusterExtension;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.DistributionMode;
+import org.apache.iceberg.FileFormat;
+import org.apache.iceberg.FileScanTask;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Table;
+import org.apache.iceberg.TableProperties;
+import org.apache.iceberg.flink.HadoopCatalogExtension;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.flink.TableLoader;
+import org.apache.iceberg.flink.TestFixtures;
+import org.apache.iceberg.flink.sink.TestBucketPartitionerUtils.TableSchemaType;
+import org.apache.iceberg.flink.source.BoundedTestSource;
+import org.apache.iceberg.flink.util.FlinkCompatibilityUtil;
+import org.apache.iceberg.io.CloseableIterable;
+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.assertj.core.api.Assertions;
+import org.jetbrains.annotations.NotNull;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+public class TestBucketPartitionerFlinkIcebergSink {
+
+ private static final int NUMBER_TASK_MANAGERS = 1;
+ private static final int SLOTS_PER_TASK_MANAGER = 8;
+
+ @RegisterExtension
+ private static final MiniClusterExtension MINI_CLUSTER_RESOURCE =
+ new MiniClusterExtension(
+ new MiniClusterResourceConfiguration.Builder()
+ .setNumberTaskManagers(NUMBER_TASK_MANAGERS)
+ .setNumberSlotsPerTaskManager(SLOTS_PER_TASK_MANAGER)
+ .setConfiguration(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .build());
+
+ @RegisterExtension
+ private static final HadoopCatalogExtension catalogExtension =
+ new HadoopCatalogExtension(DATABASE, TestFixtures.TABLE);
+
+ private static final TypeInformation<Row> ROW_TYPE_INFO =
+ new RowTypeInfo(SimpleDataUtil.FLINK_SCHEMA.getFieldTypes());
+
+ // Parallelism = 8 (parallelism > numBuckets) throughout the test suite
+ private final int parallelism = NUMBER_TASK_MANAGERS * SLOTS_PER_TASK_MANAGER;
+ private final FileFormat format = FileFormat.PARQUET;
+ private final int numBuckets = 4;
+
+ private Table table;
+ private StreamExecutionEnvironment env;
+ private TableLoader tableLoader;
+
+ private void setupEnvironment(TableSchemaType tableSchemaType) {
Review Comment:
should this use `@BeforeEach`?
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerFlinkIcebergSink.java:
##########
@@ -0,0 +1,253 @@
+/*
+ * 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.apache.iceberg.flink.MiniClusterResource.DISABLE_CLASSLOADER_CHECK_CONFIG;
+import static org.apache.iceberg.flink.TestFixtures.DATABASE;
+import static org.apache.iceberg.flink.TestFixtures.TABLE_IDENTIFIER;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.TreeMap;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.typeutils.RowTypeInfo;
+import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.data.util.DataFormatConverters;
+import org.apache.flink.test.junit5.MiniClusterExtension;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.DistributionMode;
+import org.apache.iceberg.FileFormat;
+import org.apache.iceberg.FileScanTask;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Table;
+import org.apache.iceberg.TableProperties;
+import org.apache.iceberg.flink.HadoopCatalogExtension;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.flink.TableLoader;
+import org.apache.iceberg.flink.TestFixtures;
+import org.apache.iceberg.flink.sink.TestBucketPartitionerUtils.TableSchemaType;
+import org.apache.iceberg.flink.source.BoundedTestSource;
+import org.apache.iceberg.flink.util.FlinkCompatibilityUtil;
+import org.apache.iceberg.io.CloseableIterable;
+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.assertj.core.api.Assertions;
+import org.jetbrains.annotations.NotNull;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+public class TestBucketPartitionerFlinkIcebergSink {
+
+ private static final int NUMBER_TASK_MANAGERS = 1;
+ private static final int SLOTS_PER_TASK_MANAGER = 8;
+
+ @RegisterExtension
+ private static final MiniClusterExtension MINI_CLUSTER_RESOURCE =
+ new MiniClusterExtension(
+ new MiniClusterResourceConfiguration.Builder()
+ .setNumberTaskManagers(NUMBER_TASK_MANAGERS)
+ .setNumberSlotsPerTaskManager(SLOTS_PER_TASK_MANAGER)
+ .setConfiguration(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .build());
+
+ @RegisterExtension
+ private static final HadoopCatalogExtension catalogExtension =
+ new HadoopCatalogExtension(DATABASE, TestFixtures.TABLE);
+
+ private static final TypeInformation<Row> ROW_TYPE_INFO =
+ new RowTypeInfo(SimpleDataUtil.FLINK_SCHEMA.getFieldTypes());
+
+ // Parallelism = 8 (parallelism > numBuckets) throughout the test suite
+ private final int parallelism = NUMBER_TASK_MANAGERS * SLOTS_PER_TASK_MANAGER;
+ private final FileFormat format = FileFormat.PARQUET;
+ private final int numBuckets = 4;
+
+ private Table table;
+ private StreamExecutionEnvironment env;
+ private TableLoader tableLoader;
+
+ private void setupEnvironment(TableSchemaType tableSchemaType) {
+ table = getTable(tableSchemaType);
+ env =
+ StreamExecutionEnvironment.getExecutionEnvironment(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .enableCheckpointing(100)
+ .setParallelism(parallelism)
+ .setMaxParallelism(parallelism * 2);
+ tableLoader = catalogExtension.tableLoader();
+ }
+
+ private Table getTable(TableSchemaType tableSchemaType) {
Review Comment:
I would incline to merge this method with the method above
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerFlinkIcebergSink.java:
##########
@@ -0,0 +1,253 @@
+/*
+ * 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.apache.iceberg.flink.MiniClusterResource.DISABLE_CLASSLOADER_CHECK_CONFIG;
+import static org.apache.iceberg.flink.TestFixtures.DATABASE;
+import static org.apache.iceberg.flink.TestFixtures.TABLE_IDENTIFIER;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.TreeMap;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.typeutils.RowTypeInfo;
+import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.data.util.DataFormatConverters;
+import org.apache.flink.test.junit5.MiniClusterExtension;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.DistributionMode;
+import org.apache.iceberg.FileFormat;
+import org.apache.iceberg.FileScanTask;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Table;
+import org.apache.iceberg.TableProperties;
+import org.apache.iceberg.flink.HadoopCatalogExtension;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.flink.TableLoader;
+import org.apache.iceberg.flink.TestFixtures;
+import org.apache.iceberg.flink.sink.TestBucketPartitionerUtils.TableSchemaType;
+import org.apache.iceberg.flink.source.BoundedTestSource;
+import org.apache.iceberg.flink.util.FlinkCompatibilityUtil;
+import org.apache.iceberg.io.CloseableIterable;
+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.assertj.core.api.Assertions;
+import org.jetbrains.annotations.NotNull;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+public class TestBucketPartitionerFlinkIcebergSink {
+
+ private static final int NUMBER_TASK_MANAGERS = 1;
+ private static final int SLOTS_PER_TASK_MANAGER = 8;
+
+ @RegisterExtension
+ private static final MiniClusterExtension MINI_CLUSTER_RESOURCE =
+ new MiniClusterExtension(
+ new MiniClusterResourceConfiguration.Builder()
+ .setNumberTaskManagers(NUMBER_TASK_MANAGERS)
+ .setNumberSlotsPerTaskManager(SLOTS_PER_TASK_MANAGER)
+ .setConfiguration(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .build());
+
+ @RegisterExtension
+ private static final HadoopCatalogExtension catalogExtension =
+ new HadoopCatalogExtension(DATABASE, TestFixtures.TABLE);
+
+ private static final TypeInformation<Row> ROW_TYPE_INFO =
+ new RowTypeInfo(SimpleDataUtil.FLINK_SCHEMA.getFieldTypes());
+
+ // Parallelism = 8 (parallelism > numBuckets) throughout the test suite
+ private final int parallelism = NUMBER_TASK_MANAGERS * SLOTS_PER_TASK_MANAGER;
+ private final FileFormat format = FileFormat.PARQUET;
+ private final int numBuckets = 4;
+
+ private Table table;
+ private StreamExecutionEnvironment env;
+ private TableLoader tableLoader;
+
+ private void setupEnvironment(TableSchemaType tableSchemaType) {
+ table = getTable(tableSchemaType);
+ env =
+ StreamExecutionEnvironment.getExecutionEnvironment(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .enableCheckpointing(100)
+ .setParallelism(parallelism)
+ .setMaxParallelism(parallelism * 2);
+ tableLoader = catalogExtension.tableLoader();
+ }
+
+ private Table getTable(TableSchemaType tableSchemaType) {
+ PartitionSpec partitionSpec = TableSchemaType.getPartitionSpec(tableSchemaType, numBuckets);
+
+ return catalogExtension
+ .catalog()
+ .createTable(
+ TABLE_IDENTIFIER,
+ SimpleDataUtil.SCHEMA,
+ partitionSpec,
+ ImmutableMap.of(TableProperties.DEFAULT_FILE_FORMAT, format.name()));
+ }
+
+ private void testWriteRowData(List<RowData> allRows) throws Exception {
+ DataFormatConverters.RowConverter converter =
+ new DataFormatConverters.RowConverter(SimpleDataUtil.FLINK_SCHEMA.getFieldDataTypes());
+
+ DataStream<RowData> dataStream =
+ env.addSource(
+ new BoundedTestSource<>(
+ allRows.stream().map(converter::toExternal).toArray(Row[]::new)),
Review Comment:
why do we need convert internal and external back and forth? if this is about type info, can we use `TypeInformation.of(RowData.class)`?
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionKeySelector.java:
##########
@@ -0,0 +1,65 @@
+/*
+ * 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 org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.flink.sink.TestBucketPartitionerUtils.TableSchemaType;
+import org.assertj.core.api.Assertions;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+public class TestBucketPartitionKeySelector {
+
+ @ParameterizedTest
+ @EnumSource(
+ value = TableSchemaType.class,
+ names = {"ONE_BUCKET", "IDENTITY_AND_BUCKET"})
+ public void testCorrectKeySelection(TableSchemaType tableSchemaType) {
+ final int numBuckets = 60;
Review Comment:
nit: Iceberg style doesn't use `final` for local var
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerFlinkIcebergSink.java:
##########
@@ -0,0 +1,253 @@
+/*
+ * 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.apache.iceberg.flink.MiniClusterResource.DISABLE_CLASSLOADER_CHECK_CONFIG;
+import static org.apache.iceberg.flink.TestFixtures.DATABASE;
+import static org.apache.iceberg.flink.TestFixtures.TABLE_IDENTIFIER;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.TreeMap;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.typeutils.RowTypeInfo;
+import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.data.util.DataFormatConverters;
+import org.apache.flink.test.junit5.MiniClusterExtension;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.DistributionMode;
+import org.apache.iceberg.FileFormat;
+import org.apache.iceberg.FileScanTask;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Table;
+import org.apache.iceberg.TableProperties;
+import org.apache.iceberg.flink.HadoopCatalogExtension;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.flink.TableLoader;
+import org.apache.iceberg.flink.TestFixtures;
+import org.apache.iceberg.flink.sink.TestBucketPartitionerUtils.TableSchemaType;
+import org.apache.iceberg.flink.source.BoundedTestSource;
+import org.apache.iceberg.flink.util.FlinkCompatibilityUtil;
+import org.apache.iceberg.io.CloseableIterable;
+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.assertj.core.api.Assertions;
+import org.jetbrains.annotations.NotNull;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+public class TestBucketPartitionerFlinkIcebergSink {
+
+ private static final int NUMBER_TASK_MANAGERS = 1;
+ private static final int SLOTS_PER_TASK_MANAGER = 8;
+
+ @RegisterExtension
+ private static final MiniClusterExtension MINI_CLUSTER_RESOURCE =
+ new MiniClusterExtension(
+ new MiniClusterResourceConfiguration.Builder()
+ .setNumberTaskManagers(NUMBER_TASK_MANAGERS)
+ .setNumberSlotsPerTaskManager(SLOTS_PER_TASK_MANAGER)
+ .setConfiguration(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .build());
+
+ @RegisterExtension
+ private static final HadoopCatalogExtension catalogExtension =
+ new HadoopCatalogExtension(DATABASE, TestFixtures.TABLE);
+
+ private static final TypeInformation<Row> ROW_TYPE_INFO =
+ new RowTypeInfo(SimpleDataUtil.FLINK_SCHEMA.getFieldTypes());
+
+ // Parallelism = 8 (parallelism > numBuckets) throughout the test suite
+ private final int parallelism = NUMBER_TASK_MANAGERS * SLOTS_PER_TASK_MANAGER;
+ private final FileFormat format = FileFormat.PARQUET;
+ private final int numBuckets = 4;
+
+ private Table table;
+ private StreamExecutionEnvironment env;
+ private TableLoader tableLoader;
+
+ private void setupEnvironment(TableSchemaType tableSchemaType) {
+ table = getTable(tableSchemaType);
+ env =
+ StreamExecutionEnvironment.getExecutionEnvironment(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .enableCheckpointing(100)
+ .setParallelism(parallelism)
+ .setMaxParallelism(parallelism * 2);
+ tableLoader = catalogExtension.tableLoader();
+ }
+
+ private Table getTable(TableSchemaType tableSchemaType) {
+ PartitionSpec partitionSpec = TableSchemaType.getPartitionSpec(tableSchemaType, numBuckets);
+
+ return catalogExtension
+ .catalog()
+ .createTable(
+ TABLE_IDENTIFIER,
+ SimpleDataUtil.SCHEMA,
+ partitionSpec,
+ ImmutableMap.of(TableProperties.DEFAULT_FILE_FORMAT, format.name()));
+ }
+
+ private void testWriteRowData(List<RowData> allRows) throws Exception {
+ DataFormatConverters.RowConverter converter =
+ new DataFormatConverters.RowConverter(SimpleDataUtil.FLINK_SCHEMA.getFieldDataTypes());
+
+ DataStream<RowData> dataStream =
+ env.addSource(
+ new BoundedTestSource<>(
+ allRows.stream().map(converter::toExternal).toArray(Row[]::new)),
+ ROW_TYPE_INFO)
+ .map(converter::toInternal, FlinkCompatibilityUtil.toTypeInfo(SimpleDataUtil.ROW_TYPE));
+
+ FlinkSink.forRowData(dataStream)
+ .table(table)
+ .tableLoader(tableLoader)
+ .writeParallelism(parallelism)
+ .distributionMode(DistributionMode.HASH)
+ .append();
+
+ // Execute the program.
+ env.execute("Test Iceberg DataStream");
+
+ // Assert the iceberg table's records.
+ SimpleDataUtil.assertTableRows(table, allRows);
+ }
+
+ @ParameterizedTest
+ @EnumSource(
+ value = TableSchemaType.class,
+ names = {"ONE_BUCKET", "IDENTITY_AND_BUCKET"})
+ public void testSendRecordsToAllBucketsEvenly(TableSchemaType tableSchemaType) throws Exception {
+ setupEnvironment(tableSchemaType);
+ List<RowData> rows = generateTestDataRows();
+
+ testWriteRowData(rows);
+ TableTestStats stats = extractPartitionResults(tableSchemaType);
+
+ Assertions.assertThat(stats.totalRowCount).isEqualTo(rows.size());
+ // All 4 buckets should've been written to
+ Assertions.assertThat(stats.writersPerBucket.size()).isEqualTo(numBuckets);
+ Assertions.assertThat(stats.numFilesPerBucket.size()).isEqualTo(numBuckets);
+ // Writer expectation (2 writers per bucket):
+ // - Bucket0 -> Writers [0, 4]
+ // - Bucket1 -> Writers [1, 5]
+ // - Bucket2 -> Writers [2, 6]
+ // - Bucket3 -> Writers [3, 7]
+ for (int i = 0, j = numBuckets; i < numBuckets; i++, j++) {
+ Assertions.assertThat(stats.writersPerBucket.get(i)).isEqualTo(Arrays.asList(i, j));
+ // 2 files per bucket (one file is created by each writer)
+ Assertions.assertThat((int) stats.numFilesPerBucket.get(i)).isEqualTo(2);
+ // 2 rows per file (total of 16 rows across 8 files)
+ Assertions.assertThat((long) stats.rowsPerFile.get(i)).isEqualTo(2);
+ }
+ }
+
+ /**
+ * Verifies the BucketPartitioner is not used when the PartitionSpec has more than 1 bucket, and
+ * that it should fallback to input.keyBy
+ */
+ @ParameterizedTest
+ @EnumSource(value = TableSchemaType.class, names = "TWO_BUCKETS")
+ public void testMultipleBucketsFallback(TableSchemaType tableSchemaType) throws Exception {
+ setupEnvironment(tableSchemaType);
+ List<RowData> rows = generateTestDataRows();
+
+ testWriteRowData(rows);
+ TableTestStats stats = extractPartitionResults(tableSchemaType);
+
+ Assertions.assertThat(stats.totalRowCount).isEqualTo(rows.size());
+ for (int i = 0, j = numBuckets; i < numBuckets; i++, j++) {
+ // Only 1 file per bucket will be created when falling back to input.keyBy(...)
+ Assertions.assertThat((int) stats.numFilesPerBucket.get(i)).isEqualTo(1);
+ }
+ }
+
+ /**
+ * Generating 16 rows to be sent uniformly to all writers (round-robin across 8 writers -> 4
+ * buckets)
+ */
+ @NotNull
+ private List<RowData> generateTestDataRows() {
+ int totalNumRows = parallelism * 2;
+ int numRowsPerBucket = totalNumRows / numBuckets;
+ return TestBucketPartitionerUtils.generateRowsForBucketIdRange(numRowsPerBucket, numBuckets);
+ }
+
+ private TableTestStats extractPartitionResults(TableSchemaType tableSchemaType)
+ throws IOException {
+ int totalRecordCount = 0;
+ Map<Integer, List<Integer>> writersPerBucket = Maps.newHashMap(); // <BucketId, List<WriterId>>
+ Map<Integer, Integer> filesPerBucket = Maps.newHashMap(); // <BucketId, NumFiles>
+ Map<Integer, Long> recordsPerFile = new TreeMap<>(); // <WriterId, NumRecords>
Review Comment:
this should be `recordsPerWriter`. also why is `TreeMap` needed here?
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerFlinkIcebergSink.java:
##########
@@ -0,0 +1,253 @@
+/*
+ * 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.apache.iceberg.flink.MiniClusterResource.DISABLE_CLASSLOADER_CHECK_CONFIG;
+import static org.apache.iceberg.flink.TestFixtures.DATABASE;
+import static org.apache.iceberg.flink.TestFixtures.TABLE_IDENTIFIER;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.TreeMap;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.typeutils.RowTypeInfo;
+import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.data.util.DataFormatConverters;
+import org.apache.flink.test.junit5.MiniClusterExtension;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.DistributionMode;
+import org.apache.iceberg.FileFormat;
+import org.apache.iceberg.FileScanTask;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Table;
+import org.apache.iceberg.TableProperties;
+import org.apache.iceberg.flink.HadoopCatalogExtension;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.flink.TableLoader;
+import org.apache.iceberg.flink.TestFixtures;
+import org.apache.iceberg.flink.sink.TestBucketPartitionerUtils.TableSchemaType;
+import org.apache.iceberg.flink.source.BoundedTestSource;
+import org.apache.iceberg.flink.util.FlinkCompatibilityUtil;
+import org.apache.iceberg.io.CloseableIterable;
+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.assertj.core.api.Assertions;
+import org.jetbrains.annotations.NotNull;
Review Comment:
wrong annotation. shouldn't use jetbrains. also looking at the usage. I think we can just remove the annotation usage
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitionerUtils.java:
##########
@@ -0,0 +1,117 @@
+/*
+ * 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 java.util.List;
+import java.util.Objects;
+import java.util.stream.Collectors;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+import org.apache.iceberg.transforms.PartitionSpecVisitor;
+
+final class BucketPartitionerUtils {
Review Comment:
nit: although I can find both `Util` and `Utils`, I think there are a lot more class name without the plural. so let's stay with `Util` for this and the other test class.
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerFlinkIcebergSink.java:
##########
@@ -0,0 +1,253 @@
+/*
+ * 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.apache.iceberg.flink.MiniClusterResource.DISABLE_CLASSLOADER_CHECK_CONFIG;
+import static org.apache.iceberg.flink.TestFixtures.DATABASE;
+import static org.apache.iceberg.flink.TestFixtures.TABLE_IDENTIFIER;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.TreeMap;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.typeutils.RowTypeInfo;
+import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.data.util.DataFormatConverters;
+import org.apache.flink.test.junit5.MiniClusterExtension;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.DistributionMode;
+import org.apache.iceberg.FileFormat;
+import org.apache.iceberg.FileScanTask;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Table;
+import org.apache.iceberg.TableProperties;
+import org.apache.iceberg.flink.HadoopCatalogExtension;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.flink.TableLoader;
+import org.apache.iceberg.flink.TestFixtures;
+import org.apache.iceberg.flink.sink.TestBucketPartitionerUtils.TableSchemaType;
+import org.apache.iceberg.flink.source.BoundedTestSource;
+import org.apache.iceberg.flink.util.FlinkCompatibilityUtil;
+import org.apache.iceberg.io.CloseableIterable;
+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.assertj.core.api.Assertions;
+import org.jetbrains.annotations.NotNull;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+public class TestBucketPartitionerFlinkIcebergSink {
+
+ private static final int NUMBER_TASK_MANAGERS = 1;
+ private static final int SLOTS_PER_TASK_MANAGER = 8;
+
+ @RegisterExtension
+ private static final MiniClusterExtension MINI_CLUSTER_RESOURCE =
+ new MiniClusterExtension(
+ new MiniClusterResourceConfiguration.Builder()
+ .setNumberTaskManagers(NUMBER_TASK_MANAGERS)
+ .setNumberSlotsPerTaskManager(SLOTS_PER_TASK_MANAGER)
+ .setConfiguration(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .build());
+
+ @RegisterExtension
+ private static final HadoopCatalogExtension catalogExtension =
+ new HadoopCatalogExtension(DATABASE, TestFixtures.TABLE);
+
+ private static final TypeInformation<Row> ROW_TYPE_INFO =
+ new RowTypeInfo(SimpleDataUtil.FLINK_SCHEMA.getFieldTypes());
+
+ // Parallelism = 8 (parallelism > numBuckets) throughout the test suite
+ private final int parallelism = NUMBER_TASK_MANAGERS * SLOTS_PER_TASK_MANAGER;
+ private final FileFormat format = FileFormat.PARQUET;
+ private final int numBuckets = 4;
+
+ private Table table;
+ private StreamExecutionEnvironment env;
+ private TableLoader tableLoader;
+
+ private void setupEnvironment(TableSchemaType tableSchemaType) {
+ table = getTable(tableSchemaType);
+ env =
+ StreamExecutionEnvironment.getExecutionEnvironment(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .enableCheckpointing(100)
+ .setParallelism(parallelism)
+ .setMaxParallelism(parallelism * 2);
+ tableLoader = catalogExtension.tableLoader();
+ }
+
+ private Table getTable(TableSchemaType tableSchemaType) {
+ PartitionSpec partitionSpec = TableSchemaType.getPartitionSpec(tableSchemaType, numBuckets);
+
+ return catalogExtension
+ .catalog()
+ .createTable(
+ TABLE_IDENTIFIER,
+ SimpleDataUtil.SCHEMA,
+ partitionSpec,
+ ImmutableMap.of(TableProperties.DEFAULT_FILE_FORMAT, format.name()));
+ }
+
+ private void testWriteRowData(List<RowData> allRows) throws Exception {
+ DataFormatConverters.RowConverter converter =
+ new DataFormatConverters.RowConverter(SimpleDataUtil.FLINK_SCHEMA.getFieldDataTypes());
+
+ DataStream<RowData> dataStream =
+ env.addSource(
+ new BoundedTestSource<>(
+ allRows.stream().map(converter::toExternal).toArray(Row[]::new)),
+ ROW_TYPE_INFO)
+ .map(converter::toInternal, FlinkCompatibilityUtil.toTypeInfo(SimpleDataUtil.ROW_TYPE));
+
+ FlinkSink.forRowData(dataStream)
+ .table(table)
+ .tableLoader(tableLoader)
+ .writeParallelism(parallelism)
+ .distributionMode(DistributionMode.HASH)
+ .append();
+
+ // Execute the program.
+ env.execute("Test Iceberg DataStream");
+
+ // Assert the iceberg table's records.
+ SimpleDataUtil.assertTableRows(table, allRows);
+ }
+
+ @ParameterizedTest
+ @EnumSource(
+ value = TableSchemaType.class,
+ names = {"ONE_BUCKET", "IDENTITY_AND_BUCKET"})
+ public void testSendRecordsToAllBucketsEvenly(TableSchemaType tableSchemaType) throws Exception {
+ setupEnvironment(tableSchemaType);
+ List<RowData> rows = generateTestDataRows();
+
+ testWriteRowData(rows);
+ TableTestStats stats = extractPartitionResults(tableSchemaType);
+
+ Assertions.assertThat(stats.totalRowCount).isEqualTo(rows.size());
+ // All 4 buckets should've been written to
+ Assertions.assertThat(stats.writersPerBucket.size()).isEqualTo(numBuckets);
+ Assertions.assertThat(stats.numFilesPerBucket.size()).isEqualTo(numBuckets);
+ // Writer expectation (2 writers per bucket):
+ // - Bucket0 -> Writers [0, 4]
+ // - Bucket1 -> Writers [1, 5]
+ // - Bucket2 -> Writers [2, 6]
+ // - Bucket3 -> Writers [3, 7]
+ for (int i = 0, j = numBuckets; i < numBuckets; i++, j++) {
+ Assertions.assertThat(stats.writersPerBucket.get(i)).isEqualTo(Arrays.asList(i, j));
+ // 2 files per bucket (one file is created by each writer)
+ Assertions.assertThat((int) stats.numFilesPerBucket.get(i)).isEqualTo(2);
+ // 2 rows per file (total of 16 rows across 8 files)
+ Assertions.assertThat((long) stats.rowsPerFile.get(i)).isEqualTo(2);
+ }
+ }
+
+ /**
+ * Verifies the BucketPartitioner is not used when the PartitionSpec has more than 1 bucket, and
+ * that it should fallback to input.keyBy
+ */
+ @ParameterizedTest
+ @EnumSource(value = TableSchemaType.class, names = "TWO_BUCKETS")
+ public void testMultipleBucketsFallback(TableSchemaType tableSchemaType) throws Exception {
+ setupEnvironment(tableSchemaType);
+ List<RowData> rows = generateTestDataRows();
+
+ testWriteRowData(rows);
+ TableTestStats stats = extractPartitionResults(tableSchemaType);
+
+ Assertions.assertThat(stats.totalRowCount).isEqualTo(rows.size());
+ for (int i = 0, j = numBuckets; i < numBuckets; i++, j++) {
+ // Only 1 file per bucket will be created when falling back to input.keyBy(...)
+ Assertions.assertThat((int) stats.numFilesPerBucket.get(i)).isEqualTo(1);
+ }
+ }
+
+ /**
+ * Generating 16 rows to be sent uniformly to all writers (round-robin across 8 writers -> 4
+ * buckets)
+ */
+ @NotNull
+ private List<RowData> generateTestDataRows() {
+ int totalNumRows = parallelism * 2;
+ int numRowsPerBucket = totalNumRows / numBuckets;
+ return TestBucketPartitionerUtils.generateRowsForBucketIdRange(numRowsPerBucket, numBuckets);
+ }
+
+ private TableTestStats extractPartitionResults(TableSchemaType tableSchemaType)
+ throws IOException {
+ int totalRecordCount = 0;
+ Map<Integer, List<Integer>> writersPerBucket = Maps.newHashMap(); // <BucketId, List<WriterId>>
+ Map<Integer, Integer> filesPerBucket = Maps.newHashMap(); // <BucketId, NumFiles>
+ Map<Integer, Long> recordsPerFile = new TreeMap<>(); // <WriterId, NumRecords>
+
+ try (CloseableIterable<FileScanTask> fileScanTasks = table.newScan().planFiles()) {
+ for (FileScanTask scanTask : fileScanTasks) {
+ long recordCountInFile = scanTask.file().recordCount();
+
+ String[] splitFilePath = scanTask.file().path().toString().split("/");
+ String filename = splitFilePath[splitFilePath.length - 1];
+ int writerId = Integer.parseInt(filename.split("-")[0]);
+
+ totalRecordCount += recordCountInFile;
+ int bucketId =
+ scanTask
+ .file()
+ .partition()
+ .get(TableSchemaType.bucketPartitionColumnPosition(tableSchemaType), Integer.class);
+ writersPerBucket.computeIfAbsent(bucketId, k -> Lists.newArrayList());
+ writersPerBucket.get(bucketId).add(writerId);
+ filesPerBucket.put(bucketId, filesPerBucket.getOrDefault(bucketId, 0) + 1);
+ recordsPerFile.put(writerId, recordsPerFile.getOrDefault(writerId, 0L) + recordCountInFile);
+ }
+ }
+
+ for (int k : writersPerBucket.keySet()) {
Review Comment:
sort is not needed if using `hasSameElementsAs` from assertj
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerFlinkIcebergSink.java:
##########
@@ -0,0 +1,253 @@
+/*
+ * 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.apache.iceberg.flink.MiniClusterResource.DISABLE_CLASSLOADER_CHECK_CONFIG;
+import static org.apache.iceberg.flink.TestFixtures.DATABASE;
+import static org.apache.iceberg.flink.TestFixtures.TABLE_IDENTIFIER;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.TreeMap;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.typeutils.RowTypeInfo;
+import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.data.util.DataFormatConverters;
+import org.apache.flink.test.junit5.MiniClusterExtension;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.DistributionMode;
+import org.apache.iceberg.FileFormat;
+import org.apache.iceberg.FileScanTask;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Table;
+import org.apache.iceberg.TableProperties;
+import org.apache.iceberg.flink.HadoopCatalogExtension;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.flink.TableLoader;
+import org.apache.iceberg.flink.TestFixtures;
+import org.apache.iceberg.flink.sink.TestBucketPartitionerUtils.TableSchemaType;
+import org.apache.iceberg.flink.source.BoundedTestSource;
+import org.apache.iceberg.flink.util.FlinkCompatibilityUtil;
+import org.apache.iceberg.io.CloseableIterable;
+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.assertj.core.api.Assertions;
+import org.jetbrains.annotations.NotNull;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+public class TestBucketPartitionerFlinkIcebergSink {
+
+ private static final int NUMBER_TASK_MANAGERS = 1;
+ private static final int SLOTS_PER_TASK_MANAGER = 8;
+
+ @RegisterExtension
+ private static final MiniClusterExtension MINI_CLUSTER_RESOURCE =
+ new MiniClusterExtension(
+ new MiniClusterResourceConfiguration.Builder()
+ .setNumberTaskManagers(NUMBER_TASK_MANAGERS)
+ .setNumberSlotsPerTaskManager(SLOTS_PER_TASK_MANAGER)
+ .setConfiguration(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .build());
+
+ @RegisterExtension
+ private static final HadoopCatalogExtension catalogExtension =
+ new HadoopCatalogExtension(DATABASE, TestFixtures.TABLE);
+
+ private static final TypeInformation<Row> ROW_TYPE_INFO =
+ new RowTypeInfo(SimpleDataUtil.FLINK_SCHEMA.getFieldTypes());
+
+ // Parallelism = 8 (parallelism > numBuckets) throughout the test suite
+ private final int parallelism = NUMBER_TASK_MANAGERS * SLOTS_PER_TASK_MANAGER;
+ private final FileFormat format = FileFormat.PARQUET;
+ private final int numBuckets = 4;
+
+ private Table table;
+ private StreamExecutionEnvironment env;
+ private TableLoader tableLoader;
+
+ private void setupEnvironment(TableSchemaType tableSchemaType) {
+ table = getTable(tableSchemaType);
+ env =
+ StreamExecutionEnvironment.getExecutionEnvironment(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .enableCheckpointing(100)
+ .setParallelism(parallelism)
+ .setMaxParallelism(parallelism * 2);
+ tableLoader = catalogExtension.tableLoader();
+ }
+
+ private Table getTable(TableSchemaType tableSchemaType) {
+ PartitionSpec partitionSpec = TableSchemaType.getPartitionSpec(tableSchemaType, numBuckets);
+
+ return catalogExtension
+ .catalog()
+ .createTable(
+ TABLE_IDENTIFIER,
+ SimpleDataUtil.SCHEMA,
+ partitionSpec,
+ ImmutableMap.of(TableProperties.DEFAULT_FILE_FORMAT, format.name()));
+ }
+
+ private void testWriteRowData(List<RowData> allRows) throws Exception {
+ DataFormatConverters.RowConverter converter =
+ new DataFormatConverters.RowConverter(SimpleDataUtil.FLINK_SCHEMA.getFieldDataTypes());
+
+ DataStream<RowData> dataStream =
+ env.addSource(
+ new BoundedTestSource<>(
+ allRows.stream().map(converter::toExternal).toArray(Row[]::new)),
+ ROW_TYPE_INFO)
+ .map(converter::toInternal, FlinkCompatibilityUtil.toTypeInfo(SimpleDataUtil.ROW_TYPE));
+
+ FlinkSink.forRowData(dataStream)
+ .table(table)
+ .tableLoader(tableLoader)
+ .writeParallelism(parallelism)
+ .distributionMode(DistributionMode.HASH)
+ .append();
+
+ // Execute the program.
+ env.execute("Test Iceberg DataStream");
+
+ // Assert the iceberg table's records.
+ SimpleDataUtil.assertTableRows(table, allRows);
+ }
+
+ @ParameterizedTest
+ @EnumSource(
+ value = TableSchemaType.class,
+ names = {"ONE_BUCKET", "IDENTITY_AND_BUCKET"})
+ public void testSendRecordsToAllBucketsEvenly(TableSchemaType tableSchemaType) throws Exception {
+ setupEnvironment(tableSchemaType);
+ List<RowData> rows = generateTestDataRows();
+
+ testWriteRowData(rows);
+ TableTestStats stats = extractPartitionResults(tableSchemaType);
+
+ Assertions.assertThat(stats.totalRowCount).isEqualTo(rows.size());
+ // All 4 buckets should've been written to
+ Assertions.assertThat(stats.writersPerBucket.size()).isEqualTo(numBuckets);
+ Assertions.assertThat(stats.numFilesPerBucket.size()).isEqualTo(numBuckets);
+ // Writer expectation (2 writers per bucket):
+ // - Bucket0 -> Writers [0, 4]
+ // - Bucket1 -> Writers [1, 5]
+ // - Bucket2 -> Writers [2, 6]
+ // - Bucket3 -> Writers [3, 7]
+ for (int i = 0, j = numBuckets; i < numBuckets; i++, j++) {
+ Assertions.assertThat(stats.writersPerBucket.get(i)).isEqualTo(Arrays.asList(i, j));
+ // 2 files per bucket (one file is created by each writer)
+ Assertions.assertThat((int) stats.numFilesPerBucket.get(i)).isEqualTo(2);
+ // 2 rows per file (total of 16 rows across 8 files)
+ Assertions.assertThat((long) stats.rowsPerFile.get(i)).isEqualTo(2);
+ }
+ }
+
+ /**
+ * Verifies the BucketPartitioner is not used when the PartitionSpec has more than 1 bucket, and
+ * that it should fallback to input.keyBy
+ */
+ @ParameterizedTest
+ @EnumSource(value = TableSchemaType.class, names = "TWO_BUCKETS")
+ public void testMultipleBucketsFallback(TableSchemaType tableSchemaType) throws Exception {
+ setupEnvironment(tableSchemaType);
+ List<RowData> rows = generateTestDataRows();
+
+ testWriteRowData(rows);
+ TableTestStats stats = extractPartitionResults(tableSchemaType);
+
+ Assertions.assertThat(stats.totalRowCount).isEqualTo(rows.size());
+ for (int i = 0, j = numBuckets; i < numBuckets; i++, j++) {
+ // Only 1 file per bucket will be created when falling back to input.keyBy(...)
+ Assertions.assertThat((int) stats.numFilesPerBucket.get(i)).isEqualTo(1);
+ }
+ }
+
+ /**
+ * Generating 16 rows to be sent uniformly to all writers (round-robin across 8 writers -> 4
+ * buckets)
+ */
+ @NotNull
+ private List<RowData> generateTestDataRows() {
+ int totalNumRows = parallelism * 2;
+ int numRowsPerBucket = totalNumRows / numBuckets;
+ return TestBucketPartitionerUtils.generateRowsForBucketIdRange(numRowsPerBucket, numBuckets);
+ }
+
+ private TableTestStats extractPartitionResults(TableSchemaType tableSchemaType)
+ throws IOException {
+ int totalRecordCount = 0;
+ Map<Integer, List<Integer>> writersPerBucket = Maps.newHashMap(); // <BucketId, List<WriterId>>
+ Map<Integer, Integer> filesPerBucket = Maps.newHashMap(); // <BucketId, NumFiles>
+ Map<Integer, Long> recordsPerFile = new TreeMap<>(); // <WriterId, NumRecords>
+
+ try (CloseableIterable<FileScanTask> fileScanTasks = table.newScan().planFiles()) {
+ for (FileScanTask scanTask : fileScanTasks) {
+ long recordCountInFile = scanTask.file().recordCount();
+
+ String[] splitFilePath = scanTask.file().path().toString().split("/");
Review Comment:
please add an example of the file path for easier understanding of the code.
this approach of parsing file name to extract writerId is a bit fragile (with tight coupling of file name convention in `OutputFileFactory`). But I don't have a better idea. Maybe we will live with this for now.
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/HadoopCatalogExtension.java:
##########
@@ -0,0 +1,104 @@
+/*
+ * 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;
+
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.UUID;
+import org.apache.commons.io.FileUtils;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.iceberg.CatalogProperties;
+import org.apache.iceberg.catalog.Catalog;
+import org.apache.iceberg.catalog.TableIdentifier;
+import org.apache.iceberg.hadoop.HadoopCatalog;
+import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
+import org.assertj.core.api.Assertions;
+import org.junit.jupiter.api.extension.AfterAllCallback;
+import org.junit.jupiter.api.extension.AfterEachCallback;
+import org.junit.jupiter.api.extension.BeforeAllCallback;
+import org.junit.jupiter.api.extension.BeforeEachCallback;
+import org.junit.jupiter.api.extension.ExtensionContext;
+
+public class HadoopCatalogExtension
+ implements BeforeAllCallback, BeforeEachCallback, AfterAllCallback, AfterEachCallback {
+ protected final String database;
+ protected final String tableName;
+
+ protected Path temporaryFolder;
Review Comment:
it can be passed in similar as the junit4 resource
```
public HadoopCatalogResource(TemporaryFolder temporaryFolder, String database, String tableName) {
}
```
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerFlinkIcebergSink.java:
##########
@@ -0,0 +1,253 @@
+/*
+ * 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.apache.iceberg.flink.MiniClusterResource.DISABLE_CLASSLOADER_CHECK_CONFIG;
+import static org.apache.iceberg.flink.TestFixtures.DATABASE;
+import static org.apache.iceberg.flink.TestFixtures.TABLE_IDENTIFIER;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.TreeMap;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.typeutils.RowTypeInfo;
+import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.data.util.DataFormatConverters;
+import org.apache.flink.test.junit5.MiniClusterExtension;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.DistributionMode;
+import org.apache.iceberg.FileFormat;
+import org.apache.iceberg.FileScanTask;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Table;
+import org.apache.iceberg.TableProperties;
+import org.apache.iceberg.flink.HadoopCatalogExtension;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.flink.TableLoader;
+import org.apache.iceberg.flink.TestFixtures;
+import org.apache.iceberg.flink.sink.TestBucketPartitionerUtils.TableSchemaType;
+import org.apache.iceberg.flink.source.BoundedTestSource;
+import org.apache.iceberg.flink.util.FlinkCompatibilityUtil;
+import org.apache.iceberg.io.CloseableIterable;
+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.assertj.core.api.Assertions;
+import org.jetbrains.annotations.NotNull;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+public class TestBucketPartitionerFlinkIcebergSink {
+
+ private static final int NUMBER_TASK_MANAGERS = 1;
+ private static final int SLOTS_PER_TASK_MANAGER = 8;
+
+ @RegisterExtension
+ private static final MiniClusterExtension MINI_CLUSTER_RESOURCE =
+ new MiniClusterExtension(
+ new MiniClusterResourceConfiguration.Builder()
+ .setNumberTaskManagers(NUMBER_TASK_MANAGERS)
+ .setNumberSlotsPerTaskManager(SLOTS_PER_TASK_MANAGER)
+ .setConfiguration(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .build());
+
+ @RegisterExtension
+ private static final HadoopCatalogExtension catalogExtension =
+ new HadoopCatalogExtension(DATABASE, TestFixtures.TABLE);
+
+ private static final TypeInformation<Row> ROW_TYPE_INFO =
+ new RowTypeInfo(SimpleDataUtil.FLINK_SCHEMA.getFieldTypes());
+
+ // Parallelism = 8 (parallelism > numBuckets) throughout the test suite
+ private final int parallelism = NUMBER_TASK_MANAGERS * SLOTS_PER_TASK_MANAGER;
+ private final FileFormat format = FileFormat.PARQUET;
+ private final int numBuckets = 4;
+
+ private Table table;
+ private StreamExecutionEnvironment env;
+ private TableLoader tableLoader;
+
+ private void setupEnvironment(TableSchemaType tableSchemaType) {
+ table = getTable(tableSchemaType);
+ env =
+ StreamExecutionEnvironment.getExecutionEnvironment(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .enableCheckpointing(100)
+ .setParallelism(parallelism)
+ .setMaxParallelism(parallelism * 2);
+ tableLoader = catalogExtension.tableLoader();
+ }
+
+ private Table getTable(TableSchemaType tableSchemaType) {
+ PartitionSpec partitionSpec = TableSchemaType.getPartitionSpec(tableSchemaType, numBuckets);
+
+ return catalogExtension
+ .catalog()
+ .createTable(
+ TABLE_IDENTIFIER,
+ SimpleDataUtil.SCHEMA,
+ partitionSpec,
+ ImmutableMap.of(TableProperties.DEFAULT_FILE_FORMAT, format.name()));
+ }
+
+ private void testWriteRowData(List<RowData> allRows) throws Exception {
+ DataFormatConverters.RowConverter converter =
+ new DataFormatConverters.RowConverter(SimpleDataUtil.FLINK_SCHEMA.getFieldDataTypes());
+
+ DataStream<RowData> dataStream =
+ env.addSource(
+ new BoundedTestSource<>(
+ allRows.stream().map(converter::toExternal).toArray(Row[]::new)),
+ ROW_TYPE_INFO)
+ .map(converter::toInternal, FlinkCompatibilityUtil.toTypeInfo(SimpleDataUtil.ROW_TYPE));
+
+ FlinkSink.forRowData(dataStream)
+ .table(table)
+ .tableLoader(tableLoader)
+ .writeParallelism(parallelism)
+ .distributionMode(DistributionMode.HASH)
+ .append();
+
+ // Execute the program.
Review Comment:
remove too obvious comment
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerFlinkIcebergSink.java:
##########
@@ -0,0 +1,253 @@
+/*
+ * 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.apache.iceberg.flink.MiniClusterResource.DISABLE_CLASSLOADER_CHECK_CONFIG;
+import static org.apache.iceberg.flink.TestFixtures.DATABASE;
+import static org.apache.iceberg.flink.TestFixtures.TABLE_IDENTIFIER;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.TreeMap;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.typeutils.RowTypeInfo;
+import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.data.util.DataFormatConverters;
+import org.apache.flink.test.junit5.MiniClusterExtension;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.DistributionMode;
+import org.apache.iceberg.FileFormat;
+import org.apache.iceberg.FileScanTask;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Table;
+import org.apache.iceberg.TableProperties;
+import org.apache.iceberg.flink.HadoopCatalogExtension;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.flink.TableLoader;
+import org.apache.iceberg.flink.TestFixtures;
+import org.apache.iceberg.flink.sink.TestBucketPartitionerUtils.TableSchemaType;
+import org.apache.iceberg.flink.source.BoundedTestSource;
+import org.apache.iceberg.flink.util.FlinkCompatibilityUtil;
+import org.apache.iceberg.io.CloseableIterable;
+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.assertj.core.api.Assertions;
+import org.jetbrains.annotations.NotNull;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+public class TestBucketPartitionerFlinkIcebergSink {
+
+ private static final int NUMBER_TASK_MANAGERS = 1;
+ private static final int SLOTS_PER_TASK_MANAGER = 8;
+
+ @RegisterExtension
+ private static final MiniClusterExtension MINI_CLUSTER_RESOURCE =
+ new MiniClusterExtension(
+ new MiniClusterResourceConfiguration.Builder()
+ .setNumberTaskManagers(NUMBER_TASK_MANAGERS)
+ .setNumberSlotsPerTaskManager(SLOTS_PER_TASK_MANAGER)
+ .setConfiguration(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .build());
+
+ @RegisterExtension
+ private static final HadoopCatalogExtension catalogExtension =
+ new HadoopCatalogExtension(DATABASE, TestFixtures.TABLE);
+
+ private static final TypeInformation<Row> ROW_TYPE_INFO =
+ new RowTypeInfo(SimpleDataUtil.FLINK_SCHEMA.getFieldTypes());
+
+ // Parallelism = 8 (parallelism > numBuckets) throughout the test suite
+ private final int parallelism = NUMBER_TASK_MANAGERS * SLOTS_PER_TASK_MANAGER;
+ private final FileFormat format = FileFormat.PARQUET;
+ private final int numBuckets = 4;
+
+ private Table table;
+ private StreamExecutionEnvironment env;
+ private TableLoader tableLoader;
+
+ private void setupEnvironment(TableSchemaType tableSchemaType) {
+ table = getTable(tableSchemaType);
+ env =
+ StreamExecutionEnvironment.getExecutionEnvironment(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .enableCheckpointing(100)
+ .setParallelism(parallelism)
+ .setMaxParallelism(parallelism * 2);
+ tableLoader = catalogExtension.tableLoader();
+ }
+
+ private Table getTable(TableSchemaType tableSchemaType) {
+ PartitionSpec partitionSpec = TableSchemaType.getPartitionSpec(tableSchemaType, numBuckets);
+
+ return catalogExtension
+ .catalog()
+ .createTable(
+ TABLE_IDENTIFIER,
+ SimpleDataUtil.SCHEMA,
+ partitionSpec,
+ ImmutableMap.of(TableProperties.DEFAULT_FILE_FORMAT, format.name()));
+ }
+
+ private void testWriteRowData(List<RowData> allRows) throws Exception {
Review Comment:
nit: remove `test` from method name. maybe all this `appendRowsToTable`?
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerFlinkIcebergSink.java:
##########
@@ -0,0 +1,253 @@
+/*
+ * 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.apache.iceberg.flink.MiniClusterResource.DISABLE_CLASSLOADER_CHECK_CONFIG;
+import static org.apache.iceberg.flink.TestFixtures.DATABASE;
+import static org.apache.iceberg.flink.TestFixtures.TABLE_IDENTIFIER;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.TreeMap;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.typeutils.RowTypeInfo;
+import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.data.util.DataFormatConverters;
+import org.apache.flink.test.junit5.MiniClusterExtension;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.DistributionMode;
+import org.apache.iceberg.FileFormat;
+import org.apache.iceberg.FileScanTask;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Table;
+import org.apache.iceberg.TableProperties;
+import org.apache.iceberg.flink.HadoopCatalogExtension;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.flink.TableLoader;
+import org.apache.iceberg.flink.TestFixtures;
+import org.apache.iceberg.flink.sink.TestBucketPartitionerUtils.TableSchemaType;
+import org.apache.iceberg.flink.source.BoundedTestSource;
+import org.apache.iceberg.flink.util.FlinkCompatibilityUtil;
+import org.apache.iceberg.io.CloseableIterable;
+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.assertj.core.api.Assertions;
+import org.jetbrains.annotations.NotNull;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+public class TestBucketPartitionerFlinkIcebergSink {
+
+ private static final int NUMBER_TASK_MANAGERS = 1;
+ private static final int SLOTS_PER_TASK_MANAGER = 8;
+
+ @RegisterExtension
+ private static final MiniClusterExtension MINI_CLUSTER_RESOURCE =
+ new MiniClusterExtension(
+ new MiniClusterResourceConfiguration.Builder()
+ .setNumberTaskManagers(NUMBER_TASK_MANAGERS)
+ .setNumberSlotsPerTaskManager(SLOTS_PER_TASK_MANAGER)
+ .setConfiguration(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .build());
+
+ @RegisterExtension
+ private static final HadoopCatalogExtension catalogExtension =
+ new HadoopCatalogExtension(DATABASE, TestFixtures.TABLE);
+
+ private static final TypeInformation<Row> ROW_TYPE_INFO =
+ new RowTypeInfo(SimpleDataUtil.FLINK_SCHEMA.getFieldTypes());
+
+ // Parallelism = 8 (parallelism > numBuckets) throughout the test suite
+ private final int parallelism = NUMBER_TASK_MANAGERS * SLOTS_PER_TASK_MANAGER;
+ private final FileFormat format = FileFormat.PARQUET;
+ private final int numBuckets = 4;
+
+ private Table table;
+ private StreamExecutionEnvironment env;
+ private TableLoader tableLoader;
+
+ private void setupEnvironment(TableSchemaType tableSchemaType) {
+ table = getTable(tableSchemaType);
+ env =
+ StreamExecutionEnvironment.getExecutionEnvironment(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .enableCheckpointing(100)
+ .setParallelism(parallelism)
+ .setMaxParallelism(parallelism * 2);
+ tableLoader = catalogExtension.tableLoader();
+ }
+
+ private Table getTable(TableSchemaType tableSchemaType) {
+ PartitionSpec partitionSpec = TableSchemaType.getPartitionSpec(tableSchemaType, numBuckets);
+
+ return catalogExtension
+ .catalog()
+ .createTable(
+ TABLE_IDENTIFIER,
+ SimpleDataUtil.SCHEMA,
+ partitionSpec,
+ ImmutableMap.of(TableProperties.DEFAULT_FILE_FORMAT, format.name()));
+ }
+
+ private void testWriteRowData(List<RowData> allRows) throws Exception {
+ DataFormatConverters.RowConverter converter =
+ new DataFormatConverters.RowConverter(SimpleDataUtil.FLINK_SCHEMA.getFieldDataTypes());
+
+ DataStream<RowData> dataStream =
+ env.addSource(
+ new BoundedTestSource<>(
+ allRows.stream().map(converter::toExternal).toArray(Row[]::new)),
+ ROW_TYPE_INFO)
+ .map(converter::toInternal, FlinkCompatibilityUtil.toTypeInfo(SimpleDataUtil.ROW_TYPE));
+
+ FlinkSink.forRowData(dataStream)
+ .table(table)
+ .tableLoader(tableLoader)
+ .writeParallelism(parallelism)
+ .distributionMode(DistributionMode.HASH)
+ .append();
+
+ // Execute the program.
+ env.execute("Test Iceberg DataStream");
+
+ // Assert the iceberg table's records.
Review Comment:
remove too obvious comment
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerFlinkIcebergSink.java:
##########
@@ -0,0 +1,253 @@
+/*
+ * 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.apache.iceberg.flink.MiniClusterResource.DISABLE_CLASSLOADER_CHECK_CONFIG;
+import static org.apache.iceberg.flink.TestFixtures.DATABASE;
+import static org.apache.iceberg.flink.TestFixtures.TABLE_IDENTIFIER;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.TreeMap;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.typeutils.RowTypeInfo;
+import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.data.util.DataFormatConverters;
+import org.apache.flink.test.junit5.MiniClusterExtension;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.DistributionMode;
+import org.apache.iceberg.FileFormat;
+import org.apache.iceberg.FileScanTask;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Table;
+import org.apache.iceberg.TableProperties;
+import org.apache.iceberg.flink.HadoopCatalogExtension;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.flink.TableLoader;
+import org.apache.iceberg.flink.TestFixtures;
+import org.apache.iceberg.flink.sink.TestBucketPartitionerUtils.TableSchemaType;
+import org.apache.iceberg.flink.source.BoundedTestSource;
+import org.apache.iceberg.flink.util.FlinkCompatibilityUtil;
+import org.apache.iceberg.io.CloseableIterable;
+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.assertj.core.api.Assertions;
+import org.jetbrains.annotations.NotNull;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+public class TestBucketPartitionerFlinkIcebergSink {
+
+ private static final int NUMBER_TASK_MANAGERS = 1;
+ private static final int SLOTS_PER_TASK_MANAGER = 8;
+
+ @RegisterExtension
+ private static final MiniClusterExtension MINI_CLUSTER_RESOURCE =
+ new MiniClusterExtension(
+ new MiniClusterResourceConfiguration.Builder()
+ .setNumberTaskManagers(NUMBER_TASK_MANAGERS)
+ .setNumberSlotsPerTaskManager(SLOTS_PER_TASK_MANAGER)
+ .setConfiguration(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .build());
+
+ @RegisterExtension
+ private static final HadoopCatalogExtension catalogExtension =
+ new HadoopCatalogExtension(DATABASE, TestFixtures.TABLE);
+
+ private static final TypeInformation<Row> ROW_TYPE_INFO =
+ new RowTypeInfo(SimpleDataUtil.FLINK_SCHEMA.getFieldTypes());
+
+ // Parallelism = 8 (parallelism > numBuckets) throughout the test suite
+ private final int parallelism = NUMBER_TASK_MANAGERS * SLOTS_PER_TASK_MANAGER;
+ private final FileFormat format = FileFormat.PARQUET;
+ private final int numBuckets = 4;
+
+ private Table table;
+ private StreamExecutionEnvironment env;
+ private TableLoader tableLoader;
+
+ private void setupEnvironment(TableSchemaType tableSchemaType) {
+ table = getTable(tableSchemaType);
+ env =
+ StreamExecutionEnvironment.getExecutionEnvironment(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .enableCheckpointing(100)
+ .setParallelism(parallelism)
+ .setMaxParallelism(parallelism * 2);
+ tableLoader = catalogExtension.tableLoader();
+ }
+
+ private Table getTable(TableSchemaType tableSchemaType) {
+ PartitionSpec partitionSpec = TableSchemaType.getPartitionSpec(tableSchemaType, numBuckets);
+
+ return catalogExtension
+ .catalog()
+ .createTable(
+ TABLE_IDENTIFIER,
+ SimpleDataUtil.SCHEMA,
+ partitionSpec,
+ ImmutableMap.of(TableProperties.DEFAULT_FILE_FORMAT, format.name()));
+ }
+
+ private void testWriteRowData(List<RowData> allRows) throws Exception {
+ DataFormatConverters.RowConverter converter =
+ new DataFormatConverters.RowConverter(SimpleDataUtil.FLINK_SCHEMA.getFieldDataTypes());
+
+ DataStream<RowData> dataStream =
+ env.addSource(
+ new BoundedTestSource<>(
+ allRows.stream().map(converter::toExternal).toArray(Row[]::new)),
+ ROW_TYPE_INFO)
+ .map(converter::toInternal, FlinkCompatibilityUtil.toTypeInfo(SimpleDataUtil.ROW_TYPE));
+
+ FlinkSink.forRowData(dataStream)
+ .table(table)
+ .tableLoader(tableLoader)
+ .writeParallelism(parallelism)
+ .distributionMode(DistributionMode.HASH)
+ .append();
+
+ // Execute the program.
+ env.execute("Test Iceberg DataStream");
+
+ // Assert the iceberg table's records.
+ SimpleDataUtil.assertTableRows(table, allRows);
+ }
+
+ @ParameterizedTest
+ @EnumSource(
+ value = TableSchemaType.class,
+ names = {"ONE_BUCKET", "IDENTITY_AND_BUCKET"})
+ public void testSendRecordsToAllBucketsEvenly(TableSchemaType tableSchemaType) throws Exception {
+ setupEnvironment(tableSchemaType);
+ List<RowData> rows = generateTestDataRows();
+
+ testWriteRowData(rows);
+ TableTestStats stats = extractPartitionResults(tableSchemaType);
+
+ Assertions.assertThat(stats.totalRowCount).isEqualTo(rows.size());
+ // All 4 buckets should've been written to
+ Assertions.assertThat(stats.writersPerBucket.size()).isEqualTo(numBuckets);
+ Assertions.assertThat(stats.numFilesPerBucket.size()).isEqualTo(numBuckets);
+ // Writer expectation (2 writers per bucket):
+ // - Bucket0 -> Writers [0, 4]
+ // - Bucket1 -> Writers [1, 5]
+ // - Bucket2 -> Writers [2, 6]
+ // - Bucket3 -> Writers [3, 7]
+ for (int i = 0, j = numBuckets; i < numBuckets; i++, j++) {
+ Assertions.assertThat(stats.writersPerBucket.get(i)).isEqualTo(Arrays.asList(i, j));
+ // 2 files per bucket (one file is created by each writer)
+ Assertions.assertThat((int) stats.numFilesPerBucket.get(i)).isEqualTo(2);
+ // 2 rows per file (total of 16 rows across 8 files)
+ Assertions.assertThat((long) stats.rowsPerFile.get(i)).isEqualTo(2);
+ }
+ }
+
+ /**
+ * Verifies the BucketPartitioner is not used when the PartitionSpec has more than 1 bucket, and
+ * that it should fallback to input.keyBy
+ */
+ @ParameterizedTest
+ @EnumSource(value = TableSchemaType.class, names = "TWO_BUCKETS")
+ public void testMultipleBucketsFallback(TableSchemaType tableSchemaType) throws Exception {
+ setupEnvironment(tableSchemaType);
+ List<RowData> rows = generateTestDataRows();
+
+ testWriteRowData(rows);
+ TableTestStats stats = extractPartitionResults(tableSchemaType);
+
+ Assertions.assertThat(stats.totalRowCount).isEqualTo(rows.size());
+ for (int i = 0, j = numBuckets; i < numBuckets; i++, j++) {
+ // Only 1 file per bucket will be created when falling back to input.keyBy(...)
+ Assertions.assertThat((int) stats.numFilesPerBucket.get(i)).isEqualTo(1);
+ }
+ }
+
+ /**
+ * Generating 16 rows to be sent uniformly to all writers (round-robin across 8 writers -> 4
+ * buckets)
+ */
+ @NotNull
+ private List<RowData> generateTestDataRows() {
+ int totalNumRows = parallelism * 2;
+ int numRowsPerBucket = totalNumRows / numBuckets;
+ return TestBucketPartitionerUtils.generateRowsForBucketIdRange(numRowsPerBucket, numBuckets);
+ }
+
+ private TableTestStats extractPartitionResults(TableSchemaType tableSchemaType)
+ throws IOException {
+ int totalRecordCount = 0;
+ Map<Integer, List<Integer>> writersPerBucket = Maps.newHashMap(); // <BucketId, List<WriterId>>
Review Comment:
the comments at the end of lines are unnecessary
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerFlinkIcebergSink.java:
##########
@@ -0,0 +1,253 @@
+/*
+ * 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.apache.iceberg.flink.MiniClusterResource.DISABLE_CLASSLOADER_CHECK_CONFIG;
+import static org.apache.iceberg.flink.TestFixtures.DATABASE;
+import static org.apache.iceberg.flink.TestFixtures.TABLE_IDENTIFIER;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.TreeMap;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.typeutils.RowTypeInfo;
+import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.data.util.DataFormatConverters;
+import org.apache.flink.test.junit5.MiniClusterExtension;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.DistributionMode;
+import org.apache.iceberg.FileFormat;
+import org.apache.iceberg.FileScanTask;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Table;
+import org.apache.iceberg.TableProperties;
+import org.apache.iceberg.flink.HadoopCatalogExtension;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.flink.TableLoader;
+import org.apache.iceberg.flink.TestFixtures;
+import org.apache.iceberg.flink.sink.TestBucketPartitionerUtils.TableSchemaType;
+import org.apache.iceberg.flink.source.BoundedTestSource;
+import org.apache.iceberg.flink.util.FlinkCompatibilityUtil;
+import org.apache.iceberg.io.CloseableIterable;
+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.assertj.core.api.Assertions;
+import org.jetbrains.annotations.NotNull;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+public class TestBucketPartitionerFlinkIcebergSink {
+
+ private static final int NUMBER_TASK_MANAGERS = 1;
+ private static final int SLOTS_PER_TASK_MANAGER = 8;
+
+ @RegisterExtension
+ private static final MiniClusterExtension MINI_CLUSTER_RESOURCE =
+ new MiniClusterExtension(
+ new MiniClusterResourceConfiguration.Builder()
+ .setNumberTaskManagers(NUMBER_TASK_MANAGERS)
+ .setNumberSlotsPerTaskManager(SLOTS_PER_TASK_MANAGER)
+ .setConfiguration(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .build());
+
+ @RegisterExtension
+ private static final HadoopCatalogExtension catalogExtension =
+ new HadoopCatalogExtension(DATABASE, TestFixtures.TABLE);
+
+ private static final TypeInformation<Row> ROW_TYPE_INFO =
+ new RowTypeInfo(SimpleDataUtil.FLINK_SCHEMA.getFieldTypes());
+
+ // Parallelism = 8 (parallelism > numBuckets) throughout the test suite
+ private final int parallelism = NUMBER_TASK_MANAGERS * SLOTS_PER_TASK_MANAGER;
+ private final FileFormat format = FileFormat.PARQUET;
+ private final int numBuckets = 4;
+
+ private Table table;
+ private StreamExecutionEnvironment env;
+ private TableLoader tableLoader;
+
+ private void setupEnvironment(TableSchemaType tableSchemaType) {
+ table = getTable(tableSchemaType);
+ env =
+ StreamExecutionEnvironment.getExecutionEnvironment(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .enableCheckpointing(100)
+ .setParallelism(parallelism)
+ .setMaxParallelism(parallelism * 2);
+ tableLoader = catalogExtension.tableLoader();
+ }
+
+ private Table getTable(TableSchemaType tableSchemaType) {
+ PartitionSpec partitionSpec = TableSchemaType.getPartitionSpec(tableSchemaType, numBuckets);
+
+ return catalogExtension
+ .catalog()
+ .createTable(
+ TABLE_IDENTIFIER,
+ SimpleDataUtil.SCHEMA,
+ partitionSpec,
+ ImmutableMap.of(TableProperties.DEFAULT_FILE_FORMAT, format.name()));
+ }
+
+ private void testWriteRowData(List<RowData> allRows) throws Exception {
+ DataFormatConverters.RowConverter converter =
+ new DataFormatConverters.RowConverter(SimpleDataUtil.FLINK_SCHEMA.getFieldDataTypes());
+
+ DataStream<RowData> dataStream =
+ env.addSource(
+ new BoundedTestSource<>(
+ allRows.stream().map(converter::toExternal).toArray(Row[]::new)),
+ ROW_TYPE_INFO)
+ .map(converter::toInternal, FlinkCompatibilityUtil.toTypeInfo(SimpleDataUtil.ROW_TYPE));
+
+ FlinkSink.forRowData(dataStream)
+ .table(table)
+ .tableLoader(tableLoader)
+ .writeParallelism(parallelism)
+ .distributionMode(DistributionMode.HASH)
+ .append();
+
+ // Execute the program.
+ env.execute("Test Iceberg DataStream");
+
+ // Assert the iceberg table's records.
+ SimpleDataUtil.assertTableRows(table, allRows);
+ }
+
+ @ParameterizedTest
+ @EnumSource(
+ value = TableSchemaType.class,
+ names = {"ONE_BUCKET", "IDENTITY_AND_BUCKET"})
+ public void testSendRecordsToAllBucketsEvenly(TableSchemaType tableSchemaType) throws Exception {
+ setupEnvironment(tableSchemaType);
+ List<RowData> rows = generateTestDataRows();
+
+ testWriteRowData(rows);
+ TableTestStats stats = extractPartitionResults(tableSchemaType);
+
+ Assertions.assertThat(stats.totalRowCount).isEqualTo(rows.size());
+ // All 4 buckets should've been written to
+ Assertions.assertThat(stats.writersPerBucket.size()).isEqualTo(numBuckets);
+ Assertions.assertThat(stats.numFilesPerBucket.size()).isEqualTo(numBuckets);
+ // Writer expectation (2 writers per bucket):
+ // - Bucket0 -> Writers [0, 4]
+ // - Bucket1 -> Writers [1, 5]
+ // - Bucket2 -> Writers [2, 6]
+ // - Bucket3 -> Writers [3, 7]
+ for (int i = 0, j = numBuckets; i < numBuckets; i++, j++) {
+ Assertions.assertThat(stats.writersPerBucket.get(i)).isEqualTo(Arrays.asList(i, j));
+ // 2 files per bucket (one file is created by each writer)
+ Assertions.assertThat((int) stats.numFilesPerBucket.get(i)).isEqualTo(2);
Review Comment:
why do we need the type cast?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] stevenzwu commented on a diff in pull request #7161: Flink: Custom partitioner for bucket partitions
Posted by "stevenzwu (via GitHub)" <gi...@apache.org>.
stevenzwu commented on code in PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#discussion_r1232755744
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitionKeySelector.java:
##########
@@ -0,0 +1,66 @@
+/*
+ * 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 java.util.stream.IntStream;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.types.logical.RowType;
+import org.apache.iceberg.PartitionKey;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Schema;
+import org.apache.iceberg.flink.RowDataWrapper;
+
+/**
+ * A {@link KeySelector} that extracts the bucketId from a data row's bucket partition as the key.
+ * To be used with the {@link BucketPartitioner}.
+ */
+class BucketPartitionKeySelector implements KeySelector<RowData, Integer> {
+
+ private final Schema schema;
+ private final PartitionKey partitionKey;
+ private final RowType flinkSchema;
+ private final int bucketFieldPosition;
+
+ private transient RowDataWrapper rowDataWrapper;
+
+ BucketPartitionKeySelector(PartitionSpec partitionSpec, Schema schema, RowType flinkSchema) {
+ this.schema = schema;
+ this.partitionKey = new PartitionKey(partitionSpec, schema);
+ this.flinkSchema = flinkSchema;
+ int bucketFieldId = BucketPartitionerUtil.getBucketFieldInfo(partitionSpec).f0;
+ this.bucketFieldPosition =
Review Comment:
can we combine this with the line above in a method called `getBucketFieldPosition`?
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerFlinkIcebergSink.java:
##########
@@ -0,0 +1,253 @@
+/*
+ * 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.apache.iceberg.flink.MiniClusterResource.DISABLE_CLASSLOADER_CHECK_CONFIG;
+import static org.apache.iceberg.flink.TestFixtures.DATABASE;
+import static org.apache.iceberg.flink.TestFixtures.TABLE_IDENTIFIER;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.TreeMap;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.typeutils.RowTypeInfo;
+import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.data.util.DataFormatConverters;
+import org.apache.flink.test.junit5.MiniClusterExtension;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.DistributionMode;
+import org.apache.iceberg.FileFormat;
+import org.apache.iceberg.FileScanTask;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Table;
+import org.apache.iceberg.TableProperties;
+import org.apache.iceberg.flink.HadoopCatalogExtension;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.flink.TableLoader;
+import org.apache.iceberg.flink.TestFixtures;
+import org.apache.iceberg.flink.sink.TestBucketPartitionerUtils.TableSchemaType;
+import org.apache.iceberg.flink.source.BoundedTestSource;
+import org.apache.iceberg.flink.util.FlinkCompatibilityUtil;
+import org.apache.iceberg.io.CloseableIterable;
+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.assertj.core.api.Assertions;
+import org.jetbrains.annotations.NotNull;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+public class TestBucketPartitionerFlinkIcebergSink {
+
+ private static final int NUMBER_TASK_MANAGERS = 1;
+ private static final int SLOTS_PER_TASK_MANAGER = 8;
+
+ @RegisterExtension
+ private static final MiniClusterExtension MINI_CLUSTER_RESOURCE =
+ new MiniClusterExtension(
+ new MiniClusterResourceConfiguration.Builder()
+ .setNumberTaskManagers(NUMBER_TASK_MANAGERS)
+ .setNumberSlotsPerTaskManager(SLOTS_PER_TASK_MANAGER)
+ .setConfiguration(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .build());
+
+ @RegisterExtension
+ private static final HadoopCatalogExtension catalogExtension =
+ new HadoopCatalogExtension(DATABASE, TestFixtures.TABLE);
+
+ private static final TypeInformation<Row> ROW_TYPE_INFO =
+ new RowTypeInfo(SimpleDataUtil.FLINK_SCHEMA.getFieldTypes());
+
+ // Parallelism = 8 (parallelism > numBuckets) throughout the test suite
+ private final int parallelism = NUMBER_TASK_MANAGERS * SLOTS_PER_TASK_MANAGER;
+ private final FileFormat format = FileFormat.PARQUET;
+ private final int numBuckets = 4;
+
+ private Table table;
+ private StreamExecutionEnvironment env;
+ private TableLoader tableLoader;
+
+ private void setupEnvironment(TableSchemaType tableSchemaType) {
+ table = getTable(tableSchemaType);
+ env =
+ StreamExecutionEnvironment.getExecutionEnvironment(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .enableCheckpointing(100)
+ .setParallelism(parallelism)
+ .setMaxParallelism(parallelism * 2);
+ tableLoader = catalogExtension.tableLoader();
+ }
+
+ private Table getTable(TableSchemaType tableSchemaType) {
+ PartitionSpec partitionSpec = TableSchemaType.getPartitionSpec(tableSchemaType, numBuckets);
+
+ return catalogExtension
+ .catalog()
+ .createTable(
+ TABLE_IDENTIFIER,
+ SimpleDataUtil.SCHEMA,
+ partitionSpec,
+ ImmutableMap.of(TableProperties.DEFAULT_FILE_FORMAT, format.name()));
+ }
+
+ private void testWriteRowData(List<RowData> allRows) throws Exception {
+ DataFormatConverters.RowConverter converter =
+ new DataFormatConverters.RowConverter(SimpleDataUtil.FLINK_SCHEMA.getFieldDataTypes());
+
+ DataStream<RowData> dataStream =
+ env.addSource(
+ new BoundedTestSource<>(
+ allRows.stream().map(converter::toExternal).toArray(Row[]::new)),
Review Comment:
you can just add the type info. it should take care of the serialization part.
```
DataStream<RowData> stream =
env.fromSource(
IcebergSource.<RowData>builder()
.tableLoader(tableLoader)
.assignerFactory(new SimpleSplitAssignerFactory())
.project(projected)
.build(),
WatermarkStrategy.noWatermarks(),
"testBasicRead",
TypeInformation.of(RowData.class));
```
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitioner.java:
##########
@@ -0,0 +1,107 @@
+/*
+ * 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.apache.iceberg.flink.sink.BucketPartitioner.BUCKET_GREATER_THAN_UPPER_BOUND_MESSAGE;
+import static org.apache.iceberg.flink.sink.BucketPartitioner.BUCKET_LESS_THAN_LOWER_BOUND_MESSAGE;
+import static org.apache.iceberg.flink.sink.BucketPartitioner.BUCKET_NULL_MESSAGE;
+
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.flink.sink.TestBucketPartitionerUtil.TableSchemaType;
+import org.assertj.core.api.Assertions;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.CsvSource;
+
+public class TestBucketPartitioner {
+
+ static final int DEFAULT_NUM_BUCKETS = 60;
+
+ @ParameterizedTest
+ @CsvSource({"ONE_BUCKET,50", "IDENTITY_AND_BUCKET,50", "ONE_BUCKET,60", "IDENTITY_AND_BUCKET,60"})
+ public void testPartitioningParallelismGreaterThanBuckets(
+ String schemaTypeStr, String numBucketsStr) {
+ int numPartitions = 500;
+ TableSchemaType tableSchemaType = TableSchemaType.valueOf(schemaTypeStr);
+ int numBuckets = Integer.parseInt(numBucketsStr);
+ PartitionSpec partitionSpec = tableSchemaType.getPartitionSpec(numBuckets);
+ BucketPartitioner bucketPartitioner = new BucketPartitioner(partitionSpec);
+
+ int bucketId = 0;
+ for (int expectedIdx = 0; expectedIdx < numPartitions; expectedIdx++) {
Review Comment:
nit: idx -> index
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/HadoopCatalogExtension.java:
##########
@@ -0,0 +1,104 @@
+/*
+ * 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;
+
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.UUID;
+import org.apache.commons.io.FileUtils;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.iceberg.CatalogProperties;
+import org.apache.iceberg.catalog.Catalog;
+import org.apache.iceberg.catalog.TableIdentifier;
+import org.apache.iceberg.hadoop.HadoopCatalog;
+import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
+import org.assertj.core.api.Assertions;
+import org.junit.jupiter.api.extension.AfterAllCallback;
+import org.junit.jupiter.api.extension.AfterEachCallback;
+import org.junit.jupiter.api.extension.BeforeAllCallback;
+import org.junit.jupiter.api.extension.BeforeEachCallback;
+import org.junit.jupiter.api.extension.ExtensionContext;
+
+public class HadoopCatalogExtension
+ implements BeforeAllCallback, BeforeEachCallback, AfterAllCallback, AfterEachCallback {
+ protected final String database;
+ protected final String tableName;
+
+ protected Path temporaryFolder;
+ protected Catalog catalog;
+ protected CatalogLoader catalogLoader;
+ protected String warehouse;
+ protected TableLoader tableLoader;
+
+ public HadoopCatalogExtension(String database, String tableName) {
+ this.database = database;
+ this.tableName = tableName;
+ }
+
+ @Override
+ public void beforeAll(ExtensionContext context) throws Exception {
+ this.temporaryFolder = Files.createTempDirectory("junit5_hadoop_catalog-");
+ }
+
+ @Override
+ public void afterAll(ExtensionContext context) throws Exception {
+ FileUtils.deleteDirectory(this.temporaryFolder.toFile());
Review Comment:
nit: iceberg style only uses `this.` when setting/changing values (not during read). this applies to a few other places in the main code.
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitioner.java:
##########
@@ -0,0 +1,107 @@
+/*
+ * 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 org.apache.flink.api.common.functions.Partitioner;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/**
+ * This partitioner will redirect records to writers deterministically based on the Bucket partition
+ * spec. It'll attempt to optimize the file size written depending on whether numPartitions is
+ * greater, less or equal than the maxNumBuckets. Note: The current implementation only supports ONE
+ * bucket in the partition spec.
+ */
+class BucketPartitioner implements Partitioner<Integer> {
+
+ static final String BUCKET_NULL_MESSAGE = "bucketId cannot be null";
+ static final String BUCKET_LESS_THAN_LOWER_BOUND_MESSAGE =
+ "Invalid bucket ID %s: must be non-negative.";
+ static final String BUCKET_GREATER_THAN_UPPER_BOUND_MESSAGE =
+ "Invalid bucket ID %s: must be less than bucket limit: %s.";
+
+ private final int maxNumBuckets;
+
+ // To hold the OFFSET of the next writer to use for any bucket, only used when writers > the
+ // number of buckets
+ private final int[] currentBucketWriterOffset;
+
+ BucketPartitioner(PartitionSpec partitionSpec) {
+ Tuple2<Integer, Integer> bucketFieldInfo =
+ BucketPartitionerUtil.getBucketFieldInfo(partitionSpec);
Review Comment:
it would probably be a bit cleaned to expose a method `getBucketNumber(PartitionSpec)`. then we can avoid the tuple of `.f1`
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/HadoopCatalogExtension.java:
##########
@@ -0,0 +1,104 @@
+/*
+ * 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;
+
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.UUID;
+import org.apache.commons.io.FileUtils;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.iceberg.CatalogProperties;
+import org.apache.iceberg.catalog.Catalog;
+import org.apache.iceberg.catalog.TableIdentifier;
+import org.apache.iceberg.hadoop.HadoopCatalog;
+import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
+import org.assertj.core.api.Assertions;
+import org.junit.jupiter.api.extension.AfterAllCallback;
+import org.junit.jupiter.api.extension.AfterEachCallback;
+import org.junit.jupiter.api.extension.BeforeAllCallback;
+import org.junit.jupiter.api.extension.BeforeEachCallback;
+import org.junit.jupiter.api.extension.ExtensionContext;
+
+public class HadoopCatalogExtension
+ implements BeforeAllCallback, BeforeEachCallback, AfterAllCallback, AfterEachCallback {
+ protected final String database;
+ protected final String tableName;
+
+ protected Path temporaryFolder;
Review Comment:
that's fair
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] stevenzwu commented on a diff in pull request #7161: Flink: Custom partitioner for bucket partitions
Posted by "stevenzwu (via GitHub)" <gi...@apache.org>.
stevenzwu commented on code in PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#discussion_r1233385124
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitionerUtils.java:
##########
@@ -0,0 +1,117 @@
+/*
+ * 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 java.util.List;
+import java.util.Objects;
+import java.util.stream.Collectors;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+import org.apache.iceberg.transforms.PartitionSpecVisitor;
+
+final class BucketPartitionerUtils {
+ static final String BAD_NUMBER_OF_BUCKETS_ERROR_MESSAGE =
Review Comment:
maybe we can use the same style as mentioned in the other comment
```
Preconditions.checkArgument(
numBuckets > 0, "Invalid number of buckets: %s (must be > 0)", numBuckets);
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] kengtin commented on a diff in pull request #7161: Flink: Custom partitioner for bucket partitions
Posted by "kengtin (via GitHub)" <gi...@apache.org>.
kengtin commented on code in PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#discussion_r1210920983
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerUtils.java:
##########
@@ -0,0 +1,120 @@
+/*
+ * 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 java.util.List;
+import java.util.UUID;
+import org.apache.flink.table.data.GenericRowData;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.data.StringData;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.apache.iceberg.util.BucketUtil;
+
+final class TestBucketPartitionerUtils {
+
+ enum TableSchemaType {
+ ONE_BUCKET,
+ IDENTITY_AND_BUCKET,
+ TWO_BUCKETS;
+
+ public static int bucketPartitionColumnPosition(TableSchemaType tableSchemaType) {
Review Comment:
The refactoring to non-static is done
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] stevenzwu merged pull request #7161: Flink: Custom partitioner for bucket partitions
Posted by "stevenzwu (via GitHub)" <gi...@apache.org>.
stevenzwu merged PR #7161:
URL: https://github.com/apache/iceberg/pull/7161
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] kengtin commented on a diff in pull request #7161: Flink: Custom partitioner for bucket partitions
Posted by "kengtin (via GitHub)" <gi...@apache.org>.
kengtin commented on code in PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#discussion_r1197847372
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitioner.java:
##########
@@ -0,0 +1,99 @@
+/*
+ * 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.apache.iceberg.flink.sink.BucketPartitioner.BUCKET_OUT_OF_RANGE_MESSAGE_PREFIX;
+
+import org.apache.iceberg.PartitionSpec;
+import org.assertj.core.api.Assertions;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+import org.junit.jupiter.params.provider.ValueSource;
+
+public class TestBucketPartitioner {
+
+ static final int NUM_BUCKETS = 60;
+
+ @ParameterizedTest
+ @EnumSource(
+ value = TestBucketPartitionerUtils.TableSchemaType.class,
+ names = {"ONE_BUCKET", "IDENTITY_AND_BUCKET"})
+ public void testPartitioningParallelismGreaterThanBuckets(
+ TestBucketPartitionerUtils.TableSchemaType tableSchemaType) {
+ final int numPartitions = 500;
+
+ PartitionSpec partitionSpec =
+ TestBucketPartitionerUtils.getPartitionSpec(tableSchemaType, NUM_BUCKETS);
+
+ BucketPartitioner bucketPartitioner = new BucketPartitioner(partitionSpec);
+
+ for (int expectedIdx = 0, bucketId = 0; expectedIdx < numPartitions; expectedIdx++) {
+ int actualIdx = bucketPartitioner.partition(bucketId, numPartitions);
+ Assertions.assertThat(actualIdx).isEqualTo(expectedIdx);
+ if (++bucketId == NUM_BUCKETS) {
+ bucketId = 0;
+ }
+ }
+ }
+
+ @ParameterizedTest
+ @EnumSource(
+ value = TestBucketPartitionerUtils.TableSchemaType.class,
+ names = {"ONE_BUCKET", "IDENTITY_AND_BUCKET"})
+ public void testPartitioningParallelismEqualLessThanBuckets(
+ TestBucketPartitionerUtils.TableSchemaType tableSchemaType) {
+ final int numPartitions = 30;
+
+ PartitionSpec partitionSpec =
+ TestBucketPartitionerUtils.getPartitionSpec(tableSchemaType, NUM_BUCKETS);
+
+ BucketPartitioner bucketPartitioner = new BucketPartitioner(partitionSpec);
+
+ for (int bucketId = 0; bucketId < NUM_BUCKETS; bucketId++) {
+ int actualIdx = bucketPartitioner.partition(bucketId, numPartitions);
+ Assertions.assertThat(actualIdx).isEqualTo(bucketId % numPartitions);
+ }
+ }
+
+ @ParameterizedTest
+ @EnumSource(value = TestBucketPartitionerUtils.TableSchemaType.class, names = "TWO_BUCKETS")
+ public void testPartitionerMultipleBucketsFail(
+ TestBucketPartitionerUtils.TableSchemaType tableSchemaType) {
+ PartitionSpec partitionSpec =
+ TestBucketPartitionerUtils.getPartitionSpec(tableSchemaType, NUM_BUCKETS);
+
+ Assertions.assertThatExceptionOfType(RuntimeException.class)
+ .isThrownBy(() -> new BucketPartitioner(partitionSpec))
+ .withMessageContaining(BucketPartitionerUtils.BAD_NUMBER_OF_BUCKETS_ERROR_MESSAGE);
+ }
+
+ @ParameterizedTest
Review Comment:
I'm evaluating the 2 invalid scenarios when the bucketId is out of range. It seems a bit cleaner than evaluating `.partition` twice directly.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] kengtin commented on a diff in pull request #7161: Flink: Custom partitioner for a more even data distribution when bucket partitions are present
Posted by "kengtin (via GitHub)" <gi...@apache.org>.
kengtin commented on code in PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#discussion_r1159004896
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitioner.java:
##########
@@ -0,0 +1,86 @@
+/*
+ * 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 java.util.Optional;
+import org.apache.flink.api.common.functions.Partitioner;
+import org.apache.iceberg.PartitionField;
+import org.apache.iceberg.PartitionSpec;
+
+/**
+ * This partitioner will redirect elements to writers deterministically so that each writer only
+ * targets 1 bucket. If the number of writers > number of buckets each partitioner will keep a state
+ * of multiple writers per bucket as evenly as possible, and will round-robin the requests across
+ * them.
+ */
+public class BucketPartitioner implements Partitioner<Integer> {
+
+ private final Integer maxBuckets;
+
+ private final int[] currentWriterOffset;
+
+ BucketPartitioner(PartitionSpec partitionSpec) {
+ // The current implementation redirects to writers based on the _FIRST_ bucket found
+ Optional<PartitionField> bucket =
+ partitionSpec.fields().stream()
+ .filter(f -> f.transform().dedupName().contains("bucket"))
Review Comment:
Because `org.apache.iceberg.transforms.Bucket` is inaccessible from here (package private).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Flink: Custom partitioner for bucket partitions [iceberg]
Posted by "binshuohu (via GitHub)" <gi...@apache.org>.
binshuohu commented on PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#issuecomment-2347381724
@stevenzwu Is there any plan to reapply this change to the main branch? Has there been any follow up since https://github.com/apache/iceberg/pull/8848 ?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Flink: Custom partitioner for bucket partitions [iceberg]
Posted by "stevenzwu (via GitHub)" <gi...@apache.org>.
stevenzwu commented on PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#issuecomment-2350416841
@binshuohu Currently, there is no plan to reapply this change to the main branch. We have a more general range distribution available now (guided by statistics collection): https://iceberg.apache.org/docs/nightly/flink-writes/#range-distribution-experimental. It is more general than this (bucketing only). Range distribution also handle different parallelisms and partitions well.
Range distribution has one disadvantage. It performs statistics collection and aggregation to guide the range split. That adds a little overhead. Bucketing partitioner here assumes traffic are evenly distributed across buckets, which should be true (hash % nBuckets).
cc @pvary
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] stevenzwu commented on pull request #7161: Flink: Custom partitioner for bucket partitions
Posted by "stevenzwu (via GitHub)" <gi...@apache.org>.
stevenzwu commented on PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#issuecomment-1703073520
@advancedxy thanks for a lot for raising the multi-dimensional bucketing. if there is enough use cases, we should be able to adapt the implementation to support it under the assumption that data are evenly distributed. Let's use an example with 2 bucketing columns (m and n). The assumption would be data are evenly distributed in the `m x n` combinational bucket tuples.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] Flink: Custom partitioner for bucket partitions [iceberg]
Posted by "stevenzwu (via GitHub)" <gi...@apache.org>.
stevenzwu commented on PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#issuecomment-1862079413
@bendevera here is our presentation: https://www.youtube.com/watch?v=GJplmOO7ULA&t=18s. here is the design doc: https://docs.google.com/document/d/13N8cMqPi-ZPSKbkXGOBMPOzbv2Fua59j8bIjjtxLWqo/edit#heading=h.o4q8a61sahkq
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] kengtin commented on a diff in pull request #7161: Flink: Custom partitioner for bucket partitions
Posted by "kengtin (via GitHub)" <gi...@apache.org>.
kengtin commented on code in PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#discussion_r1197860233
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitioner.java:
##########
@@ -0,0 +1,99 @@
+/*
+ * 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.apache.iceberg.flink.sink.BucketPartitioner.BUCKET_OUT_OF_RANGE_MESSAGE_PREFIX;
+
+import org.apache.iceberg.PartitionSpec;
+import org.assertj.core.api.Assertions;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+import org.junit.jupiter.params.provider.ValueSource;
+
+public class TestBucketPartitioner {
+
+ static final int NUM_BUCKETS = 60;
+
+ @ParameterizedTest
+ @EnumSource(
+ value = TestBucketPartitionerUtils.TableSchemaType.class,
+ names = {"ONE_BUCKET", "IDENTITY_AND_BUCKET"})
+ public void testPartitioningParallelismGreaterThanBuckets(
+ TestBucketPartitionerUtils.TableSchemaType tableSchemaType) {
+ final int numPartitions = 500;
+
+ PartitionSpec partitionSpec =
+ TestBucketPartitionerUtils.getPartitionSpec(tableSchemaType, NUM_BUCKETS);
+
+ BucketPartitioner bucketPartitioner = new BucketPartitioner(partitionSpec);
+
+ for (int expectedIdx = 0, bucketId = 0; expectedIdx < numPartitions; expectedIdx++) {
+ int actualIdx = bucketPartitioner.partition(bucketId, numPartitions);
+ Assertions.assertThat(actualIdx).isEqualTo(expectedIdx);
+ if (++bucketId == NUM_BUCKETS) {
+ bucketId = 0;
+ }
+ }
+ }
+
+ @ParameterizedTest
+ @EnumSource(
+ value = TestBucketPartitionerUtils.TableSchemaType.class,
+ names = {"ONE_BUCKET", "IDENTITY_AND_BUCKET"})
+ public void testPartitioningParallelismEqualLessThanBuckets(
+ TestBucketPartitionerUtils.TableSchemaType tableSchemaType) {
+ final int numPartitions = 30;
+
+ PartitionSpec partitionSpec =
+ TestBucketPartitionerUtils.getPartitionSpec(tableSchemaType, NUM_BUCKETS);
+
+ BucketPartitioner bucketPartitioner = new BucketPartitioner(partitionSpec);
+
+ for (int bucketId = 0; bucketId < NUM_BUCKETS; bucketId++) {
+ int actualIdx = bucketPartitioner.partition(bucketId, numPartitions);
+ Assertions.assertThat(actualIdx).isEqualTo(bucketId % numPartitions);
+ }
+ }
+
+ @ParameterizedTest
+ @EnumSource(value = TestBucketPartitionerUtils.TableSchemaType.class, names = "TWO_BUCKETS")
+ public void testPartitionerMultipleBucketsFail(
+ TestBucketPartitionerUtils.TableSchemaType tableSchemaType) {
+ PartitionSpec partitionSpec =
+ TestBucketPartitionerUtils.getPartitionSpec(tableSchemaType, NUM_BUCKETS);
+
+ Assertions.assertThatExceptionOfType(RuntimeException.class)
+ .isThrownBy(() -> new BucketPartitioner(partitionSpec))
+ .withMessageContaining(BucketPartitionerUtils.BAD_NUMBER_OF_BUCKETS_ERROR_MESSAGE);
+ }
+
+ @ParameterizedTest
Review Comment:
I ended up modifying it to better address the comment below
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] stevenzwu commented on a diff in pull request #7161: Flink: Custom partitioner for bucket partitions
Posted by "stevenzwu (via GitHub)" <gi...@apache.org>.
stevenzwu commented on code in PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#discussion_r1186975564
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitioner.java:
##########
@@ -0,0 +1,108 @@
+/*
+ * 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 org.apache.flink.api.common.functions.Partitioner;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/**
+ * This partitioner will redirect records to writers deterministically based on the Bucket partition
+ * spec. It'll attempt to optimize the file size written depending on whether numPartitions is
+ * greater, less or equal than the maxNumBuckets. Note: The current implementation only supports ONE
+ * bucket in the partition spec.
+ */
+class BucketPartitioner implements Partitioner<Integer> {
+
+ static final String BUCKET_OUT_OF_RANGE_MESSAGE_PREFIX = "bucketId out of range: ";
+
+ private final int maxNumBuckets;
+
+ // To hold the OFFSET of the next writer to use for any bucket
+ private final int[] currentBucketWriterOffset;
+
+ BucketPartitioner(PartitionSpec partitionSpec) {
+ Tuple2<Integer, Integer> bucketFieldInfo =
+ BucketPartitionerUtils.getBucketFieldInfo(partitionSpec);
+
+ this.maxNumBuckets = bucketFieldInfo.f1;
+ this.currentBucketWriterOffset = new int[this.maxNumBuckets];
+ }
+
+ /**
+ * Determine the partition id based on the following criteria: If the number of writers <= the
+ * number of buckets, an evenly distributed number of buckets will be assigned to each writer (one
+ * writer -> many buckets). Conversely, if the number of writers > the number of buckets the logic
+ * is handled by the {@link #getPartitionWritersGreaterThanBuckets
+ * getPartitionWritersGreaterThanBuckets} method.
+ *
+ * @param bucketId the bucketId for each request
+ * @param numPartitions the total number of partitions
+ * @return the partition id (writer) to use for each request
+ */
+ @Override
+ public int partition(Integer bucketId, int numPartitions) {
+ Preconditions.checkArgument(
+ bucketId >= 0, BUCKET_OUT_OF_RANGE_MESSAGE_PREFIX + bucketId + " (must be >= 0)");
Review Comment:
Nit: error msg could be `Invalid bucket ID: %s. Must be non-negative.`. note that `Preconditions` supports arg format.
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitionKeySelector.java:
##########
@@ -0,0 +1,74 @@
+/*
+ * 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 java.util.stream.IntStream;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.types.logical.RowType;
+import org.apache.iceberg.PartitionKey;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Schema;
+import org.apache.iceberg.flink.RowDataWrapper;
+
+/**
+ * A {@link KeySelector} that extracts the bucketId from a data row's bucket partition as the key.
+ * To be used with the {@link BucketPartitioner}.
+ */
+class BucketPartitionKeySelector implements KeySelector<RowData, Integer> {
+
+ private final Schema schema;
+ private final PartitionKey partitionKey;
+ private final RowType flinkSchema;
+ private final int bucketFieldPosition;
+
+ private transient RowDataWrapper rowDataWrapper;
+
+ BucketPartitionKeySelector(PartitionSpec partitionSpec, Schema schema, RowType flinkSchema) {
+ Tuple2<Integer, Integer> bucketFieldInfo =
+ BucketPartitionerUtils.getBucketFieldInfo(partitionSpec);
+
+ int bucketFieldId = bucketFieldInfo.f0;
+ this.schema = schema;
+ this.partitionKey = new PartitionKey(partitionSpec, schema);
+ this.flinkSchema = flinkSchema;
+ this.bucketFieldPosition =
+ IntStream.range(0, partitionSpec.fields().size())
+ .filter(i -> partitionSpec.fields().get(i).fieldId() == bucketFieldId)
+ .toArray()[0];
+ }
+
+ /**
Review Comment:
nit: this comment is probably not necessary. also tries to avoid use words like `we`
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerUtils.java:
##########
@@ -0,0 +1,121 @@
+/*
+ * 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 java.util.List;
+import java.util.UUID;
+import org.apache.flink.table.data.util.DataFormatConverters;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.apache.iceberg.util.BucketUtil;
+
+final class TestBucketPartitionerUtils {
+
+ enum TableSchemaType {
+ ONE_BUCKET,
+ IDENTITY_AND_BUCKET,
+ TWO_BUCKETS;
+ }
+
+ private TestBucketPartitionerUtils() {}
+
+ static final DataFormatConverters.RowConverter CONVERTER =
+ new DataFormatConverters.RowConverter(SimpleDataUtil.FLINK_SCHEMA.getFieldDataTypes());
+
+ static PartitionSpec getPartitionSpec(TableSchemaType tableSchemaType, int numBuckets) {
+ PartitionSpec partitionSpec = null;
+
+ switch (tableSchemaType) {
+ case ONE_BUCKET:
+ partitionSpec =
+ PartitionSpec.builderFor(SimpleDataUtil.SCHEMA).bucket("data", numBuckets).build();
+ break;
+ case IDENTITY_AND_BUCKET:
+ partitionSpec =
+ PartitionSpec.builderFor(SimpleDataUtil.SCHEMA)
+ .identity("id")
+ .bucket("data", numBuckets)
+ .build();
+ break;
+ case TWO_BUCKETS:
+ partitionSpec =
+ PartitionSpec.builderFor(SimpleDataUtil.SCHEMA)
+ .bucket("id", numBuckets)
+ .bucket("data", numBuckets)
+ .build();
+ break;
+ }
+
+ Preconditions.checkNotNull(
+ partitionSpec, "Invalid tableSchemaType provided: " + tableSchemaType);
+ return partitionSpec;
+ }
+
+ /**
+ * Utility method to generate rows whose values will "hash" to a range of bucketIds (from 0 to
+ * numBuckets - 1)
+ *
+ * @param numRowsPerBucket how many different rows should be generated per bucket
+ * @param numBuckets max number of buckets to consider
+ * @return the list of rows whose data "hashes" to the desired bucketId
+ */
+ static List<Row> generateRowsForBucketIdRange(int numRowsPerBucket, int numBuckets) {
+ List<Row> rows = Lists.newArrayListWithCapacity(numBuckets * numRowsPerBucket);
+ // For some of our tests, this order of the generated rows matters
+ for (int i = 0; i < numRowsPerBucket; i++) {
+ for (int bucketId = 0; bucketId < numBuckets; bucketId++) {
+ String value = generateValueForBucketId(bucketId, numBuckets);
+ rows.add(Row.of(1, value));
+ }
+ }
+ return rows;
+ }
+
+ /**
+ * Utility method to generate a UUID string that will "hash" to a desired bucketId
+ *
+ * @param bucketId the desired bucketId
+ * @return the string data that "hashes" to the desired bucketId
+ */
+ private static String generateValueForBucketId(int bucketId, int numBuckets) {
+ String value = "";
+ while (true) {
+ String uuid = UUID.randomUUID().toString();
+ if (computeBucketId(numBuckets, uuid) == bucketId) {
+ value = uuid;
+ break;
+ }
+ }
+ return value;
+ }
+
+ /**
+ * Utility that performs the same hashing/bucketing mechanism used by Bucket.java
Review Comment:
looking at Bucket.java. we also need to add test coverage for null bucket id.
```
public Integer apply(T value) {
if (value == null) {
return null;
}
return (hash(value) & Integer.MAX_VALUE) % numBuckets;
}
```
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitioner.java:
##########
@@ -0,0 +1,108 @@
+/*
+ * 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 org.apache.flink.api.common.functions.Partitioner;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/**
+ * This partitioner will redirect records to writers deterministically based on the Bucket partition
+ * spec. It'll attempt to optimize the file size written depending on whether numPartitions is
+ * greater, less or equal than the maxNumBuckets. Note: The current implementation only supports ONE
+ * bucket in the partition spec.
+ */
+class BucketPartitioner implements Partitioner<Integer> {
+
+ static final String BUCKET_OUT_OF_RANGE_MESSAGE_PREFIX = "bucketId out of range: ";
+
+ private final int maxNumBuckets;
+
+ // To hold the OFFSET of the next writer to use for any bucket
+ private final int[] currentBucketWriterOffset;
+
+ BucketPartitioner(PartitionSpec partitionSpec) {
+ Tuple2<Integer, Integer> bucketFieldInfo =
+ BucketPartitionerUtils.getBucketFieldInfo(partitionSpec);
+
+ this.maxNumBuckets = bucketFieldInfo.f1;
+ this.currentBucketWriterOffset = new int[this.maxNumBuckets];
+ }
+
+ /**
+ * Determine the partition id based on the following criteria: If the number of writers <= the
+ * number of buckets, an evenly distributed number of buckets will be assigned to each writer (one
+ * writer -> many buckets). Conversely, if the number of writers > the number of buckets the logic
+ * is handled by the {@link #getPartitionWritersGreaterThanBuckets
+ * getPartitionWritersGreaterThanBuckets} method.
+ *
+ * @param bucketId the bucketId for each request
+ * @param numPartitions the total number of partitions
+ * @return the partition id (writer) to use for each request
+ */
+ @Override
+ public int partition(Integer bucketId, int numPartitions) {
+ Preconditions.checkArgument(
+ bucketId >= 0, BUCKET_OUT_OF_RANGE_MESSAGE_PREFIX + bucketId + " (must be >= 0)");
+ Preconditions.checkArgument(
+ bucketId < maxNumBuckets,
+ BUCKET_OUT_OF_RANGE_MESSAGE_PREFIX
Review Comment:
nite: error msg could be `Invalid bucket ID: %s. Must be less than bucket limit: %s`
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitioner.java:
##########
@@ -0,0 +1,99 @@
+/*
+ * 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.apache.iceberg.flink.sink.BucketPartitioner.BUCKET_OUT_OF_RANGE_MESSAGE_PREFIX;
+
+import org.apache.iceberg.PartitionSpec;
+import org.assertj.core.api.Assertions;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+import org.junit.jupiter.params.provider.ValueSource;
+
+public class TestBucketPartitioner {
+
+ static final int NUM_BUCKETS = 60;
+
+ @ParameterizedTest
+ @EnumSource(
+ value = TestBucketPartitionerUtils.TableSchemaType.class,
+ names = {"ONE_BUCKET", "IDENTITY_AND_BUCKET"})
+ public void testPartitioningParallelismGreaterThanBuckets(
+ TestBucketPartitionerUtils.TableSchemaType tableSchemaType) {
+ final int numPartitions = 500;
+
+ PartitionSpec partitionSpec =
+ TestBucketPartitionerUtils.getPartitionSpec(tableSchemaType, NUM_BUCKETS);
+
+ BucketPartitioner bucketPartitioner = new BucketPartitioner(partitionSpec);
+
+ for (int expectedIdx = 0, bucketId = 0; expectedIdx < numPartitions; expectedIdx++) {
+ int actualIdx = bucketPartitioner.partition(bucketId, numPartitions);
+ Assertions.assertThat(actualIdx).isEqualTo(expectedIdx);
+ if (++bucketId == NUM_BUCKETS) {
+ bucketId = 0;
+ }
+ }
+ }
+
+ @ParameterizedTest
+ @EnumSource(
+ value = TestBucketPartitionerUtils.TableSchemaType.class,
+ names = {"ONE_BUCKET", "IDENTITY_AND_BUCKET"})
+ public void testPartitioningParallelismEqualLessThanBuckets(
Review Comment:
similarly, test dividable and non-dividable
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitioner.java:
##########
@@ -0,0 +1,99 @@
+/*
+ * 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.apache.iceberg.flink.sink.BucketPartitioner.BUCKET_OUT_OF_RANGE_MESSAGE_PREFIX;
+
+import org.apache.iceberg.PartitionSpec;
+import org.assertj.core.api.Assertions;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+import org.junit.jupiter.params.provider.ValueSource;
+
+public class TestBucketPartitioner {
+
+ static final int NUM_BUCKETS = 60;
+
+ @ParameterizedTest
+ @EnumSource(
+ value = TestBucketPartitionerUtils.TableSchemaType.class,
+ names = {"ONE_BUCKET", "IDENTITY_AND_BUCKET"})
+ public void testPartitioningParallelismGreaterThanBuckets(
+ TestBucketPartitionerUtils.TableSchemaType tableSchemaType) {
+ final int numPartitions = 500;
+
+ PartitionSpec partitionSpec =
+ TestBucketPartitionerUtils.getPartitionSpec(tableSchemaType, NUM_BUCKETS);
+
+ BucketPartitioner bucketPartitioner = new BucketPartitioner(partitionSpec);
+
+ for (int expectedIdx = 0, bucketId = 0; expectedIdx < numPartitions; expectedIdx++) {
+ int actualIdx = bucketPartitioner.partition(bucketId, numPartitions);
+ Assertions.assertThat(actualIdx).isEqualTo(expectedIdx);
+ if (++bucketId == NUM_BUCKETS) {
+ bucketId = 0;
+ }
+ }
+ }
+
+ @ParameterizedTest
+ @EnumSource(
+ value = TestBucketPartitionerUtils.TableSchemaType.class,
+ names = {"ONE_BUCKET", "IDENTITY_AND_BUCKET"})
+ public void testPartitioningParallelismEqualLessThanBuckets(
+ TestBucketPartitionerUtils.TableSchemaType tableSchemaType) {
+ final int numPartitions = 30;
+
+ PartitionSpec partitionSpec =
+ TestBucketPartitionerUtils.getPartitionSpec(tableSchemaType, NUM_BUCKETS);
+
+ BucketPartitioner bucketPartitioner = new BucketPartitioner(partitionSpec);
+
+ for (int bucketId = 0; bucketId < NUM_BUCKETS; bucketId++) {
+ int actualIdx = bucketPartitioner.partition(bucketId, numPartitions);
+ Assertions.assertThat(actualIdx).isEqualTo(bucketId % numPartitions);
+ }
+ }
+
+ @ParameterizedTest
Review Comment:
parameterized test seems unnecessary here
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerFlinkIcebergSink.java:
##########
@@ -0,0 +1,256 @@
+/*
+ * 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.apache.iceberg.flink.MiniClusterResource.DISABLE_CLASSLOADER_CHECK_CONFIG;
+import static org.apache.iceberg.flink.TestFixtures.DATABASE;
+import static org.apache.iceberg.flink.TestFixtures.TABLE_IDENTIFIER;
+import static org.apache.iceberg.flink.sink.TestBucketPartitionerUtils.CONVERTER;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.TreeMap;
+import java.util.stream.Collectors;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.typeutils.RowTypeInfo;
+import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.test.junit5.MiniClusterExtension;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.DistributionMode;
+import org.apache.iceberg.FileFormat;
+import org.apache.iceberg.FileScanTask;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Table;
+import org.apache.iceberg.TableProperties;
+import org.apache.iceberg.flink.HadoopCatalogExtension;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.flink.TableLoader;
+import org.apache.iceberg.flink.TestFixtures;
+import org.apache.iceberg.flink.sink.TestBucketPartitionerUtils.TableSchemaType;
+import org.apache.iceberg.flink.source.BoundedTestSource;
+import org.apache.iceberg.flink.util.FlinkCompatibilityUtil;
+import org.apache.iceberg.io.CloseableIterable;
+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.assertj.core.api.Assertions;
+import org.jetbrains.annotations.NotNull;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+public class TestBucketPartitionerFlinkIcebergSink {
+
+ private static final int NUMBER_TASK_MANAGERS = 1;
+ private static final int SLOTS_PER_TASK_MANAGER = 8;
+
+ @RegisterExtension
+ private static final MiniClusterExtension MINI_CLUSTER_RESOURCE =
+ new MiniClusterExtension(
+ new MiniClusterResourceConfiguration.Builder()
+ .setNumberTaskManagers(NUMBER_TASK_MANAGERS)
+ .setNumberSlotsPerTaskManager(SLOTS_PER_TASK_MANAGER)
+ .setConfiguration(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .build());
+
+ @RegisterExtension
+ private static final HadoopCatalogExtension catalogExtension =
+ new HadoopCatalogExtension(DATABASE, TestFixtures.TABLE);
+
+ private static final TypeInformation<Row> ROW_TYPE_INFO =
+ new RowTypeInfo(SimpleDataUtil.FLINK_SCHEMA.getFieldTypes());
+
+ // Parallelism = 8 (parallelism > numBuckets) throughout the test suite
+ private final int parallelism = NUMBER_TASK_MANAGERS * SLOTS_PER_TASK_MANAGER;
+ private final FileFormat format = FileFormat.PARQUET;
+ private final int numBuckets = 4;
+
+ private Table table;
+ private StreamExecutionEnvironment env;
+ private TableLoader tableLoader;
+
+ private void setupEnvironment(TableSchemaType tableSchemaType) {
+ table = getTable(tableSchemaType);
+ env =
+ StreamExecutionEnvironment.getExecutionEnvironment(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .enableCheckpointing(100)
+ .setParallelism(parallelism)
+ .setMaxParallelism(parallelism * 2);
+ tableLoader = catalogExtension.tableLoader();
+ }
+
+ private Table getTable(TableSchemaType tableSchemaType) {
+ PartitionSpec partitionSpec =
+ TestBucketPartitionerUtils.getPartitionSpec(tableSchemaType, numBuckets);
+
+ return catalogExtension
+ .catalog()
+ .createTable(
+ TABLE_IDENTIFIER,
+ SimpleDataUtil.SCHEMA,
+ partitionSpec,
+ ImmutableMap.of(TableProperties.DEFAULT_FILE_FORMAT, format.name()));
+ }
+
+ private List<RowData> convertToRowData(List<Row> rows) {
+ return rows.stream().map(CONVERTER::toInternal).collect(Collectors.toList());
+ }
+
+ private BoundedTestSource<Row> createBoundedSource(List<Row> rows) {
+ return new BoundedTestSource<>(rows.toArray(new Row[0]));
+ }
+
+ private TableTestStats extractTableTestStats(TableSchemaType tableSchemaType) throws IOException {
+ int totalRecordCount = 0;
+ Map<Integer, List<Integer>> writersPerBucket = Maps.newHashMap(); // <BucketId, List<WriterId>>
+ Map<Integer, Integer> filesPerBucket = Maps.newHashMap(); // <BucketId, NumFiles>
+ Map<Integer, Long> recordsPerFile = new TreeMap<>(); // <WriterId, NumRecords>
+
+ try (CloseableIterable<FileScanTask> fileScanTasks = table.newScan().planFiles()) {
+ for (FileScanTask scanTask : fileScanTasks) {
+ long recordCountInFile = scanTask.file().recordCount();
+
+ String[] splitFilePath = scanTask.file().path().toString().split("/");
+ String filename = splitFilePath[splitFilePath.length - 1];
+ int writerId = Integer.parseInt(filename.split("-")[0]);
+
+ totalRecordCount += recordCountInFile;
+ int bucketId =
+ scanTask
+ .file()
+ .partition()
+ .get(tableSchemaType == TableSchemaType.ONE_BUCKET ? 0 : 1, Integer.class);
Review Comment:
> tableSchemaType == TableSchemaType.ONE_BUCKET ? 0 : 1
this code is not intuitive. you can move it inside the enum class. maybe expose a method like `bucketPartitionColumnPosition()`
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerFlinkIcebergSink.java:
##########
@@ -0,0 +1,256 @@
+/*
+ * 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.apache.iceberg.flink.MiniClusterResource.DISABLE_CLASSLOADER_CHECK_CONFIG;
+import static org.apache.iceberg.flink.TestFixtures.DATABASE;
+import static org.apache.iceberg.flink.TestFixtures.TABLE_IDENTIFIER;
+import static org.apache.iceberg.flink.sink.TestBucketPartitionerUtils.CONVERTER;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.TreeMap;
+import java.util.stream.Collectors;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.typeutils.RowTypeInfo;
+import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.test.junit5.MiniClusterExtension;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.DistributionMode;
+import org.apache.iceberg.FileFormat;
+import org.apache.iceberg.FileScanTask;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Table;
+import org.apache.iceberg.TableProperties;
+import org.apache.iceberg.flink.HadoopCatalogExtension;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.flink.TableLoader;
+import org.apache.iceberg.flink.TestFixtures;
+import org.apache.iceberg.flink.sink.TestBucketPartitionerUtils.TableSchemaType;
+import org.apache.iceberg.flink.source.BoundedTestSource;
+import org.apache.iceberg.flink.util.FlinkCompatibilityUtil;
+import org.apache.iceberg.io.CloseableIterable;
+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.assertj.core.api.Assertions;
+import org.jetbrains.annotations.NotNull;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+public class TestBucketPartitionerFlinkIcebergSink {
+
+ private static final int NUMBER_TASK_MANAGERS = 1;
+ private static final int SLOTS_PER_TASK_MANAGER = 8;
+
+ @RegisterExtension
+ private static final MiniClusterExtension MINI_CLUSTER_RESOURCE =
+ new MiniClusterExtension(
+ new MiniClusterResourceConfiguration.Builder()
+ .setNumberTaskManagers(NUMBER_TASK_MANAGERS)
+ .setNumberSlotsPerTaskManager(SLOTS_PER_TASK_MANAGER)
+ .setConfiguration(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .build());
+
+ @RegisterExtension
+ private static final HadoopCatalogExtension catalogExtension =
+ new HadoopCatalogExtension(DATABASE, TestFixtures.TABLE);
+
+ private static final TypeInformation<Row> ROW_TYPE_INFO =
+ new RowTypeInfo(SimpleDataUtil.FLINK_SCHEMA.getFieldTypes());
+
+ // Parallelism = 8 (parallelism > numBuckets) throughout the test suite
+ private final int parallelism = NUMBER_TASK_MANAGERS * SLOTS_PER_TASK_MANAGER;
+ private final FileFormat format = FileFormat.PARQUET;
+ private final int numBuckets = 4;
+
+ private Table table;
+ private StreamExecutionEnvironment env;
+ private TableLoader tableLoader;
+
+ private void setupEnvironment(TableSchemaType tableSchemaType) {
+ table = getTable(tableSchemaType);
+ env =
+ StreamExecutionEnvironment.getExecutionEnvironment(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .enableCheckpointing(100)
+ .setParallelism(parallelism)
+ .setMaxParallelism(parallelism * 2);
+ tableLoader = catalogExtension.tableLoader();
+ }
+
+ private Table getTable(TableSchemaType tableSchemaType) {
+ PartitionSpec partitionSpec =
+ TestBucketPartitionerUtils.getPartitionSpec(tableSchemaType, numBuckets);
+
+ return catalogExtension
+ .catalog()
+ .createTable(
+ TABLE_IDENTIFIER,
+ SimpleDataUtil.SCHEMA,
+ partitionSpec,
+ ImmutableMap.of(TableProperties.DEFAULT_FILE_FORMAT, format.name()));
+ }
+
+ private List<RowData> convertToRowData(List<Row> rows) {
+ return rows.stream().map(CONVERTER::toInternal).collect(Collectors.toList());
+ }
+
+ private BoundedTestSource<Row> createBoundedSource(List<Row> rows) {
+ return new BoundedTestSource<>(rows.toArray(new Row[0]));
+ }
+
+ private TableTestStats extractTableTestStats(TableSchemaType tableSchemaType) throws IOException {
Review Comment:
move the util method after the `@Test` methods
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/HadoopCatalogExtension.java:
##########
@@ -0,0 +1,104 @@
+/*
+ * 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;
+
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.UUID;
+import org.apache.commons.io.FileUtils;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.iceberg.CatalogProperties;
+import org.apache.iceberg.catalog.Catalog;
+import org.apache.iceberg.catalog.TableIdentifier;
+import org.apache.iceberg.hadoop.HadoopCatalog;
+import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
+import org.assertj.core.api.Assertions;
+import org.junit.jupiter.api.extension.AfterAllCallback;
+import org.junit.jupiter.api.extension.AfterEachCallback;
+import org.junit.jupiter.api.extension.BeforeAllCallback;
+import org.junit.jupiter.api.extension.BeforeEachCallback;
+import org.junit.jupiter.api.extension.ExtensionContext;
+
+public class HadoopCatalogExtension
+ implements BeforeAllCallback, BeforeEachCallback, AfterAllCallback, AfterEachCallback {
+ protected final String database;
+ protected final String tableName;
+
+ protected Path temporaryFolder;
Review Comment:
this can be replaced with `@TempDir private File tempDir;`. then we can remove the beforeAll and afterAll callbacks
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitioner.java:
##########
@@ -0,0 +1,108 @@
+/*
+ * 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 org.apache.flink.api.common.functions.Partitioner;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/**
+ * This partitioner will redirect records to writers deterministically based on the Bucket partition
+ * spec. It'll attempt to optimize the file size written depending on whether numPartitions is
+ * greater, less or equal than the maxNumBuckets. Note: The current implementation only supports ONE
+ * bucket in the partition spec.
+ */
+class BucketPartitioner implements Partitioner<Integer> {
+
+ static final String BUCKET_OUT_OF_RANGE_MESSAGE_PREFIX = "bucketId out of range: ";
+
+ private final int maxNumBuckets;
+
+ // To hold the OFFSET of the next writer to use for any bucket
+ private final int[] currentBucketWriterOffset;
+
+ BucketPartitioner(PartitionSpec partitionSpec) {
+ Tuple2<Integer, Integer> bucketFieldInfo =
+ BucketPartitionerUtils.getBucketFieldInfo(partitionSpec);
+
+ this.maxNumBuckets = bucketFieldInfo.f1;
+ this.currentBucketWriterOffset = new int[this.maxNumBuckets];
+ }
+
+ /**
+ * Determine the partition id based on the following criteria: If the number of writers <= the
+ * number of buckets, an evenly distributed number of buckets will be assigned to each writer (one
+ * writer -> many buckets). Conversely, if the number of writers > the number of buckets the logic
+ * is handled by the {@link #getPartitionWritersGreaterThanBuckets
+ * getPartitionWritersGreaterThanBuckets} method.
+ *
+ * @param bucketId the bucketId for each request
+ * @param numPartitions the total number of partitions
+ * @return the partition id (writer) to use for each request
+ */
+ @Override
+ public int partition(Integer bucketId, int numPartitions) {
+ Preconditions.checkArgument(
+ bucketId >= 0, BUCKET_OUT_OF_RANGE_MESSAGE_PREFIX + bucketId + " (must be >= 0)");
+ Preconditions.checkArgument(
+ bucketId < maxNumBuckets,
+ BUCKET_OUT_OF_RANGE_MESSAGE_PREFIX
+ + bucketId
+ + " (must be >= 0), maxNumBuckets: "
+ + maxNumBuckets);
+
+ if (numPartitions <= maxNumBuckets) {
+ return bucketId % numPartitions;
+ } else {
+ return getPartitionWritersGreaterThanBuckets(bucketId, numPartitions);
+ }
+ }
+
+ /*-
+ * If the number of writers > the number of buckets each partitioner will keep a state of multiple
+ * writers per bucket as evenly as possible, and will round-robin the requests across them, in this
+ * case each writer will target only one bucket at all times (many writers -> one bucket). Example:
+ * Configuration: numPartitions (writers) = 5, maxBuckets = 2
+ * Expected behavior:
+ * - Records for Bucket 0 will be "round robin" between Writers 0, 2 and 4
+ * - Records for Bucket 1 will always use Writer 1 and 3
+ * Notes:
+ * - maxNumWritersPerBucket determines when to reset the currentBucketWriterOffset to 0 for this bucketId
+ * - When numPartitions is not evenly divisible by maxBuckets, some buckets will have one more writer (extraWriter).
+ * In this example Bucket 0 has an "extra writer" to consider before resetting its offset to 0.
+ *
+ * @param bucketId the bucketId for each request
+ * @param numPartitions the total number of partitions
+ * @return the partition index (writer) to use for each request
Review Comment:
nit: `@return the destination partition index (writer subtask id)`.
`to use for each request` is a bit inaccurate
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerUtils.java:
##########
@@ -0,0 +1,121 @@
+/*
+ * 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 java.util.List;
+import java.util.UUID;
+import org.apache.flink.table.data.util.DataFormatConverters;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.apache.iceberg.util.BucketUtil;
+
+final class TestBucketPartitionerUtils {
+
+ enum TableSchemaType {
+ ONE_BUCKET,
+ IDENTITY_AND_BUCKET,
+ TWO_BUCKETS;
+ }
+
+ private TestBucketPartitionerUtils() {}
+
+ static final DataFormatConverters.RowConverter CONVERTER =
+ new DataFormatConverters.RowConverter(SimpleDataUtil.FLINK_SCHEMA.getFieldDataTypes());
+
+ static PartitionSpec getPartitionSpec(TableSchemaType tableSchemaType, int numBuckets) {
Review Comment:
it is probably bit cleaner to move it into the enum class
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitioner.java:
##########
@@ -0,0 +1,108 @@
+/*
+ * 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 org.apache.flink.api.common.functions.Partitioner;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/**
+ * This partitioner will redirect records to writers deterministically based on the Bucket partition
+ * spec. It'll attempt to optimize the file size written depending on whether numPartitions is
+ * greater, less or equal than the maxNumBuckets. Note: The current implementation only supports ONE
+ * bucket in the partition spec.
+ */
+class BucketPartitioner implements Partitioner<Integer> {
+
+ static final String BUCKET_OUT_OF_RANGE_MESSAGE_PREFIX = "bucketId out of range: ";
+
+ private final int maxNumBuckets;
+
+ // To hold the OFFSET of the next writer to use for any bucket
+ private final int[] currentBucketWriterOffset;
+
+ BucketPartitioner(PartitionSpec partitionSpec) {
+ Tuple2<Integer, Integer> bucketFieldInfo =
+ BucketPartitionerUtils.getBucketFieldInfo(partitionSpec);
+
+ this.maxNumBuckets = bucketFieldInfo.f1;
+ this.currentBucketWriterOffset = new int[this.maxNumBuckets];
+ }
+
+ /**
+ * Determine the partition id based on the following criteria: If the number of writers <= the
+ * number of buckets, an evenly distributed number of buckets will be assigned to each writer (one
+ * writer -> many buckets). Conversely, if the number of writers > the number of buckets the logic
+ * is handled by the {@link #getPartitionWritersGreaterThanBuckets
+ * getPartitionWritersGreaterThanBuckets} method.
+ *
+ * @param bucketId the bucketId for each request
+ * @param numPartitions the total number of partitions
+ * @return the partition id (writer) to use for each request
+ */
+ @Override
+ public int partition(Integer bucketId, int numPartitions) {
+ Preconditions.checkArgument(
+ bucketId >= 0, BUCKET_OUT_OF_RANGE_MESSAGE_PREFIX + bucketId + " (must be >= 0)");
+ Preconditions.checkArgument(
+ bucketId < maxNumBuckets,
+ BUCKET_OUT_OF_RANGE_MESSAGE_PREFIX
+ + bucketId
+ + " (must be >= 0), maxNumBuckets: "
+ + maxNumBuckets);
+
+ if (numPartitions <= maxNumBuckets) {
+ return bucketId % numPartitions;
+ } else {
+ return getPartitionWritersGreaterThanBuckets(bucketId, numPartitions);
+ }
+ }
+
+ /*-
+ * If the number of writers > the number of buckets each partitioner will keep a state of multiple
+ * writers per bucket as evenly as possible, and will round-robin the requests across them, in this
+ * case each writer will target only one bucket at all times (many writers -> one bucket). Example:
+ * Configuration: numPartitions (writers) = 5, maxBuckets = 2
+ * Expected behavior:
+ * - Records for Bucket 0 will be "round robin" between Writers 0, 2 and 4
+ * - Records for Bucket 1 will always use Writer 1 and 3
+ * Notes:
+ * - maxNumWritersPerBucket determines when to reset the currentBucketWriterOffset to 0 for this bucketId
+ * - When numPartitions is not evenly divisible by maxBuckets, some buckets will have one more writer (extraWriter).
+ * In this example Bucket 0 has an "extra writer" to consider before resetting its offset to 0.
+ *
+ * @param bucketId the bucketId for each request
Review Comment:
it seems to me that the two `@param` are not useful. maybe remove them.
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitioner.java:
##########
@@ -0,0 +1,108 @@
+/*
+ * 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 org.apache.flink.api.common.functions.Partitioner;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/**
+ * This partitioner will redirect records to writers deterministically based on the Bucket partition
+ * spec. It'll attempt to optimize the file size written depending on whether numPartitions is
+ * greater, less or equal than the maxNumBuckets. Note: The current implementation only supports ONE
+ * bucket in the partition spec.
+ */
+class BucketPartitioner implements Partitioner<Integer> {
+
+ static final String BUCKET_OUT_OF_RANGE_MESSAGE_PREFIX = "bucketId out of range: ";
+
+ private final int maxNumBuckets;
+
+ // To hold the OFFSET of the next writer to use for any bucket
+ private final int[] currentBucketWriterOffset;
+
+ BucketPartitioner(PartitionSpec partitionSpec) {
+ Tuple2<Integer, Integer> bucketFieldInfo =
+ BucketPartitionerUtils.getBucketFieldInfo(partitionSpec);
+
+ this.maxNumBuckets = bucketFieldInfo.f1;
+ this.currentBucketWriterOffset = new int[this.maxNumBuckets];
+ }
+
+ /**
+ * Determine the partition id based on the following criteria: If the number of writers <= the
+ * number of buckets, an evenly distributed number of buckets will be assigned to each writer (one
+ * writer -> many buckets). Conversely, if the number of writers > the number of buckets the logic
+ * is handled by the {@link #getPartitionWritersGreaterThanBuckets
+ * getPartitionWritersGreaterThanBuckets} method.
+ *
+ * @param bucketId the bucketId for each request
+ * @param numPartitions the total number of partitions
+ * @return the partition id (writer) to use for each request
+ */
+ @Override
+ public int partition(Integer bucketId, int numPartitions) {
+ Preconditions.checkArgument(
+ bucketId >= 0, BUCKET_OUT_OF_RANGE_MESSAGE_PREFIX + bucketId + " (must be >= 0)");
+ Preconditions.checkArgument(
+ bucketId < maxNumBuckets,
+ BUCKET_OUT_OF_RANGE_MESSAGE_PREFIX
+ + bucketId
+ + " (must be >= 0), maxNumBuckets: "
+ + maxNumBuckets);
+
+ if (numPartitions <= maxNumBuckets) {
+ return bucketId % numPartitions;
+ } else {
+ return getPartitionWritersGreaterThanBuckets(bucketId, numPartitions);
+ }
+ }
+
+ /*-
+ * If the number of writers > the number of buckets each partitioner will keep a state of multiple
Review Comment:
very nice explanation here
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerUtils.java:
##########
@@ -0,0 +1,121 @@
+/*
+ * 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 java.util.List;
+import java.util.UUID;
+import org.apache.flink.table.data.util.DataFormatConverters;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.apache.iceberg.util.BucketUtil;
+
+final class TestBucketPartitionerUtils {
+
+ enum TableSchemaType {
+ ONE_BUCKET,
+ IDENTITY_AND_BUCKET,
+ TWO_BUCKETS;
+ }
+
+ private TestBucketPartitionerUtils() {}
+
+ static final DataFormatConverters.RowConverter CONVERTER =
+ new DataFormatConverters.RowConverter(SimpleDataUtil.FLINK_SCHEMA.getFieldDataTypes());
+
+ static PartitionSpec getPartitionSpec(TableSchemaType tableSchemaType, int numBuckets) {
+ PartitionSpec partitionSpec = null;
+
+ switch (tableSchemaType) {
+ case ONE_BUCKET:
+ partitionSpec =
+ PartitionSpec.builderFor(SimpleDataUtil.SCHEMA).bucket("data", numBuckets).build();
+ break;
+ case IDENTITY_AND_BUCKET:
+ partitionSpec =
+ PartitionSpec.builderFor(SimpleDataUtil.SCHEMA)
+ .identity("id")
+ .bucket("data", numBuckets)
+ .build();
+ break;
+ case TWO_BUCKETS:
+ partitionSpec =
+ PartitionSpec.builderFor(SimpleDataUtil.SCHEMA)
+ .bucket("id", numBuckets)
+ .bucket("data", numBuckets)
+ .build();
+ break;
+ }
+
+ Preconditions.checkNotNull(
+ partitionSpec, "Invalid tableSchemaType provided: " + tableSchemaType);
+ return partitionSpec;
+ }
+
+ /**
+ * Utility method to generate rows whose values will "hash" to a range of bucketIds (from 0 to
+ * numBuckets - 1)
+ *
+ * @param numRowsPerBucket how many different rows should be generated per bucket
+ * @param numBuckets max number of buckets to consider
+ * @return the list of rows whose data "hashes" to the desired bucketId
+ */
+ static List<Row> generateRowsForBucketIdRange(int numRowsPerBucket, int numBuckets) {
Review Comment:
it should return `RowData` right? You can use `GenericRowData`. Then we don't need `CONVERTER `
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitioner.java:
##########
@@ -0,0 +1,108 @@
+/*
+ * 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 org.apache.flink.api.common.functions.Partitioner;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/**
+ * This partitioner will redirect records to writers deterministically based on the Bucket partition
+ * spec. It'll attempt to optimize the file size written depending on whether numPartitions is
+ * greater, less or equal than the maxNumBuckets. Note: The current implementation only supports ONE
+ * bucket in the partition spec.
+ */
+class BucketPartitioner implements Partitioner<Integer> {
+
+ static final String BUCKET_OUT_OF_RANGE_MESSAGE_PREFIX = "bucketId out of range: ";
+
+ private final int maxNumBuckets;
+
+ // To hold the OFFSET of the next writer to use for any bucket
+ private final int[] currentBucketWriterOffset;
+
+ BucketPartitioner(PartitionSpec partitionSpec) {
+ Tuple2<Integer, Integer> bucketFieldInfo =
+ BucketPartitionerUtils.getBucketFieldInfo(partitionSpec);
+
+ this.maxNumBuckets = bucketFieldInfo.f1;
+ this.currentBucketWriterOffset = new int[this.maxNumBuckets];
+ }
+
+ /**
+ * Determine the partition id based on the following criteria: If the number of writers <= the
+ * number of buckets, an evenly distributed number of buckets will be assigned to each writer (one
+ * writer -> many buckets). Conversely, if the number of writers > the number of buckets the logic
+ * is handled by the {@link #getPartitionWritersGreaterThanBuckets
+ * getPartitionWritersGreaterThanBuckets} method.
+ *
+ * @param bucketId the bucketId for each request
+ * @param numPartitions the total number of partitions
+ * @return the partition id (writer) to use for each request
+ */
+ @Override
+ public int partition(Integer bucketId, int numPartitions) {
+ Preconditions.checkArgument(
+ bucketId >= 0, BUCKET_OUT_OF_RANGE_MESSAGE_PREFIX + bucketId + " (must be >= 0)");
+ Preconditions.checkArgument(
+ bucketId < maxNumBuckets,
+ BUCKET_OUT_OF_RANGE_MESSAGE_PREFIX
+ + bucketId
+ + " (must be >= 0), maxNumBuckets: "
+ + maxNumBuckets);
+
+ if (numPartitions <= maxNumBuckets) {
+ return bucketId % numPartitions;
+ } else {
+ return getPartitionWritersGreaterThanBuckets(bucketId, numPartitions);
+ }
+ }
+
+ /*-
+ * If the number of writers > the number of buckets each partitioner will keep a state of multiple
+ * writers per bucket as evenly as possible, and will round-robin the requests across them, in this
+ * case each writer will target only one bucket at all times (many writers -> one bucket). Example:
+ * Configuration: numPartitions (writers) = 5, maxBuckets = 2
+ * Expected behavior:
+ * - Records for Bucket 0 will be "round robin" between Writers 0, 2 and 4
+ * - Records for Bucket 1 will always use Writer 1 and 3
+ * Notes:
+ * - maxNumWritersPerBucket determines when to reset the currentBucketWriterOffset to 0 for this bucketId
+ * - When numPartitions is not evenly divisible by maxBuckets, some buckets will have one more writer (extraWriter).
+ * In this example Bucket 0 has an "extra writer" to consider before resetting its offset to 0.
+ *
+ * @param bucketId the bucketId for each request
+ * @param numPartitions the total number of partitions
+ * @return the partition index (writer) to use for each request
+ */
+ private int getPartitionWritersGreaterThanBuckets(int bucketId, int numPartitions) {
Review Comment:
nit: I would call this `partitionWithMoreWritersThanBuckets`
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitioner.java:
##########
@@ -0,0 +1,108 @@
+/*
+ * 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 org.apache.flink.api.common.functions.Partitioner;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/**
+ * This partitioner will redirect records to writers deterministically based on the Bucket partition
+ * spec. It'll attempt to optimize the file size written depending on whether numPartitions is
+ * greater, less or equal than the maxNumBuckets. Note: The current implementation only supports ONE
+ * bucket in the partition spec.
+ */
+class BucketPartitioner implements Partitioner<Integer> {
+
+ static final String BUCKET_OUT_OF_RANGE_MESSAGE_PREFIX = "bucketId out of range: ";
+
+ private final int maxNumBuckets;
+
+ // To hold the OFFSET of the next writer to use for any bucket
Review Comment:
nit: add to the comment that this is only used for the case where the number of buckets is less than the writer parallelism
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerUtils.java:
##########
@@ -0,0 +1,121 @@
+/*
+ * 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 java.util.List;
+import java.util.UUID;
+import org.apache.flink.table.data.util.DataFormatConverters;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.apache.iceberg.util.BucketUtil;
+
+final class TestBucketPartitionerUtils {
+
+ enum TableSchemaType {
+ ONE_BUCKET,
+ IDENTITY_AND_BUCKET,
+ TWO_BUCKETS;
+ }
+
+ private TestBucketPartitionerUtils() {}
+
+ static final DataFormatConverters.RowConverter CONVERTER =
+ new DataFormatConverters.RowConverter(SimpleDataUtil.FLINK_SCHEMA.getFieldDataTypes());
+
+ static PartitionSpec getPartitionSpec(TableSchemaType tableSchemaType, int numBuckets) {
+ PartitionSpec partitionSpec = null;
+
+ switch (tableSchemaType) {
+ case ONE_BUCKET:
+ partitionSpec =
+ PartitionSpec.builderFor(SimpleDataUtil.SCHEMA).bucket("data", numBuckets).build();
+ break;
+ case IDENTITY_AND_BUCKET:
+ partitionSpec =
+ PartitionSpec.builderFor(SimpleDataUtil.SCHEMA)
+ .identity("id")
+ .bucket("data", numBuckets)
+ .build();
+ break;
+ case TWO_BUCKETS:
+ partitionSpec =
+ PartitionSpec.builderFor(SimpleDataUtil.SCHEMA)
+ .bucket("id", numBuckets)
+ .bucket("data", numBuckets)
+ .build();
+ break;
+ }
+
+ Preconditions.checkNotNull(
+ partitionSpec, "Invalid tableSchemaType provided: " + tableSchemaType);
+ return partitionSpec;
+ }
+
+ /**
+ * Utility method to generate rows whose values will "hash" to a range of bucketIds (from 0 to
+ * numBuckets - 1)
+ *
+ * @param numRowsPerBucket how many different rows should be generated per bucket
+ * @param numBuckets max number of buckets to consider
+ * @return the list of rows whose data "hashes" to the desired bucketId
+ */
+ static List<Row> generateRowsForBucketIdRange(int numRowsPerBucket, int numBuckets) {
+ List<Row> rows = Lists.newArrayListWithCapacity(numBuckets * numRowsPerBucket);
+ // For some of our tests, this order of the generated rows matters
+ for (int i = 0; i < numRowsPerBucket; i++) {
+ for (int bucketId = 0; bucketId < numBuckets; bucketId++) {
+ String value = generateValueForBucketId(bucketId, numBuckets);
+ rows.add(Row.of(1, value));
+ }
+ }
+ return rows;
+ }
+
+ /**
+ * Utility method to generate a UUID string that will "hash" to a desired bucketId
+ *
+ * @param bucketId the desired bucketId
+ * @return the string data that "hashes" to the desired bucketId
+ */
+ private static String generateValueForBucketId(int bucketId, int numBuckets) {
+ String value = "";
+ while (true) {
+ String uuid = UUID.randomUUID().toString();
+ if (computeBucketId(numBuckets, uuid) == bucketId) {
+ value = uuid;
+ break;
Review Comment:
nit: just return the uuid? avoided the need of value variable.
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerUtils.java:
##########
@@ -0,0 +1,121 @@
+/*
+ * 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 java.util.List;
+import java.util.UUID;
+import org.apache.flink.table.data.util.DataFormatConverters;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.apache.iceberg.util.BucketUtil;
+
+final class TestBucketPartitionerUtils {
+
+ enum TableSchemaType {
+ ONE_BUCKET,
+ IDENTITY_AND_BUCKET,
+ TWO_BUCKETS;
+ }
+
+ private TestBucketPartitionerUtils() {}
+
+ static final DataFormatConverters.RowConverter CONVERTER =
+ new DataFormatConverters.RowConverter(SimpleDataUtil.FLINK_SCHEMA.getFieldDataTypes());
+
+ static PartitionSpec getPartitionSpec(TableSchemaType tableSchemaType, int numBuckets) {
+ PartitionSpec partitionSpec = null;
+
+ switch (tableSchemaType) {
+ case ONE_BUCKET:
+ partitionSpec =
+ PartitionSpec.builderFor(SimpleDataUtil.SCHEMA).bucket("data", numBuckets).build();
+ break;
+ case IDENTITY_AND_BUCKET:
+ partitionSpec =
+ PartitionSpec.builderFor(SimpleDataUtil.SCHEMA)
+ .identity("id")
+ .bucket("data", numBuckets)
+ .build();
+ break;
+ case TWO_BUCKETS:
+ partitionSpec =
+ PartitionSpec.builderFor(SimpleDataUtil.SCHEMA)
+ .bucket("id", numBuckets)
+ .bucket("data", numBuckets)
+ .build();
+ break;
+ }
+
+ Preconditions.checkNotNull(
+ partitionSpec, "Invalid tableSchemaType provided: " + tableSchemaType);
+ return partitionSpec;
+ }
+
+ /**
+ * Utility method to generate rows whose values will "hash" to a range of bucketIds (from 0 to
+ * numBuckets - 1)
+ *
+ * @param numRowsPerBucket how many different rows should be generated per bucket
+ * @param numBuckets max number of buckets to consider
+ * @return the list of rows whose data "hashes" to the desired bucketId
+ */
+ static List<Row> generateRowsForBucketIdRange(int numRowsPerBucket, int numBuckets) {
+ List<Row> rows = Lists.newArrayListWithCapacity(numBuckets * numRowsPerBucket);
+ // For some of our tests, this order of the generated rows matters
+ for (int i = 0; i < numRowsPerBucket; i++) {
+ for (int bucketId = 0; bucketId < numBuckets; bucketId++) {
+ String value = generateValueForBucketId(bucketId, numBuckets);
+ rows.add(Row.of(1, value));
+ }
+ }
+ return rows;
+ }
+
+ /**
+ * Utility method to generate a UUID string that will "hash" to a desired bucketId
+ *
+ * @param bucketId the desired bucketId
+ * @return the string data that "hashes" to the desired bucketId
+ */
+ private static String generateValueForBucketId(int bucketId, int numBuckets) {
+ String value = "";
+ while (true) {
+ String uuid = UUID.randomUUID().toString();
+ if (computeBucketId(numBuckets, uuid) == bucketId) {
+ value = uuid;
+ break;
+ }
+ }
+ return value;
Review Comment:
nit: Iceberg style adds an empty line after a control block `}`
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitioner.java:
##########
@@ -0,0 +1,99 @@
+/*
+ * 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.apache.iceberg.flink.sink.BucketPartitioner.BUCKET_OUT_OF_RANGE_MESSAGE_PREFIX;
+
+import org.apache.iceberg.PartitionSpec;
+import org.assertj.core.api.Assertions;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+import org.junit.jupiter.params.provider.ValueSource;
+
+public class TestBucketPartitioner {
+
+ static final int NUM_BUCKETS = 60;
+
+ @ParameterizedTest
+ @EnumSource(
+ value = TestBucketPartitionerUtils.TableSchemaType.class,
+ names = {"ONE_BUCKET", "IDENTITY_AND_BUCKET"})
+ public void testPartitioningParallelismGreaterThanBuckets(
+ TestBucketPartitionerUtils.TableSchemaType tableSchemaType) {
+ final int numPartitions = 500;
+
+ PartitionSpec partitionSpec =
+ TestBucketPartitionerUtils.getPartitionSpec(tableSchemaType, NUM_BUCKETS);
+
+ BucketPartitioner bucketPartitioner = new BucketPartitioner(partitionSpec);
+
+ for (int expectedIdx = 0, bucketId = 0; expectedIdx < numPartitions; expectedIdx++) {
Review Comment:
move `bucketId` initialization before the for loop. would make it easier to see the for loop
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitioner.java:
##########
@@ -0,0 +1,99 @@
+/*
+ * 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.apache.iceberg.flink.sink.BucketPartitioner.BUCKET_OUT_OF_RANGE_MESSAGE_PREFIX;
+
+import org.apache.iceberg.PartitionSpec;
+import org.assertj.core.api.Assertions;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+import org.junit.jupiter.params.provider.ValueSource;
+
+public class TestBucketPartitioner {
+
+ static final int NUM_BUCKETS = 60;
+
+ @ParameterizedTest
+ @EnumSource(
+ value = TestBucketPartitionerUtils.TableSchemaType.class,
+ names = {"ONE_BUCKET", "IDENTITY_AND_BUCKET"})
+ public void testPartitioningParallelismGreaterThanBuckets(
Review Comment:
I feel we need to test two separate scenarios: dividable or non-dividable
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitioner.java:
##########
@@ -0,0 +1,99 @@
+/*
+ * 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.apache.iceberg.flink.sink.BucketPartitioner.BUCKET_OUT_OF_RANGE_MESSAGE_PREFIX;
+
+import org.apache.iceberg.PartitionSpec;
+import org.assertj.core.api.Assertions;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+import org.junit.jupiter.params.provider.ValueSource;
+
+public class TestBucketPartitioner {
+
+ static final int NUM_BUCKETS = 60;
+
+ @ParameterizedTest
+ @EnumSource(
+ value = TestBucketPartitionerUtils.TableSchemaType.class,
+ names = {"ONE_BUCKET", "IDENTITY_AND_BUCKET"})
+ public void testPartitioningParallelismGreaterThanBuckets(
+ TestBucketPartitionerUtils.TableSchemaType tableSchemaType) {
+ final int numPartitions = 500;
+
+ PartitionSpec partitionSpec =
+ TestBucketPartitionerUtils.getPartitionSpec(tableSchemaType, NUM_BUCKETS);
+
Review Comment:
nit: empty lines in line 40 and 43 seems unnecessary
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitioner.java:
##########
@@ -0,0 +1,99 @@
+/*
+ * 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.apache.iceberg.flink.sink.BucketPartitioner.BUCKET_OUT_OF_RANGE_MESSAGE_PREFIX;
+
+import org.apache.iceberg.PartitionSpec;
+import org.assertj.core.api.Assertions;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+import org.junit.jupiter.params.provider.ValueSource;
+
+public class TestBucketPartitioner {
+
+ static final int NUM_BUCKETS = 60;
+
+ @ParameterizedTest
+ @EnumSource(
+ value = TestBucketPartitionerUtils.TableSchemaType.class,
+ names = {"ONE_BUCKET", "IDENTITY_AND_BUCKET"})
+ public void testPartitioningParallelismGreaterThanBuckets(
+ TestBucketPartitionerUtils.TableSchemaType tableSchemaType) {
+ final int numPartitions = 500;
+
+ PartitionSpec partitionSpec =
+ TestBucketPartitionerUtils.getPartitionSpec(tableSchemaType, NUM_BUCKETS);
+
+ BucketPartitioner bucketPartitioner = new BucketPartitioner(partitionSpec);
+
+ for (int expectedIdx = 0, bucketId = 0; expectedIdx < numPartitions; expectedIdx++) {
+ int actualIdx = bucketPartitioner.partition(bucketId, numPartitions);
Review Comment:
unless very obvious, Iceberg style doesn't use acronym. so maybe change `actualIdx` to `actualPartitionIndex`.
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitioner.java:
##########
@@ -0,0 +1,99 @@
+/*
+ * 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.apache.iceberg.flink.sink.BucketPartitioner.BUCKET_OUT_OF_RANGE_MESSAGE_PREFIX;
+
+import org.apache.iceberg.PartitionSpec;
+import org.assertj.core.api.Assertions;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+import org.junit.jupiter.params.provider.ValueSource;
+
+public class TestBucketPartitioner {
+
+ static final int NUM_BUCKETS = 60;
+
+ @ParameterizedTest
+ @EnumSource(
+ value = TestBucketPartitionerUtils.TableSchemaType.class,
+ names = {"ONE_BUCKET", "IDENTITY_AND_BUCKET"})
+ public void testPartitioningParallelismGreaterThanBuckets(
+ TestBucketPartitionerUtils.TableSchemaType tableSchemaType) {
+ final int numPartitions = 500;
+
+ PartitionSpec partitionSpec =
+ TestBucketPartitionerUtils.getPartitionSpec(tableSchemaType, NUM_BUCKETS);
+
+ BucketPartitioner bucketPartitioner = new BucketPartitioner(partitionSpec);
+
+ for (int expectedIdx = 0, bucketId = 0; expectedIdx < numPartitions; expectedIdx++) {
+ int actualIdx = bucketPartitioner.partition(bucketId, numPartitions);
+ Assertions.assertThat(actualIdx).isEqualTo(expectedIdx);
+ if (++bucketId == NUM_BUCKETS) {
+ bucketId = 0;
+ }
+ }
+ }
+
+ @ParameterizedTest
+ @EnumSource(
+ value = TestBucketPartitionerUtils.TableSchemaType.class,
+ names = {"ONE_BUCKET", "IDENTITY_AND_BUCKET"})
+ public void testPartitioningParallelismEqualLessThanBuckets(
+ TestBucketPartitionerUtils.TableSchemaType tableSchemaType) {
+ final int numPartitions = 30;
+
+ PartitionSpec partitionSpec =
+ TestBucketPartitionerUtils.getPartitionSpec(tableSchemaType, NUM_BUCKETS);
+
+ BucketPartitioner bucketPartitioner = new BucketPartitioner(partitionSpec);
+
+ for (int bucketId = 0; bucketId < NUM_BUCKETS; bucketId++) {
+ int actualIdx = bucketPartitioner.partition(bucketId, numPartitions);
+ Assertions.assertThat(actualIdx).isEqualTo(bucketId % numPartitions);
+ }
+ }
+
+ @ParameterizedTest
+ @EnumSource(value = TestBucketPartitionerUtils.TableSchemaType.class, names = "TWO_BUCKETS")
+ public void testPartitionerMultipleBucketsFail(
+ TestBucketPartitionerUtils.TableSchemaType tableSchemaType) {
+ PartitionSpec partitionSpec =
+ TestBucketPartitionerUtils.getPartitionSpec(tableSchemaType, NUM_BUCKETS);
+
+ Assertions.assertThatExceptionOfType(RuntimeException.class)
+ .isThrownBy(() -> new BucketPartitioner(partitionSpec))
+ .withMessageContaining(BucketPartitionerUtils.BAD_NUMBER_OF_BUCKETS_ERROR_MESSAGE);
+ }
+
+ @ParameterizedTest
Review Comment:
parameterized test seems unnecessary here
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitioner.java:
##########
@@ -0,0 +1,99 @@
+/*
+ * 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.apache.iceberg.flink.sink.BucketPartitioner.BUCKET_OUT_OF_RANGE_MESSAGE_PREFIX;
+
+import org.apache.iceberg.PartitionSpec;
+import org.assertj.core.api.Assertions;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+import org.junit.jupiter.params.provider.ValueSource;
+
+public class TestBucketPartitioner {
+
+ static final int NUM_BUCKETS = 60;
+
+ @ParameterizedTest
+ @EnumSource(
+ value = TestBucketPartitionerUtils.TableSchemaType.class,
+ names = {"ONE_BUCKET", "IDENTITY_AND_BUCKET"})
+ public void testPartitioningParallelismGreaterThanBuckets(
+ TestBucketPartitionerUtils.TableSchemaType tableSchemaType) {
+ final int numPartitions = 500;
+
+ PartitionSpec partitionSpec =
+ TestBucketPartitionerUtils.getPartitionSpec(tableSchemaType, NUM_BUCKETS);
+
+ BucketPartitioner bucketPartitioner = new BucketPartitioner(partitionSpec);
+
+ for (int expectedIdx = 0, bucketId = 0; expectedIdx < numPartitions; expectedIdx++) {
+ int actualIdx = bucketPartitioner.partition(bucketId, numPartitions);
+ Assertions.assertThat(actualIdx).isEqualTo(expectedIdx);
+ if (++bucketId == NUM_BUCKETS) {
+ bucketId = 0;
+ }
+ }
+ }
+
+ @ParameterizedTest
+ @EnumSource(
+ value = TestBucketPartitionerUtils.TableSchemaType.class,
+ names = {"ONE_BUCKET", "IDENTITY_AND_BUCKET"})
+ public void testPartitioningParallelismEqualLessThanBuckets(
+ TestBucketPartitionerUtils.TableSchemaType tableSchemaType) {
+ final int numPartitions = 30;
+
+ PartitionSpec partitionSpec =
+ TestBucketPartitionerUtils.getPartitionSpec(tableSchemaType, NUM_BUCKETS);
+
+ BucketPartitioner bucketPartitioner = new BucketPartitioner(partitionSpec);
+
+ for (int bucketId = 0; bucketId < NUM_BUCKETS; bucketId++) {
+ int actualIdx = bucketPartitioner.partition(bucketId, numPartitions);
+ Assertions.assertThat(actualIdx).isEqualTo(bucketId % numPartitions);
+ }
+ }
+
+ @ParameterizedTest
+ @EnumSource(value = TestBucketPartitionerUtils.TableSchemaType.class, names = "TWO_BUCKETS")
+ public void testPartitionerMultipleBucketsFail(
+ TestBucketPartitionerUtils.TableSchemaType tableSchemaType) {
+ PartitionSpec partitionSpec =
+ TestBucketPartitionerUtils.getPartitionSpec(tableSchemaType, NUM_BUCKETS);
+
+ Assertions.assertThatExceptionOfType(RuntimeException.class)
+ .isThrownBy(() -> new BucketPartitioner(partitionSpec))
+ .withMessageContaining(BucketPartitionerUtils.BAD_NUMBER_OF_BUCKETS_ERROR_MESSAGE);
+ }
+
+ @ParameterizedTest
+ @ValueSource(ints = {-1, NUM_BUCKETS})
+ public void testPartitionerBucketIdOutOfRangeFail(int bucketId) {
+ PartitionSpec partitionSpec =
+ TestBucketPartitionerUtils.getPartitionSpec(
+ TestBucketPartitionerUtils.TableSchemaType.ONE_BUCKET, NUM_BUCKETS);
+
+ BucketPartitioner bucketPartitioner = new BucketPartitioner(partitionSpec);
+
+ Assertions.assertThatExceptionOfType(IllegalArgumentException.class)
+ .isThrownBy(() -> bucketPartitioner.partition(bucketId, 1))
+ .withMessageContaining(BUCKET_OUT_OF_RANGE_MESSAGE_PREFIX);
Review Comment:
nit: I haven't seen Iceberg code use constant string for unit test error message. maybe also assert the exact error msg.
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitioner.java:
##########
@@ -0,0 +1,99 @@
+/*
+ * 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.apache.iceberg.flink.sink.BucketPartitioner.BUCKET_OUT_OF_RANGE_MESSAGE_PREFIX;
+
+import org.apache.iceberg.PartitionSpec;
+import org.assertj.core.api.Assertions;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+import org.junit.jupiter.params.provider.ValueSource;
+
+public class TestBucketPartitioner {
+
+ static final int NUM_BUCKETS = 60;
+
+ @ParameterizedTest
+ @EnumSource(
+ value = TestBucketPartitionerUtils.TableSchemaType.class,
+ names = {"ONE_BUCKET", "IDENTITY_AND_BUCKET"})
+ public void testPartitioningParallelismGreaterThanBuckets(
+ TestBucketPartitionerUtils.TableSchemaType tableSchemaType) {
+ final int numPartitions = 500;
+
+ PartitionSpec partitionSpec =
+ TestBucketPartitionerUtils.getPartitionSpec(tableSchemaType, NUM_BUCKETS);
+
+ BucketPartitioner bucketPartitioner = new BucketPartitioner(partitionSpec);
+
+ for (int expectedIdx = 0, bucketId = 0; expectedIdx < numPartitions; expectedIdx++) {
+ int actualIdx = bucketPartitioner.partition(bucketId, numPartitions);
+ Assertions.assertThat(actualIdx).isEqualTo(expectedIdx);
+ if (++bucketId == NUM_BUCKETS) {
Review Comment:
Iceberg style doesn't use `++` with side effect. you can use the ternary `? : ` operator here
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerFlinkIcebergSink.java:
##########
@@ -0,0 +1,256 @@
+/*
+ * 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.apache.iceberg.flink.MiniClusterResource.DISABLE_CLASSLOADER_CHECK_CONFIG;
+import static org.apache.iceberg.flink.TestFixtures.DATABASE;
+import static org.apache.iceberg.flink.TestFixtures.TABLE_IDENTIFIER;
+import static org.apache.iceberg.flink.sink.TestBucketPartitionerUtils.CONVERTER;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.TreeMap;
+import java.util.stream.Collectors;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.typeutils.RowTypeInfo;
+import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.test.junit5.MiniClusterExtension;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.DistributionMode;
+import org.apache.iceberg.FileFormat;
+import org.apache.iceberg.FileScanTask;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Table;
+import org.apache.iceberg.TableProperties;
+import org.apache.iceberg.flink.HadoopCatalogExtension;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.flink.TableLoader;
+import org.apache.iceberg.flink.TestFixtures;
+import org.apache.iceberg.flink.sink.TestBucketPartitionerUtils.TableSchemaType;
+import org.apache.iceberg.flink.source.BoundedTestSource;
+import org.apache.iceberg.flink.util.FlinkCompatibilityUtil;
+import org.apache.iceberg.io.CloseableIterable;
+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.assertj.core.api.Assertions;
+import org.jetbrains.annotations.NotNull;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+public class TestBucketPartitionerFlinkIcebergSink {
+
+ private static final int NUMBER_TASK_MANAGERS = 1;
+ private static final int SLOTS_PER_TASK_MANAGER = 8;
+
+ @RegisterExtension
+ private static final MiniClusterExtension MINI_CLUSTER_RESOURCE =
+ new MiniClusterExtension(
+ new MiniClusterResourceConfiguration.Builder()
+ .setNumberTaskManagers(NUMBER_TASK_MANAGERS)
+ .setNumberSlotsPerTaskManager(SLOTS_PER_TASK_MANAGER)
+ .setConfiguration(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .build());
+
+ @RegisterExtension
+ private static final HadoopCatalogExtension catalogExtension =
+ new HadoopCatalogExtension(DATABASE, TestFixtures.TABLE);
+
+ private static final TypeInformation<Row> ROW_TYPE_INFO =
+ new RowTypeInfo(SimpleDataUtil.FLINK_SCHEMA.getFieldTypes());
+
+ // Parallelism = 8 (parallelism > numBuckets) throughout the test suite
+ private final int parallelism = NUMBER_TASK_MANAGERS * SLOTS_PER_TASK_MANAGER;
+ private final FileFormat format = FileFormat.PARQUET;
+ private final int numBuckets = 4;
+
+ private Table table;
+ private StreamExecutionEnvironment env;
+ private TableLoader tableLoader;
+
+ private void setupEnvironment(TableSchemaType tableSchemaType) {
+ table = getTable(tableSchemaType);
+ env =
+ StreamExecutionEnvironment.getExecutionEnvironment(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .enableCheckpointing(100)
+ .setParallelism(parallelism)
+ .setMaxParallelism(parallelism * 2);
+ tableLoader = catalogExtension.tableLoader();
+ }
+
+ private Table getTable(TableSchemaType tableSchemaType) {
+ PartitionSpec partitionSpec =
+ TestBucketPartitionerUtils.getPartitionSpec(tableSchemaType, numBuckets);
+
+ return catalogExtension
+ .catalog()
+ .createTable(
+ TABLE_IDENTIFIER,
+ SimpleDataUtil.SCHEMA,
+ partitionSpec,
+ ImmutableMap.of(TableProperties.DEFAULT_FILE_FORMAT, format.name()));
+ }
+
+ private List<RowData> convertToRowData(List<Row> rows) {
+ return rows.stream().map(CONVERTER::toInternal).collect(Collectors.toList());
+ }
+
+ private BoundedTestSource<Row> createBoundedSource(List<Row> rows) {
Review Comment:
this util method seems not necessary. it is better to stay with `RowData`
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerFlinkIcebergSink.java:
##########
@@ -0,0 +1,256 @@
+/*
+ * 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.apache.iceberg.flink.MiniClusterResource.DISABLE_CLASSLOADER_CHECK_CONFIG;
+import static org.apache.iceberg.flink.TestFixtures.DATABASE;
+import static org.apache.iceberg.flink.TestFixtures.TABLE_IDENTIFIER;
+import static org.apache.iceberg.flink.sink.TestBucketPartitionerUtils.CONVERTER;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.TreeMap;
+import java.util.stream.Collectors;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.typeutils.RowTypeInfo;
+import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.test.junit5.MiniClusterExtension;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.DistributionMode;
+import org.apache.iceberg.FileFormat;
+import org.apache.iceberg.FileScanTask;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Table;
+import org.apache.iceberg.TableProperties;
+import org.apache.iceberg.flink.HadoopCatalogExtension;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.flink.TableLoader;
+import org.apache.iceberg.flink.TestFixtures;
+import org.apache.iceberg.flink.sink.TestBucketPartitionerUtils.TableSchemaType;
+import org.apache.iceberg.flink.source.BoundedTestSource;
+import org.apache.iceberg.flink.util.FlinkCompatibilityUtil;
+import org.apache.iceberg.io.CloseableIterable;
+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.assertj.core.api.Assertions;
+import org.jetbrains.annotations.NotNull;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+public class TestBucketPartitionerFlinkIcebergSink {
+
+ private static final int NUMBER_TASK_MANAGERS = 1;
+ private static final int SLOTS_PER_TASK_MANAGER = 8;
+
+ @RegisterExtension
+ private static final MiniClusterExtension MINI_CLUSTER_RESOURCE =
+ new MiniClusterExtension(
+ new MiniClusterResourceConfiguration.Builder()
+ .setNumberTaskManagers(NUMBER_TASK_MANAGERS)
+ .setNumberSlotsPerTaskManager(SLOTS_PER_TASK_MANAGER)
+ .setConfiguration(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .build());
+
+ @RegisterExtension
+ private static final HadoopCatalogExtension catalogExtension =
+ new HadoopCatalogExtension(DATABASE, TestFixtures.TABLE);
+
+ private static final TypeInformation<Row> ROW_TYPE_INFO =
+ new RowTypeInfo(SimpleDataUtil.FLINK_SCHEMA.getFieldTypes());
+
+ // Parallelism = 8 (parallelism > numBuckets) throughout the test suite
+ private final int parallelism = NUMBER_TASK_MANAGERS * SLOTS_PER_TASK_MANAGER;
+ private final FileFormat format = FileFormat.PARQUET;
+ private final int numBuckets = 4;
+
+ private Table table;
+ private StreamExecutionEnvironment env;
+ private TableLoader tableLoader;
+
+ private void setupEnvironment(TableSchemaType tableSchemaType) {
+ table = getTable(tableSchemaType);
+ env =
+ StreamExecutionEnvironment.getExecutionEnvironment(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .enableCheckpointing(100)
+ .setParallelism(parallelism)
+ .setMaxParallelism(parallelism * 2);
+ tableLoader = catalogExtension.tableLoader();
+ }
+
+ private Table getTable(TableSchemaType tableSchemaType) {
+ PartitionSpec partitionSpec =
+ TestBucketPartitionerUtils.getPartitionSpec(tableSchemaType, numBuckets);
+
+ return catalogExtension
+ .catalog()
+ .createTable(
+ TABLE_IDENTIFIER,
+ SimpleDataUtil.SCHEMA,
+ partitionSpec,
+ ImmutableMap.of(TableProperties.DEFAULT_FILE_FORMAT, format.name()));
+ }
+
+ private List<RowData> convertToRowData(List<Row> rows) {
+ return rows.stream().map(CONVERTER::toInternal).collect(Collectors.toList());
+ }
+
+ private BoundedTestSource<Row> createBoundedSource(List<Row> rows) {
+ return new BoundedTestSource<>(rows.toArray(new Row[0]));
+ }
+
+ private TableTestStats extractTableTestStats(TableSchemaType tableSchemaType) throws IOException {
+ int totalRecordCount = 0;
+ Map<Integer, List<Integer>> writersPerBucket = Maps.newHashMap(); // <BucketId, List<WriterId>>
+ Map<Integer, Integer> filesPerBucket = Maps.newHashMap(); // <BucketId, NumFiles>
+ Map<Integer, Long> recordsPerFile = new TreeMap<>(); // <WriterId, NumRecords>
+
+ try (CloseableIterable<FileScanTask> fileScanTasks = table.newScan().planFiles()) {
+ for (FileScanTask scanTask : fileScanTasks) {
+ long recordCountInFile = scanTask.file().recordCount();
+
+ String[] splitFilePath = scanTask.file().path().toString().split("/");
+ String filename = splitFilePath[splitFilePath.length - 1];
+ int writerId = Integer.parseInt(filename.split("-")[0]);
Review Comment:
We can also disable checkpoint so that there is only one flush/commit in the test.
case 1: number of buckets (like 2) is less than the number of writers (like 5). with keyBy, there will be two files committed. with bucket partitioner, there should be 5 files.
case 2: number of buckets (like 4) is less than the number of writers (like 2). both keyBy and bucket partition will write 4 files. this is the case where we need the extra complexity of paring the writer id. hopefully/likely, the keyBy will result in unbalanced assignment. e.g. one writer handles 3 buckets while the other handles 1 bucket. that will demonstrate the benefit and expected behavior of bucket partitioner
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerFlinkIcebergSink.java:
##########
@@ -0,0 +1,256 @@
+/*
+ * 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.apache.iceberg.flink.MiniClusterResource.DISABLE_CLASSLOADER_CHECK_CONFIG;
+import static org.apache.iceberg.flink.TestFixtures.DATABASE;
+import static org.apache.iceberg.flink.TestFixtures.TABLE_IDENTIFIER;
+import static org.apache.iceberg.flink.sink.TestBucketPartitionerUtils.CONVERTER;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.TreeMap;
+import java.util.stream.Collectors;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.typeutils.RowTypeInfo;
+import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.test.junit5.MiniClusterExtension;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.DistributionMode;
+import org.apache.iceberg.FileFormat;
+import org.apache.iceberg.FileScanTask;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Table;
+import org.apache.iceberg.TableProperties;
+import org.apache.iceberg.flink.HadoopCatalogExtension;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.flink.TableLoader;
+import org.apache.iceberg.flink.TestFixtures;
+import org.apache.iceberg.flink.sink.TestBucketPartitionerUtils.TableSchemaType;
+import org.apache.iceberg.flink.source.BoundedTestSource;
+import org.apache.iceberg.flink.util.FlinkCompatibilityUtil;
+import org.apache.iceberg.io.CloseableIterable;
+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.assertj.core.api.Assertions;
+import org.jetbrains.annotations.NotNull;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+public class TestBucketPartitionerFlinkIcebergSink {
+
+ private static final int NUMBER_TASK_MANAGERS = 1;
+ private static final int SLOTS_PER_TASK_MANAGER = 8;
+
+ @RegisterExtension
+ private static final MiniClusterExtension MINI_CLUSTER_RESOURCE =
+ new MiniClusterExtension(
+ new MiniClusterResourceConfiguration.Builder()
+ .setNumberTaskManagers(NUMBER_TASK_MANAGERS)
+ .setNumberSlotsPerTaskManager(SLOTS_PER_TASK_MANAGER)
+ .setConfiguration(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .build());
+
+ @RegisterExtension
+ private static final HadoopCatalogExtension catalogExtension =
+ new HadoopCatalogExtension(DATABASE, TestFixtures.TABLE);
+
+ private static final TypeInformation<Row> ROW_TYPE_INFO =
+ new RowTypeInfo(SimpleDataUtil.FLINK_SCHEMA.getFieldTypes());
+
+ // Parallelism = 8 (parallelism > numBuckets) throughout the test suite
+ private final int parallelism = NUMBER_TASK_MANAGERS * SLOTS_PER_TASK_MANAGER;
+ private final FileFormat format = FileFormat.PARQUET;
+ private final int numBuckets = 4;
+
+ private Table table;
+ private StreamExecutionEnvironment env;
+ private TableLoader tableLoader;
+
+ private void setupEnvironment(TableSchemaType tableSchemaType) {
+ table = getTable(tableSchemaType);
+ env =
+ StreamExecutionEnvironment.getExecutionEnvironment(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .enableCheckpointing(100)
+ .setParallelism(parallelism)
+ .setMaxParallelism(parallelism * 2);
+ tableLoader = catalogExtension.tableLoader();
+ }
+
+ private Table getTable(TableSchemaType tableSchemaType) {
+ PartitionSpec partitionSpec =
+ TestBucketPartitionerUtils.getPartitionSpec(tableSchemaType, numBuckets);
+
+ return catalogExtension
+ .catalog()
+ .createTable(
+ TABLE_IDENTIFIER,
+ SimpleDataUtil.SCHEMA,
+ partitionSpec,
+ ImmutableMap.of(TableProperties.DEFAULT_FILE_FORMAT, format.name()));
+ }
+
+ private List<RowData> convertToRowData(List<Row> rows) {
+ return rows.stream().map(CONVERTER::toInternal).collect(Collectors.toList());
+ }
+
+ private BoundedTestSource<Row> createBoundedSource(List<Row> rows) {
+ return new BoundedTestSource<>(rows.toArray(new Row[0]));
+ }
+
+ private TableTestStats extractTableTestStats(TableSchemaType tableSchemaType) throws IOException {
Review Comment:
also the method name is not intuitive. it seems like `PartitionResult` to me
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] kengtin commented on a diff in pull request #7161: Flink: Custom partitioner for a more even data distribution when bucket partitions are present
Posted by "kengtin (via GitHub)" <gi...@apache.org>.
kengtin commented on code in PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#discussion_r1159051792
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitionKeySelector.java:
##########
@@ -0,0 +1,85 @@
+/*
+ * 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 java.util.Optional;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.types.logical.RowType;
+import org.apache.iceberg.PartitionKey;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Schema;
+import org.apache.iceberg.flink.RowDataWrapper;
+
+/**
+ * Very similar to the {@link PartitionKeySelector}, but optimized to extract and return an Integer
+ * bucketId as the key.
+ */
+public class BucketPartitionKeySelector implements KeySelector<RowData, Integer> {
+
+ private final Schema schema;
+ private final PartitionKey partitionKey;
+ private final RowType flinkSchema;
+ private static final Pattern intPattern = Pattern.compile("\\d+");
+
+ private transient RowDataWrapper rowDataWrapper;
+
+ BucketPartitionKeySelector(PartitionSpec spec, Schema schema, RowType flinkSchema) {
+ this.schema = schema;
+ this.partitionKey = new PartitionKey(spec, schema);
+ this.flinkSchema = flinkSchema;
+ }
+
+ /**
+ * Construct the {@link RowDataWrapper} lazily here because few members in it are not
+ * serializable. In this way, we don't have to serialize them with forcing.
+ */
+ private RowDataWrapper lazyRowDataWrapper() {
+ if (rowDataWrapper == null) {
+ rowDataWrapper = new RowDataWrapper(flinkSchema, schema.asStruct());
+ }
+ return rowDataWrapper;
+ }
+
+ @Override
+ public Integer getKey(RowData rowData) throws Exception {
+ partitionKey.partition(lazyRowDataWrapper().wrap(rowData));
+ Optional<Integer> bucketId = BucketPartitionKeySelector.extractInteger(partitionKey.toPath());
+
+ if (!bucketId.isPresent()) {
+ throw new IllegalStateException("Unable to extract bucket in partition key: " + partitionKey);
+ }
+
+ return bucketId.get();
+ }
+
+ /**
+ * Utility method to extract a bucketId from a string value. Input examples: bucket[5] or
+ * bucket_data[11], would return 5 and 11 bucket Ids respectively.
+ *
+ * @param value the string with the Bucket Id
+ * @return the Bucket Id as an Integer or Empty if not found
+ */
+ public static Optional<Integer> extractInteger(String value) {
Review Comment:
Yes, `partitionKey.toPath()` already returns a string so I'm just extracting precomputed bucket from there. Not sure I follow the tuples suggestion but we can sync offline.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] kengtin commented on a diff in pull request #7161: Flink: Custom partitioner for bucket partitions
Posted by "kengtin (via GitHub)" <gi...@apache.org>.
kengtin commented on code in PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#discussion_r1199503493
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerFlinkIcebergSink.java:
##########
@@ -0,0 +1,256 @@
+/*
+ * 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.apache.iceberg.flink.MiniClusterResource.DISABLE_CLASSLOADER_CHECK_CONFIG;
+import static org.apache.iceberg.flink.TestFixtures.DATABASE;
+import static org.apache.iceberg.flink.TestFixtures.TABLE_IDENTIFIER;
+import static org.apache.iceberg.flink.sink.TestBucketPartitionerUtils.CONVERTER;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.TreeMap;
+import java.util.stream.Collectors;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.typeutils.RowTypeInfo;
+import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.test.junit5.MiniClusterExtension;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.DistributionMode;
+import org.apache.iceberg.FileFormat;
+import org.apache.iceberg.FileScanTask;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Table;
+import org.apache.iceberg.TableProperties;
+import org.apache.iceberg.flink.HadoopCatalogExtension;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.flink.TableLoader;
+import org.apache.iceberg.flink.TestFixtures;
+import org.apache.iceberg.flink.sink.TestBucketPartitionerUtils.TableSchemaType;
+import org.apache.iceberg.flink.source.BoundedTestSource;
+import org.apache.iceberg.flink.util.FlinkCompatibilityUtil;
+import org.apache.iceberg.io.CloseableIterable;
+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.assertj.core.api.Assertions;
+import org.jetbrains.annotations.NotNull;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+public class TestBucketPartitionerFlinkIcebergSink {
+
+ private static final int NUMBER_TASK_MANAGERS = 1;
+ private static final int SLOTS_PER_TASK_MANAGER = 8;
+
+ @RegisterExtension
+ private static final MiniClusterExtension MINI_CLUSTER_RESOURCE =
+ new MiniClusterExtension(
+ new MiniClusterResourceConfiguration.Builder()
+ .setNumberTaskManagers(NUMBER_TASK_MANAGERS)
+ .setNumberSlotsPerTaskManager(SLOTS_PER_TASK_MANAGER)
+ .setConfiguration(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .build());
+
+ @RegisterExtension
+ private static final HadoopCatalogExtension catalogExtension =
+ new HadoopCatalogExtension(DATABASE, TestFixtures.TABLE);
+
+ private static final TypeInformation<Row> ROW_TYPE_INFO =
+ new RowTypeInfo(SimpleDataUtil.FLINK_SCHEMA.getFieldTypes());
+
+ // Parallelism = 8 (parallelism > numBuckets) throughout the test suite
+ private final int parallelism = NUMBER_TASK_MANAGERS * SLOTS_PER_TASK_MANAGER;
+ private final FileFormat format = FileFormat.PARQUET;
+ private final int numBuckets = 4;
+
+ private Table table;
+ private StreamExecutionEnvironment env;
+ private TableLoader tableLoader;
+
+ private void setupEnvironment(TableSchemaType tableSchemaType) {
+ table = getTable(tableSchemaType);
+ env =
+ StreamExecutionEnvironment.getExecutionEnvironment(DISABLE_CLASSLOADER_CHECK_CONFIG)
+ .enableCheckpointing(100)
+ .setParallelism(parallelism)
+ .setMaxParallelism(parallelism * 2);
+ tableLoader = catalogExtension.tableLoader();
+ }
+
+ private Table getTable(TableSchemaType tableSchemaType) {
+ PartitionSpec partitionSpec =
+ TestBucketPartitionerUtils.getPartitionSpec(tableSchemaType, numBuckets);
+
+ return catalogExtension
+ .catalog()
+ .createTable(
+ TABLE_IDENTIFIER,
+ SimpleDataUtil.SCHEMA,
+ partitionSpec,
+ ImmutableMap.of(TableProperties.DEFAULT_FILE_FORMAT, format.name()));
+ }
+
+ private List<RowData> convertToRowData(List<Row> rows) {
+ return rows.stream().map(CONVERTER::toInternal).collect(Collectors.toList());
+ }
+
+ private BoundedTestSource<Row> createBoundedSource(List<Row> rows) {
+ return new BoundedTestSource<>(rows.toArray(new Row[0]));
+ }
+
+ private TableTestStats extractTableTestStats(TableSchemaType tableSchemaType) throws IOException {
+ int totalRecordCount = 0;
+ Map<Integer, List<Integer>> writersPerBucket = Maps.newHashMap(); // <BucketId, List<WriterId>>
+ Map<Integer, Integer> filesPerBucket = Maps.newHashMap(); // <BucketId, NumFiles>
+ Map<Integer, Long> recordsPerFile = new TreeMap<>(); // <WriterId, NumRecords>
+
+ try (CloseableIterable<FileScanTask> fileScanTasks = table.newScan().planFiles()) {
+ for (FileScanTask scanTask : fileScanTasks) {
+ long recordCountInFile = scanTask.file().recordCount();
+
+ String[] splitFilePath = scanTask.file().path().toString().split("/");
+ String filename = splitFilePath[splitFilePath.length - 1];
+ int writerId = Integer.parseInt(filename.split("-")[0]);
+
+ totalRecordCount += recordCountInFile;
+ int bucketId =
+ scanTask
+ .file()
+ .partition()
+ .get(tableSchemaType == TableSchemaType.ONE_BUCKET ? 0 : 1, Integer.class);
Review Comment:
I was going to rely on the partitionSpec definition but this is more straightforward
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] stevenzwu commented on a diff in pull request #7161: Flink: Custom partitioner for bucket partitions
Posted by "stevenzwu (via GitHub)" <gi...@apache.org>.
stevenzwu commented on code in PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#discussion_r1288988277
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitionKeySelector.java:
##########
@@ -0,0 +1,69 @@
+/*
+ * 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 java.util.stream.IntStream;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.types.logical.RowType;
+import org.apache.iceberg.PartitionKey;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Schema;
+import org.apache.iceberg.flink.RowDataWrapper;
+
+/**
+ * A {@link KeySelector} that extracts the bucketId from a data row's bucket partition as the key.
+ * To be used with the {@link BucketPartitioner}.
+ */
+class BucketPartitionKeySelector implements KeySelector<RowData, Integer> {
+
+ private final Schema schema;
+ private final PartitionKey partitionKey;
+ private final RowType flinkSchema;
+ private final int bucketFieldPosition;
+
+ private transient RowDataWrapper rowDataWrapper;
+
+ BucketPartitionKeySelector(PartitionSpec partitionSpec, Schema schema, RowType flinkSchema) {
+ this.schema = schema;
+ this.partitionKey = new PartitionKey(partitionSpec, schema);
+ this.flinkSchema = flinkSchema;
+ this.bucketFieldPosition = getBucketFieldPosition(partitionSpec);
+ }
+
+ private int getBucketFieldPosition(PartitionSpec partitionSpec) {
+ int bucketFieldId = BucketPartitionerUtil.getBucketFieldId(partitionSpec);
+ return IntStream.range(0, partitionSpec.fields().size())
+ .filter(i -> partitionSpec.fields().get(i).fieldId() == bucketFieldId)
+ .toArray()[0];
+ }
+
+ private RowDataWrapper lazyRowDataWrapper() {
+ if (rowDataWrapper == null) {
+ rowDataWrapper = new RowDataWrapper(flinkSchema, schema.asStruct());
+ }
+ return rowDataWrapper;
Review Comment:
nit: iceberg style adds an empty line after the end of a control block `}`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] advancedxy commented on pull request #7161: Flink: Custom partitioner for bucket partitions
Posted by "advancedxy (via GitHub)" <gi...@apache.org>.
advancedxy commented on PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#issuecomment-1689300761
> Perhaps we could allow multi-bucket/key definitions but I'd still restrict the routing/distribution to be based only on the first bucket (or something like that) to not overcomplicate it. Thoughts?
Ah, yeah. it's definitely much simpler to routing/distribution by one bucket partition(first one or the one has the largest number of bucket, etc., a.k.a the most significant bucket transform). When multi-arg bucket transform is not supported, perhaps we may allow multiple bucket transforms in the table but choose only one bucket transform. Once multi-arg bucket transform is supported, the multi-arg bucket transform should probably to be selected first.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] kengtin commented on a diff in pull request #7161: Flink: Custom partitioner for a more even data distribution when bucket partitions are present
Posted by "kengtin (via GitHub)" <gi...@apache.org>.
kengtin commented on code in PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#discussion_r1174613739
##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/BucketPartitionerUtils.java:
##########
@@ -0,0 +1,97 @@
+/*
+ * 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 java.util.List;
+import java.util.Objects;
+import java.util.stream.Collectors;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.transforms.PartitionSpecVisitor;
+
+final class BucketPartitionerUtils {
+ static final String BAD_NUMBER_OF_BUCKETS_ERROR_MESSAGE =
+ "Expected 1 Bucket transform in the provided PartitionSpec, found: ";
+
+ private BucketPartitionerUtils() {}
+
+ public static boolean hasOneBucketField(PartitionSpec partitionSpec) {
+ List<Tuple2<Integer, Integer>> bucketFields = getBucketFields(partitionSpec);
+ return bucketFields != null && bucketFields.size() == 1;
+ }
+
+ public static List<Tuple2<Integer, Integer>> getBucketFields(PartitionSpec spec) {
+ return PartitionSpecVisitor.visit(
+ spec,
+ new PartitionSpecVisitor<Tuple2<Integer, Integer>>() {
Review Comment:
Sure, I was just following the same convention of the `Partitioning` implementation
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
[GitHub] [iceberg] stevenzwu commented on a diff in pull request #7161: Flink: Custom partitioner for bucket partitions
Posted by "stevenzwu (via GitHub)" <gi...@apache.org>.
stevenzwu commented on code in PR #7161:
URL: https://github.com/apache/iceberg/pull/7161#discussion_r1263978026
##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/TestBucketPartitionerUtil.java:
##########
@@ -0,0 +1,118 @@
+/*
+ * 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 java.util.List;
+import java.util.UUID;
+import org.apache.flink.table.data.GenericRowData;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.data.StringData;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.apache.iceberg.util.BucketUtil;
+
+final class TestBucketPartitionerUtil {
+
+ enum TableSchemaType {
+ ONE_BUCKET {
+ @Override
+ public int bucketPartitionColumnPosition() {
+ return 0;
+ }
+
+ @Override
+ public PartitionSpec getPartitionSpec(int numBuckets) {
+ return PartitionSpec.builderFor(SimpleDataUtil.SCHEMA).bucket("data", numBuckets).build();
+ }
+ },
+ IDENTITY_AND_BUCKET {
+ @Override
+ public PartitionSpec getPartitionSpec(int numBuckets) {
+ return PartitionSpec.builderFor(SimpleDataUtil.SCHEMA)
+ .identity("id")
+ .bucket("data", numBuckets)
+ .build();
+ }
+ },
+ TWO_BUCKETS {
+ @Override
+ public PartitionSpec getPartitionSpec(int numBuckets) {
+ return PartitionSpec.builderFor(SimpleDataUtil.SCHEMA)
+ .bucket("id", numBuckets)
+ .bucket("data", numBuckets)
+ .build();
+ }
+ };
+
+ public int bucketPartitionColumnPosition() {
Review Comment:
yes, it will add a couple more lines. but I think it is more readable. it is not intuitive to define this in the base without defining the partition spec in the base. I would rather to keep both closely together.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org