This is an automated email from the ASF dual-hosted git repository.
szehon-ho pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/iceberg.git
The following commit(s) were added to refs/heads/main by this push:
new fe30749129 Parquet: Track avg_value_size_in_bytes for Geo types
(#17333)
fe30749129 is described below
commit fe30749129b5a629cdf5b02d43d6bad759b1d193
Author: Xin Huang <[email protected]>
AuthorDate: Fri Jul 31 10:40:04 2026 -0700
Parquet: Track avg_value_size_in_bytes for Geo types (#17333)
---
.../main/java/org/apache/iceberg/FieldMetrics.java | 19 ++++++
.../java/org/apache/iceberg/FieldStatsStruct.java | 2 +-
.../org/apache/iceberg/ValueSizeFieldMetrics.java | 51 ++++++++++++++++
.../org/apache/iceberg/TestFieldStatsStruct.java | 4 +-
.../apache/iceberg/TestValueSizeFieldMetrics.java | 59 ++++++++++++++++++
.../iceberg/data/parquet/BaseParquetWriter.java | 4 +-
.../org/apache/iceberg/parquet/ParquetMetrics.java | 9 ++-
.../iceberg/parquet/ParquetValueWriters.java | 29 ++++++++-
.../iceberg/parquet/TestParquetValueWriters.java | 63 +++++++++++++++++++
.../iceberg/spark/data/SparkParquetWriters.java | 37 +++++++++--
.../iceberg/spark/data/TestSparkParquetWriter.java | 71 ++++++++++++++++++++++
11 files changed, 334 insertions(+), 14 deletions(-)
diff --git a/core/src/main/java/org/apache/iceberg/FieldMetrics.java
b/core/src/main/java/org/apache/iceberg/FieldMetrics.java
index 422339568a..d037e71047 100644
--- a/core/src/main/java/org/apache/iceberg/FieldMetrics.java
+++ b/core/src/main/java/org/apache/iceberg/FieldMetrics.java
@@ -29,6 +29,7 @@ public class FieldMetrics<T> {
private final T lowerBound;
private final T upperBound;
private final Type originalType;
+ private final Integer avgValueSizeInBytes;
public FieldMetrics(int id, long valueCount, long nullValueCount) {
this(id, valueCount, nullValueCount, -1L, null, null, null);
@@ -65,6 +66,18 @@ public class FieldMetrics<T> {
T lowerBound,
T upperBound,
Type originalType) {
+ this(id, valueCount, nullValueCount, nanValueCount, lowerBound,
upperBound, originalType, null);
+ }
+
+ public FieldMetrics(
+ int id,
+ long valueCount,
+ long nullValueCount,
+ long nanValueCount,
+ T lowerBound,
+ T upperBound,
+ Type originalType,
+ Integer avgValueSizeInBytes) {
this.id = id;
this.valueCount = valueCount;
this.nullValueCount = nullValueCount;
@@ -72,6 +85,7 @@ public class FieldMetrics<T> {
this.lowerBound = lowerBound;
this.upperBound = upperBound;
this.originalType = originalType;
+ this.avgValueSizeInBytes = avgValueSizeInBytes;
}
/** Returns the id of the field that the metrics within this class are
associated with. */
@@ -112,6 +126,11 @@ public class FieldMetrics<T> {
return originalType;
}
+ /** Returns the average size in bytes over non-null values, or null if it is
not known. */
+ public Integer avgValueSizeInBytes() {
+ return avgValueSizeInBytes;
+ }
+
/** Returns if the metrics has bounds (i.e. there is at least non-null value
for this field) */
public boolean hasBounds() {
return upperBound != null;
diff --git a/core/src/main/java/org/apache/iceberg/FieldStatsStruct.java
b/core/src/main/java/org/apache/iceberg/FieldStatsStruct.java
index 44c2e2009c..93de05a2e6 100644
--- a/core/src/main/java/org/apache/iceberg/FieldStatsStruct.java
+++ b/core/src/main/java/org/apache/iceberg/FieldStatsStruct.java
@@ -98,7 +98,7 @@ class FieldStatsStruct<T> implements FieldStats<T>,
StructLike, Serializable {
this.valueCount = fieldMetrics.valueCount();
this.nullValueCount = fieldMetrics.nullValueCount() < 0 ? null :
fieldMetrics.nullValueCount();
this.nanValueCount = fieldMetrics.nanValueCount() < 0 ? null :
fieldMetrics.nanValueCount();
- this.avgValueSize = null;
+ this.avgValueSize = fieldMetrics.avgValueSizeInBytes();
}
private boolean isBinary() {
diff --git a/core/src/main/java/org/apache/iceberg/ValueSizeFieldMetrics.java
b/core/src/main/java/org/apache/iceberg/ValueSizeFieldMetrics.java
new file mode 100644
index 0000000000..dbd2ae7f50
--- /dev/null
+++ b/core/src/main/java/org/apache/iceberg/ValueSizeFieldMetrics.java
@@ -0,0 +1,51 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg;
+
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/** Field-level metrics for tracking the average size of variable-length
values. */
+public class ValueSizeFieldMetrics extends FieldMetrics<Object> {
+
+ private ValueSizeFieldMetrics(int id, long valueCount, Integer
avgValueSizeInBytes) {
+ super(id, valueCount, 0L, -1L, null, null, null, avgValueSizeInBytes);
+ }
+
+ public static class Builder {
+ private final int id;
+ private long valueCount = 0;
+ private long totalValueSizeInBytes = 0;
+
+ public Builder(int id) {
+ this.id = id;
+ }
+
+ public void addValueSize(int sizeInBytes) {
+ Preconditions.checkArgument(sizeInBytes >= 0, "Invalid value size: %s",
sizeInBytes);
+ this.valueCount += 1;
+ this.totalValueSizeInBytes += sizeInBytes;
+ }
+
+ public ValueSizeFieldMetrics build() {
+ Integer avgValueSizeInBytes =
+ valueCount > 0 ? Math.toIntExact(totalValueSizeInBytes / valueCount)
: null;
+ return new ValueSizeFieldMetrics(id, valueCount, avgValueSizeInBytes);
+ }
+ }
+}
diff --git a/core/src/test/java/org/apache/iceberg/TestFieldStatsStruct.java
b/core/src/test/java/org/apache/iceberg/TestFieldStatsStruct.java
index 936fc43c77..847d05ba41 100644
--- a/core/src/test/java/org/apache/iceberg/TestFieldStatsStruct.java
+++ b/core/src/test/java/org/apache/iceberg/TestFieldStatsStruct.java
@@ -150,14 +150,14 @@ public class TestFieldStatsStruct {
public void fromFieldMetricsString() {
FieldStatsStruct<String> stats = new FieldStatsStruct<>(STRING_STATS);
- stats.fromFieldMetrics(new FieldMetrics<>(100, 28, 2, "a", "z"));
+ stats.fromFieldMetrics(new FieldMetrics<>(100, 28, 2, -1, "a", "z", null,
4));
assertThat(stats.lowerBound()).isEqualTo("a");
assertThat(stats.upperBound()).isEqualTo("z");
assertThat(stats.tightBounds()).isFalse();
assertThat(stats.valueCount()).isEqualTo(28L);
assertThat(stats.nullValueCount()).isEqualTo(2L);
- assertThat(stats.avgValueSizeInBytes()).isNull(); // unknown
+ assertThat(stats.avgValueSizeInBytes()).isEqualTo(4);
}
@Test
diff --git
a/core/src/test/java/org/apache/iceberg/TestValueSizeFieldMetrics.java
b/core/src/test/java/org/apache/iceberg/TestValueSizeFieldMetrics.java
new file mode 100644
index 0000000000..67d4a081ab
--- /dev/null
+++ b/core/src/test/java/org/apache/iceberg/TestValueSizeFieldMetrics.java
@@ -0,0 +1,59 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+import org.junit.jupiter.api.Test;
+
+class TestValueSizeFieldMetrics {
+
+ @Test
+ void averageValueSize() {
+ ValueSizeFieldMetrics.Builder builder = new
ValueSizeFieldMetrics.Builder(2);
+ builder.addValueSize(21);
+ builder.addValueSize(42);
+
+ FieldMetrics<?> metrics = builder.build();
+
+ assertThat(metrics.id()).isEqualTo(2);
+ assertThat(metrics.valueCount()).isEqualTo(2);
+ assertThat(metrics.nullValueCount()).isZero();
+ assertThat(metrics.nanValueCount()).isEqualTo(-1);
+ assertThat(metrics.avgValueSizeInBytes()).isEqualTo(31);
+ }
+
+ @Test
+ void noValues() {
+ FieldMetrics<?> metrics = new ValueSizeFieldMetrics.Builder(2).build();
+
+ assertThat(metrics.valueCount()).isZero();
+ assertThat(metrics.avgValueSizeInBytes()).isNull();
+ }
+
+ @Test
+ void rejectsNegativeValueSize() {
+ ValueSizeFieldMetrics.Builder builder = new
ValueSizeFieldMetrics.Builder(2);
+
+ assertThatThrownBy(() -> builder.addValueSize(-1))
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessage("Invalid value size: -1");
+ }
+}
diff --git
a/parquet/src/main/java/org/apache/iceberg/data/parquet/BaseParquetWriter.java
b/parquet/src/main/java/org/apache/iceberg/data/parquet/BaseParquetWriter.java
index dcc93f939d..78b924db5b 100644
---
a/parquet/src/main/java/org/apache/iceberg/data/parquet/BaseParquetWriter.java
+++
b/parquet/src/main/java/org/apache/iceberg/data/parquet/BaseParquetWriter.java
@@ -269,14 +269,14 @@ abstract class BaseParquetWriter<T> {
public Optional<ParquetValueWriter<?>> visit(
LogicalTypeAnnotation.GeometryLogicalTypeAnnotation geometryType) {
// geometry values are pure WKB stored in a BINARY column
- return Optional.of(ParquetValueWriters.byteBuffers(desc));
+ return Optional.of(ParquetValueWriters.geospatial(desc));
}
@Override
public Optional<ParquetValueWriter<?>> visit(
LogicalTypeAnnotation.GeographyLogicalTypeAnnotation geographyType) {
// geography values are pure WKB stored in a BINARY column
- return Optional.of(ParquetValueWriters.byteBuffers(desc));
+ return Optional.of(ParquetValueWriters.geospatial(desc));
}
@Override
diff --git
a/parquet/src/main/java/org/apache/iceberg/parquet/ParquetMetrics.java
b/parquet/src/main/java/org/apache/iceberg/parquet/ParquetMetrics.java
index 67d6a0bc73..0412ebc695 100644
--- a/parquet/src/main/java/org/apache/iceberg/parquet/ParquetMetrics.java
+++ b/parquet/src/main/java/org/apache/iceberg/parquet/ParquetMetrics.java
@@ -264,7 +264,11 @@ class ParquetMetrics {
fieldMetrics.id(),
fieldMetrics.valueCount(),
fieldMetrics.nullValueCount(),
- fieldMetrics.nanValueCount());
+ fieldMetrics.nanValueCount(),
+ null,
+ null,
+ null,
+ fieldMetrics.avgValueSizeInBytes());
} else {
T lowerBound = truncateLowerBound(icebergType,
fieldMetrics.lowerBound(), truncateLength);
T upperBound = truncateUpperBound(icebergType,
fieldMetrics.upperBound(), truncateLength);
@@ -275,7 +279,8 @@ class ParquetMetrics {
fieldMetrics.nanValueCount(),
lowerBound,
upperBound,
- icebergType);
+ icebergType,
+ fieldMetrics.avgValueSizeInBytes());
}
}
diff --git
a/parquet/src/main/java/org/apache/iceberg/parquet/ParquetValueWriters.java
b/parquet/src/main/java/org/apache/iceberg/parquet/ParquetValueWriters.java
index 5554c86b46..298ffa1215 100644
--- a/parquet/src/main/java/org/apache/iceberg/parquet/ParquetValueWriters.java
+++ b/parquet/src/main/java/org/apache/iceberg/parquet/ParquetValueWriters.java
@@ -37,6 +37,7 @@ import org.apache.iceberg.DoubleFieldMetrics;
import org.apache.iceberg.FieldMetrics;
import org.apache.iceberg.FloatFieldMetrics;
import org.apache.iceberg.StructLike;
+import org.apache.iceberg.ValueSizeFieldMetrics;
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;
@@ -120,6 +121,10 @@ public class ParquetValueWriters {
return new BytesWriter(desc);
}
+ public static PrimitiveWriter<ByteBuffer> geospatial(ColumnDescriptor desc) {
+ return new GeospatialWriter(desc);
+ }
+
public static PrimitiveWriter<ByteBuffer> fixedBuffers(ColumnDescriptor
desc) {
return new FixedBufferWriter(desc);
}
@@ -336,6 +341,27 @@ public class ParquetValueWriters {
}
}
+ private static class GeospatialWriter extends PrimitiveWriter<ByteBuffer> {
+ private final ValueSizeFieldMetrics.Builder metricsBuilder;
+
+ private GeospatialWriter(ColumnDescriptor desc) {
+ super(desc);
+ this.metricsBuilder =
+ new
ValueSizeFieldMetrics.Builder(desc.getPrimitiveType().getId().intValue());
+ }
+
+ @Override
+ public void write(int repetitionLevel, ByteBuffer buffer) {
+ metricsBuilder.addValueSize(buffer.remaining());
+ column.writeBinary(repetitionLevel, Binary.fromReusedByteBuffer(buffer));
+ }
+
+ @Override
+ public Stream<FieldMetrics<?>> metrics() {
+ return Stream.of(metricsBuilder.build());
+ }
+ }
+
private static class FixedBufferWriter extends PrimitiveWriter<ByteBuffer> {
private final int length;
@@ -457,7 +483,8 @@ public class ParquetValueWriters {
metrics.nanValueCount(),
metrics.lowerBound(),
metrics.upperBound(),
- metrics.originalType()));
+ metrics.originalType(),
+ metrics.avgValueSizeInBytes()));
} else {
throw new IllegalStateException(
String.format(
diff --git
a/parquet/src/test/java/org/apache/iceberg/parquet/TestParquetValueWriters.java
b/parquet/src/test/java/org/apache/iceberg/parquet/TestParquetValueWriters.java
new file mode 100644
index 0000000000..33c4044a0a
--- /dev/null
+++
b/parquet/src/test/java/org/apache/iceberg/parquet/TestParquetValueWriters.java
@@ -0,0 +1,63 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.parquet;
+
+import static org.apache.iceberg.types.Types.NestedField.optional;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import java.nio.ByteBuffer;
+import org.apache.iceberg.FieldMetrics;
+import org.apache.iceberg.Schema;
+import org.apache.iceberg.types.Types;
+import org.apache.parquet.column.ColumnDescriptor;
+import org.apache.parquet.column.ColumnWriteStore;
+import org.apache.parquet.column.ColumnWriter;
+import org.apache.parquet.schema.MessageType;
+import org.apache.parquet.schema.Type;
+import org.junit.jupiter.api.Test;
+
+class TestParquetValueWriters {
+
+ @Test
+ void geospatialValueSizeMetricsExcludeNulls() {
+ Schema schema = new Schema(optional(2, "geom",
Types.GeometryType.crs84()));
+ MessageType parquetSchema = ParquetSchemaUtil.convert(schema, "table");
+ Type parquetType = parquetSchema.getType("geom");
+ ColumnDescriptor desc = parquetSchema.getColumnDescription(new String[]
{"geom"});
+ ParquetValueWriter<ByteBuffer> writer =
+ ParquetValueWriters.option(
+ parquetType,
+ parquetSchema.getMaxDefinitionLevel(new String[] {"geom"}),
+ ParquetValueWriters.geospatial(desc));
+
+ ColumnWriteStore columnStore = mock(ColumnWriteStore.class);
+
when(columnStore.getColumnWriter(desc)).thenReturn(mock(ColumnWriter.class));
+ writer.setColumnStore(columnStore);
+ writer.write(0, ByteBuffer.allocate(21));
+ writer.write(0, ByteBuffer.allocate(42));
+ writer.write(0, null);
+
+ FieldMetrics<?> metrics = writer.metrics().findFirst().orElseThrow();
+ assertThat(metrics.valueCount()).isEqualTo(3);
+ assertThat(metrics.nullValueCount()).isEqualTo(1);
+ assertThat(metrics.avgValueSizeInBytes()).isEqualTo(31);
+ }
+}
diff --git
a/spark/v4.1/spark/src/main/java/org/apache/iceberg/spark/data/SparkParquetWriters.java
b/spark/v4.1/spark/src/main/java/org/apache/iceberg/spark/data/SparkParquetWriters.java
index 24fabf5b63..a31ec560df 100644
---
a/spark/v4.1/spark/src/main/java/org/apache/iceberg/spark/data/SparkParquetWriters.java
+++
b/spark/v4.1/spark/src/main/java/org/apache/iceberg/spark/data/SparkParquetWriters.java
@@ -31,6 +31,7 @@ import java.util.stream.IntStream;
import java.util.stream.Stream;
import org.apache.iceberg.FieldMetrics;
import org.apache.iceberg.Schema;
+import org.apache.iceberg.ValueSizeFieldMetrics;
import org.apache.iceberg.parquet.ParquetValueReaders.ReusableEntry;
import org.apache.iceberg.parquet.ParquetValueWriter;
import org.apache.iceberg.parquet.ParquetValueWriters;
@@ -488,29 +489,53 @@ public class SparkParquetWriters {
}
}
+ private abstract static class GeospatialWriter<T> extends PrimitiveWriter<T>
{
+ private final ValueSizeFieldMetrics.Builder metricsBuilder;
+
+ private GeospatialWriter(ColumnDescriptor desc) {
+ super(desc);
+ this.metricsBuilder =
+ new
ValueSizeFieldMetrics.Builder(desc.getPrimitiveType().getId().intValue());
+ }
+
+ @Override
+ public void write(int repetitionLevel, T value) {
+ byte[] wkb = toWkb(value);
+ metricsBuilder.addValueSize(wkb.length);
+ column.writeBinary(repetitionLevel, Binary.fromReusedByteArray(wkb));
+ }
+
+ @Override
+ public Stream<FieldMetrics<?>> metrics() {
+ return Stream.of(metricsBuilder.build());
+ }
+
+ protected abstract byte[] toWkb(T value);
+ }
+
/** Writes a Spark {@link GeometryVal} as its WKB bytes into a BINARY
column. */
- private static class GeometryWriter extends PrimitiveWriter<GeometryVal> {
+ private static class GeometryWriter extends GeospatialWriter<GeometryVal> {
private GeometryWriter(ColumnDescriptor desc) {
super(desc);
}
@Override
- public void write(int repetitionLevel, GeometryVal value) {
+ protected byte[] toWkb(GeometryVal value) {
// Spark stores geometry as [SRID | WKB]; Iceberg stores pure WKB, so
strip the SRID header.
- column.writeBinary(repetitionLevel,
Binary.fromReusedByteArray(STUtils.stAsBinary(value)));
+ return STUtils.stAsBinary(value);
}
}
/** Writes a Spark {@link GeographyVal} as its WKB bytes into a BINARY
column. */
- private static class GeographyWriter extends PrimitiveWriter<GeographyVal> {
+ private static class GeographyWriter extends GeospatialWriter<GeographyVal> {
private GeographyWriter(ColumnDescriptor desc) {
super(desc);
}
@Override
- public void write(int repetitionLevel, GeographyVal value) {
+ protected byte[] toWkb(GeographyVal value) {
// Spark stores geography as [SRID | WKB]; Iceberg stores pure WKB, so
strip the SRID header.
- column.writeBinary(repetitionLevel,
Binary.fromReusedByteArray(STUtils.stAsBinary(value)));
+ return STUtils.stAsBinary(value);
}
}
diff --git
a/spark/v4.1/spark/src/test/java/org/apache/iceberg/spark/data/TestSparkParquetWriter.java
b/spark/v4.1/spark/src/test/java/org/apache/iceberg/spark/data/TestSparkParquetWriter.java
index 86a61b90d8..6a5bbe8d39 100644
---
a/spark/v4.1/spark/src/test/java/org/apache/iceberg/spark/data/TestSparkParquetWriter.java
+++
b/spark/v4.1/spark/src/test/java/org/apache/iceberg/spark/data/TestSparkParquetWriter.java
@@ -24,6 +24,9 @@ import static
org.apache.iceberg.TableProperties.PARQUET_BLOOM_FILTER_COLUMN_NDV
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.io.File;
import java.io.IOException;
@@ -31,17 +34,24 @@ import java.lang.reflect.Field;
import java.nio.file.Path;
import java.util.Iterator;
import java.util.List;
+import java.util.Map;
import java.util.OptionalLong;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+import org.apache.iceberg.FieldMetrics;
import org.apache.iceberg.Files;
import org.apache.iceberg.Schema;
import org.apache.iceberg.io.CloseableIterable;
import org.apache.iceberg.io.FileAppender;
import org.apache.iceberg.parquet.Parquet;
import org.apache.iceberg.parquet.ParquetSchemaUtil;
+import org.apache.iceberg.parquet.ParquetValueWriter;
import org.apache.iceberg.relocated.com.google.common.collect.Lists;
import org.apache.iceberg.spark.SparkSchemaUtil;
import org.apache.iceberg.types.Types;
import org.apache.parquet.column.ColumnDescriptor;
+import org.apache.parquet.column.ColumnWriteStore;
+import org.apache.parquet.column.ColumnWriter;
import org.apache.parquet.column.ParquetProperties;
import org.apache.parquet.schema.MessageType;
import org.apache.spark.sql.catalyst.InternalRow;
@@ -182,6 +192,67 @@ public class TestSparkParquetWriter {
}
}
+ @Test
+ public void testGeospatialAvgValueSizeMetrics() throws IOException {
+ Schema geoSchema =
+ new Schema(
+ required(1, "id", Types.LongType.get()),
+ optional(2, "geom", Types.GeometryType.crs84()),
+ optional(3, "geog", Types.GeographyType.crs84()));
+
+ // WKB payloads of 3 and 5 bytes for geometry (avg 4), 7 bytes for
geography.
+ byte[] geomWkbSmall = new byte[] {0x01, 0x02, 0x03};
+ byte[] geomWkbLarge = new byte[] {0x01, 0x02, 0x03, 0x04, 0x05};
+ byte[] geogWkb = new byte[] {0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a};
+
+ InternalRow first = new GenericInternalRow(3);
+ first.update(0, 1L);
+ // Spark's GeometryVal/GeographyVal wrap [SRID | WKB]; build them from the
pure WKB.
+ first.update(1, STUtils.stGeomFromWKB(geomWkbSmall));
+ first.update(2, STUtils.stGeogFromWKB(geogWkb));
+ InternalRow second = new GenericInternalRow(3);
+ second.update(0, 2L);
+ second.update(1, STUtils.stGeomFromWKB(geomWkbLarge));
+ // geography left null on the second row, so it must not affect the
average.
+
+ File testFile = File.createTempFile("junit", null, temp.toFile());
+ assertThat(testFile.delete()).as("Delete should succeed").isTrue();
+
+ MessageType parquetSchema = ParquetSchemaUtil.convert(geoSchema, "table");
+ ParquetValueWriter<InternalRow> writer =
+ SparkParquetWriters.buildWriter(SparkSchemaUtil.convert(geoSchema),
parquetSchema);
+
+ ColumnWriteStore columnStore = mock(ColumnWriteStore.class);
+
when(columnStore.getColumnWriter(any())).thenReturn(mock(ColumnWriter.class));
+ writer.setColumnStore(columnStore);
+ writer.write(0, first);
+ writer.write(0, second);
+
+ Map<Integer, FieldMetrics<?>> metricsById =
+ writer.metrics().collect(Collectors.toMap(FieldMetrics::id,
Function.identity()));
+
+ int geomId = fieldId(parquetSchema, "geom");
+ int geogId = fieldId(parquetSchema, "geog");
+
+ FieldMetrics<?> geomMetrics = metricsById.get(geomId);
+ assertThat(geomMetrics.valueCount()).isEqualTo(2);
+ assertThat(geomMetrics.nullValueCount()).isZero();
+ assertThat(geomMetrics.avgValueSizeInBytes()).isEqualTo(4);
+
+ FieldMetrics<?> geogMetrics = metricsById.get(geogId);
+ assertThat(geogMetrics.valueCount()).isEqualTo(2);
+ assertThat(geogMetrics.nullValueCount()).isEqualTo(1);
+ assertThat(geogMetrics.avgValueSizeInBytes()).isEqualTo(7);
+ }
+
+ private static int fieldId(MessageType parquetSchema, String column) {
+ return parquetSchema
+ .getColumnDescription(new String[] {column})
+ .getPrimitiveType()
+ .getId()
+ .intValue();
+ }
+
@Test
public void testFpp() throws IOException, NoSuchFieldException,
IllegalAccessException {
File testFile = File.createTempFile("junit", null, temp.toFile());