Parquet: Fix null counts for nested optional fields of a null struct

data/src/test/java/org/apache/iceberg/parquet/TestParquetMetrics.java

@@ -18,26 +18,42 @@
  */
 package org.apache.iceberg.parquet;
 
+import static org.apache.iceberg.types.Types.NestedField.optional;
+import static org.apache.iceberg.types.Types.NestedField.required;
+import static org.assertj.core.api.Assertions.assertThat;
+
 import java.io.File;
 import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
 import java.util.Map;
 import java.util.UUID;
+import org.apache.iceberg.DataFile;
+import org.apache.iceberg.DataFiles;
 import org.apache.iceberg.FileFormat;
 import org.apache.iceberg.Files;
 import org.apache.iceberg.Metrics;
 import org.apache.iceberg.MetricsConfig;
 import org.apache.iceberg.ParameterizedTestExtension;
+import org.apache.iceberg.PartitionSpec;
 import org.apache.iceberg.Schema;
 import org.apache.iceberg.TableProperties;
 import org.apache.iceberg.TestMetrics;
+import org.apache.iceberg.data.GenericRecord;
 import org.apache.iceberg.data.Record;
 import org.apache.iceberg.data.parquet.GenericParquetWriter;
 import org.apache.iceberg.io.FileAppender;
 import org.apache.iceberg.io.InputFile;
 import org.apache.iceberg.io.OutputFile;
 import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
 import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.apache.iceberg.types.Types.DoubleType;
+import org.apache.iceberg.types.Types.FloatType;
+import org.apache.iceberg.types.Types.GeometryType;
+import org.apache.iceberg.types.Types.LongType;
+import org.apache.iceberg.types.Types.StructType;
 import org.apache.parquet.hadoop.ParquetFileReader;
+import org.junit.jupiter.api.TestTemplate;
 import org.junit.jupiter.api.extension.ExtendWith;
 
 /** Test Metrics for Parquet. */
