Parquet: Fix null counts for nested 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,60 @@ 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. A required
+    // leaf is included because a null struct makes even a required field null.
+    StructType struct =
+        StructType.of(
+            optional(2, "optDouble", DoubleType.get()),
+            required(3, "reqFloat", 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("reqFloat", 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);
+  }
+
+  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

@@ -388,7 +388,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()));
@@ -474,17 +474,7 @@
           // 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()));
+          return Stream.of(withNullValues(fieldMetricsFromWriter.get(0)));
         } else {
           throw new IllegalStateException(
               String.format(
@@ -494,9 +484,26 @@
         }
       }
 
-      // 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 their metrics.
+      return writer.metrics().map(this::withNullValues);
+    }
+
+    /** Adds the nulls counted by this writer to metrics produced by a descendant column. */
+    private FieldMetrics<?> withNullValues(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 +731,7 @@
 
     @Override
     protected Object get(PositionDelete<R> delete, int index) {
       switch (index) {
         case 0:
           return pathTransformFunc.apply(delete.path());
         case 1:

parquet/src/test/java/org/apache/iceberg/parquet/TestParquetValueWriters.java

@@ -19,13 +19,21 @@
 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 static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
 import java.nio.ByteBuffer;
+import java.util.Map;
+import java.util.function.Function;
+import java.util.stream.Collectors;
 import org.apache.iceberg.FieldMetrics;
 import org.apache.iceberg.Schema;
+import org.apache.iceberg.data.GenericRecord;
+import org.apache.iceberg.data.Record;
+import org.apache.iceberg.data.parquet.InternalWriter;
 import org.apache.iceberg.types.Types;
 import org.apache.parquet.column.ColumnDescriptor;
 import org.apache.parquet.column.ColumnWriteStore;
@@ -60,4 +68,117 @@ void geospatialValueSizeMetricsExcludeNulls() {
     assertThat(metrics.nullValueCount()).isEqualTo(1);
     assertThat(metrics.avgValueSizeInBytes()).isEqualTo(31);
   }
+
+  @Test
+  void nullStructCountsNullsForNestedFields() {
+    // a null struct is also null for the fields it contains, but those columns are written by the
+    // struct's writer and never see the value, so the struct must count the nulls for them
+    Types.StructType struct =
+        Types.StructType.of(
+            optional(2, "d", Types.DoubleType.get()), required(3, "f", Types.FloatType.get()));
+    Schema schema = new Schema(optional(1, "s", struct));
+
+    ParquetValueWriter<Record> writer = writerFor(schema);
+    Record inner = GenericRecord.create(struct);
+    inner.set(0, 2.0D);
+    inner.set(1, 1.0F);
+
+    writer.write(0, record(schema, inner));
+    writer.write(0, record(schema, null));
+    writer.write(0, record(schema, null));
+
+    Map<Integer, FieldMetrics<?>> metrics = metricsById(writer);
+    // both fields have one non-null value and two nulls from the null structs
+    assertThat(metrics.get(2).nullValueCount()).isEqualTo(2);
+    assertThat(metrics.get(2).valueCount()).isEqualTo(3);
+    assertThat(metrics.get(3).nullValueCount()).isEqualTo(2);
+    assertThat(metrics.get(3).valueCount()).isEqualTo(3);
+  }
+
+  @Test
+  void nullStructAddsToNullsCountedByNestedField() {
+    Types.StructType struct = Types.StructType.of(optional(2, "d", Types.DoubleType.get()));
+    Schema schema = new Schema(optional(1, "s", struct));
+
+    ParquetValueWriter<Record> writer = writerFor(schema);
+    Record present = GenericRecord.create(struct);
+    present.set(0, 2.0D);
+    Record nullField = GenericRecord.create(struct);
+    nullField.set(0, null);
+
+    writer.write(0, record(schema, present));
+    // the field is null while the struct is present, so the field's own writer counts it
+    writer.write(0, record(schema, nullField));
+    writer.write(0, record(schema, null));
+
+    Map<Integer, FieldMetrics<?>> metrics = metricsById(writer);
+    assertThat(metrics.get(2).nullValueCount()).isEqualTo(2);
+    assertThat(metrics.get(2).valueCount()).isEqualTo(3);
+  }
+
+  @Test
+  void nullStructCountsNullsForDeeplyNestedFields() {
+    Types.StructType inner = Types.StructType.of(optional(3, "d", Types.DoubleType.get()));
+    Types.StructType outer = Types.StructType.of(optional(2, "inner", inner));
+    Schema schema = new Schema(optional(1, "s", outer));
+
+    ParquetValueWriter<Record> writer = writerFor(schema);
+    Record innerRecord = GenericRecord.create(inner);
+    innerRecord.set(0, 2.0D);
+    Record withInner = GenericRecord.create(outer);
+    withInner.set(0, innerRecord);
+    Record withoutInner = GenericRecord.create(outer);
+    withoutInner.set(0, null);
+
+    writer.write(0, record(schema, withInner));
+    // a null at either level is a null for the leaf
+    writer.write(0, record(schema, withoutInner));
+    writer.write(0, record(schema, null));
+
+    Map<Integer, FieldMetrics<?>> metrics = metricsById(writer);
+    assertThat(metrics.get(3).nullValueCount()).isEqualTo(2);
+    assertThat(metrics.get(3).valueCount()).isEqualTo(3);
+  }
+
+  @Test
+  void nullStructCountsNullsForNestedGeospatialField() {
+    // geospatial writers also report metrics, so they are affected in the same way
+    Types.StructType struct = Types.StructType.of(optional(2, "g", Types.GeometryType.crs84()));
+    Schema schema = new Schema(optional(1, "s", struct));
+
+    ParquetValueWriter<Record> writer = writerFor(schema);
+    Record present = GenericRecord.create(struct);
+    present.set(0, ByteBuffer.allocate(21));
+
+    writer.write(0, record(schema, present));
+    writer.write(0, record(schema, null));
+
+    Map<Integer, FieldMetrics<?>> metrics = metricsById(writer);
+    assertThat(metrics.get(2).nullValueCount()).isEqualTo(1);
+    assertThat(metrics.get(2).valueCount()).isEqualTo(2);
+    // the size of the one non-null value is still reported
+    assertThat(metrics.get(2).avgValueSizeInBytes()).isEqualTo(21);
+  }
+
+  private static Record record(Schema schema, Record struct) {
+    Record record = GenericRecord.create(schema);
+    record.set(0, struct);
+    return record;
+  }
+
+  /** Returns a writer for the given schema, with a mocked column store. */
+  private static ParquetValueWriter<Record> writerFor(Schema schema) {
+    MessageType parquetSchema = ParquetSchemaUtil.convert(schema, "table");
+    ParquetValueWriter<Record> writer = InternalWriter.createWriter(schema, parquetSchema);
+
+    ColumnWriteStore columnStore = mock(ColumnWriteStore.class);
+    when(columnStore.getColumnWriter(any())).thenReturn(mock(ColumnWriter.class));
+    writer.setColumnStore(columnStore);
+
+    return writer;
+  }
+
+  private static Map<Integer, FieldMetrics<?>> metricsById(ParquetValueWriter<?> writer) {
+    return writer.metrics().collect(Collectors.toMap(FieldMetrics::id, Function.identity()));
+  }
 }