kevinjqliu commented on code in PR #16882:
URL: https://github.com/apache/iceberg/pull/16882#discussion_r3444906213
##########
spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/actions/TestRewriteDataFilesAction.java:
##########
@@ -2340,123 +2327,111 @@ protected Table createTable() {
* @return the created table
*/
protected Table createTable(int files) {
- String key =
- String.format(
- "unpartitioned|fv=%d|rowGroup=%d|files=%d|rows=%d",
- formatVersion, INPUT_PARQUET_ROW_GROUP_SIZE_BYTES, files, SCALE);
- List<DataFile> inputFiles =
- cachedInputFiles(
- key,
- () -> {
- Table golden = createTable();
- writeRecords(files, SCALE);
- golden.refresh();
- return golden;
- });
- Table table = createTable();
- appendInputFiles(table, inputFiles);
+ PartitionSpec spec = PartitionSpec.unpartitioned();
+ Map<String, String> options = unpartitionedTableOptions();
+ Table table = createTable(this.tableLocation, spec, options);
+ writeRecords(table, spec, options, files, SCALE);
table.refresh();
return table;
}
protected Table createTablePartitioned(
- int partitions, int files, int numRecords, Map<String, String> options) {
+ int partitionCount, int files, int numRecords, Map<String, String>
options) {
PartitionSpec spec =
PartitionSpec.builderFor(SCHEMA).identity("c1").truncate("c2", 2).build();
- String key =
- String.format(
- "partitioned|fv=%d|spec=%s|opts=%s|files=%d|rows=%d|partitions=%d",
- formatVersion, spec, new TreeMap<>(options), files, numRecords,
partitions);
- List<DataFile> inputFiles =
- cachedInputFiles(
- key,
- () -> {
- Table golden = TABLES.create(SCHEMA, spec, options,
tableLocation);
- assertThat(golden.currentSnapshot()).as("Table must be
empty").isNull();
- writeRecords(files, numRecords, partitions);
- golden.refresh();
- return golden;
- });
- Table table = TABLES.create(SCHEMA, spec, options, tableLocation);
- assertThat(table.currentSnapshot()).as("Table must be empty").isNull();
- appendInputFiles(table, inputFiles);
+ Table table = createTable(this.tableLocation, spec, options);
+
+ writeRecords(table, spec, options, files, numRecords, partitionCount);
table.refresh();
return table;
}
- /**
- * Returns the input data files for a given table shape, materializing them
with Spark exactly
- * once per JVM fork and reusing them afterwards. On a cache miss the {@code
goldenBuilder} writes
- * the data into a stable cache location (kept alive for the whole class via
a static {@link
- * TempDir}); on a hit the cached {@link DataFile}s are returned and
re-appended to a fresh table
- * by {@link #appendInputFiles}. The data is deterministic (fixed RNG seed)
so reuse is
- * byte-identical to regenerating it.
- */
- private List<DataFile> cachedInputFiles(String key, Supplier<Table>
goldenBuilder) {
- List<DataFile> cached = INPUT_FILE_CACHE.get(key);
- if (cached != null) {
- return cached;
- }
- // Serialize builds per key: concurrent callers requesting the same table
shape block on the
- // first build and then reuse its result, instead of materializing
identical input twice. The
- // heavy Spark write happens outside any map lock, so distinct shapes can
still build in
- // parallel.
- Object lock = INPUT_CACHE_LOCKS.computeIfAbsent(key, ignored -> new
Object());
- synchronized (lock) {
- List<DataFile> existing = INPUT_FILE_CACHE.get(key);
- if (existing != null) {
- return existing;
- }
- String savedLocation = this.tableLocation;
- try {
- this.tableLocation =
- inputCacheDir.resolve("input-" +
INPUT_CACHE_SEQ.incrementAndGet()).toUri().toString();
- Table golden = goldenBuilder.get();
- // includeColumnStats() is required: a plain scan drops lower/upper
bounds and
- // value counts, and re-appending stat-less files breaks tests that
read bounds.
- // planFiles() returns a CloseableIterable holding manifest readers
open, so close it via
- // try-with-resources; otherwise every cache miss leaks file
descriptors and can leave
- // manifest files locked on Windows.
- List<DataFile> built;
- try (CloseableIterable<FileScanTask> tasks =
- golden.newScan().includeColumnStats().planFiles()) {
- built =
- Streams.stream(tasks)
- .map(FileScanTask::file)
- .map(DataFile::copy)
- .collect(ImmutableList.toImmutableList());
- } catch (IOException e) {
- throw new UncheckedIOException("Failed to plan cached input files",
e);
- }
- INPUT_FILE_CACHE.put(key, built);
- return built;
- } finally {
- this.tableLocation = savedLocation;
- }
- }
- }
-
- private static void appendInputFiles(Table table, List<DataFile> inputFiles)
{
- AppendFiles append = table.newAppend();
- inputFiles.forEach(append::appendFile);
- append.commit();
- }
-
- protected Table createTablePartitioned(int partitions, int files) {
+ protected Table createTablePartitioned(int partitionCount, int files) {
return createTablePartitioned(
- partitions,
+ partitionCount,
files,
SCALE,
ImmutableMap.of(TableProperties.FORMAT_VERSION,
String.valueOf(formatVersion)));
}
- protected Table createTablePartitioned(int partitions, int files, int
numRecords) {
+ protected Table createTablePartitioned(int partitionCount, int files, int
numRecords) {
return createTablePartitioned(
- partitions,
+ partitionCount,
files,
numRecords,
ImmutableMap.of(TableProperties.FORMAT_VERSION,
String.valueOf(formatVersion)));
}
+ // Reuse cached input files when available; otherwise materialize them under
the cache dir.
+ private void writeRecords(
+ Table targetTable,
+ PartitionSpec spec,
+ Map<String, String> tableOptions,
+ int files,
+ int numRecords) {
+ writeRecords(targetTable, spec, tableOptions, files, numRecords, 0);
+ }
+
+ private void writeRecords(
+ Table targetTable,
+ PartitionSpec spec,
+ Map<String, String> tableOptions,
+ int files,
+ int numRecords,
+ int partitionCount) {
+ if (spec.isUnpartitioned()) {
+ assertThat(partitionCount).as("Unpartitioned input partition
count").isZero();
+ } else {
+ assertThat(partitionCount).as("Partitioned input partition
count").isPositive();
+ }
+
+ CachedDataFilesKey key =
+ new CachedDataFilesKey(spec, tableOptions, files, numRecords,
partitionCount);
+ List<DataFile> cachedDataFiles = CACHED_DATA_FILES.get(key);
+ if (cachedDataFiles == null) {
+ cachedDataFiles = cacheDataFiles(key, spec, tableOptions, files,
numRecords, partitionCount);
+ }
+
+ appendDataFiles(targetTable, spec, cachedDataFiles);
+ }
+
+ private List<DataFile> cacheDataFiles(
+ CachedDataFilesKey key,
+ PartitionSpec spec,
+ Map<String, String> tableOptions,
+ int files,
+ int numRecords,
+ int partitionCount) {
+ String cacheLocation = new File(inputCacheDir,
UUID.randomUUID().toString()).toURI().toString();
+ Table cacheTable = createTable(cacheLocation, spec, tableOptions);
+
+ writeRecords(files, numRecords, partitionCount, cacheTable.location());
+ cacheTable.refresh();
+
+ try (CloseableIterable<FileScanTask> tasks =
+ cacheTable.newScan().includeColumnStats().planFiles()) {
+ List<DataFile> dataFiles =
+ Streams.stream(tasks)
+ .map(FileScanTask::file)
+ // Clear v3 first_row_id so each fresh target table assigns its
own on re-append.
+ .map(file ->
DataFiles.builder(spec).copy(file).withFirstRowId(null).build())
+ .collect(Collectors.toList());
+
+ List<DataFile> existingDataFiles = CACHED_DATA_FILES.putIfAbsent(key,
dataFiles);
+ return existingDataFiles != null ? existingDataFiles : dataFiles;
+ } catch (IOException e) {
+ throw new UncheckedIOException("Failed to plan cached data files", e);
+ }
+ }
+
+ private static void appendDataFiles(
+ Table targetTable, PartitionSpec spec, List<DataFile> dataFiles) {
+ AppendFiles append = targetTable.newAppend();
+ dataFiles.stream()
+ .map(file ->
DataFiles.builder(spec).copy(file).withFirstRowId(null).build())
Review Comment:
decided to keep this to be explicit, added a 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]