@@ -107,4 +123,96 @@ public int splitCount(InputFile inputFile) throws IOException {
   public boolean supportsSmallRowGroups() {
     return true;
   }
+
+  @TestTemplate
+  public void testMetricsForNullStructWithFloatingAndGeoLeaves() throws IOException {
+    // float, double, geometry and geography are the only types whose writers report metrics, so
+    // they are the ones whose nested null counts could be dropped when a struct is null. Null
+    // counts are only tracked for optional fields.
+    StructType struct =
+        StructType.of(
+            optional(2, "optDouble", DoubleType.get()),
+            optional(3, "optFloat", FloatType.get()),
+            optional(4, "geom", GeometryType.crs84()),
+            optional(5, "optLong", LongType.get()));
+    Schema schema = new Schema(optional(1, "struct", struct));
+
+    Record inner = GenericRecord.create(struct);
+    inner.setField("optDouble", 1.5D);
+    inner.setField("optFloat", 2.5F);
+    inner.setField("geom", wkbPoint(30, 10));
+    inner.setField("optLong", 10L);
+    Record withStruct = GenericRecord.create(schema);
+    withStruct.setField("struct", inner);
+    Record nullStruct = GenericRecord.create(schema);
+    nullStruct.setField("struct", null);
+
+    Metrics metrics = getMetrics(schema, withStruct, nullStruct, nullStruct);
+
+    assertThat(metrics.recordCount()).isEqualTo(3L);
+    // each leaf has one value from the populated struct and two nulls from the null structs
+    assertCounts(2, 3L, 2L, 0L, metrics);
+    assertCounts(3, 3L, 2L, 0L, metrics);
+    assertCounts(4, 3L, 2L, metrics);
+    // a type without writer metrics was already correct via the footer; included as a control
+    assertCounts(5, 3L, 2L, metrics);
+
+    // the counts also reach a data file built from these metrics
+    DataFile dataFile =
+        DataFiles.builder(PartitionSpec.unpartitioned())
+            .withPath("/path/to/file.parquet")
+            .withFileSizeInBytes(1024)
+            .withFormat(FileFormat.PARQUET)
+            .withMetrics(metrics)
+            .build();
+    assertThat(dataFile.nullValueCounts()).containsEntry(2, 2L).containsEntry(3, 2L);
+  }
+
+  @TestTemplate
+  public void testMetricsForRequiredNestedFieldInNullStruct() throws IOException {
+    // null counts are only tracked for optional fields, so a required float/double leaf keeps its
+    // prior behavior when the struct is null: the null count is not recorded.
+    StructType struct =
+        StructType.of(
+            required(2, "reqDouble", DoubleType.get()), required(3, "reqLong", LongType.get()));
+    Schema schema = new Schema(optional(1, "struct", struct));
+
+    Record inner = GenericRecord.create(struct);
+    inner.setField("reqDouble", 1.5D);
+    inner.setField("reqLong", 10L);
+    Record withStruct = GenericRecord.create(schema);
+    withStruct.setField("struct", inner);
+    Record nullStruct = GenericRecord.create(schema);
+    nullStruct.setField("struct", null);
+
+    Metrics metrics = getMetrics(schema, withStruct, nullStruct, nullStruct);
+
+    assertThat(metrics.recordCount()).isEqualTo(3L);
+    // reqDouble uses writer metrics, which do not count nulls for required fields: only the
+    // populated struct's value is seen
+    assertCounts(2, 1L, 0L, 0L, metrics);
+    // reqLong uses footer stats, so its counts are unaffected
+    assertCounts(3, 3L, 2L, metrics);
+
+    DataFile dataFile =
+        DataFiles.builder(PartitionSpec.unpartitioned())
+            .withPath("/path/to/file.parquet")
+            .withFileSizeInBytes(1024)
+            .withFormat(FileFormat.PARQUET)
+            .withMetrics(metrics)
+            .build();
+    assertThat(dataFile.nullValueCounts()).containsEntry(2, 0L).containsEntry(3, 2L);
+  }
+
+  private static ByteBuffer wkbPoint(double xCoord, double yCoord) {
+    // little-endian WKB encoding of a point
+    return ByteBuffer.wrap(
+        ByteBuffer.allocate(21)
+            .order(ByteOrder.LITTLE_ENDIAN)
+            .put((byte) 1) // byte order: little endian
+            .putInt(1) // WKB geometry type: Point
+            .putDouble(xCoord)
+            .putDouble(yCoord)
+            .array());
+  }
 }

parquet/src/main/java/org/apache/iceberg/parquet/ParquetValueWriters.java

@@ -27,6 +27,7 @@
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.UUID;
 import java.util.function.Function;
 import java.util.stream.Collectors;
@@ -41,6 +42,7 @@
 import org.apache.iceberg.deletes.PositionDelete;
 import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
 import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList;
+import org.apache.iceberg.relocated.com.google.common.collect.Sets;
 import org.apache.iceberg.types.TypeUtil;
 import org.apache.iceberg.types.Types;
 import org.apache.iceberg.util.DecimalUtil;
@@ -56,12 +58,38 @@
   public static <T> ParquetValueWriter<T> option(
       Type type, int definitionLevel, ParquetValueWriter<T> writer) {
     if (type.isRepetition(Type.Repetition.OPTIONAL)) {
-      return new OptionWriter<>(definitionLevel, writer);
+      return new OptionWriter<>(definitionLevel, writer, optionalLeafIds(type));
     }
 
     return writer;
   }
 
