From e21da73857be74900f0ceefa75d73c72e0c5c27b Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 26 Aug 2026 09:17:15 -0400 Subject: [PATCH 1/2] AddFiles: tighten file schemas with footer null counts, folded per declared schema Parquet writers commonly declare every column optional whatever the data holds. Taken at face value, a file with no nulls in a column the table has as required would ask the pre-pass to relax that column for no reason. Tightening uses the footer's null counts to drop such requests. The evidence travels beside the schema, not inside it. ReadFooterSchema emits (canonical declared-schema JSON, proven-null-free paths) and CollectDistinctSchemas folds the paths per distinct schema by intersection: a column stays proven only when every file proved it, equivalently a relaxation is requested when any file needs it. The commit side will rebuild the group's most conservative member with FileSchemas.markRequired before classifying, which produces the same table schema as tightening each file individually would. union adds new columns as optional whatever the file declares. Tightening only affects columns the table already has as required: the union adds new columns as optional whatever the file declares. --- .../io/iceberg/CollectDistinctSchemas.java | 230 +++++++++++++--- .../beam/sdk/io/iceberg/FileSchemas.java | 187 ++++++++++++- .../beam/sdk/io/iceberg/ReadFooterSchema.java | 38 ++- .../iceberg/CollectDistinctSchemasTest.java | 107 ++++++-- .../beam/sdk/io/iceberg/FileSchemasTest.java | 256 ++++++++++++++++++ .../sdk/io/iceberg/ReadFooterSchemaTest.java | 49 +++- 6 files changed, 777 insertions(+), 90 deletions(-) diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/CollectDistinctSchemas.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/CollectDistinctSchemas.java index 1b81e008f19d..16e877619b10 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/CollectDistinctSchemas.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/CollectDistinctSchemas.java @@ -17,87 +17,251 @@ */ package org.apache.beam.sdk.io.iceberg; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; import java.util.ArrayList; import java.util.List; import java.util.Map; +import java.util.Objects; import java.util.TreeMap; +import java.util.TreeSet; import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.coders.CoderRegistry; -import org.apache.beam.sdk.coders.KvCoder; +import org.apache.beam.sdk.coders.CustomCoder; import org.apache.beam.sdk.coders.ListCoder; import org.apache.beam.sdk.coders.MapCoder; import org.apache.beam.sdk.coders.StringUtf8Coder; import org.apache.beam.sdk.coders.VarLongCoder; import org.apache.beam.sdk.transforms.Combine; -import org.apache.beam.sdk.values.KV; +import org.checkerframework.checker.nullness.qual.Nullable; /** - * Collects the distinct schemas among canonical file schema JSONs (see {@link FileSchemas}), with - * the number of files per schema, most common first (ties broken by JSON). The commit side applies - * schemas in this order, so the schema covering the most files wins a conflict. - * - *

Inputs are compared as strings, so they must already be canonical. + * One output entry per distinct schema: its file count and the columns EVERY file carrying it + * proved free of nulls; one file with a null in "name" forces "name" to relax, however many clean + * files sit next to it. Entries come out most common first (ties broken by the JSON text) because + * the commit side applies schemas in that order and the most common schema should win a conflict. + * Schemas are compared as strings, so inputs must already be canonical. */ class CollectDistinctSchemas - extends Combine.CombineFn, List>> { + extends Combine.CombineFn< + CollectDistinctSchemas.SchemaGroup, + Map, + List> { + + /** Mutable accumulator counterpart of {@link SchemaGroup}. */ + static final class Group { + long files; + TreeSet nullFreeColumns; + + Group(long files, TreeSet nullFreeColumns) { + this.files = files; + this.nullFreeColumns = nullFreeColumns; + } + + @Override + public boolean equals(@Nullable Object other) { + if (!(other instanceof Group)) { + return false; + } + Group that = (Group) other; + return files == that.files && nullFreeColumns.equals(that.nullFreeColumns); + } + + @Override + public int hashCode() { + return Objects.hash(files, nullFreeColumns); + } + } + + /** + * A schema, how many files carry it, and the columns all of them proved free of nulls. + * ReadFooterSchema emits one per file ({@code files} = 1); this combiner merges them. + */ + static final class SchemaGroup { + final String schemaJson; + final long files; + final List nullFreeColumns; + + SchemaGroup(String schemaJson, long files, List nullFreeColumns) { + this.schemaJson = schemaJson; + this.files = files; + this.nullFreeColumns = nullFreeColumns; + } + + @Override + public boolean equals(@Nullable Object other) { + if (!(other instanceof SchemaGroup)) { + return false; + } + SchemaGroup that = (SchemaGroup) other; + return files == that.files + && schemaJson.equals(that.schemaJson) + && nullFreeColumns.equals(that.nullFreeColumns); + } + + @Override + public int hashCode() { + return Objects.hash(schemaJson, files, nullFreeColumns); + } + + @Override + public String toString() { + return files + " file(s), null-free in " + nullFreeColumns + ", schema " + schemaJson; + } + } @Override - public Map createAccumulator() { + public Map createAccumulator() { return new TreeMap<>(); } @Override - public Map addInput(Map accumulator, String schemaJson) { - add(accumulator, schemaJson, 1L); + public Map addInput(Map accumulator, SchemaGroup file) { + add(accumulator, file.schemaJson, file.files, file.nullFreeColumns); return accumulator; } @Override - public Map mergeAccumulators(Iterable> accumulators) { - Map merged = createAccumulator(); - for (Map accumulator : accumulators) { - for (Map.Entry entry : accumulator.entrySet()) { - add(merged, entry.getKey(), entry.getValue()); + public Map mergeAccumulators(Iterable> accumulators) { + Map merged = createAccumulator(); + for (Map accumulator : accumulators) { + for (Map.Entry entry : accumulator.entrySet()) { + add(merged, entry.getKey(), entry.getValue().files, entry.getValue().nullFreeColumns); } } return merged; } @Override - public List> extractOutput(Map accumulator) { - List> schemas = new ArrayList<>(); - for (Map.Entry entry : accumulator.entrySet()) { - schemas.add(KV.of(entry.getKey(), entry.getValue())); + public List extractOutput(Map accumulator) { + List schemas = new ArrayList<>(); + for (Map.Entry entry : accumulator.entrySet()) { + schemas.add( + new SchemaGroup( + entry.getKey(), + entry.getValue().files, + new ArrayList<>(entry.getValue().nullFreeColumns))); } schemas.sort( (a, b) -> { - int byCount = Long.compare(b.getValue(), a.getValue()); + int byCount = Long.compare(b.files, a.files); if (byCount != 0) { return byCount; } - return a.getKey().compareTo(b.getKey()); + return a.schemaJson.compareTo(b.schemaJson); }); return schemas; } @Override - public Coder> getAccumulatorCoder( - CoderRegistry registry, Coder inputCoder) { - return MapCoder.of(StringUtf8Coder.of(), VarLongCoder.of()); + public Coder> getAccumulatorCoder( + CoderRegistry registry, Coder inputCoder) { + return MapCoder.of(StringUtf8Coder.of(), GroupCoder.INSTANCE); } @Override - public Coder>> getDefaultOutputCoder( - CoderRegistry registry, Coder inputCoder) { - return ListCoder.of(KvCoder.of(StringUtf8Coder.of(), VarLongCoder.of())); + public Coder> getDefaultOutputCoder( + CoderRegistry registry, Coder inputCoder) { + return outputCoder(); } - private static void add(Map accumulator, String schemaJson, long count) { - Long existing = accumulator.get(schemaJson); + static Coder groupCoder() { + return SchemaGroupCoder.INSTANCE; + } + + static Coder> outputCoder() { + return ListCoder.of(SchemaGroupCoder.INSTANCE); + } + + private static final Coder> COLUMNS_CODER = ListCoder.of(StringUtf8Coder.of()); + + /** Singletons with class equality, so repeated mentions compare equal; deterministic encoding. */ + private static class GroupCoder extends CustomCoder { + static final GroupCoder INSTANCE = new GroupCoder(); + + private GroupCoder() {} + + @Override + public void verifyDeterministic() {} + + @Override + public boolean equals(@Nullable Object other) { + return other instanceof GroupCoder; + } + + @Override + public int hashCode() { + return getClass().hashCode(); + } + + @Override + public void encode(Group value, OutputStream out) throws IOException { + VarLongCoder.of().encode(value.files, out); + COLUMNS_CODER.encode(new ArrayList<>(value.nullFreeColumns), out); + } + + @Override + public Group decode(InputStream in) throws IOException { + long files = VarLongCoder.of().decode(in); + return new Group(files, new TreeSet<>(COLUMNS_CODER.decode(in))); + } + } + + private static class SchemaGroupCoder extends CustomCoder { + static final SchemaGroupCoder INSTANCE = new SchemaGroupCoder(); + + private SchemaGroupCoder() {} + + @Override + public void verifyDeterministic() {} + + @Override + public boolean equals(@Nullable Object other) { + return other instanceof SchemaGroupCoder; + } + + @Override + public int hashCode() { + return getClass().hashCode(); + } + + @Override + public void encode(SchemaGroup value, OutputStream out) throws IOException { + StringUtf8Coder.of().encode(value.schemaJson, out); + VarLongCoder.of().encode(value.files, out); + COLUMNS_CODER.encode(value.nullFreeColumns, out); + } + + @Override + public SchemaGroup decode(InputStream in) throws IOException { + String schemaJson = StringUtf8Coder.of().decode(in); + long files = VarLongCoder.of().decode(in); + return new SchemaGroup(schemaJson, files, COLUMNS_CODER.decode(in)); + } + } + + private static void add( + Map accumulator, + String schemaJson, + long files, + Iterable nullFreeColumns) { + Group existing = accumulator.get(schemaJson); if (existing == null) { - accumulator.put(schemaJson, count); - } else { - accumulator.put(schemaJson, existing + count); + TreeSet copy = new TreeSet<>(); + for (String column : nullFreeColumns) { + copy.add(column); + } + accumulator.put(schemaJson, new Group(files, copy)); + return; + } + existing.files += files; + TreeSet stillNullFree = new TreeSet<>(); + for (String column : nullFreeColumns) { + if (existing.nullFreeColumns.contains(column)) { + stillNullFree.add(column); + } } + existing.nullFreeColumns = stillNullFree; } } diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/FileSchemas.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/FileSchemas.java index 592e11e8c767..2d824f9f6670 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/FileSchemas.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/FileSchemas.java @@ -18,29 +18,208 @@ package org.apache.beam.sdk.io.iceberg; import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collection; +import java.util.Collections; +import java.util.HashSet; import java.util.List; +import java.util.Set; import org.apache.iceberg.Schema; import org.apache.iceberg.SchemaParser; import org.apache.iceberg.parquet.ParquetSchemaUtil; import org.apache.iceberg.types.Type; import org.apache.iceberg.types.TypeUtil; import org.apache.iceberg.types.Types; +import org.apache.parquet.column.statistics.Statistics; +import org.apache.parquet.hadoop.metadata.BlockMetaData; +import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData; import org.apache.parquet.hadoop.metadata.ParquetMetadata; /** - * Derives the schema a file contributes to schema inference. The canonical form sorts struct fields - * by name at every level and renumbers ids in deterministic order, so files that differ only in - * column order produce identical JSON. Ids are positional and meaningless: the commit side - * reconciles columns by name. + * What a file contributes to schema inference: the canonical form of the schema it declares, and + * the columns its footer proves free of nulls. + * + *

The schema half depends only on the declared schema, never on the data, so files written by + * the same job dedup to one entry no matter where their nulls fall. The null evidence is combined + * per schema by {@link CollectDistinctSchemas} and reapplied by the commit side via {@link + * #markRequired}. + * + *

The canonical form sorts struct fields by name at every level and renumbers ids. Ids are + * positional and meaningless (the commit side reconciles columns by name), so never diff two file + * schemas by id. Other field attributes (doc, defaults) are preserved, matching what SchemaDelta + * compares. Column paths are dotted, like pins. */ final class FileSchemas { private FileSchemas() {} + /** Canonical JSON of the schema the file declares. */ static String canonicalJson(ParquetMetadata footer) { Schema converted = ParquetSchemaUtil.convert(footer.getFileMetaData().getSchema()); return SchemaParser.toJson(canonical(converted)); } + /** This file as a schema group of one: its declared schema and its null-free columns. */ + static CollectDistinctSchemas.SchemaGroup schemaGroup(ParquetMetadata footer) { + Schema converted = ParquetSchemaUtil.convert(footer.getFileMetaData().getSchema()); + Schema tightened = tighten(converted, footer); + return new CollectDistinctSchemas.SchemaGroup( + SchemaParser.toJson(canonical(converted)), 1, changedToRequired(converted, tightened)); + } + + /** + * Marks a declared-optional column required when every row group has a null count of zero for it, + * so the file does not request a relaxation it does not need; an absent count is not proof. A + * struct is null-free when any leaf under it is (a null struct nulls all its leaves). Nothing + * under lists or maps is tightened: a zero count there would be valid evidence too, but mapping + * physical chunk paths (writer-dependent names like {@code list.element}, {@code array}) onto the + * converted schema is not worth it. A file with no row groups has no rows and proves every + * column, matching how the pin check treats empty files. + */ + static Schema tighten(Schema schema, ParquetMetadata footer) { + if (footer.getBlocks().isEmpty()) { + return new Schema(tightenAll(schema.asStruct()).fields()); + } + Set> zeroNullLeaves = leafPathsWithZeroNullCounts(footer); + if (zeroNullLeaves.isEmpty()) { + return schema; + } + return new Schema( + tightenStruct(schema.asStruct(), new ArrayList<>(), zeroNullLeaves).struct.fields()); + } + + /** With no rows, nothing can hold a null: every leaf and struct outside lists and maps. */ + private static Types.StructType tightenAll(Types.StructType struct) { + List fields = new ArrayList<>(); + for (Types.NestedField field : struct.fields()) { + Type type = field.type(); + if (type.isStructType()) { + fields.add(withOptionality(field, tightenAll(type.asStructType()), false)); + } else if (type.isPrimitiveType()) { + fields.add(withOptionality(field, type, false)); + } else { + fields.add(field); + } + } + return Types.StructType.of(fields); + } + + /** + * Returns the schema with the given dotted column paths made required. The commit side parses a + * group's schema JSON (optionality as the writer declared it) and applies the group's null-free + * columns with this before classifying, so only relaxations some file actually needs remain. + */ + static Schema markRequired(Schema declared, Collection columns) { + if (columns.isEmpty()) { + return declared; + } + Types.StructType required = markRequiredStruct(declared.asStruct(), "", new HashSet<>(columns)); + return new Schema(required.fields()); + } + + private static Types.StructType markRequiredStruct( + Types.StructType struct, String prefix, Set columns) { + List fields = new ArrayList<>(); + for (Types.NestedField field : struct.fields()) { + String path = prefix + field.name(); + Type type = field.type(); + if (type.isStructType()) { + type = markRequiredStruct(type.asStructType(), path + ".", columns); + } + boolean required = !field.isRequired() && columns.contains(path); + fields.add(withOptionality(field, type, field.isOptional() && !required)); + } + return Types.StructType.of(fields); + } + + /** Dotted paths of fields the tightened schema made required, sorted. */ + private static List changedToRequired(Schema declared, Schema tightened) { + List paths = new ArrayList<>(); + collectChangedToRequired(declared.asStruct(), tightened.asStruct(), "", paths); + Collections.sort(paths); + return paths; + } + + private static void collectChangedToRequired( + Types.StructType declared, Types.StructType tightened, String prefix, List out) { + for (int i = 0; i < declared.fields().size(); i++) { + Types.NestedField before = declared.fields().get(i); + Types.NestedField after = tightened.fields().get(i); + String path = prefix + before.name(); + if (before.isOptional() && after.isRequired()) { + out.add(path); + } + if (before.type().isStructType()) { + collectChangedToRequired( + before.type().asStructType(), after.type().asStructType(), path + ".", out); + } + } + } + + /** Leaf paths proven null-free in every row group (intersection over blocks). */ + private static Set> leafPathsWithZeroNullCounts(ParquetMetadata footer) { + Set> proven = null; + for (BlockMetaData block : footer.getBlocks()) { + Set> provenHere = new HashSet<>(); + for (ColumnChunkMetaData chunk : block.getColumns()) { + // A zero-row row group proves trivially: zero rows hold zero nulls. + Statistics stats = chunk.getStatistics(); + if (stats != null && stats.isNumNullsSet() && stats.getNumNulls() == 0) { + provenHere.add(Arrays.asList(chunk.getPath().toArray())); + } + } + if (proven == null) { + proven = provenHere; + } else { + proven.retainAll(provenHere); + } + } + if (proven == null) { + return new HashSet<>(); + } + return proven; + } + + private static final class Tightened { + final Types.StructType struct; + + /** Some leaf below, not under a list or map, is proven: the struct itself was never null. */ + final boolean hasNullFreeLeaf; + + Tightened(Types.StructType struct, boolean hasNullFreeLeaf) { + this.struct = struct; + this.hasNullFreeLeaf = hasNullFreeLeaf; + } + } + + private static Tightened tightenStruct( + Types.StructType struct, List path, Set> zeroNulls) { + List fields = new ArrayList<>(); + boolean hasNullFreeLeaf = false; + for (Types.NestedField field : struct.fields()) { + path.add(field.name()); + if (field.type().isPrimitiveType()) { + boolean nullFreeLeaf = zeroNulls.contains(path); + fields.add(nullFreeLeaf ? withOptionality(field, field.type(), false) : field); + hasNullFreeLeaf |= nullFreeLeaf; + } else if (field.type().isStructType()) { + Tightened child = tightenStruct(field.type().asStructType(), path, zeroNulls); + boolean optional = field.isOptional() && !child.hasNullFreeLeaf; + fields.add(withOptionality(field, child.struct, optional)); + hasNullFreeLeaf |= child.hasNullFreeLeaf; + } else { + fields.add(field); + } + path.remove(path.size() - 1); + } + return new Tightened(Types.StructType.of(fields), hasNullFreeLeaf); + } + + /** Copies every attribute (id, name, doc, defaults), replacing only type and optionality. */ + private static Types.NestedField withOptionality( + Types.NestedField field, Type type, boolean optional) { + return Types.NestedField.from(field).ofType(type).isOptional(optional).build(); + } + static Schema canonical(Schema schema) { Type sorted = TypeUtil.visit(schema.asStruct(), new SortFields()); int[] nextId = {0}; diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/ReadFooterSchema.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/ReadFooterSchema.java index 571c0b552441..e559feab3512 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/ReadFooterSchema.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/ReadFooterSchema.java @@ -35,10 +35,10 @@ import org.slf4j.LoggerFactory; /** - * Emits the canonical schema (see {@link FileSchemas}) of every readable Parquet file as JSON. - * Unreadable or non-Parquet files contribute nothing. + * Emits one {@link CollectDistinctSchemas.SchemaGroup} of one file per readable Parquet file: its + * declared schema and its null-free columns. Unreadable or non-Parquet files contribute nothing. */ -class ReadFooterSchema extends DoFn { +class ReadFooterSchema extends DoFn { private static final Logger LOG = LoggerFactory.getLogger(ReadFooterSchema.class); static final int DEFAULT_THREAD_POOL_SIZE = 10; @@ -65,24 +65,21 @@ class ReadFooterSchema extends DoFn { this.maxInFlightTasks = maxInFlightTasks; } - /** - * {@code schemaJson} is null when the file contributes no schema. Counters are updated when the - * result is delivered, on the processing thread: metrics touched from the executor are lost. - */ + /** Counters are updated on the processing thread: metrics touched from the executor are lost. */ private static class ReadResult { - final @Nullable String schemaJson; + final CollectDistinctSchemas.@Nullable SchemaGroup schema; final boolean footerError; final Instant timestamp; final BoundedWindow window; final PaneInfo paneInfo; ReadResult( - @Nullable String schemaJson, + CollectDistinctSchemas.@Nullable SchemaGroup schema, boolean footerError, Instant timestamp, BoundedWindow window, PaneInfo paneInfo) { - this.schemaJson = schemaJson; + this.schema = schema; this.footerError = footerError; this.timestamp = timestamp; this.window = window; @@ -114,7 +111,7 @@ public void process( @Timestamp Instant timestamp, BoundedWindow window, PaneInfo paneInfo, - OutputReceiver output) + OutputReceiver output) throws Exception { numFilesRead.inc(); Callable task = createReadTask(filePath, timestamp, window, paneInfo); @@ -128,24 +125,22 @@ public void finishBundle(FinishBundleContext context) throws Exception { private static void outputAtFinish(ReadResult result, FinishBundleContext context) { count(result); - if (result.schemaJson != null) { - context.output(result.schemaJson, result.timestamp, result.window); + if (result.schema != null) { + context.output(result.schema, result.timestamp, result.window); } } - private static void outputResult(ReadResult result, OutputReceiver output) { + private static void outputResult( + ReadResult result, OutputReceiver output) { count(result); - if (result.schemaJson != null) { + if (result.schema != null) { output.outputWindowedValue( - result.schemaJson, - result.timestamp, - Collections.singleton(result.window), - result.paneInfo); + result.schema, result.timestamp, Collections.singleton(result.window), result.paneInfo); } } private static void count(ReadResult result) { - if (result.schemaJson != null) { + if (result.schema != null) { numSchemasEmitted.inc(); } if (result.footerError) { @@ -167,8 +162,7 @@ private static Callable createReadTask( } try { ParquetMetadata footer = ParquetFooters.read(filePath); - return new ReadResult( - FileSchemas.canonicalJson(footer), false, timestamp, window, paneInfo); + return new ReadResult(FileSchemas.schemaGroup(footer), false, timestamp, window, paneInfo); } catch (Exception e) { LOG.warn( "Could not read the footer of {}; the file will not contribute to schema inference: {}", diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/CollectDistinctSchemasTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/CollectDistinctSchemasTest.java index 206e13acd0a6..476954f0fb4d 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/CollectDistinctSchemasTest.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/CollectDistinctSchemasTest.java @@ -22,15 +22,17 @@ import static org.junit.Assert.assertEquals; import java.util.Arrays; +import java.util.Collections; import java.util.List; import java.util.Map; import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.io.iceberg.CollectDistinctSchemas.Group; +import org.apache.beam.sdk.io.iceberg.CollectDistinctSchemas.SchemaGroup; import org.apache.beam.sdk.testing.CoderProperties; import org.apache.beam.sdk.testing.PAssert; import org.apache.beam.sdk.testing.TestPipeline; import org.apache.beam.sdk.transforms.Combine; import org.apache.beam.sdk.transforms.Create; -import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.PCollection; import org.apache.iceberg.Schema; import org.apache.iceberg.SchemaParser; @@ -60,38 +62,72 @@ public class CollectDistinctSchemasTest { required(1, "id", Types.LongType.get()), optional(2, "name", Types.StringType.get()))); + private static final List NONE = Collections.emptyList(); + private final CollectDistinctSchemas fn = new CollectDistinctSchemas(); @Test public void testDedupsIdenticalSchemas() { - assertEquals(Arrays.asList(KV.of(ID_NAME, 3L)), combine(ID_NAME, ID_NAME, ID_NAME)); + assertEquals( + Arrays.asList(group(ID_NAME, 3, NONE)), + combine(group(ID_NAME, 1, NONE), group(ID_NAME, 1, NONE), group(ID_NAME, 1, NONE))); } - /** Inputs are compared as strings; canonicalization is ReadFooterSchema's job. */ + /** Schemas are compared as strings; canonicalization is ReadFooterSchema's job. */ @Test public void testDifferentStringsAreDistinct() { - List> out = combine(ID_NAME, NAME_ID, ID_LONG_NAME); + List out = + combine(group(ID_NAME, 1, NONE), group(NAME_ID, 1, NONE), group(ID_LONG_NAME, 1, NONE)); assertEquals(3, out.size()); - for (KV entry : out) { - assertEquals(Long.valueOf(1L), entry.getValue()); + for (SchemaGroup entry : out) { + assertEquals(1L, entry.files); } } @Test public void testMostCommonFirstThenJson() { - List> out = combine(NAME_ID, ID_LONG_NAME, ID_NAME, ID_LONG_NAME, NAME_ID); + List out = + combine( + group(NAME_ID, 1, NONE), + group(ID_LONG_NAME, 1, NONE), + group(ID_NAME, 1, NONE), + group(ID_LONG_NAME, 1, NONE), + group(NAME_ID, 1, NONE)); + assertEquals( + Arrays.asList( + group(ID_LONG_NAME, 2, NONE), group(NAME_ID, 2, NONE), group(ID_NAME, 1, NONE)), + out); + } + + /** + * A column counts as proven for the group only if every file proved it: one file with nulls in a + * column is enough to make the table relax that column. + */ + @Test + public void testNullFreeColumnsIntersect() { + List out = + combine( + group(ID_NAME, 1, Arrays.asList("id", "name")), + group(ID_NAME, 1, Arrays.asList("id")), + group(NAME_ID, 1, Arrays.asList("name"))); assertEquals( - Arrays.asList(KV.of(ID_LONG_NAME, 2L), KV.of(NAME_ID, 2L), KV.of(ID_NAME, 1L)), out); + Arrays.asList( + group(ID_NAME, 2, Arrays.asList("id")), group(NAME_ID, 1, Arrays.asList("name"))), + out); } @Test - public void testMergeSumsCounts() { - Map first = fn.addInput(fn.createAccumulator(), ID_NAME); - Map second = fn.addInput(fn.createAccumulator(), ID_NAME); - second = fn.addInput(second, NAME_ID); - List> out = - fn.extractOutput(fn.mergeAccumulators(Arrays.asList(first, second))); - assertEquals(Arrays.asList(KV.of(ID_NAME, 2L), KV.of(NAME_ID, 1L)), out); + public void testNullFreeColumnsIntersectAcrossMergedAccumulators() { + Map first = + fn.addInput(fn.createAccumulator(), group(ID_NAME, 1, Arrays.asList("id", "name"))); + Map second = + fn.addInput(fn.createAccumulator(), group(ID_NAME, 1, Arrays.asList("name"))); + second = fn.addInput(second, group(NAME_ID, 1, Arrays.asList("id"))); + List out = fn.extractOutput(fn.mergeAccumulators(Arrays.asList(first, second))); + assertEquals( + Arrays.asList( + group(ID_NAME, 2, Arrays.asList("name")), group(NAME_ID, 1, Arrays.asList("id"))), + out); } @Test @@ -101,30 +137,51 @@ public void testEmptyInput() { @Test public void testAccumulatorCoderRoundTrip() throws Exception { - Coder> coder = fn.getAccumulatorCoder(null, null); - Map accumulator = fn.addInput(fn.createAccumulator(), ID_NAME); - accumulator = fn.addInput(accumulator, NAME_ID); + Coder> coder = fn.getAccumulatorCoder(null, null); + Map accumulator = + fn.addInput(fn.createAccumulator(), group(ID_NAME, 1, Arrays.asList("id"))); + accumulator = fn.addInput(accumulator, group(NAME_ID, 1, NONE)); CoderProperties.coderDecodeEncodeEqual(coder, accumulator); } + /** Coders from separate calls must compare equal, or coder inference treats them as different. */ + @Test + public void testCodersFromSeparateCallsAreEqual() throws Exception { + assertEquals(CollectDistinctSchemas.outputCoder(), CollectDistinctSchemas.outputCoder()); + assertEquals(fn.getAccumulatorCoder(null, null), fn.getAccumulatorCoder(null, null)); + // The output coder is deterministic; the accumulator coder is not required to be (MapCoder). + CollectDistinctSchemas.outputCoder().verifyDeterministic(); + } + @Test public void testPipeline() { - PCollection>> out = + PCollection> out = pipeline - .apply(Create.of(ID_NAME, NAME_ID, ID_NAME)) + .apply( + Create.of( + group(ID_NAME, 1, Arrays.asList("id", "name")), + group(NAME_ID, 1, NONE), + group(ID_NAME, 1, Arrays.asList("id"))) + .withCoder(CollectDistinctSchemas.groupCoder())) .apply(Combine.globally(new CollectDistinctSchemas())); - PAssert.that(out).containsInAnyOrder(Arrays.asList(KV.of(ID_NAME, 2L), KV.of(NAME_ID, 1L))); + PAssert.that(out) + .containsInAnyOrder( + Arrays.asList(group(ID_NAME, 2, Arrays.asList("id")), group(NAME_ID, 1, NONE))); pipeline.run(); } - private List> combine(String... schemaJsons) { - Map accumulator = fn.createAccumulator(); - for (String schemaJson : schemaJsons) { - accumulator = fn.addInput(accumulator, schemaJson); + private List combine(SchemaGroup... files) { + Map accumulator = fn.createAccumulator(); + for (SchemaGroup file : files) { + accumulator = fn.addInput(accumulator, file); } return fn.extractOutput(accumulator); } + private static SchemaGroup group(String json, long files, List proven) { + return new SchemaGroup(json, files, proven); + } + private static String json(Schema schema) { return SchemaParser.toJson(schema); } diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/FileSchemasTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/FileSchemasTest.java index 9723695186be..7358e8d50bd2 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/FileSchemasTest.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/FileSchemasTest.java @@ -20,17 +20,273 @@ import static org.apache.iceberg.types.Types.NestedField.optional; import static org.apache.iceberg.types.Types.NestedField.required; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertTrue; +import java.io.File; +import java.io.IOException; +import java.util.function.IntPredicate; +import org.apache.hadoop.fs.Path; import org.apache.iceberg.Schema; import org.apache.iceberg.SchemaParser; +import org.apache.iceberg.parquet.ParquetSchemaUtil; import org.apache.iceberg.types.Types; +import org.apache.parquet.example.data.Group; +import org.apache.parquet.example.data.simple.SimpleGroupFactory; +import org.apache.parquet.hadoop.ParquetWriter; +import org.apache.parquet.hadoop.example.ExampleParquetWriter; +import org.apache.parquet.hadoop.metadata.ParquetMetadata; +import org.apache.parquet.schema.LogicalTypeAnnotation; +import org.apache.parquet.schema.MessageType; +import org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName; +import org.apache.parquet.schema.Type.Repetition; +import org.junit.Rule; import org.junit.Test; +import org.junit.rules.TemporaryFolder; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @RunWith(JUnit4.class) public class FileSchemasTest { + @Rule public final TemporaryFolder tmp = new TemporaryFolder(); + + // ---- tightening + + // root: required id, optional name, optional struct address {city, zip}, optional list tags + private static final MessageType MIXED = + org.apache.parquet.schema.Types.buildMessage() + .required(PrimitiveTypeName.INT64) + .named("id") + .optional(PrimitiveTypeName.BINARY) + .as(LogicalTypeAnnotation.stringType()) + .named("name") + .addField( + org.apache.parquet.schema.Types.buildGroup(Repetition.OPTIONAL) + .optional(PrimitiveTypeName.BINARY) + .as(LogicalTypeAnnotation.stringType()) + .named("city") + .optional(PrimitiveTypeName.INT32) + .named("zip") + .named("address")) + .addField( + org.apache.parquet.schema.Types.buildGroup(Repetition.OPTIONAL) + .as(LogicalTypeAnnotation.listType()) + .addField( + org.apache.parquet.schema.Types.repeatedGroup() + .optional(PrimitiveTypeName.BINARY) + .as(LogicalTypeAnnotation.stringType()) + .named("element") + .named("list")) + .named("tags")) + .named("root"); + + private static final class Nulls { + final IntPredicate name; + final IntPredicate address; + final IntPredicate city; + final IntPredicate zip; + + Nulls(IntPredicate name, IntPredicate address, IntPredicate city, IntPredicate zip) { + this.name = name; + this.address = address; + this.city = city; + this.zip = zip; + } + + static final Nulls NONE = new Nulls(r -> false, r -> false, r -> false, r -> false); + } + + private ParquetMetadata write(int rows, int rowGroups, boolean stats, Nulls nulls) + throws IOException { + File file = new File(tmp.getRoot(), "t" + System.nanoTime() + ".parquet"); + ExampleParquetWriter.Builder builder = + ExampleParquetWriter.builder(new Path(file.getAbsolutePath())) + .withType(MIXED) + .withStatisticsEnabled(stats); + if (rowGroups > 1) { + int rowsPerGroup = rows / rowGroups; + builder = + builder + .withRowGroupSize(1L) + .withMinRowCountForPageSizeCheck(rowsPerGroup) + .withMaxRowCountForPageSizeCheck(rowsPerGroup); + } + SimpleGroupFactory factory = new SimpleGroupFactory(MIXED); + try (ParquetWriter writer = builder.build()) { + for (int row = 0; row < rows; row++) { + Group group = factory.newGroup(); + group.add("id", (long) row); + if (!nulls.name.test(row)) { + group.add("name", "n" + row); + } + if (!nulls.address.test(row)) { + Group address = group.addGroup("address"); + if (!nulls.city.test(row)) { + address.add("city", "c" + row); + } + if (!nulls.zip.test(row)) { + address.add("zip", row); + } + } + Group tags = group.addGroup("tags"); + tags.addGroup("list").add("element", "t" + row); + writer.write(group); + } + } + return ParquetFooters.read(file.getAbsolutePath()); + } + + private static Schema tightened(ParquetMetadata footer) { + return FileSchemas.tighten( + ParquetSchemaUtil.convert(footer.getFileMetaData().getSchema()), footer); + } + + private static boolean isRequired(Schema schema, String path) { + return schema.findField(path).isRequired(); + } + + @Test + public void testProvenNullFreeColumnsBecomeRequired() throws IOException { + Schema schema = tightened(write(10, 1, true, Nulls.NONE)); + assertTrue(isRequired(schema, "id")); + assertTrue(isRequired(schema, "name")); + assertTrue(isRequired(schema, "address")); + assertTrue(isRequired(schema, "address.city")); + assertTrue(isRequired(schema, "address.zip")); + } + + @Test + public void testListAndElementStayAsDeclared() throws IOException { + Schema schema = tightened(write(10, 1, true, Nulls.NONE)); + assertFalse(isRequired(schema, "tags")); + assertFalse(schema.findField("tags").type().asListType().isElementRequired()); + } + + @Test + public void testSomeNullsStayOptional() throws IOException { + Schema schema = + tightened(write(10, 1, true, new Nulls(r -> r == 3, r -> false, r -> false, r -> false))); + assertFalse(isRequired(schema, "name")); + assertTrue(isRequired(schema, "address.city")); + } + + @Test + public void testAllNullsStayOptional() throws IOException { + Schema schema = + tightened(write(10, 1, true, new Nulls(r -> true, r -> false, r -> false, r -> false))); + assertFalse(isRequired(schema, "name")); + } + + @Test + public void testStatsDisabledStaysOptional() throws IOException { + Schema schema = tightened(write(10, 1, false, Nulls.NONE)); + assertTrue(isRequired(schema, "id")); + assertFalse(isRequired(schema, "name")); + assertFalse(isRequired(schema, "address")); + assertFalse(isRequired(schema, "address.city")); + } + + @Test + public void testOneRowGroupWithNullsSpoilsTheProof() throws IOException { + Schema schema = + tightened(write(100, 4, true, new Nulls(r -> r == 60, r -> false, r -> false, r -> false))); + assertFalse(isRequired(schema, "name")); + assertTrue(isRequired(schema, "address.zip")); + } + + /** With no rows nothing can violate a required column, so every column counts as proven. */ + @Test + public void testZeroRowsProveEverything() throws IOException { + Schema schema = tightened(write(0, 1, true, Nulls.NONE)); + assertTrue(isRequired(schema, "id")); + assertTrue(isRequired(schema, "name")); + assertTrue(isRequired(schema, "address")); + assertTrue(isRequired(schema, "address.city")); + assertFalse(isRequired(schema, "tags")); + } + + @Test + public void testCanonicalWithNullFreeColumnsReportsChangedOnly() throws IOException { + ParquetMetadata footer = + write(10, 1, true, new Nulls(r -> false, r -> false, r -> false, r -> r == 4)); + CollectDistinctSchemas.SchemaGroup group = FileSchemas.schemaGroup(footer); + // id is declared required already; zip has a null; the rest flipped + assertEquals(java.util.Arrays.asList("address", "address.city", "name"), group.nullFreeColumns); + Schema declared = SchemaParser.fromJson(group.schemaJson); + assertFalse(isRequired(declared, "name")); + } + + @Test + public void testMarkRequiredFlipsOnlyNamedColumns() { + Schema declared = + new Schema( + optional(1, "name", Types.StringType.get()), + optional( + 2, + "address", + Types.StructType.of( + optional(3, "city", Types.StringType.get()), + optional(4, "zip", Types.IntegerType.get())))); + Schema required = + FileSchemas.markRequired( + declared, java.util.Arrays.asList("address", "address.city", "not_a_column")); + assertFalse(isRequired(required, "name")); + assertTrue(isRequired(required, "address")); + assertTrue(isRequired(required, "address.city")); + assertFalse(isRequired(required, "address.zip")); + assertEquals( + declared.asStruct(), + FileSchemas.markRequired(declared, java.util.Arrays.asList()).asStruct()); + } + + @Test + public void testNullStructKeepsStructAndLeavesOptional() throws IOException { + Schema schema = + tightened(write(10, 1, true, new Nulls(r -> false, r -> r == 5, r -> false, r -> false))); + assertFalse(isRequired(schema, "address")); + assertFalse(isRequired(schema, "address.city")); + assertFalse(isRequired(schema, "address.zip")); + } + + @Test + public void testOneProvenLeafProvesTheStruct() throws IOException { + Schema schema = + tightened(write(10, 1, true, new Nulls(r -> false, r -> false, r -> r == 2, r -> false))); + assertTrue(isRequired(schema, "address")); + assertFalse(isRequired(schema, "address.city")); + assertTrue(isRequired(schema, "address.zip")); + } + + @Test + public void testTightenPreservesIdsNamesAndTypes() throws IOException { + ParquetMetadata footer = write(10, 1, true, Nulls.NONE); + Schema converted = ParquetSchemaUtil.convert(footer.getFileMetaData().getSchema()); + Schema schema = FileSchemas.tighten(converted, footer); + assertEquals(converted.columns().size(), schema.columns().size()); + for (Types.NestedField field : converted.columns()) { + Types.NestedField after = schema.findField(field.fieldId()); + assertEquals(field.name(), after.name()); + assertEquals(field.type().typeId(), after.type().typeId()); + } + } + + @Test + public void testTightenAndCanonicalPreserveDocAndDefaults() { + Types.NestedField withAttributes = + Types.NestedField.optional("b") + .withId(2) + .ofType(Types.LongType.get()) + .withDoc("the b") + .withWriteDefault(org.apache.iceberg.expressions.Literal.of(7L)) + .build(); + Schema schema = new Schema(withAttributes, required(1, "a", Types.StringType.get())); + Schema canonical = FileSchemas.canonical(schema); + Types.NestedField b = canonical.findField("b"); + assertEquals("the b", b.doc()); + assertEquals(7L, b.writeDefault()); + } + + // ---- canonicalization @Test public void testSortsTopLevelFieldsAndRenumbers() { diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/ReadFooterSchemaTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/ReadFooterSchemaTest.java index 6b3679913918..bc4c0ae81507 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/ReadFooterSchemaTest.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/ReadFooterSchemaTest.java @@ -164,6 +164,38 @@ public void testNonParquetEmitsNothing() throws IOException { pipeline.run(); } + /** + * The schema is emitted exactly as declared; which columns proved null-free travels beside it. + */ + @Test + public void testNullFreeColumnsAreReported() throws IOException { + Record full = GenericRecord.create(FLAT_SCHEMA).copy("id", 1, "name", "a"); + String clean = writeParquet("clean.parquet", FLAT_SCHEMA, full); + PCollection out = run(clean); + assertSchemas(out, FLAT_SCHEMA); + assertNullFreeColumns(out, Arrays.asList("name")); + pipeline.run(); + } + + @Test + public void testColumnWithNullsIsNotReportedNullFree() throws IOException { + String withNull = writeParquet("null.parquet", FLAT_SCHEMA, record(FLAT_SCHEMA, "id", 1)); + PCollection out = run(withNull); + assertSchemas(out, FLAT_SCHEMA); + assertNullFreeColumns(out, Arrays.asList()); + pipeline.run(); + } + + private static void assertNullFreeColumns( + PCollection out, List expected) { + PAssert.thatSingleton(out) + .satisfies( + group -> { + assertEquals(expected, group.nullFreeColumns); + return null; + }); + } + @Test public void testPermutedColumnsProduceIdenticalSchema() throws IOException { Schema permuted = @@ -177,7 +209,7 @@ public void testPermutedColumnsProduceIdenticalSchema() throws IOException { .satisfies( actual -> { List jsons = new ArrayList<>(); - actual.forEach(jsons::add); + actual.forEach(group -> jsons.add(group.schemaJson)); assertEquals(2, jsons.size()); assertEquals(jsons.get(0), jsons.get(1)); return null; @@ -218,12 +250,16 @@ private static long counter(PipelineResult result, String name) { return total; } - private PCollection run(String... paths) { - return pipeline.apply(Create.of(Arrays.asList(paths))).apply(ParDo.of(new ReadFooterSchema())); + private PCollection run(String... paths) { + return pipeline + .apply(Create.of(Arrays.asList(paths))) + .apply(ParDo.of(new ReadFooterSchema())) + .setCoder(CollectDistinctSchemas.groupCoder()); } - /** Asserts the emitted schemas equal the canonical forms of {@code expected}, in any order. */ - private static void assertSchemas(PCollection out, Schema... expected) { + /** Asserts the emitted declared schemas equal the canonical forms of {@code expected}. */ + private static void assertSchemas( + PCollection out, Schema... expected) { List expectedJson = new ArrayList<>(); for (Schema schema : expected) { expectedJson.add(SchemaParser.toJson(FileSchemas.canonical(schema))); @@ -232,7 +268,8 @@ private static void assertSchemas(PCollection out, Schema... expected) { .satisfies( actual -> { List remaining = new ArrayList<>(expectedJson); - for (String json : actual) { + for (CollectDistinctSchemas.SchemaGroup group : actual) { + String json = group.schemaJson; Schema schema = SchemaParser.fromJson(json); boolean matched = false; for (int i = 0; i < remaining.size(); i++) { From bc7707aa189831ae04258440fd142522b635d008 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 3 Sep 2026 13:11:26 -0400 Subject: [PATCH 2/2] comments --- .../io/iceberg/CollectDistinctSchemas.java | 128 +++++++----------- .../beam/sdk/io/iceberg/FileSchemas.java | 26 +++- .../iceberg/CollectDistinctSchemasTest.java | 4 +- .../beam/sdk/io/iceberg/FileSchemasTest.java | 99 +++++++++++++- .../sdk/io/iceberg/ReadFooterSchemaTest.java | 6 +- 5 files changed, 167 insertions(+), 96 deletions(-) diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/CollectDistinctSchemas.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/CollectDistinctSchemas.java index 16e877619b10..5c36967cfeba 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/CollectDistinctSchemas.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/CollectDistinctSchemas.java @@ -17,6 +17,7 @@ */ package org.apache.beam.sdk.io.iceberg; +import com.google.auto.value.AutoValue; import java.io.IOException; import java.io.InputStream; import java.io.OutputStream; @@ -26,14 +27,21 @@ import java.util.Objects; import java.util.TreeMap; import java.util.TreeSet; +import org.apache.beam.sdk.coders.AtomicCoder; import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.coders.CoderRegistry; -import org.apache.beam.sdk.coders.CustomCoder; import org.apache.beam.sdk.coders.ListCoder; import org.apache.beam.sdk.coders.MapCoder; import org.apache.beam.sdk.coders.StringUtf8Coder; import org.apache.beam.sdk.coders.VarLongCoder; +import org.apache.beam.sdk.schemas.AutoValueSchema; +import org.apache.beam.sdk.schemas.NoSuchSchemaException; +import org.apache.beam.sdk.schemas.SchemaCoder; +import org.apache.beam.sdk.schemas.SchemaRegistry; +import org.apache.beam.sdk.schemas.annotations.DefaultSchema; +import org.apache.beam.sdk.schemas.annotations.SchemaFieldNumber; import org.apache.beam.sdk.transforms.Combine; +import org.checkerframework.checker.nullness.qual.MonotonicNonNull; import org.checkerframework.checker.nullness.qual.Nullable; /** @@ -78,36 +86,42 @@ public int hashCode() { * A schema, how many files carry it, and the columns all of them proved free of nulls. * ReadFooterSchema emits one per file ({@code files} = 1); this combiner merges them. */ - static final class SchemaGroup { - final String schemaJson; - final long files; - final List nullFreeColumns; - - SchemaGroup(String schemaJson, long files, List nullFreeColumns) { - this.schemaJson = schemaJson; - this.files = files; - this.nullFreeColumns = nullFreeColumns; - } - - @Override - public boolean equals(@Nullable Object other) { - if (!(other instanceof SchemaGroup)) { - return false; + @DefaultSchema(AutoValueSchema.class) + @AutoValue + abstract static class SchemaGroup { + private static @MonotonicNonNull SchemaCoder coder; + + static SchemaGroup of(String schemaJson, long files, List nullFreeColumns) { + return new AutoValue_CollectDistinctSchemas_SchemaGroup(schemaJson, files, nullFreeColumns); + } + + static SchemaCoder getCoder() { + if (coder == null) { + try { + coder = SchemaRegistry.createDefault().getSchemaCoder(SchemaGroup.class); + } catch (NoSuchSchemaException e) { + throw new RuntimeException(e); + } } - SchemaGroup that = (SchemaGroup) other; - return files == that.files - && schemaJson.equals(that.schemaJson) - && nullFreeColumns.equals(that.nullFreeColumns); + return coder; } - @Override - public int hashCode() { - return Objects.hash(schemaJson, files, nullFreeColumns); - } + @SchemaFieldNumber("0") + abstract String getSchemaJson(); + + @SchemaFieldNumber("1") + abstract long getFiles(); + + @SchemaFieldNumber("2") + abstract List getNullFreeColumns(); @Override - public String toString() { - return files + " file(s), null-free in " + nullFreeColumns + ", schema " + schemaJson; + public final String toString() { + return getFiles() + + " file(s), null-free in " + + getNullFreeColumns() + + ", schema " + + getSchemaJson(); } } @@ -118,7 +132,7 @@ public Map createAccumulator() { @Override public Map addInput(Map accumulator, SchemaGroup file) { - add(accumulator, file.schemaJson, file.files, file.nullFreeColumns); + add(accumulator, file.getSchemaJson(), file.getFiles(), file.getNullFreeColumns()); return accumulator; } @@ -138,18 +152,18 @@ public List extractOutput(Map accumulator) { List schemas = new ArrayList<>(); for (Map.Entry entry : accumulator.entrySet()) { schemas.add( - new SchemaGroup( + SchemaGroup.of( entry.getKey(), entry.getValue().files, new ArrayList<>(entry.getValue().nullFreeColumns))); } schemas.sort( (a, b) -> { - int byCount = Long.compare(b.files, a.files); + int byCount = Long.compare(b.getFiles(), a.getFiles()); if (byCount != 0) { return byCount; } - return a.schemaJson.compareTo(b.schemaJson); + return a.getSchemaJson().compareTo(b.getSchemaJson()); }); return schemas; } @@ -167,34 +181,21 @@ public Coder> getDefaultOutputCoder( } static Coder groupCoder() { - return SchemaGroupCoder.INSTANCE; + return SchemaGroup.getCoder(); } static Coder> outputCoder() { - return ListCoder.of(SchemaGroupCoder.INSTANCE); + return ListCoder.of(SchemaGroup.getCoder()); } private static final Coder> COLUMNS_CODER = ListCoder.of(StringUtf8Coder.of()); - /** Singletons with class equality, so repeated mentions compare equal; deterministic encoding. */ - private static class GroupCoder extends CustomCoder { + /** Sorted columns, so the encoding is deterministic. */ + private static class GroupCoder extends AtomicCoder { static final GroupCoder INSTANCE = new GroupCoder(); private GroupCoder() {} - @Override - public void verifyDeterministic() {} - - @Override - public boolean equals(@Nullable Object other) { - return other instanceof GroupCoder; - } - - @Override - public int hashCode() { - return getClass().hashCode(); - } - @Override public void encode(Group value, OutputStream out) throws IOException { VarLongCoder.of().encode(value.files, out); @@ -208,39 +209,6 @@ public Group decode(InputStream in) throws IOException { } } - private static class SchemaGroupCoder extends CustomCoder { - static final SchemaGroupCoder INSTANCE = new SchemaGroupCoder(); - - private SchemaGroupCoder() {} - - @Override - public void verifyDeterministic() {} - - @Override - public boolean equals(@Nullable Object other) { - return other instanceof SchemaGroupCoder; - } - - @Override - public int hashCode() { - return getClass().hashCode(); - } - - @Override - public void encode(SchemaGroup value, OutputStream out) throws IOException { - StringUtf8Coder.of().encode(value.schemaJson, out); - VarLongCoder.of().encode(value.files, out); - COLUMNS_CODER.encode(value.nullFreeColumns, out); - } - - @Override - public SchemaGroup decode(InputStream in) throws IOException { - String schemaJson = StringUtf8Coder.of().decode(in); - long files = VarLongCoder.of().decode(in); - return new SchemaGroup(schemaJson, files, COLUMNS_CODER.decode(in)); - } - } - private static void add( Map accumulator, String schemaJson, diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/FileSchemas.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/FileSchemas.java index 2d824f9f6670..d95321153d5e 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/FileSchemas.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/FileSchemas.java @@ -62,7 +62,7 @@ static String canonicalJson(ParquetMetadata footer) { static CollectDistinctSchemas.SchemaGroup schemaGroup(ParquetMetadata footer) { Schema converted = ParquetSchemaUtil.convert(footer.getFileMetaData().getSchema()); Schema tightened = tighten(converted, footer); - return new CollectDistinctSchemas.SchemaGroup( + return CollectDistinctSchemas.SchemaGroup.of( SchemaParser.toJson(canonical(converted)), 1, changedToRequired(converted, tightened)); } @@ -72,11 +72,12 @@ static CollectDistinctSchemas.SchemaGroup schemaGroup(ParquetMetadata footer) { * struct is null-free when any leaf under it is (a null struct nulls all its leaves). Nothing * under lists or maps is tightened: a zero count there would be valid evidence too, but mapping * physical chunk paths (writer-dependent names like {@code list.element}, {@code array}) onto the - * converted schema is not worth it. A file with no row groups has no rows and proves every - * column, matching how the pin check treats empty files. + * converted schema is not worth it. A file with no rows (no row groups, or only empty ones, as + * pyarrow writes an empty table) proves every column, matching how the pin check treats empty + * files. */ static Schema tighten(Schema schema, ParquetMetadata footer) { - if (footer.getBlocks().isEmpty()) { + if (rowCount(footer) == 0) { return new Schema(tightenAll(schema.asStruct()).fields()); } Set> zeroNullLeaves = leafPathsWithZeroNullCounts(footer); @@ -155,13 +156,26 @@ private static void collectChangedToRequired( } } - /** Leaf paths proven null-free in every row group (intersection over blocks). */ + private static long rowCount(ParquetMetadata footer) { + long rows = 0; + for (BlockMetaData block : footer.getBlocks()) { + rows += block.getRowCount(); + } + return rows; + } + + /** + * Leaf paths proven null-free in every row group that has rows (intersection over blocks). An + * empty row group holds no nulls whatever its statistics say, so it constrains nothing. + */ private static Set> leafPathsWithZeroNullCounts(ParquetMetadata footer) { Set> proven = null; for (BlockMetaData block : footer.getBlocks()) { + if (block.getRowCount() == 0) { + continue; + } Set> provenHere = new HashSet<>(); for (ColumnChunkMetaData chunk : block.getColumns()) { - // A zero-row row group proves trivially: zero rows hold zero nulls. Statistics stats = chunk.getStatistics(); if (stats != null && stats.isNumNullsSet() && stats.getNumNulls() == 0) { provenHere.add(Arrays.asList(chunk.getPath().toArray())); diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/CollectDistinctSchemasTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/CollectDistinctSchemasTest.java index 476954f0fb4d..7b4a22b2618e 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/CollectDistinctSchemasTest.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/CollectDistinctSchemasTest.java @@ -80,7 +80,7 @@ public void testDifferentStringsAreDistinct() { combine(group(ID_NAME, 1, NONE), group(NAME_ID, 1, NONE), group(ID_LONG_NAME, 1, NONE)); assertEquals(3, out.size()); for (SchemaGroup entry : out) { - assertEquals(1L, entry.files); + assertEquals(1L, entry.getFiles()); } } @@ -179,7 +179,7 @@ private List combine(SchemaGroup... files) { } private static SchemaGroup group(String json, long files, List proven) { - return new SchemaGroup(json, files, proven); + return SchemaGroup.of(json, files, proven); } private static String json(Schema schema) { diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/FileSchemasTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/FileSchemasTest.java index 7358e8d50bd2..6a6bef483b73 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/FileSchemasTest.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/FileSchemasTest.java @@ -25,16 +25,26 @@ import java.io.File; import java.io.IOException; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; import java.util.function.IntPredicate; import org.apache.hadoop.fs.Path; import org.apache.iceberg.Schema; import org.apache.iceberg.SchemaParser; import org.apache.iceberg.parquet.ParquetSchemaUtil; import org.apache.iceberg.types.Types; +import org.apache.parquet.column.ColumnDescriptor; +import org.apache.parquet.column.EncodingStats; +import org.apache.parquet.column.statistics.Statistics; import org.apache.parquet.example.data.Group; import org.apache.parquet.example.data.simple.SimpleGroupFactory; import org.apache.parquet.hadoop.ParquetWriter; import org.apache.parquet.hadoop.example.ExampleParquetWriter; +import org.apache.parquet.hadoop.metadata.BlockMetaData; +import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData; +import org.apache.parquet.hadoop.metadata.ColumnPath; +import org.apache.parquet.hadoop.metadata.CompressionCodecName; import org.apache.parquet.hadoop.metadata.ParquetMetadata; import org.apache.parquet.schema.LogicalTypeAnnotation; import org.apache.parquet.schema.MessageType; @@ -188,8 +198,12 @@ public void testStatsDisabledStaysOptional() throws IOException { @Test public void testOneRowGroupWithNullsSpoilsTheProof() throws IOException { - Schema schema = - tightened(write(100, 4, true, new Nulls(r -> r == 60, r -> false, r -> false, r -> false))); + ParquetMetadata footer = + write(100, 4, true, new Nulls(r -> r == 60, r -> false, r -> false, r -> false)); + assertTrue("expected several row groups", footer.getBlocks().size() > 1); + + Schema schema = tightened(footer); + assertFalse(isRequired(schema, "name")); assertTrue(isRequired(schema, "address.zip")); } @@ -197,7 +211,11 @@ public void testOneRowGroupWithNullsSpoilsTheProof() throws IOException { /** With no rows nothing can violate a required column, so every column counts as proven. */ @Test public void testZeroRowsProveEverything() throws IOException { - Schema schema = tightened(write(0, 1, true, Nulls.NONE)); + ParquetMetadata footer = write(0, 1, true, Nulls.NONE); + assertEquals("parquet-mr writes no row group for zero rows", 0, footer.getBlocks().size()); + + Schema schema = tightened(footer); + assertTrue(isRequired(schema, "id")); assertTrue(isRequired(schema, "name")); assertTrue(isRequired(schema, "address")); @@ -205,14 +223,85 @@ public void testZeroRowsProveEverything() throws IOException { assertFalse(isRequired(schema, "tags")); } + /** + * pyarrow writes an empty table as one row group with zero rows and no statistics. parquet-mr + * never produces that shape, so the footer is assembled by hand. + */ + @Test + public void testEmptyRowGroupWithoutStatsProvesEverything() throws IOException { + ParquetMetadata footer = withBlocks(write(0, 1, true, Nulls.NONE), emptyBlockWithoutStats()); + + Schema schema = tightened(footer); + + assertTrue(isRequired(schema, "id")); + assertTrue(isRequired(schema, "name")); + assertTrue(isRequired(schema, "address")); + assertTrue(isRequired(schema, "address.city")); + assertFalse(isRequired(schema, "tags")); + } + + @Test + public void testEmptyRowGroupDoesNotSpoilTheProof() throws IOException { + ParquetMetadata written = write(10, 1, true, Nulls.NONE); + ParquetMetadata footer = + withBlocks(written, written.getBlocks().get(0), emptyBlockWithoutStats()); + + Schema schema = tightened(footer); + + assertTrue(isRequired(schema, "name")); + assertTrue(isRequired(schema, "address.city")); + } + + @Test + public void testEmptyRowGroupDoesNotHideNullsElsewhere() throws IOException { + ParquetMetadata written = + write(10, 1, true, new Nulls(r -> r == 3, r -> false, r -> false, r -> false)); + ParquetMetadata footer = + withBlocks(written, emptyBlockWithoutStats(), written.getBlocks().get(0)); + + Schema schema = tightened(footer); + + assertFalse(isRequired(schema, "name")); + assertTrue(isRequired(schema, "address.city")); + } + + private static ParquetMetadata withBlocks(ParquetMetadata footer, BlockMetaData... blocks) { + List list = new ArrayList<>(); + Collections.addAll(list, blocks); + return new ParquetMetadata(footer.getFileMetaData(), list); + } + + /** One chunk per leaf of {@link #MIXED}, zero rows, statistics absent as a reader sees them. */ + private static BlockMetaData emptyBlockWithoutStats() { + BlockMetaData block = new BlockMetaData(); + block.setRowCount(0); + for (ColumnDescriptor column : MIXED.getColumns()) { + block.addColumn( + ColumnChunkMetaData.get( + ColumnPath.get(column.getPath()), + column.getPrimitiveType(), + CompressionCodecName.UNCOMPRESSED, + new EncodingStats.Builder().build(), + Collections.emptySet(), + Statistics.getBuilderForReading(column.getPrimitiveType()).build(), + 0, + 0, + 0, + 0, + 0)); + } + return block; + } + @Test public void testCanonicalWithNullFreeColumnsReportsChangedOnly() throws IOException { ParquetMetadata footer = write(10, 1, true, new Nulls(r -> false, r -> false, r -> false, r -> r == 4)); CollectDistinctSchemas.SchemaGroup group = FileSchemas.schemaGroup(footer); // id is declared required already; zip has a null; the rest flipped - assertEquals(java.util.Arrays.asList("address", "address.city", "name"), group.nullFreeColumns); - Schema declared = SchemaParser.fromJson(group.schemaJson); + assertEquals( + java.util.Arrays.asList("address", "address.city", "name"), group.getNullFreeColumns()); + Schema declared = SchemaParser.fromJson(group.getSchemaJson()); assertFalse(isRequired(declared, "name")); } diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/ReadFooterSchemaTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/ReadFooterSchemaTest.java index bc4c0ae81507..5cac095201db 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/ReadFooterSchemaTest.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/ReadFooterSchemaTest.java @@ -191,7 +191,7 @@ private static void assertNullFreeColumns( PAssert.thatSingleton(out) .satisfies( group -> { - assertEquals(expected, group.nullFreeColumns); + assertEquals(expected, group.getNullFreeColumns()); return null; }); } @@ -209,7 +209,7 @@ public void testPermutedColumnsProduceIdenticalSchema() throws IOException { .satisfies( actual -> { List jsons = new ArrayList<>(); - actual.forEach(group -> jsons.add(group.schemaJson)); + actual.forEach(group -> jsons.add(group.getSchemaJson())); assertEquals(2, jsons.size()); assertEquals(jsons.get(0), jsons.get(1)); return null; @@ -269,7 +269,7 @@ private static void assertSchemas( actual -> { List remaining = new ArrayList<>(expectedJson); for (CollectDistinctSchemas.SchemaGroup group : actual) { - String json = group.schemaJson; + String json = group.getSchemaJson(); Schema schema = SchemaParser.fromJson(json); boolean matched = false; for (int i = 0; i < remaining.size(); i++) {