Baunsgaard commented on code in PR #16937:
URL: https://github.com/apache/iceberg/pull/16937#discussion_r3463555511
##########
spark/v4.1/spark/src/test/java/org/apache/iceberg/spark/actions/TestRewriteDataFilesAction.java:
##########
@@ -2466,6 +2509,101 @@ private void writeRecords(int files, int numRecords,
int partitions) {
writeDF(df);
}
+ private List<Record> generateRecords(int numRecords, int partitions) {
+ int rowDimension = (int) Math.ceil(Math.sqrt(numRecords));
+ List<Pair<Integer, Integer>> data =
+ IntStream.range(0, rowDimension)
+ .boxed()
+ .flatMap(x -> IntStream.range(0, rowDimension).boxed().map(y ->
Pair.of(x, y)))
+ .collect(Collectors.toList());
+ Collections.shuffle(data, new Random(42));
+ List<Record> records = Lists.newArrayListWithExpectedSize(data.size());
+ for (Pair<Integer, Integer> pair : data) {
+ GenericRecord record = GenericRecord.create(SCHEMA);
+ record.setField("c1", partitions > 0 ? pair.first() % partitions :
pair.first());
+ record.setField("c2", "foo" + pair.first());
+ record.setField("c3", "bar" + pair.second());
+ records.add(record);
+ }
+ return records;
+ }
+
+ /**
+ * Append records as Iceberg data files without going through Spark, writing
{@code files} data
+ * files per table partition. This mirrors the file layout that {@code
df.repartition(files)}
+ * would produce: the suite asserts on file counts and row content, not on
the row-to-file
+ * assignment, so a deterministic contiguous split is equivalent while
avoiding a Spark job,
+ * shuffle, sort, and DataFrame conversion per write.
+ */
+ private void appendRecords(List<Record> records, int files) {
+ for (Record record : records) {
+ writtenRecords.add(
+ new Object[] {record.getField("c1"), record.getField("c2"),
record.getField("c3")});
+ }
+
+ Table table = TABLES.load(tableLocation);
+ PartitionSpec spec = table.spec();
+ AppendFiles append = table.newAppend();
+
+ if (spec.isUnpartitioned()) {
+ for (List<Record> chunk : split(records, files)) {
+ append.appendFile(writeDataFileNative(table, null, chunk));
+ }
+ } else {
+ Map<String, PartitionKey> keys = Maps.newLinkedHashMap();
+ Map<String, List<Record>> groups = Maps.newLinkedHashMap();
+ PartitionKey key = new PartitionKey(spec, table.schema());
+ for (Record record : records) {
+ key.partition(record);
+ String path = key.toPath();
+ keys.computeIfAbsent(path, p -> key.copy());
+ groups.computeIfAbsent(path, p -> Lists.newArrayList()).add(record);
+ }
+ groups.forEach(
+ (path, groupRecords) -> {
+ for (List<Record> chunk : split(groupRecords, files)) {
+ append.appendFile(writeDataFileNative(table, keys.get(path),
chunk));
+ }
+ });
+ }
+
+ append.commit();
+ }
+
+ private DataFile writeDataFileNative(Table table, StructLike partition,
List<Record> records) {
+ OutputFile outputFile =
+ table
+ .io()
+ .newOutputFile(
+ table
+ .locationProvider()
+ .newDataLocation(
+
FileFormat.PARQUET.addExtension(UUID.randomUUID().toString())));
+ try {
+ return FileHelpers.writeDataFile(table, outputFile, partition, records);
+ } catch (IOException e) {
+ throw new UncheckedIOException(e);
+ }
+ }
+
+ /** Split a list into at most {@code parts} contiguous, non-empty sublists.
*/
+ private static List<List<Record>> split(List<Record> records, int parts) {
+ int count = Math.min(parts, records.size());
+ List<List<Record>> chunks = Lists.newArrayListWithCapacity(Math.max(count,
0));
Review Comment:
valid point, i simplified it in accordance with your comment!
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