+  /**
+   * Collects the ids of optional primitive leaves within a type. A null written for an optional
+   * value is also a null for these fields, but not for required fields, whose null count is left as
+   * it was before (per the spec, null counts are only tracked for optional fields).
+   */
+  private static Set<Integer> optionalLeafIds(Type type) {
+    Set<Integer> ids = Sets.newHashSet();
+    collectOptionalLeafIds(type, ids);
+    return ids;
+  }
+
+  private static void collectOptionalLeafIds(Type type, Set<Integer> ids) {
+    if (type.isRepetition(Type.Repetition.REPEATED)) {
+      // repeated (list and map) fields are not propagated to from an enclosing option
+      return;
+    } else if (type.isPrimitive()) {
+      if (type.isRepetition(Type.Repetition.OPTIONAL) && type.getId() != null) {
+        ids.add(type.getId().intValue());
+      }
+    } else {
+      for (Type field : type.asGroupType().getFields()) {
+        collectOptionalLeafIds(field, ids);
+      }
+    }
+  }
+
   public static UnboxedWriter<Boolean> booleans(ColumnDescriptor desc) {
     return new UnboxedWriter<>(desc);
   }
@@ -388,7 +416,7 @@
 
     @Override
     public void write(int repetitionLevel, CharSequence value) {
       if (value instanceof Utf8) {
         Utf8 utf8 = (Utf8) value;
         column.writeBinary(
             repetitionLevel, Binary.fromReusedByteArray(utf8.getBytes(), 0, utf8.getByteLength()));
@@ -433,12 +461,14 @@
     private final int definitionLevel;
     private final ParquetValueWriter<T> writer;
     private final List<TripleWriter<?>> children;
+    private final Set<Integer> optionalFieldIds;
     private long nullValueCount = 0;
 
-    OptionWriter(int definitionLevel, ParquetValueWriter<T> writer) {
+    OptionWriter(int definitionLevel, ParquetValueWriter<T> writer, Set<Integer> optionalFieldIds) {
       this.definitionLevel = definitionLevel;
       this.writer = writer;
       this.children = writer.columns();
+      this.optionalFieldIds = optionalFieldIds;
     }
 
     @Override
@@ -474,17 +504,8 @@
           // we are not tracking field metrics for this type ourselves
           return Stream.empty();
         } else if (fieldMetricsFromWriter.size() == 1) {
-          FieldMetrics<?> metrics = fieldMetricsFromWriter.get(0);
-          return Stream.of(
-              new FieldMetrics<>(
-                  metrics.id(),
-                  metrics.valueCount() + nullValueCount,
-                  nullValueCount,
-                  metrics.nanValueCount(),
-                  metrics.lowerBound(),
-                  metrics.upperBound(),
-                  metrics.originalType(),
-                  metrics.avgValueSizeInBytes()));
+          // this option directly wraps an optional primitive, so its null count applies
+          return Stream.of(addNullValues(fieldMetricsFromWriter.get(0)));
         } else {
           throw new IllegalStateException(
               String.format(
@@ -494,9 +515,31 @@
         }
       }
 
-      // skipping updating null stats for non-primitive types since we don't use them today, to
-      // avoid unnecessary work
-      return writer.metrics();
+      // A null value here is also null for every descendant column, but those columns are written
+      // directly and never see it, so their writers cannot count it. Add it to the metrics of
+      // optional descendants; null counts are not tracked for required fields.
+      return writer
+          .metrics()
+          .map(
+              metrics ->
+                  optionalFieldIds.contains(metrics.id()) ? addNullValues(metrics) : metrics);
+    }
+
+    /** Adds the nulls counted by this writer to metrics produced by a descendant column. */
+    private FieldMetrics<?> addNullValues(FieldMetrics<?> metrics) {
+      if (nullValueCount == 0) {
+        return metrics;
+      }
+
+      return new FieldMetrics<>(
+          metrics.id(),
+          metrics.valueCount() + nullValueCount,
+          metrics.nullValueCount() + nullValueCount,
+          metrics.nanValueCount(),
+          metrics.lowerBound(),
+          metrics.upperBound(),
+          metrics.originalType(),
+          metrics.avgValueSizeInBytes());
     }
   }
 
@@ -724,7 +767,7 @@
 
     @Override
     protected Object get(PositionDelete<R> delete, int index) {
       switch (index) {
         case 0:
           return pathTransformFunc.apply(delete.path());
         case 1: