diff --git a/hudi-spark-datasource/hudi-spark/src/main/java/org/apache/hudi/QuickstartUtils.java b/hudi-spark-datasource/hudi-spark/src/main/java/org/apache/hudi/QuickstartUtils.java index 2ba1e286481cc..47b30173cfed0 100644 --- a/hudi-spark-datasource/hudi-spark/src/main/java/org/apache/hudi/QuickstartUtils.java +++ b/hudi-spark-datasource/hudi-spark/src/main/java/org/apache/hudi/QuickstartUtils.java @@ -23,6 +23,7 @@ import org.apache.hudi.common.model.HoodieRecord; import org.apache.hudi.common.model.OverwriteWithLatestAvroPayload; import org.apache.hudi.common.util.Option; +import org.apache.hudi.exception.HoodieException; import org.apache.hudi.exception.HoodieIOException; import org.apache.avro.Schema; @@ -31,7 +32,6 @@ import org.apache.spark.sql.Row; import java.io.IOException; -import java.util.ArrayList; import java.util.Arrays; import java.util.HashMap; import java.util.List; @@ -175,15 +175,18 @@ public HoodieRecord generateUpdateRecord(HoodieKey key, String randomString) thr * @param n Number of updates (including dups) * @return list of hoodie record updates */ - public List generateUpdates(Integer n) throws IOException { - String randomString = generateRandomString(); - List updates = new ArrayList<>(); - for (int i = 0; i < n; i++) { - HoodieKey key = existingKeys.get(rand.nextInt(numExistingKeys)); - HoodieRecord record = generateUpdateRecord(key, randomString); - updates.add(record); + public List generateUpdates(Integer n) { + if (numExistingKeys == 0) { + throw new HoodieException("Data must have been written before performing the update operation"); } - return updates; + String randomString = generateRandomString(); + return IntStream.range(0, n).boxed().map(x -> { + try { + return generateUpdateRecord(existingKeys.get(rand.nextInt(numExistingKeys)), randomString); + } catch (IOException e) { + throw new HoodieIOException(e.getMessage(), e); + } + }).collect(Collectors.toList()); } /** diff --git a/hudi-spark-datasource/hudi-spark/src/test/java/org/apache/hudi/TestQuickstartUtils.java b/hudi-spark-datasource/hudi-spark/src/test/java/org/apache/hudi/TestQuickstartUtils.java new file mode 100644 index 0000000000000..2042249ecd788 --- /dev/null +++ b/hudi-spark-datasource/hudi-spark/src/test/java/org/apache/hudi/TestQuickstartUtils.java @@ -0,0 +1,45 @@ +/* + * 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; + +import org.apache.hudi.exception.HoodieException; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.junit.jupiter.MockitoExtension; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +@ExtendWith(MockitoExtension.class) +public class TestQuickstartUtils { + + @Test + public void testGenerateUpdates() throws Exception { + QuickstartUtils.DataGenerator dataGenerator = new QuickstartUtils.DataGenerator(); + + //Call generateUpdates directly then throws HoodieException + assertEquals(Assertions.assertThrows(HoodieException.class, () -> { + dataGenerator.generateUpdates(10); + }).getMessage(), "Data must have been written before performing the update operation"); + + //Execute generateInserts before executing generateUpdates then no exception + assertEquals(dataGenerator.generateInserts(10).size(), 10); + assertEquals(dataGenerator.generateUpdates(10).size(), 10); + } +}