szehon-ho commented on code in PR #16937:
URL: https://github.com/apache/iceberg/pull/16937#discussion_r3462619440
##########
spark/v4.1/spark/src/test/java/org/apache/iceberg/spark/actions/TestRewriteDataFilesAction.java:
##########
@@ -164,6 +167,12 @@ protected static List<Integer> parameters() {
private final ScanTaskSetManager manager = ScanTaskSetManager.get();
private String tableLocation = null;
+ // Mirror of the rows written natively via appendRecords, in the {c1, c2,
c3} shape that
+ // currentData() returns. Lets tests that only compact (which preserves row
content) assert
+ // against the known input via expectedData() instead of paying a second
Spark read+sort+collect
+ // just to re-materialize what they wrote. Reset per test in
setupTableLocation().
+ private final List<Object[]> writtenRecords = Lists.newArrayList();
Review Comment:
Nit (optional): this field comment (plus the `expectedData` /
`writeRecordsViaSpark` / `appendRecords` javadocs) is a bit verbose. It's
accurate and the two-path subtlety is worth a note, but it could be trimmed to
just the "why" (e.g. that compaction-only tests assert against this to avoid a
second Spark read).
##########
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) {
Review Comment:
Nit (optional): `generateRecords` duplicates the record-generation logic in
`writeRecordsViaSpark` (the IntStream grid + seed-42 shuffle + `c1 = first %
partitions` / `c2` / `c3`). Consider sharing a single `List<Pair<Integer,
Integer>>` generator and mapping it to `Record` here vs `ThreeColumnRecord`
there. Minor, since `writeRecordsViaSpark` has only two callers.
##########
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:
Nit (optional): `Math.max(count, 0)` is redundant here. `count =
Math.min(parts, records.size())` is always >= 0, and the `count <= 0` early
return on the next line already handles the empty case, so
`Lists.newArrayListWithCapacity(count)` reads cleaner.
--
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]
