diff --git a/hudi-common/src/main/java/org/apache/hudi/common/util/collection/CloseableConcatenatingIterator.java b/hudi-common/src/main/java/org/apache/hudi/common/util/collection/CloseableConcatenatingIterator.java index f0ea1cd911413..23fd5721fd23f 100644 --- a/hudi-common/src/main/java/org/apache/hudi/common/util/collection/CloseableConcatenatingIterator.java +++ b/hudi-common/src/main/java/org/apache/hudi/common/util/collection/CloseableConcatenatingIterator.java @@ -25,7 +25,8 @@ * Provides closeable iterator interface over list of iterators. Consumes all records from first iterator element * before moving to next iterator in the list. That is concatenating elements across multiple iterators. */ -public class CloseableConcatenatingIterator extends ConcatenatingIterator { +public class CloseableConcatenatingIterator extends ConcatenatingIterator + implements ClosableIterator { public CloseableConcatenatingIterator(List> iterators) { super(iterators); @@ -37,4 +38,12 @@ protected Iterator advanceIterator() { previous.close(); return previous; } + + @Override + public void close() { + Iterator iterator; + while ((iterator = super.advanceIterator()) != null) { + ((ClosableIterator) iterator).close(); + } + } } diff --git a/hudi-common/src/test/java/org/apache/hudi/common/util/collection/TestConcatenatingIterator.java b/hudi-common/src/test/java/org/apache/hudi/common/util/collection/TestConcatenatingIterator.java index 71838f290748c..c62acd6fa9419 100644 --- a/hudi-common/src/test/java/org/apache/hudi/common/util/collection/TestConcatenatingIterator.java +++ b/hudi-common/src/test/java/org/apache/hudi/common/util/collection/TestConcatenatingIterator.java @@ -29,6 +29,9 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.fail; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; public class TestConcatenatingIterator { @@ -62,4 +65,21 @@ public void testConcatError() { // } } -} \ No newline at end of file + + @Test + public void testCloseableConcatClosesRemainingIterators() { + ClosableIterator first = mock(ClosableIterator.class); + ClosableIterator second = mock(ClosableIterator.class); + when(first.hasNext()).thenReturn(true); + when(first.next()).thenReturn(1); + + CloseableConcatenatingIterator iterator = + new CloseableConcatenatingIterator<>(Arrays.asList(first, second)); + assertEquals(1, iterator.next()); + + iterator.close(); + + verify(first).close(); + verify(second).close(); + } +} diff --git a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/bulk/BulkInsertWriterHelper.java b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/bulk/BulkInsertWriterHelper.java index 30a19e1423aa6..212910bb612a7 100644 --- a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/bulk/BulkInsertWriterHelper.java +++ b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/bulk/BulkInsertWriterHelper.java @@ -59,7 +59,7 @@ * Helper class for bulk insert used by Flink. */ @Slf4j -public class BulkInsertWriterHelper { +public class BulkInsertWriterHelper implements AutoCloseable { @Getter protected final String instantTime; @@ -170,6 +170,7 @@ private HoodieRowDataCreateHandle getRowCreateHandle(String partitionPath) throw return handles.get(partitionPath); } + @Override public void close() throws IOException { if (handles.isEmpty()) { return; @@ -228,4 +229,3 @@ private WriteStatus closeWriteHandle(HoodieRowDataCreateHandle rowCreateHandle) } } - diff --git a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/clustering/ClusteringOperator.java b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/clustering/ClusteringOperator.java index 21297a1b5b313..9c1b1e6d7c1ca 100644 --- a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/clustering/ClusteringOperator.java +++ b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/clustering/ClusteringOperator.java @@ -87,7 +87,6 @@ import java.io.IOException; import java.util.ArrayList; import java.util.Collections; -import java.util.Iterator; import java.util.List; import java.util.stream.Collectors; @@ -229,50 +228,55 @@ public void endInput() { private void doClustering(String instantTime, List clusteringOperations) throws Exception { clusteringMetrics.startClustering(); - BulkInsertWriterHelper writerHelper = new BulkInsertWriterHelper(this.conf, this.table, this.writeConfig, + try (BulkInsertWriterHelper writerHelper = new BulkInsertWriterHelper(this.conf, this.table, this.writeConfig, instantTime, this.taskID, RuntimeContextUtils.getNumberOfParallelSubtasks(getRuntimeContext()), - RuntimeContextUtils.getAttemptNumber(getRuntimeContext()), this.rowType, true); - - Iterator iterator; - if (clusteringOperations.stream().anyMatch(operation -> CollectionUtils.nonEmpty(operation.getDeltaFilePaths()))) { - // if there are log files, we read all records into memory for a file group and apply updates. - iterator = readRecordsForGroupWithLogs(clusteringOperations, instantTime); - } else { - // We want to optimize reading records for case there are no log files. - iterator = readRecordsForGroupBaseFiles(clusteringOperations); - } - - if (this.sortClusteringEnabled) { - RowDataSerializer rowDataSerializer = new RowDataSerializer(rowType); - BinaryExternalSorter sorter = initSorter(); - while (iterator.hasNext()) { - RowData rowData = iterator.next(); - BinaryRowData binaryRowData = rowDataSerializer.toBinaryRow(rowData).copy(); - sorter.write(binaryRowData); + RuntimeContextUtils.getAttemptNumber(getRuntimeContext()), this.rowType, true)) { + ClosableIterator iterator; + if (clusteringOperations.stream().anyMatch(operation -> CollectionUtils.nonEmpty(operation.getDeltaFilePaths()))) { + // if there are log files, we read all records into memory for a file group and apply updates. + iterator = readRecordsForGroupWithLogs(clusteringOperations, instantTime); + } else { + // We want to optimize reading records for case there are no log files. + iterator = readRecordsForGroupBaseFiles(clusteringOperations); } - BinaryRowData row = binarySerializer.createInstance(); - while ((row = sorter.getIterator().next(row)) != null) { - writerHelper.write(row); + try (ClosableIterator closeableIterator = iterator) { + if (this.sortClusteringEnabled) { + RowDataSerializer rowDataSerializer = new RowDataSerializer(rowType); + BinaryExternalSorter sorter = initSorter(); + try { + while (closeableIterator.hasNext()) { + RowData rowData = closeableIterator.next(); + BinaryRowData binaryRowData = rowDataSerializer.toBinaryRow(rowData).copy(); + sorter.write(binaryRowData); + } + + BinaryRowData row = binarySerializer.createInstance(); + while ((row = sorter.getIterator().next(row)) != null) { + writerHelper.write(row); + } + } finally { + sorter.close(); + } + } else { + while (closeableIterator.hasNext()) { + writerHelper.write(closeableIterator.next()); + } + } } - sorter.close(); - } else { - while (iterator.hasNext()) { - writerHelper.write(iterator.next()); - } - } - List writeStatuses = writerHelper.getWriteStatuses(this.taskID); - clusteringMetrics.endClustering(); - collector.collect(new ClusteringCommitEvent(instantTime, getFileIds(clusteringOperations), writeStatuses, this.taskID)); - writerHelper.close(); + List writeStatuses = writerHelper.getWriteStatuses(this.taskID); + clusteringMetrics.endClustering(); + collector.collect(new ClusteringCommitEvent(instantTime, getFileIds(clusteringOperations), writeStatuses, this.taskID)); + } } /** * Read records from baseFiles, apply updates and convert to Iterator. */ @SuppressWarnings("unchecked") - private Iterator readRecordsForGroupWithLogs(List clusteringOps, String instantTime) { + private ClosableIterator readRecordsForGroupWithLogs( + List clusteringOps, String instantTime) { List> recordIterators = new ArrayList<>(); long maxMemoryPerCompaction = MergeUtils.getMaxMemoryPerCompaction(new FlinkTaskContextSupplier(null), writeConfig); log.info("MaxMemoryPerCompaction run as part of clustering => {}", maxMemoryPerCompaction); @@ -304,7 +308,7 @@ private ClosableIterator getRecordIterator(ClusteringOperation clusterO /** * Read records from baseFiles and get iterator. */ - private Iterator readRecordsForGroupBaseFiles(List clusteringOps) { + private ClosableIterator readRecordsForGroupBaseFiles(List clusteringOps) { List> iteratorsForPartition = clusteringOps.stream().map(clusteringOp -> { try { HoodieFileReaderFactory fileReaderFactory = HoodieIOFactory.getIOFactory(table.getStorage()) diff --git a/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/clustering/TestClusteringCommitSink.java b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/clustering/TestClusteringCommitSink.java new file mode 100644 index 0000000000000..e3486f3d3cc50 --- /dev/null +++ b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/clustering/TestClusteringCommitSink.java @@ -0,0 +1,252 @@ +/* + * 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.hudi.sink.clustering; + +import org.apache.hudi.avro.model.HoodieClusteringGroup; +import org.apache.hudi.avro.model.HoodieClusteringPlan; +import org.apache.hudi.client.HoodieFlinkWriteClient; +import org.apache.hudi.client.WriteStatus; +import org.apache.hudi.common.model.HoodieCommitMetadata; +import org.apache.hudi.common.model.HoodieFileGroupId; +import org.apache.hudi.common.model.HoodieWriteStat; +import org.apache.hudi.common.model.TableServiceType; +import org.apache.hudi.common.table.HoodieTableMetaClient; +import org.apache.hudi.common.table.timeline.HoodieActiveTimeline; +import org.apache.hudi.common.table.timeline.HoodieInstant; +import org.apache.hudi.common.table.timeline.InstantGenerator; +import org.apache.hudi.common.util.ClusteringUtils; +import org.apache.hudi.common.util.Option; +import org.apache.hudi.common.util.collection.Pair; +import org.apache.hudi.config.HoodieWriteConfig; +import org.apache.hudi.configuration.FlinkOptions; +import org.apache.hudi.sink.utils.MockStreamingRuntimeContext; +import org.apache.hudi.table.HoodieFlinkTable; +import org.apache.hudi.util.ClusteringUtil; +import org.apache.hudi.util.FlinkWriteClients; + +import org.apache.flink.configuration.Configuration; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.MockedStatic; + +import java.util.Arrays; +import java.util.List; +import java.util.stream.Stream; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.ArgumentMatchers.same; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.mockStatic; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +/** + * Tests {@link ClusteringCommitSink}. + */ +class TestClusteringCommitSink { + + private static final String INSTANT = "20240101000000000"; + + private Configuration conf; + private HoodieFlinkWriteClient writeClient; + private HoodieFlinkTable table; + private HoodieTableMetaClient metaClient; + private HoodieActiveTimeline activeTimeline; + private InstantGenerator instantGenerator; + private HoodieInstant inflightInstant; + + @BeforeEach + @SuppressWarnings("unchecked") + void setUp() { + conf = new Configuration(); + conf.set(FlinkOptions.CLEAN_ASYNC_ENABLED, false); + writeClient = mock(HoodieFlinkWriteClient.class); + table = mock(HoodieFlinkTable.class); + metaClient = mock(HoodieTableMetaClient.class); + activeTimeline = mock(HoodieActiveTimeline.class); + instantGenerator = mock(InstantGenerator.class); + inflightInstant = mock(HoodieInstant.class); + + when(writeClient.getHoodieTable()).thenReturn(table); + when(table.getMetaClient()).thenReturn(metaClient); + when(table.getInstantGenerator()).thenReturn(instantGenerator); + when(metaClient.getActiveTimeline()).thenReturn(activeTimeline); + } + + @Test + void testIgnoreEventsWhoseClusteringPlanNoLongerExists() throws Exception { + ClusteringCommitSink sink = openSink(); + try (MockedStatic clusteringUtils = mockStatic(ClusteringUtils.class); + MockedStatic clusteringUtil = mockStatic(ClusteringUtil.class)) { + clusteringUtils.when(() -> ClusteringUtils.getInflightClusteringInstant( + INSTANT, activeTimeline, instantGenerator)).thenReturn(Option.empty()); + + sink.invoke(new ClusteringCommitEvent(INSTANT, "file-1", 0), null); + + clusteringUtil.verifyNoInteractions(); + verify(writeClient, never()).completeTableService( + any(), any(), any(), any()); + } + } + + @Test + void testWaitForEveryGroupThenRollbackFailedClustering() throws Exception { + ClusteringCommitSink sink = openSink(); + HoodieClusteringPlan plan = planWithGroups(2); + try (MockedStatic clusteringUtils = mockStatic(ClusteringUtils.class); + MockedStatic clusteringUtil = mockStatic(ClusteringUtil.class)) { + stubClusteringPlan(clusteringUtils, plan); + + sink.invoke(successEvent("file-1", new WriteStatus()), null); + clusteringUtil.verifyNoInteractions(); + + sink.invoke(new ClusteringCommitEvent(INSTANT, "file-2", 1), null); + clusteringUtil.verify( + () -> ClusteringUtil.rollbackClustering(table, writeClient, INSTANT), times(1)); + + // A reset sink starts buffering from scratch and reloads the plan. + sink.invoke(successEvent("file-1", new WriteStatus()), null); + clusteringUtils.verify( + () -> ClusteringUtils.getClusteringPlan(metaClient, inflightInstant), times(2)); + } + } + + @Test + void testRollbackWriteStatusesWithErrorsUnlessConfiguredToIgnore() throws Exception { + ClusteringCommitSink sink = openSink(); + HoodieClusteringPlan plan = planWithGroups(1); + WriteStatus failedStatus = new WriteStatus(); + failedStatus.setTotalErrorRecords(3); + + try (MockedStatic clusteringUtils = mockStatic(ClusteringUtils.class); + MockedStatic clusteringUtil = mockStatic(ClusteringUtil.class)) { + stubClusteringPlan(clusteringUtils, plan); + + sink.invoke(successEvent("file-1", failedStatus), null); + + clusteringUtil.verify( + () -> ClusteringUtil.rollbackClustering(table, writeClient, INSTANT)); + verify(writeClient, never()).completeTableService( + any(), any(), any(), any()); + } + } + + @Test + void testCommitWriteStatusesWithErrorsWhenConfiguredToIgnore() throws Exception { + conf.set(FlinkOptions.IGNORE_FAILED, true); + ClusteringCommitSink sink = openSink(); + HoodieClusteringPlan plan = planWithGroups(1); + HoodieWriteStat writeStat = new HoodieWriteStat(); + writeStat.setPartitionPath("partition"); + writeStat.setFileId("new-file"); + writeStat.setNumWrites(1); + WriteStatus failedStatus = new WriteStatus(); + failedStatus.setStat(writeStat); + failedStatus.setTotalErrorRecords(3); + HoodieWriteConfig writeConfig = mock(HoodieWriteConfig.class); + when(writeConfig.getSchema()).thenReturn("{}"); + when(writeClient.getConfig()).thenReturn(writeConfig); + + try (MockedStatic clusteringUtils = mockStatic(ClusteringUtils.class); + MockedStatic clusteringUtil = mockStatic(ClusteringUtil.class)) { + stubClusteringPlan(clusteringUtils, plan); + clusteringUtils.when(() -> ClusteringUtils.getFileGroupsFromClusteringPlan(plan)) + .thenReturn(Stream.of(new HoodieFileGroupId("partition", "old-file"))); + + sink.invoke(successEvent("file-1", failedStatus), null); + + clusteringUtil.verifyNoInteractions(); + verify(writeClient).completeTableService( + eq(TableServiceType.CLUSTER), + any(HoodieCommitMetadata.class), + same(table), + eq(INSTANT)); + } + } + + @Test + void testCommitSuccessfulClusteringAndCleanInline() throws Exception { + ClusteringCommitSink sink = openSink(); + HoodieClusteringPlan plan = planWithGroups(1); + HoodieWriteStat writeStat = new HoodieWriteStat(); + writeStat.setPartitionPath("partition"); + writeStat.setFileId("new-file"); + writeStat.setNumWrites(1); + WriteStatus status = new WriteStatus(); + status.setStat(writeStat); + HoodieWriteConfig writeConfig = mock(HoodieWriteConfig.class); + when(writeConfig.getSchema()).thenReturn("{}"); + when(writeClient.getConfig()).thenReturn(writeConfig); + + try (MockedStatic clusteringUtils = mockStatic(ClusteringUtils.class)) { + stubClusteringPlan(clusteringUtils, plan); + clusteringUtils.when(() -> ClusteringUtils.getFileGroupsFromClusteringPlan(plan)) + .thenReturn(Stream.of( + new HoodieFileGroupId("partition", "old-file"), + new HoodieFileGroupId("partition", "new-file"))); + + sink.invoke(successEvent("file-1", status), null); + } + + verify(metaClient).reloadActiveTimeline(); + verify(writeClient).completeTableService( + eq(TableServiceType.CLUSTER), + any(HoodieCommitMetadata.class), + same(table), + eq(INSTANT)); + verify(writeClient).clean(); + } + + private ClusteringCommitSink openSink() throws Exception { + ClusteringCommitSink sink = new ClusteringCommitSink(conf); + MockStreamingRuntimeContext runtimeContext = new MockStreamingRuntimeContext(false, 1, 0); + sink.setRuntimeContext(runtimeContext); + try (MockedStatic writeClients = mockStatic(FlinkWriteClients.class)) { + writeClients.when(() -> FlinkWriteClients.createWriteClient(conf, runtimeContext)) + .thenReturn(writeClient); + sink.open(conf); + } + return sink; + } + + private void stubClusteringPlan( + MockedStatic clusteringUtils, + HoodieClusteringPlan plan) { + clusteringUtils.when(() -> ClusteringUtils.getInflightClusteringInstant( + INSTANT, activeTimeline, instantGenerator)).thenReturn(Option.of(inflightInstant)); + clusteringUtils.when(() -> ClusteringUtils.getClusteringPlan(metaClient, inflightInstant)) + .thenReturn(Option.of(Pair.of(inflightInstant, plan))); + } + + private HoodieClusteringPlan planWithGroups(int numGroups) { + HoodieClusteringPlan plan = mock(HoodieClusteringPlan.class); + HoodieClusteringGroup[] groups = new HoodieClusteringGroup[numGroups]; + Arrays.setAll(groups, ignored -> mock(HoodieClusteringGroup.class)); + when(plan.getInputGroups()).thenReturn(Arrays.asList(groups)); + return plan; + } + + private ClusteringCommitEvent successEvent(String fileId, WriteStatus... statuses) { + List statusList = Arrays.asList(statuses); + return new ClusteringCommitEvent(INSTANT, fileId, statusList, 0); + } +} diff --git a/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/clustering/TestClusteringOperator.java b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/clustering/TestClusteringOperator.java new file mode 100644 index 0000000000000..1d05302cb76f3 --- /dev/null +++ b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/clustering/TestClusteringOperator.java @@ -0,0 +1,426 @@ +/* + * 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.hudi.sink.clustering; + +import org.apache.hudi.adapter.Utils; +import org.apache.hudi.client.HoodieFlinkWriteClient; +import org.apache.hudi.common.model.ClusteringGroupInfo; +import org.apache.hudi.common.model.ClusteringOperation; +import org.apache.hudi.common.model.HoodieRecord; +import org.apache.hudi.common.model.WriteOperationType; +import org.apache.hudi.common.schema.HoodieSchema; +import org.apache.hudi.common.table.HoodieTableMetaClient; +import org.apache.hudi.common.util.collection.ClosableIterator; +import org.apache.hudi.config.HoodieWriteConfig; +import org.apache.hudi.configuration.FlinkOptions; +import org.apache.hudi.core.io.storage.HoodieFileReaderFactory; +import org.apache.hudi.core.io.storage.HoodieIOFactory; +import org.apache.hudi.io.MergeUtils; +import org.apache.hudi.sink.bulk.BulkInsertWriterHelper; +import org.apache.hudi.sink.utils.NonThrownExecutor; +import org.apache.hudi.storage.HoodieStorage; +import org.apache.hudi.storage.StoragePath; +import org.apache.hudi.table.HoodieFlinkTable; +import org.apache.hudi.table.format.FormatUtils; +import org.apache.hudi.table.format.HoodieRowDataParquetReader; +import org.apache.hudi.util.DataTypeUtils; +import org.apache.hudi.util.FlinkWriteClients; +import org.apache.hudi.utils.TestConfigurations; + +import org.apache.flink.api.common.functions.RuntimeContext; +import org.apache.flink.configuration.Configuration; +import org.apache.flink.streaming.runtime.streamrecord.StreamRecord; +import org.apache.flink.streaming.util.OneInputStreamOperatorTestHarness; +import org.apache.flink.table.data.GenericRowData; +import org.apache.flink.table.data.RowData; +import org.apache.flink.table.data.binary.BinaryRowData; +import org.apache.flink.table.runtime.operators.sort.BinaryExternalSorter; +import org.apache.flink.util.MutableObjectIterator; +import org.apache.flink.util.function.ThrowingRunnable; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; +import org.mockito.MockedConstruction; +import org.mockito.MockedStatic; + +import java.io.File; +import java.io.IOException; +import java.util.Collections; +import java.util.List; +import java.util.Map; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.ArgumentMatchers.same; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.mockConstruction; +import static org.mockito.Mockito.mockStatic; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +/** + * Tests {@link ClusteringOperator}. + */ +class TestClusteringOperator { + + @TempDir + File tempDir; + + @Test + @SuppressWarnings("unchecked") + void testProcessClusteringPlanSynchronously() throws Exception { + Configuration conf = TestConfigurations.getDefaultConf(tempDir.getAbsolutePath()); + conf.set(FlinkOptions.CLUSTERING_PLAN_STRATEGY_SMALL_FILE_LIMIT, 2048L); + HoodieFlinkWriteClient writeClient = mock(HoodieFlinkWriteClient.class); + HoodieFlinkTable table = mock(HoodieFlinkTable.class); + HoodieWriteConfig writeConfig = mock(HoodieWriteConfig.class); + when(writeClient.getHoodieTable()).thenReturn(table); + + ClusteringOperator operator = new ClusteringOperator(conf, TestConfigurations.ROW_TYPE); + assertEquals(2048L, conf.get(FlinkOptions.CLUSTERING_PLAN_STRATEGY_SMALL_FILE_LIMIT)); + ClusteringPlanEvent planEvent = event("001"); + + try (MockedStatic writeClients = mockStatic(FlinkWriteClients.class); + MockedConstruction writerHelpers = + mockConstruction(BulkInsertWriterHelper.class, (writerHelper, context) -> + when(writerHelper.getWriteStatuses(anyInt())).thenReturn(Collections.emptyList())); + OneInputStreamOperatorTestHarness harness = + new OneInputStreamOperatorTestHarness<>(operator, 1, 1, 0)) { + writeClients.when(() -> FlinkWriteClients.getHoodieClientConfig( + any(Configuration.class), eq(false), eq(false))).thenReturn(writeConfig); + writeClients.when(() -> FlinkWriteClients.createWriteClient( + any(Configuration.class), any(RuntimeContext.class))).thenReturn(writeClient); + + harness.open(); + harness.processElement(new StreamRecord<>(planEvent)); + operator.endInput(); + + StreamRecord output = + (StreamRecord) harness.getOutput().poll(); + assertEquals("001", output.getValue().getInstant()); + assertEquals("", output.getValue().getFileIds()); + assertEquals(0, output.getValue().getTaskID()); + assertFalse(output.getValue().isFailed()); + assertEquals(1, writerHelpers.constructed().size()); + verify(writerHelpers.constructed().get(0)).close(); + } + + verify(writeClient).close(); + } + + @Test + @SuppressWarnings("unchecked") + void testAsyncFailureProducesFailedCommitEvent() throws Exception { + Configuration conf = TestConfigurations.getDefaultConf(tempDir.getAbsolutePath()); + conf.set(FlinkOptions.OPERATION, WriteOperationType.INSERT.value()); + conf.set(FlinkOptions.CLUSTERING_ASYNC_ENABLED, true); + HoodieFlinkWriteClient writeClient = mock(HoodieFlinkWriteClient.class); + HoodieFlinkTable table = mock(HoodieFlinkTable.class); + HoodieWriteConfig writeConfig = mock(HoodieWriteConfig.class); + when(writeClient.getHoodieTable()).thenReturn(table); + + ClusteringOperator operator = new ClusteringOperator(conf, TestConfigurations.ROW_TYPE); + try (MockedStatic writeClients = mockStatic(FlinkWriteClients.class); + MockedConstruction executors = + mockConstruction(NonThrownExecutor.class, (executor, context) -> + doAnswer(invocation -> { + NonThrownExecutor.ExceptionHook hook = invocation.getArgument(1); + hook.apply("expected failure", new RuntimeException("test")); + return null; + }).when(executor).execute( + any(ThrowingRunnable.class), + any(NonThrownExecutor.ExceptionHook.class), + anyString(), + anyString(), + anyInt())); + OneInputStreamOperatorTestHarness harness = + new OneInputStreamOperatorTestHarness<>(operator, 1, 1, 0)) { + writeClients.when(() -> FlinkWriteClients.getHoodieClientConfig( + any(Configuration.class), eq(false), eq(false))).thenReturn(writeConfig); + writeClients.when(() -> FlinkWriteClients.createWriteClient( + any(Configuration.class), any(RuntimeContext.class))).thenReturn(writeClient); + + harness.open(); + harness.processElement(new StreamRecord<>(event("002"))); + + StreamRecord output = + (StreamRecord) harness.getOutput().poll(); + assertEquals("002", output.getValue().getInstant()); + assertEquals("", output.getValue().getFileIds()); + assertTrue(output.getValue().isFailed()); + assertEquals(1, executors.constructed().size()); + } + + verify(writeClient).close(); + } + + @Test + @SuppressWarnings("unchecked") + void testReadAndWriteRecordsFromBaseFiles() throws Exception { + Configuration conf = TestConfigurations.getDefaultConf(tempDir.getAbsolutePath()); + HoodieFlinkWriteClient writeClient = mock(HoodieFlinkWriteClient.class); + HoodieFlinkTable table = mock(HoodieFlinkTable.class); + HoodieWriteConfig writeConfig = mock(HoodieWriteConfig.class); + HoodieStorage storage = mock(HoodieStorage.class); + HoodieIOFactory ioFactory = mock(HoodieIOFactory.class); + HoodieFileReaderFactory readerFactory = mock(HoodieFileReaderFactory.class); + HoodieRowDataParquetReader fileReader = mock(HoodieRowDataParquetReader.class); + HoodieRecord hoodieRecord = mock(HoodieRecord.class); + RowData row = mock(RowData.class); + when(writeClient.getHoodieTable()).thenReturn(table); + when(table.getStorage()).thenReturn(storage); + when(table.getConfig()).thenReturn(writeConfig); + when(ioFactory.getReaderFactory(HoodieRecord.HoodieRecordType.FLINK)).thenReturn(readerFactory); + when(readerFactory.getFileReader( + same(writeConfig), any(StoragePath.class))).thenReturn(fileReader); + when(hoodieRecord.getData()).thenReturn(row); + when(fileReader.getRecordIterator(any(HoodieSchema.class))).thenReturn( + ClosableIterator.wrap(Collections.singletonList(hoodieRecord).iterator())); + + ClusteringOperation operation = new ClusteringOperation( + tempDir.toPath().resolve("base.parquet").toString(), + Collections.emptyList(), + "old-file", + "partition", + null, + 0); + ClusteringOperator operator = new ClusteringOperator(conf, TestConfigurations.ROW_TYPE); + + try (MockedStatic writeClients = mockStatic(FlinkWriteClients.class); + MockedStatic ioFactories = mockStatic(HoodieIOFactory.class); + MockedConstruction writerHelpers = + mockConstruction(BulkInsertWriterHelper.class, (writerHelper, context) -> + when(writerHelper.getWriteStatuses(anyInt())).thenReturn(Collections.emptyList())); + OneInputStreamOperatorTestHarness harness = + new OneInputStreamOperatorTestHarness<>(operator, 1, 1, 0)) { + writeClients.when(() -> FlinkWriteClients.getHoodieClientConfig( + any(Configuration.class), eq(false), eq(false))).thenReturn(writeConfig); + writeClients.when(() -> FlinkWriteClients.createWriteClient( + any(Configuration.class), any(RuntimeContext.class))).thenReturn(writeClient); + ioFactories.when(() -> HoodieIOFactory.getIOFactory(storage)).thenReturn(ioFactory); + + harness.open(); + harness.processElement(new StreamRecord<>(event("003", Collections.singletonList(operation)))); + + verify(writerHelpers.constructed().get(0)).write(row); + StreamRecord output = + (StreamRecord) harness.getOutput().poll(); + assertEquals("old-file", output.getValue().getFileIds()); + assertFalse(output.getValue().isFailed()); + } + } + + @Test + @SuppressWarnings("unchecked") + void testCloseReaderAndWriterWhenWritingFails() throws Exception { + Configuration conf = TestConfigurations.getDefaultConf(tempDir.getAbsolutePath()); + HoodieFlinkWriteClient writeClient = mock(HoodieFlinkWriteClient.class); + HoodieFlinkTable table = mock(HoodieFlinkTable.class); + HoodieWriteConfig writeConfig = mock(HoodieWriteConfig.class); + HoodieStorage storage = mock(HoodieStorage.class); + HoodieIOFactory ioFactory = mock(HoodieIOFactory.class); + HoodieFileReaderFactory readerFactory = mock(HoodieFileReaderFactory.class); + HoodieRowDataParquetReader fileReader = mock(HoodieRowDataParquetReader.class); + ClosableIterator> recordIterator = mock(ClosableIterator.class); + HoodieRecord hoodieRecord = mock(HoodieRecord.class); + RowData row = mock(RowData.class); + when(writeClient.getHoodieTable()).thenReturn(table); + when(table.getStorage()).thenReturn(storage); + when(table.getConfig()).thenReturn(writeConfig); + when(ioFactory.getReaderFactory(HoodieRecord.HoodieRecordType.FLINK)).thenReturn(readerFactory); + when(readerFactory.getFileReader( + same(writeConfig), any(StoragePath.class))).thenReturn(fileReader); + when(fileReader.getRecordIterator(any(HoodieSchema.class))).thenReturn(recordIterator); + when(recordIterator.hasNext()).thenReturn(true); + when(recordIterator.next()).thenReturn(hoodieRecord); + when(hoodieRecord.getData()).thenReturn(row); + + ClusteringOperation operation = new ClusteringOperation( + tempDir.toPath().resolve("base.parquet").toString(), + Collections.emptyList(), + "old-file", + "partition", + null, + 0); + ClusteringOperator operator = new ClusteringOperator(conf, TestConfigurations.ROW_TYPE); + + try (MockedStatic writeClients = mockStatic(FlinkWriteClients.class); + MockedStatic ioFactories = mockStatic(HoodieIOFactory.class); + MockedConstruction writerHelpers = + mockConstruction(BulkInsertWriterHelper.class, (writerHelper, context) -> + doThrow(new IOException("expected failure")).when(writerHelper).write(row)); + OneInputStreamOperatorTestHarness harness = + new OneInputStreamOperatorTestHarness<>(operator, 1, 1, 0)) { + writeClients.when(() -> FlinkWriteClients.getHoodieClientConfig( + any(Configuration.class), eq(false), eq(false))).thenReturn(writeConfig); + writeClients.when(() -> FlinkWriteClients.createWriteClient( + any(Configuration.class), any(RuntimeContext.class))).thenReturn(writeClient); + ioFactories.when(() -> HoodieIOFactory.getIOFactory(storage)).thenReturn(ioFactory); + + harness.open(); + assertThrows(IOException.class, () -> + harness.processElement(new StreamRecord<>(event("004", Collections.singletonList(operation))))); + + verify(recordIterator).close(); + verify(writerHelpers.constructed().get(0)).close(); + } + } + + @Test + @SuppressWarnings("unchecked") + void testReadAndWriteRecordsFromFileSliceWithLogs() throws Exception { + Configuration conf = TestConfigurations.getDefaultConf(tempDir.getAbsolutePath()); + HoodieFlinkWriteClient writeClient = mock(HoodieFlinkWriteClient.class); + HoodieFlinkTable table = mock(HoodieFlinkTable.class); + HoodieTableMetaClient metaClient = mock(HoodieTableMetaClient.class); + HoodieWriteConfig writeConfig = mock(HoodieWriteConfig.class); + org.apache.hudi.common.table.read.HoodieRecordReader recordReader = + mock(org.apache.hudi.common.table.read.HoodieRecordReader.class); + RowData row = mock(RowData.class); + when(writeClient.getHoodieTable()).thenReturn(table); + when(table.getMetaClient()).thenReturn(metaClient); + when(recordReader.getClosableIterator()).thenReturn( + ClosableIterator.wrap(Collections.singletonList(row).iterator())); + + ClusteringOperation operation = new ClusteringOperation( + tempDir.toPath().resolve("old-file_0-0-0_001.parquet").toString(), + Collections.singletonList(tempDir.toPath().resolve(".old-file_001.log.1_0-0-0").toString()), + "old-file", + "partition", + null, + 0); + ClusteringOperator operator = new ClusteringOperator(conf, TestConfigurations.ROW_TYPE); + + try (MockedStatic writeClients = mockStatic(FlinkWriteClients.class); + MockedStatic mergeUtils = mockStatic(MergeUtils.class); + MockedStatic formatUtils = mockStatic(FormatUtils.class); + MockedConstruction writerHelpers = + mockConstruction(BulkInsertWriterHelper.class, (writerHelper, context) -> + when(writerHelper.getWriteStatuses(anyInt())).thenReturn(Collections.emptyList())); + OneInputStreamOperatorTestHarness harness = + new OneInputStreamOperatorTestHarness<>(operator, 1, 1, 0)) { + writeClients.when(() -> FlinkWriteClients.getHoodieClientConfig( + any(Configuration.class), eq(false), eq(false))).thenReturn(writeConfig); + writeClients.when(() -> FlinkWriteClients.createWriteClient( + any(Configuration.class), any(RuntimeContext.class))).thenReturn(writeClient); + mergeUtils.when(() -> MergeUtils.getMaxMemoryPerCompaction(any(), same(writeConfig))) + .thenReturn(1024L); + formatUtils.when(() -> FormatUtils.createRecordReader( + same(metaClient), same(writeConfig), any(), any(), any(), any(), eq("004"), + eq(FlinkOptions.REALTIME_PAYLOAD_COMBINE), eq(false), eq(Collections.emptyList()), any())) + .thenReturn(recordReader); + + harness.open(); + harness.processElement(new StreamRecord<>(event("004", Collections.singletonList(operation)))); + + verify(recordReader).getClosableIterator(); + verify(writerHelpers.constructed().get(0)).write(row); + StreamRecord output = + (StreamRecord) harness.getOutput().poll(); + assertEquals("old-file", output.getValue().getFileIds()); + assertFalse(output.getValue().isFailed()); + } + } + + @Test + @SuppressWarnings("unchecked") + void testSortClustering() throws Exception { + Configuration conf = TestConfigurations.getDefaultConf(tempDir.getAbsolutePath()); + conf.set(FlinkOptions.CLUSTERING_SORT_COLUMNS, TestConfigurations.ROW_TYPE.getFieldNames().get(0)); + HoodieFlinkWriteClient writeClient = mock(HoodieFlinkWriteClient.class); + HoodieFlinkTable table = mock(HoodieFlinkTable.class); + HoodieWriteConfig writeConfig = mock(HoodieWriteConfig.class); + HoodieStorage storage = mock(HoodieStorage.class); + HoodieIOFactory ioFactory = mock(HoodieIOFactory.class); + HoodieFileReaderFactory readerFactory = mock(HoodieFileReaderFactory.class); + HoodieRowDataParquetReader fileReader = mock(HoodieRowDataParquetReader.class); + HoodieRecord hoodieRecord = mock(HoodieRecord.class); + GenericRowData inputRow = new GenericRowData( + DataTypeUtils.addMetadataFields(TestConfigurations.ROW_TYPE, false).getFieldCount()); + BinaryExternalSorter sorter = mock(BinaryExternalSorter.class); + MutableObjectIterator iterator = mock(MutableObjectIterator.class); + BinaryRowData sortedRow = mock(BinaryRowData.class); + when(writeClient.getHoodieTable()).thenReturn(table); + when(table.getStorage()).thenReturn(storage); + when(table.getConfig()).thenReturn(writeConfig); + when(ioFactory.getReaderFactory(HoodieRecord.HoodieRecordType.FLINK)).thenReturn(readerFactory); + when(readerFactory.getFileReader( + same(writeConfig), any(StoragePath.class))).thenReturn(fileReader); + when(hoodieRecord.getData()).thenReturn(inputRow); + when(fileReader.getRecordIterator(any(HoodieSchema.class))).thenReturn( + ClosableIterator.wrap(Collections.singletonList(hoodieRecord).iterator())); + when(sorter.getIterator()).thenReturn(iterator); + when(iterator.next(any(BinaryRowData.class))).thenReturn(sortedRow).thenReturn(null); + + ClusteringOperation operation = new ClusteringOperation( + tempDir.toPath().resolve("base.parquet").toString(), + Collections.emptyList(), + "old-file", + "partition", + null, + 0); + ClusteringOperator operator = new ClusteringOperator(conf, TestConfigurations.ROW_TYPE); + try (MockedStatic writeClients = mockStatic(FlinkWriteClients.class); + MockedStatic utils = mockStatic(Utils.class); + MockedStatic ioFactories = mockStatic(HoodieIOFactory.class); + MockedConstruction writerHelpers = + mockConstruction(BulkInsertWriterHelper.class, (writerHelper, context) -> + when(writerHelper.getWriteStatuses(anyInt())).thenReturn(Collections.emptyList())); + OneInputStreamOperatorTestHarness harness = + new OneInputStreamOperatorTestHarness<>(operator, 1, 1, 0)) { + writeClients.when(() -> FlinkWriteClients.getHoodieClientConfig( + any(Configuration.class), eq(false), eq(false))).thenReturn(writeConfig); + writeClients.when(() -> FlinkWriteClients.createWriteClient( + any(Configuration.class), any(RuntimeContext.class))).thenReturn(writeClient); + utils.when(() -> Utils.getBinaryExternalSorter( + any(), any(), anyLong(), any(), any(), any(), any(), any(), any(Configuration.class))) + .thenReturn(sorter); + ioFactories.when(() -> HoodieIOFactory.getIOFactory(storage)).thenReturn(ioFactory); + + harness.open(); + harness.processElement(new StreamRecord<>(event("005", Collections.singletonList(operation)))); + + verify(sorter).startThreads(); + verify(sorter).write(any(BinaryRowData.class)); + verify(writerHelpers.constructed().get(0)).write(same(sortedRow)); + verify(sorter).close(); + assertEquals("005", + ((StreamRecord) harness.getOutput().poll()).getValue().getInstant()); + } + } + + private ClusteringPlanEvent event(String instant) { + return event(instant, Collections.emptyList()); + } + + private ClusteringPlanEvent event(String instant, List operations) { + ClusteringGroupInfo groupInfo = new ClusteringGroupInfo(); + groupInfo.setOperations(operations); + groupInfo.setNumOutputGroups(1); + return new ClusteringPlanEvent(instant, groupInfo, Map.of()); + } +} diff --git a/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/clustering/update/strategy/TestFlinkConsistentBucketUpdateStrategy.java b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/clustering/update/strategy/TestFlinkConsistentBucketUpdateStrategy.java new file mode 100644 index 0000000000000..c1169c2bf36a1 --- /dev/null +++ b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/clustering/update/strategy/TestFlinkConsistentBucketUpdateStrategy.java @@ -0,0 +1,176 @@ +/* + * 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.hudi.sink.clustering.update.strategy; + +import org.apache.hudi.client.HoodieFlinkWriteClient; +import org.apache.hudi.common.engine.HoodieEngineContext; +import org.apache.hudi.common.fs.FSUtils; +import org.apache.hudi.common.model.ConsistentHashingNode; +import org.apache.hudi.common.model.HoodieAvroRecord; +import org.apache.hudi.common.model.HoodieConsistentHashingMetadata; +import org.apache.hudi.common.model.HoodieFileGroupId; +import org.apache.hudi.common.model.HoodieKey; +import org.apache.hudi.common.model.HoodieRecord; +import org.apache.hudi.common.model.HoodieRecordLocation; +import org.apache.hudi.common.model.HoodieRecordPayload; +import org.apache.hudi.common.table.HoodieTableMetaClient; +import org.apache.hudi.common.table.timeline.HoodieInstant; +import org.apache.hudi.common.table.view.TableFileSystemView; +import org.apache.hudi.common.util.ClusteringUtils; +import org.apache.hudi.common.util.collection.Pair; +import org.apache.hudi.index.bucket.ConsistentBucketIdentifier; +import org.apache.hudi.table.HoodieFlinkTable; +import org.apache.hudi.table.action.cluster.util.ConsistentHashingUpdateStrategyUtils; + +import org.junit.jupiter.api.Test; +import org.mockito.MockedStatic; + +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.stream.Stream; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.anySet; +import static org.mockito.ArgumentMatchers.same; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.mockStatic; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.when; + +/** + * Tests {@link FlinkConsistentBucketUpdateStrategy}. + */ +class TestFlinkConsistentBucketUpdateStrategy { + + private static final String PARTITION = "partition"; + private static final String OLD_FILE_ID = "old-file-id"; + private static final String CLUSTERING_INSTANT = "002"; + private static final String NEW_FILE_ID_PREFIX = "new-file"; + + @Test + void testInitializationIsRequiredAndRegularUpdatesAreNotDuplicated() { + TestContext context = new TestContext(); + FlinkConsistentBucketUpdateStrategy strategy = + new FlinkConsistentBucketUpdateStrategy<>(context.writeClient, Collections.emptyList()); + List, String>> records = records("003"); + + assertThrows(IllegalArgumentException.class, () -> strategy.handleUpdate(records)); + + try (MockedStatic clusteringUtils = mockStatic(ClusteringUtils.class)) { + clusteringUtils.when(() -> ClusteringUtils.getPendingClusteringInstantTimes(context.metaClient)) + .thenReturn(Collections.emptyList()); + + strategy.initialize(context.writeClient); + strategy.initialize(context.writeClient); + } + + Pair, String>>, Set> result = + strategy.handleUpdate(records); + assertSame(records, result.getLeft()); + assertEquals(Collections.singleton(new HoodieFileGroupId(PARTITION, OLD_FILE_ID)), result.getRight()); + assertFalse(strategy.needDualWrite(new HoodieFileGroupId(PARTITION, OLD_FILE_ID))); + + assertThrows(IllegalArgumentException.class, () -> strategy.handleUpdate(Collections.emptyList())); + strategy.reset(); + assertThrows(IllegalArgumentException.class, () -> strategy.handleUpdate(records)); + } + + @Test + void testUpdatesToPendingFileGroupAreRoutedToOldAndNewBuckets() { + TestContext context = new TestContext(); + HoodieFileGroupId oldFileGroup = new HoodieFileGroupId(PARTITION, OLD_FILE_ID); + HoodieInstant pendingInstant = mock(HoodieInstant.class); + when(pendingInstant.requestedTime()).thenReturn(CLUSTERING_INSTANT); + when(context.fileSystemView.getFileGroupsInPendingClustering()) + .thenReturn(Stream.of(Pair.of(oldFileGroup, pendingInstant))); + + ConsistentHashingNode newNode = new ConsistentHashingNode(Integer.MAX_VALUE, NEW_FILE_ID_PREFIX); + HoodieConsistentHashingMetadata hashingMetadata = new HoodieConsistentHashingMetadata( + (short) 0, PARTITION, CLUSTERING_INSTANT, 1, 0, Collections.singletonList(newNode)); + ConsistentBucketIdentifier identifier = new ConsistentBucketIdentifier(hashingMetadata); + Map> identifiers = + Collections.singletonMap(PARTITION, Pair.of(CLUSTERING_INSTANT, identifier)); + + FlinkConsistentBucketUpdateStrategy strategy = + new FlinkConsistentBucketUpdateStrategy<>(context.writeClient, Collections.emptyList()); + try (MockedStatic clusteringUtils = mockStatic(ClusteringUtils.class); + MockedStatic updateStrategyUtils = + mockStatic(ConsistentHashingUpdateStrategyUtils.class)) { + clusteringUtils.when(() -> ClusteringUtils.getPendingClusteringInstantTimes(context.metaClient)) + .thenReturn(Collections.singletonList(pendingInstant)); + updateStrategyUtils.when(() -> ConsistentHashingUpdateStrategyUtils.constructPartitionToIdentifier( + anySet(), same(context.table))).thenReturn(identifiers); + + strategy.initialize(context.writeClient); + assertTrue(strategy.needDualWrite(oldFileGroup)); + + List, String>> originalRecords = records("003"); + Pair, String>>, Set> result = + strategy.handleUpdate(originalRecords); + strategy.handleUpdate(records("004")); + + assertEquals(2, result.getLeft().size()); + Pair, String> duplicateRecords = result.getLeft().get(0); + assertEquals(CLUSTERING_INSTANT, duplicateRecords.getRight()); + assertEquals(FSUtils.createNewFileId(NEW_FILE_ID_PREFIX, 0), + duplicateRecords.getLeft().get(0).getCurrentLocation().getFileId()); + assertSame(originalRecords.get(0), result.getLeft().get(1)); + assertEquals(2, result.getRight().size()); + assertTrue(result.getRight().contains(oldFileGroup)); + assertTrue(result.getRight().contains(new HoodieFileGroupId( + PARTITION, FSUtils.createNewFileId(NEW_FILE_ID_PREFIX, 0)))); + + updateStrategyUtils.verify(() -> ConsistentHashingUpdateStrategyUtils.constructPartitionToIdentifier( + anySet(), same(context.table)), times(1)); + } + } + + @SuppressWarnings({"rawtypes", "unchecked"}) + private static List, String>> records(String instant) { + HoodieRecordPayload payload = mock(HoodieRecordPayload.class); + HoodieRecord record = new HoodieAvroRecord<>(new HoodieKey("record-key", PARTITION), payload); + record.setCurrentLocation(new HoodieRecordLocation("001", OLD_FILE_ID)); + return Collections.singletonList(Pair.of(Collections.singletonList(record), instant)); + } + + private static class TestContext { + private final HoodieFlinkWriteClient writeClient; + private final HoodieFlinkTable table; + private final HoodieTableMetaClient metaClient; + private final TableFileSystemView fileSystemView; + + @SuppressWarnings("unchecked") + private TestContext() { + writeClient = mock(HoodieFlinkWriteClient.class); + table = mock(HoodieFlinkTable.class); + metaClient = mock(HoodieTableMetaClient.class); + fileSystemView = mock(TableFileSystemView.class); + when(writeClient.getEngineContext()).thenReturn(mock(HoodieEngineContext.class)); + when(writeClient.getHoodieTable()).thenReturn(table); + when(table.getMetaClient()).thenReturn(metaClient); + when(table.getFileSystemView()).thenReturn(fileSystemView); + } + } +} diff --git a/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/compact/TestCompactionPlanSourceFunction.java b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/compact/TestCompactionPlanSourceFunction.java new file mode 100644 index 0000000000000..6204dfcc74309 --- /dev/null +++ b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/compact/TestCompactionPlanSourceFunction.java @@ -0,0 +1,94 @@ +/* + * 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.hudi.sink.compact; + +import org.apache.hudi.adapter.SourceFunctionAdapter; +import org.apache.hudi.avro.model.HoodieCompactionOperation; +import org.apache.hudi.avro.model.HoodieCompactionPlan; +import org.apache.hudi.common.table.HoodieTableMetaClient; +import org.apache.hudi.common.table.timeline.HoodieActiveTimeline; +import org.apache.hudi.common.table.timeline.HoodieTimeline; +import org.apache.hudi.common.util.collection.Pair; +import org.apache.hudi.util.StreamerUtil; + +import org.apache.flink.configuration.Configuration; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; +import org.mockito.MockedStatic; + +import java.util.Arrays; +import java.util.Collections; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.mockStatic; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +/** + * Tests {@link CompactionPlanSourceFunction}. + */ +class TestCompactionPlanSourceFunction { + + @Test + @SuppressWarnings("unchecked") + void testEmitOperationsOnlyForPendingCompactionInstants() throws Exception { + Configuration conf = new Configuration(); + HoodieCompactionPlan missingPlan = mock(HoodieCompactionPlan.class); + HoodieCompactionPlan pendingPlan = mock(HoodieCompactionPlan.class); + HoodieCompactionOperation operation = mock(HoodieCompactionOperation.class); + when(operation.getBaseInstantTime()).thenReturn("000"); + when(operation.getDataFilePath()).thenReturn(null); + when(operation.getDeltaFilePaths()).thenReturn(Collections.singletonList("file.log.1")); + when(operation.getPartitionPath()).thenReturn("partition"); + when(operation.getFileId()).thenReturn("file-id"); + when(operation.getMetrics()).thenReturn(Collections.emptyMap()); + when(operation.getBootstrapFilePath()).thenReturn(null); + when(pendingPlan.getOperations()).thenReturn(Collections.singletonList(operation)); + + HoodieTableMetaClient metaClient = mock(HoodieTableMetaClient.class); + HoodieActiveTimeline activeTimeline = mock(HoodieActiveTimeline.class); + HoodieTimeline pendingTimeline = mock(HoodieTimeline.class); + when(metaClient.getActiveTimeline()).thenReturn(activeTimeline); + when(activeTimeline.filterPendingCompactionTimeline()).thenReturn(pendingTimeline); + when(pendingTimeline.containsInstant("001")).thenReturn(true); + + SourceFunctionAdapter.SourceContext sourceContext = + mock(SourceFunctionAdapter.SourceContext.class); + CompactionPlanSourceFunction function = new CompactionPlanSourceFunction( + Arrays.asList(Pair.of("000", missingPlan), Pair.of("001", pendingPlan)), conf); + + try (MockedStatic streamerUtil = mockStatic(StreamerUtil.class)) { + streamerUtil.when(() -> StreamerUtil.createMetaClient(conf)).thenReturn(metaClient); + + function.open(conf); + function.run(sourceContext); + function.cancel(); + function.close(); + } + + ArgumentCaptor eventCaptor = + ArgumentCaptor.forClass(CompactionPlanEvent.class); + verify(sourceContext).collect(eventCaptor.capture()); + CompactionPlanEvent event = eventCaptor.getValue(); + assertEquals("001", event.getCompactionInstantTime()); + assertEquals("partition", event.getOperation().getPartitionPath()); + assertEquals("file-id", event.getOperation().getFileId()); + } +}