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..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,87 +17,219 @@ */ 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; 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.AtomicCoder; 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.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.apache.beam.sdk.values.KV; +import org.checkerframework.checker.nullness.qual.MonotonicNonNull; +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. + */ + @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); + } + } + return coder; + } + + @SchemaFieldNumber("0") + abstract String getSchemaJson(); + + @SchemaFieldNumber("1") + abstract long getFiles(); + + @SchemaFieldNumber("2") + abstract List getNullFreeColumns(); + + @Override + public final String toString() { + return getFiles() + + " file(s), null-free in " + + getNullFreeColumns() + + ", schema " + + getSchemaJson(); + } + } @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.getSchemaJson(), file.getFiles(), file.getNullFreeColumns()); 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( + SchemaGroup.of( + 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.getFiles(), a.getFiles()); if (byCount != 0) { return byCount; } - return a.getKey().compareTo(b.getKey()); + return a.getSchemaJson().compareTo(b.getSchemaJson()); }); 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(); + } + + static Coder groupCoder() { + return SchemaGroup.getCoder(); + } + + static Coder> outputCoder() { + return ListCoder.of(SchemaGroup.getCoder()); + } + + private static final Coder> COLUMNS_CODER = ListCoder.of(StringUtf8Coder.of()); + + /** Sorted columns, so the encoding is deterministic. */ + private static class GroupCoder extends AtomicCoder { + static final GroupCoder INSTANCE = new GroupCoder(); + + private GroupCoder() {} + + @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 void add(Map accumulator, String schemaJson, long count) { - Long existing = accumulator.get(schemaJson); + 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..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 @@ -18,29 +18,222 @@ 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 CollectDistinctSchemas.SchemaGroup.of( + 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 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 (rowCount(footer) == 0) { + 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); + } + } + } + + 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()) { + 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..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 @@ -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.getFiles()); } } @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 SchemaGroup.of(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..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 @@ -20,17 +20,362 @@ 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.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; +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 { + 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")); + } + + /** With no rows nothing can violate a required column, so every column counts as proven. */ + @Test + public void testZeroRowsProveEverything() throws IOException { + 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")); + assertTrue(isRequired(schema, "address.city")); + 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.getNullFreeColumns()); + Schema declared = SchemaParser.fromJson(group.getSchemaJson()); + 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..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 @@ -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.getNullFreeColumns()); + 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.getSchemaJson())); 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.getSchemaJson(); Schema schema = SchemaParser.fromJson(json); boolean matched = false; for (int i = 0; i < remaining.size(); i++) {