laskoviymishka commented on code in PR #16882:
URL: https://github.com/apache/iceberg/pull/16882#discussion_r3449084279
##########
spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/actions/TestRewriteDataFilesAction.java:
##########
@@ -2340,123 +2330,113 @@ 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();
+
+ // includeColumnStats preserves bounds and value counts needed by tests
that inspect file stats.
+ 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(ImmutableList.toImmutableList());
+
+ 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()
+ // Ensure reused files don't carry first_row_id from the cache table
into the target table.
+ .map(file ->
DataFiles.builder(spec).copy(file).withFirstRowId(null).build())
Review Comment:
Same "cost of the cache" note, not a blocker: we clear `first_row_id` twice.
`cacheDataFiles` already stores files through `withFirstRowId(null)`, so by the
time we're here it's already null and this rebuild is a no-op clear plus a
fresh builder/object per file on every hit (the common path). The comment also
reads as if this is the load-bearing clear, but the one in `cacheDataFiles` is.
If the cache stays, I'd drop the rebuild here and just
`dataFiles.forEach(append::appendFile)`, keeping the explanatory comment up at
the `cacheDataFiles` clear. (Same in v4.0/v4.1.)
##########
spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/actions/TestRewriteDataFilesAction.java:
##########
@@ -2340,123 +2330,113 @@ 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);
Review Comment:
Flagging this as a real cost of the cache rather than a blocker, given the
open question above about whether we keep it at all.
This `get()`-then-`cacheDataFiles()` is check-then-act, and the per-key lock
the old code had was doing real work. Two threads can both see `null`, both run
the full SCALE=400K Spark write into separate UUID dirs, and `putIfAbsent`
keeps one — the loser's dir is orphaned under `inputCacheDir` for the class
lifetime, and `clearCachedDataFiles()` never reclaims it. Parameterized
formatVersion runs don't collide today (formatVersion is in the key), but
concurrent test methods or a forkCount=0 profile would.
If the cache stays, `CACHED_DATA_FILES.computeIfAbsent(key, k -> ...)` makes
the heavy build run once per key and keeps the "build each shape once per fork"
guarantee. (Same in v4.0/v4.1.)
--
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]
