This is an automated email from the ASF dual-hosted git repository.
zeroshade pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/iceberg-go.git
The following commit(s) were added to refs/heads/main by this push:
new 0e6b8d2d4 feat(table): record NaN value counts for float columns on
write (#2420)
0e6b8d2d4 is described below
commit 0e6b8d2d4fa387d2aaeb1cf43a2cc72a0466512d
Author: Revanth <[email protected]>
AuthorDate: Wed Nov 4 11:21:22 2026 -0600
feat(table): record NaN value counts for float columns on write (#2420)
---
table/internal/parquet_files.go | 128 +++++++++++++++++
table/internal/parquet_files_nan_counts_test.go | 178 ++++++++++++++++++++++++
2 files changed, 306 insertions(+)
diff --git a/table/internal/parquet_files.go b/table/internal/parquet_files.go
index c1cc61c37..29bf5609a 100644
--- a/table/internal/parquet_files.go
+++ b/table/internal/parquet_files.go
@@ -517,6 +517,8 @@ type ParquetFileWriter struct {
geoCols []geoColumn
geoNormalizeCols []int
geoAccs map[int]*geoBoundsAccumulator
+ nanCols []nanColumn
+ nanCounts map[int]int64
arrowSchema *arrow.Schema
rowGroupBytes int64
}
@@ -671,6 +673,12 @@ func (p parquetFormat) NewFileWriter(ctx context.Context,
fs iceio.WriteFileIO,
geoAccs[gc.fieldID] = newGeoBoundsAccumulator(isGeog)
}
+ nanCols := collectNaNColumns(arrowSchema, colMapping, info.StatsCols)
+ nanCounts := make(map[int]int64, len(nanCols))
+ for _, nc := range nanCols {
+ nanCounts[nc.fieldID] = 0
+ }
+
return &ParquetFileWriter{
pqWriter: writer,
counter: counter,
@@ -685,6 +693,8 @@ func (p parquetFormat) NewFileWriter(ctx context.Context,
fs iceio.WriteFileIO,
geoCols: geoCols,
geoNormalizeCols: geoNormalizeCols,
geoAccs: geoAccs,
+ nanCols: nanCols,
+ nanCounts: nanCounts,
arrowSchema: arrowSchema,
rowGroupBytes: rowGroupTargetSizeBytes,
}, nil
@@ -856,6 +866,10 @@ func (w *ParquetFileWriter) Write(batch arrow.RecordBatch)
error {
return err
}
+ if err := w.countNaNs(writeBatch); err != nil {
+ return err
+ }
+
// Rotate before writing the next non-empty batch. Rotating immediately
after
// crossing the target would leave an empty trailing row group when the
file
// is closed without another write.
@@ -1306,6 +1320,118 @@ func accumulateWKBValues(acc *geoBoundsAccumulator, arr
arrow.Array, lo, hi int)
return nil
}
+// nanColumn locates a float or double leaf whose NaN values are counted during
+// the write: the position of its top-level column in the Arrow record batch,
+// the struct child index taken at each struct on the way down, and the Iceberg
+// field ID the count is recorded under.
+type nanColumn struct {
+ colIdx int
+ path []int
+ fieldID int
+}
+
+// collectNaNColumns finds the float and double leaves that get a NaN value
+// count. Parquet column statistics carry no NaN count, so it comes from the
+// written values. Like Java's MetricsUtil.createNanValueCounts, this skips
+// leaves inside a list or map and columns whose metrics mode is none.
+func collectNaNColumns(sc *arrow.Schema, colMapping map[string]int, statsCols
map[int]StatisticsCollector) []nanColumn {
+ var result []nanColumn
+ for i, f := range sc.Fields() {
+ result = appendNaNColumns(result, f.Type, f.Name, i, nil,
colMapping, statsCols)
+ }
+
+ return result
+}
+
+// appendNaNColumns walks dt through structs only and appends a nanColumn for
+// each float or double leaf. name is the Parquet column path used as the
+// colMapping key.
+func appendNaNColumns(result []nanColumn, dt arrow.DataType, name string,
colIdx int, path []int, colMapping map[string]int, statsCols
map[int]StatisticsCollector) []nanColumn {
+ switch t := dt.(type) {
+ case *arrow.Float32Type, *arrow.Float64Type:
+ fieldID, ok := colMapping[name]
+ if !ok {
+ return result
+ }
+ if sc, ok := statsCols[fieldID]; !ok || sc.Mode.Typ ==
MetricModeNone {
+ return result
+ }
+ result = append(result, nanColumn{colIdx: colIdx, path:
slices.Clone(path), fieldID: fieldID})
+ case *arrow.StructType:
+ for i, f := range t.Fields() {
+ result = appendNaNColumns(result, f.Type,
name+"."+f.Name, colIdx, append(path, i), colMapping, statsCols)
+ }
+ }
+
+ return result
+}
+
+// countNaNs adds the NaN values in this batch to the per-field counts. A value
+// under a null struct is not counted, since the row has no value there.
+func (w *ParquetFileWriter) countNaNs(batch arrow.RecordBatch) error {
+ for _, nc := range w.nanCols {
+ if nc.colIdx >= int(batch.NumCols()) {
+ return fmt.Errorf("float field %d: column %d missing
from batch with %d columns",
+ nc.fieldID, nc.colIdx, batch.NumCols())
+ }
+
+ col := batch.Column(nc.colIdx)
+ n, err := countNaNValues(col, nc.path, 0, col.Len())
+ if err != nil {
+ return fmt.Errorf("float field %d: %w", nc.fieldID, err)
+ }
+ w.nanCounts[nc.fieldID] += n
+ }
+
+ return nil
+}
+
+func countNaNValues(arr arrow.Array, path []int, lo, hi int) (int64, error) {
+ var n int64
+ switch arr := arr.(type) {
+ case *array.Struct:
+ if len(path) == 0 || path[0] >= arr.NumField() {
+ return 0, fmt.Errorf("struct array %s does not match
the float column path", arr.DataType())
+ }
+ child := arr.Field(path[0])
+ if arr.NullN() == 0 {
+ return countNaNValues(child, path[1:], lo, hi)
+ }
+ for i := lo; i < hi; i++ {
+ if arr.IsNull(i) {
+ continue
+ }
+ c, err := countNaNValues(child, path[1:], i, i+1)
+ if err != nil {
+ return 0, err
+ }
+ n += c
+ }
+ case *array.Float32:
+ if len(path) != 0 {
+ return 0, fmt.Errorf("%s array does not match the float
column path", arr.DataType())
+ }
+ for i := lo; i < hi; i++ {
+ if arr.IsValid(i) && math.IsNaN(float64(arr.Value(i))) {
+ n++
+ }
+ }
+ case *array.Float64:
+ if len(path) != 0 {
+ return 0, fmt.Errorf("%s array does not match the float
column path", arr.DataType())
+ }
+ for i := lo; i < hi; i++ {
+ if arr.IsValid(i) && math.IsNaN(arr.Value(i)) {
+ n++
+ }
+ }
+ default:
+ return 0, fmt.Errorf("unexpected %s array for a float column",
arr.DataType())
+ }
+
+ return n, nil
+}
+
// BytesWritten returns the number of bytes flushed to the output so far.
func (w *ParquetFileWriter) BytesWritten() int64 {
return w.counter.Count
@@ -1332,6 +1458,8 @@ func (w *ParquetFileWriter) Close() (_ iceberg.DataFile,
err error) {
return nil, err
}
+ maps.Copy(stats.NanValueCounts, w.nanCounts)
+
opts := DataFileOpts{
Schema: w.info.FileSchema,
Spec: w.info.Spec,
diff --git a/table/internal/parquet_files_nan_counts_test.go
b/table/internal/parquet_files_nan_counts_test.go
new file mode 100644
index 000000000..f089c02ab
--- /dev/null
+++ b/table/internal/parquet_files_nan_counts_test.go
@@ -0,0 +1,178 @@
+// 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 internal_test
+
+import (
+ "context"
+ "math"
+ "slices"
+ "testing"
+
+ "github.com/apache/arrow-go/v18/arrow"
+ "github.com/apache/arrow-go/v18/arrow/array"
+ "github.com/apache/arrow-go/v18/arrow/bitutil"
+ "github.com/apache/arrow-go/v18/arrow/memory"
+ "github.com/apache/iceberg-go"
+ iceio "github.com/apache/iceberg-go/io"
+ "github.com/apache/iceberg-go/table"
+ "github.com/apache/iceberg-go/table/internal"
+ "github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
+)
+
+// TestWriteDataFileRecordsNaNValueCounts writes float and double columns
+// through the full ParquetFileWriter path. Parquet statistics carry no NaN
+// count, so the writer counts NaN values itself; like Java, it records a
count,
+// zero included, for each float leaf outside a list or map whose metrics mode
+// isn't none.
+func TestWriteDataFileRecordsNaNValueCounts(t *testing.T) {
+ iceSchema := iceberg.NewSchema(0,
+ iceberg.NestedField{ID: 1, Name: "d", Type:
iceberg.PrimitiveTypes.Float64},
+ iceberg.NestedField{ID: 2, Name: "f", Type:
iceberg.PrimitiveTypes.Float32},
+ iceberg.NestedField{ID: 3, Name: "s", Type:
&iceberg.StructType{FieldList: []iceberg.NestedField{
+ {ID: 4, Name: "x", Type:
iceberg.PrimitiveTypes.Float64},
+ }}},
+ iceberg.NestedField{ID: 5, Name: "l", Type: &iceberg.ListType{
+ ElementID: 6, Element: iceberg.PrimitiveTypes.Float64,
+ }},
+ iceberg.NestedField{ID: 7, Name: "n", Type:
iceberg.PrimitiveTypes.Float64},
+ iceberg.NestedField{ID: 8, Name: "z", Type:
iceberg.PrimitiveTypes.Float64},
+ iceberg.NestedField{ID: 9, Name: "i", Type:
iceberg.PrimitiveTypes.Int32},
+ )
+ arrowSchema, err := table.SchemaToArrowSchema(iceSchema, nil, true,
false)
+ require.NoError(t, err)
+
+ nan := math.NaN()
+ type row struct {
+ d, x, n, z float64
+ f float32
+ dNull bool
+ sNull bool
+ list []float64
+ }
+ batch := func(rows []row) arrow.RecordBatch {
+ bldr := array.NewRecordBuilder(memory.DefaultAllocator,
arrowSchema)
+ defer bldr.Release()
+
+ for _, r := range rows {
+ if r.dNull {
+ bldr.Field(0).AppendNull()
+ } else {
+
bldr.Field(0).(*array.Float64Builder).Append(r.d)
+ }
+ bldr.Field(1).(*array.Float32Builder).Append(r.f)
+ sb := bldr.Field(2).(*array.StructBuilder)
+ sb.Append(true)
+ sb.FieldBuilder(0).(*array.Float64Builder).Append(r.x)
+ lb := bldr.Field(3).(*array.ListBuilder)
+ lb.Append(true)
+ for _, v := range r.list {
+
lb.ValueBuilder().(*array.Float64Builder).Append(v)
+ }
+ bldr.Field(4).(*array.Float64Builder).Append(r.n)
+ bldr.Field(5).(*array.Float64Builder).Append(r.z)
+ bldr.Field(6).(*array.Int32Builder).Append(1)
+ }
+
+ rec := bldr.NewRecordBatch()
+ defer rec.Release()
+
+ // StructBuilder.Append(false) appends a null child too, so
null struct
+ // rows are set afterwards over the built child: a NaN then
sits under a
+ // null struct row without belonging to it.
+ valid := make([]byte, bitutil.BytesForBits(int64(len(rows))))
+ nulls := 0
+ for i, r := range rows {
+ if r.sNull {
+ nulls++
+ } else {
+ bitutil.SetBit(valid, i)
+ }
+ }
+ structData := rec.Column(2).Data()
+ withNulls := array.NewData(structData.DataType(),
structData.Len(),
+ []*memory.Buffer{memory.NewBufferBytes(valid)},
structData.Children(), nulls, 0)
+ defer withNulls.Release()
+ structArr := array.MakeFromData(withNulls)
+ defer structArr.Release()
+
+ cols := slices.Clone(rec.Columns())
+ cols[2] = structArr
+
+ return array.NewRecordBatch(arrowSchema, cols, rec.NumRows())
+ }
+
+ first := batch([]row{
+ {d: 1, f: float32(nan), x: nan, list: []float64{nan}, n: nan,
z: 1},
+ {d: nan, f: 2, x: nan, sNull: true, list: []float64{nan, 1}, n:
nan, z: 2},
+ })
+ defer first.Release()
+ second := batch([]row{
+ {dNull: true, f: 3, x: 1, list: []float64{nan}, n: nan, z: 3},
+ {d: nan, f: 4, x: nan, n: nan, z: 4},
+ })
+ defer second.Release()
+
+ // The NaN under the null struct row must really be in the child array,
+ // or the s.x count below would pass without the null check.
+ structCol := first.Column(2).(*array.Struct)
+ require.True(t, structCol.IsNull(1))
+ require.True(t,
math.IsNaN(structCol.Field(0).(*array.Float64).Value(1)))
+
+ full := internal.MetricsMode{Typ: internal.MetricModeFull}
+ counts := internal.MetricsMode{Typ: internal.MetricModeCounts}
+ statsCols := map[int]internal.StatisticsCollector{
+ 1: {FieldID: 1, IcebergTyp: iceberg.PrimitiveTypes.Float64,
ColName: "d", Mode: full},
+ 2: {FieldID: 2, IcebergTyp: iceberg.PrimitiveTypes.Float32,
ColName: "f", Mode: counts},
+ 4: {FieldID: 4, IcebergTyp: iceberg.PrimitiveTypes.Float64,
ColName: "s.x", Mode: counts},
+ 6: {FieldID: 6, IcebergTyp: iceberg.PrimitiveTypes.Float64,
ColName: "l.element", Mode: counts},
+ 7: {FieldID: 7, IcebergTyp: iceberg.PrimitiveTypes.Float64,
ColName: "n", Mode: internal.MetricsMode{Typ: internal.MetricModeNone}},
+ 8: {FieldID: 8, IcebergTyp: iceberg.PrimitiveTypes.Float64,
ColName: "z", Mode: full},
+ 9: {FieldID: 9, IcebergTyp: iceberg.PrimitiveTypes.Int32,
ColName: "i", Mode: full},
+ }
+
+ fm := internal.GetFileFormat(iceberg.ParquetFile)
+ df, err := fm.WriteDataFile(context.Background(), iceio.NewMemFS(),
nil, internal.WriteFileInfo{
+ FileSchema: iceSchema,
+ Spec: *iceberg.UnpartitionedSpec,
+ FileName: "nan.parquet",
+ StatsCols: statsCols,
+ WriteProps: fm.GetWriteProperties(iceberg.Properties{}),
+ Content: iceberg.EntryContentData,
+ }, []arrow.RecordBatch{first, second})
+ require.NoError(t, err)
+
+ // d: NaN in both batches, and the NULL row is not a NaN. f: counts mode
+ // still records NaN counts. s.x: three NaN child values, one of them
+ // under the null struct row. z: no NaN still records a zero, which is
+ // what lets a metrics evaluator rule NaN out. l.element is in a list
and
+ // n has metrics mode none, so neither is recorded, nor is the int
column.
+ assert.Equal(t, map[int]int64{1: 2, 2: 1, 4: 2, 8: 0},
df.NaNValueCounts())
+
+ // Parquet already leaves NaN out of float min/max, so the bounds hold
+ // only the finite values.
+ lower, upper := df.LowerBoundValues(), df.UpperBoundValues()
+ d := func(b []byte) float64 {
+ lit, err :=
iceberg.LiteralFromBytes(iceberg.PrimitiveTypes.Float64, b)
+ require.NoError(t, err)
+
+ return float64(lit.(iceberg.Float64Literal))
+ }
+ assert.Equal(t, 1.0, d(lower[1]))
+ assert.Equal(t, 1.0, d(upper[1]))
+}
