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/arrow-go.git
The following commit(s) were added to refs/heads/main by this push:
new ea989eb2 perf(parquet/pqarrow): build nullable fixed-size list
children directly (#1195)
ea989eb2 is described below
commit ea989eb2eba113b70a342815d0022c8bba649e4f
Author: Minh Vu <[email protected]>
AuthorDate: Wed Aug 19 00:25:47 2026 +0200
perf(parquet/pqarrow): build nullable fixed-size list children directly
(#1195)
## Summary
- keep the existing zero-null fixed-size-list path unchanged
- validate nullable parent spans once
- use the existing slice concatenation for a small number of validity
runs
- build one child array with a typed Take for highly fragmented validity
- avoid one temporary Arrow array per validity run
## Why
Parquet does not write child values for null fixed-size-list parents.
Arrow still needs `list_size` child slots for every parent, so the
reader has to add null child values.
The old nullable path creates one temporary array for each validity run
and concatenates them. Alternating validity can therefore create one
temporary array per parent row.
## Benchmark
Apple M1 Pro, 65,536 `FixedSizeList<int32, 4>` parents:
| validity | old | new | old allocs | new allocs |
| --- | ---: | ---: | ---: | ---: |
| 10% nulls | ~4.4 ms | ~1.3 ms | 85,220 | 46 |
| alternating nulls | ~23–30 ms | ~1.1 ms | 458,770 | 46 |
| one clustered null run | ~0.3 ms | ~0.3 ms | 28 | 27 |
The low-run path is kept for the clustered case, where creating a single
pair of slices is cheaper than building an index array.
## Tests
- `go vet ./parquet/pqarrow`
- `go test -race ./parquet/pqarrow -run
'TestBuildFixedSizeListArray|TestParquetArrowIO/TestFixedSizeList'`
- `PARQUET_TEST_DATA=parquet-testing/data
ARROW_TEST_DATA=arrow-testing/data go test ./...`
---
parquet/pqarrow/column_readers.go | 109 ++++++++++++++++++++++--------
parquet/pqarrow/column_readers_test.go | 117 +++++++++++++++++++++++++++++++++
2 files changed, 198 insertions(+), 28 deletions(-)
diff --git a/parquet/pqarrow/column_readers.go
b/parquet/pqarrow/column_readers.go
index 1c175391..3456fccb 100644
--- a/parquet/pqarrow/column_readers.go
+++ b/parquet/pqarrow/column_readers.go
@@ -573,44 +573,97 @@ func (lr *listReader) buildFixedSizeListArray(length int,
offsets []int32, valid
return arrow.NewChunked(lr.field.Type, []arrow.Array{out}), nil
}
- // Each validity run becomes one piece. Alternating valid and null
parents
- // create O(length) temporary arrays before concatenation.
- pieces := make([]arrow.Array, 0, length)
- defer func() { releaseArrays(pieces) }()
-
- for i := 0; i < length; {
- valid := !lr.field.Nullable ||
bitutil.BitIsSet(validityBuffer.Bytes(), i)
- end := i + 1
- for end < length {
- nextValid := !lr.field.Nullable ||
bitutil.BitIsSet(validityBuffer.Bytes(), end)
- if nextValid != valid {
- break
+ validity := validityBuffer.Bytes()
+ runCount := 0
+ previousValid := false
+ for idx := 0; idx < length; idx++ {
+ valid := !lr.field.Nullable || bitutil.BitIsSet(validity, idx)
+ if valid {
+ if size := offsets[idx+1] - offsets[idx]; size !=
int32(listSize) {
+ return nil, fmt.Errorf("expected all lists to
be of size=%d, but index %d had size=%d", listSize, idx, size)
+ }
+ } else {
+ if size := offsets[idx+1] - offsets[idx]; size != 0 {
+ return nil, fmt.Errorf("null fixed-size list at
index %d consumed %d child values", idx, size)
}
- end++
}
+ if idx == 0 || valid != previousValid {
+ runCount++
+ }
+ previousValid = valid
+ }
- if valid {
- for idx := i; idx < end; idx++ {
- if size := offsets[idx+1] - offsets[idx]; size
!= int32(listSize) {
- return nil, fmt.Errorf("expected all
lists to be of size=%d, but index %d had size=%d", listSize, idx, size)
+ var child arrow.Array
+ var err error
+ // For a small number of runs, concatenating slices avoids materializing
+ // indices. Once the run count is high, the temporary arrays dominate.
+ const minRunsForTake = 1024
+ if runCount < minRunsForTake {
+ pieces := make([]arrow.Array, 0, runCount)
+ defer func() { releaseArrays(pieces) }()
+ for i := 0; i < length; {
+ valid := !lr.field.Nullable ||
bitutil.BitIsSet(validity, i)
+ end := i + 1
+ for end < length {
+ nextValid := !lr.field.Nullable ||
bitutil.BitIsSet(validity, end)
+ if nextValid != valid {
+ break
}
+ end++
}
- pieces = append(pieces, array.NewSlice(item,
int64(offsets[i]), int64(offsets[end])))
- } else {
- for idx := i; idx < end; idx++ {
- if size := offsets[idx+1] - offsets[idx]; size
!= 0 {
- return nil, fmt.Errorf("null fixed-size
list at index %d consumed %d child values", idx, size)
+ if valid {
+ pieces = append(pieces, array.NewSlice(item,
int64(offsets[i]), int64(offsets[end])))
+ } else {
+ pieces = append(pieces,
array.MakeArrayOfNull(lr.rctx.mem, listType.Elem(), (end-i)*listSize))
+ }
+ i = end
+ }
+ if len(pieces) == 0 {
+ pieces = append(pieces,
array.MakeArrayOfNull(lr.rctx.mem, listType.Elem(), 0))
+ }
+ child, err = array.Concatenate(pieces, lr.rctx.mem)
+ } else {
+ childLength := length * listSize
+ indicesBuffer := memory.NewResizableBuffer(lr.rctx.mem)
+ defer indicesBuffer.Release()
+
indicesBuffer.Resize(arrow.Int32Traits.BytesRequired(childLength))
+ indices :=
arrow.Int32Traits.CastFromBytes(indicesBuffer.Bytes())
+
+ indicesValidity := memory.NewResizableBuffer(lr.rctx.mem)
+ defer indicesValidity.Release()
+
indicesValidity.Resize(int(bitutil.BytesForBits(int64(childLength))))
+ clear(indicesValidity.Bytes())
+
+ for i := 0; i < length; {
+ valid := !lr.field.Nullable ||
bitutil.BitIsSet(validity, i)
+ end := i + 1
+ for end < length {
+ nextValid := !lr.field.Nullable ||
bitutil.BitIsSet(validity, end)
+ if nextValid != valid {
+ break
}
+ end++
}
- pieces = append(pieces,
array.MakeArrayOfNull(lr.rctx.mem, listType.Elem(), (end-i)*listSize))
+ if valid {
+ childStart := i * listSize
+ runChildLength := (end - i) * listSize
+ bitutil.SetBitsTo(indicesValidity.Bytes(),
int64(childStart), int64(runChildLength), true)
+ for idx := 0; idx < runChildLength; idx++ {
+ indices[childStart+idx] = offsets[i] +
int32(idx)
+ }
+ }
+ i = end
}
- i = end
- }
- if len(pieces) == 0 {
- pieces = append(pieces, array.MakeArrayOfNull(lr.rctx.mem,
listType.Elem(), 0))
+ indicesData := array.NewData(arrow.PrimitiveTypes.Int32,
childLength,
+ []*memory.Buffer{indicesValidity, indicesBuffer}, nil,
int(nullCount)*listSize, 0)
+ defer indicesData.Release()
+ indicesArr := array.NewInt32Data(indicesData)
+ defer indicesArr.Release()
+
+ ctx := compute.WithAllocator(context.Background(), lr.rctx.mem)
+ child, err = compute.TakeArrayOpts(ctx, item, indicesArr,
compute.TakeOptions{BoundsCheck: false})
}
- child, err := array.Concatenate(pieces, lr.rctx.mem)
if err != nil {
return nil, err
}
diff --git a/parquet/pqarrow/column_readers_test.go
b/parquet/pqarrow/column_readers_test.go
index 1371d6f9..f3aba2d7 100644
--- a/parquet/pqarrow/column_readers_test.go
+++ b/parquet/pqarrow/column_readers_test.go
@@ -19,11 +19,13 @@ package pqarrow
import (
"bytes"
"context"
+ "fmt"
"io"
"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/extensions"
"github.com/apache/arrow-go/v18/arrow/memory"
"github.com/apache/arrow-go/v18/parquet"
@@ -266,6 +268,121 @@ func
TestBuildFixedSizeListArrayConcatenatesSpecializedChildren(t *testing.T) {
})
}
+func TestBuildFixedSizeListArrayDirectChildren(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ const (
+ length = 2048
+ listSize = 2
+ )
+ validityBytes := make([]byte, bitutil.BytesForBits(length))
+ offsets := make([]int32, length+1)
+ validCount := 0
+ for i := 0; i < length; i++ {
+ offsets[i] = int32(validCount * listSize)
+ if i%2 == 0 {
+ bitutil.SetBit(validityBytes, i)
+ validCount++
+ }
+ }
+ offsets[length] = int32(validCount * listSize)
+
+ builder := array.NewInt32Builder(mem)
+ values := make([]int32, validCount*listSize)
+ for i := range values {
+ values[i] = int32(i)
+ }
+ builder.AppendValues(values, nil)
+ item := builder.NewArray()
+ builder.Release()
+ defer item.Release()
+
+ validity := memory.NewBufferBytes(validityBytes)
+ defer validity.Release()
+ listType := arrow.FixedSizeListOf(listSize, arrow.PrimitiveTypes.Int32)
+ field := arrow.Field{Type: listType, Nullable: true}
+ lr := &listReader{rctx: &readerCtx{mem: mem}, field: &field}
+
+ out, err := lr.buildFixedSizeListArray(length, offsets, validity,
length/2, item)
+ require.NoError(t, err)
+ defer out.Release()
+
+ list := out.Chunk(0).(*array.FixedSizeList)
+ valuesArray := list.ListValues().(*array.Int32)
+ require.Equal(t, length*listSize, valuesArray.Len())
+ for i := 0; i < length; i++ {
+ if i%2 == 0 {
+ assert.True(t, list.IsValid(i))
+ assert.Equal(t, int32((i/2)*listSize),
valuesArray.Value(i*listSize))
+ } else {
+ assert.True(t, list.IsNull(i))
+ assert.True(t, valuesArray.IsNull(i*listSize))
+ }
+ }
+}
+
+func BenchmarkBuildFixedSizeListArray(b *testing.B) {
+ const (
+ length = 1 << 16
+ listSize = 4
+ )
+
+ tests := []struct {
+ name string
+ valid func(int) bool
+ }{
+ {name: "no_nulls", valid: func(int) bool { return true }},
+ {name: "ten_percent_nulls", valid: func(i int) bool { return
i%10 != 0 }},
+ {name: "clustered_half_null", valid: func(i int) bool { return
i < length/2 }},
+ {name: "alternating", valid: func(i int) bool { return i%2 == 0
}},
+ }
+
+ for _, tt := range tests {
+ b.Run(fmt.Sprintf("%s/length=%d", tt.name, length), func(b
*testing.B) {
+ mem := memory.NewGoAllocator()
+ offsets := make([]int32, length+1)
+ validityBytes := make([]byte,
bitutil.BytesForBits(length))
+ validCount := 0
+ for i := 0; i < length; i++ {
+ offsets[i] = int32(validCount * listSize)
+ if tt.valid(i) {
+ bitutil.SetBit(validityBytes, i)
+ validCount++
+ }
+ }
+ offsets[length] = int32(validCount * listSize)
+
+ builder := array.NewInt32Builder(mem)
+ builder.AppendValues(make([]int32,
validCount*listSize), nil)
+ item := builder.NewArray()
+ builder.Release()
+ defer item.Release()
+
+ var validity *memory.Buffer
+ nullCount := int64(length - validCount)
+ if nullCount > 0 {
+ validity = memory.NewBufferBytes(validityBytes)
+ defer validity.Release()
+ }
+
+ listType := arrow.FixedSizeListOf(listSize,
arrow.PrimitiveTypes.Int32)
+ field := arrow.Field{Type: listType, Nullable: true}
+ lr := &listReader{rctx: &readerCtx{mem: mem}, field:
&field}
+
+ b.ReportAllocs()
+ b.ResetTimer()
+ for i := 0; i < b.N; i++ {
+ out, err := lr.buildFixedSizeListArray(length,
offsets, validity, nullCount, item)
+ if err != nil {
+ b.Fatal(err)
+ }
+ out.Release()
+ }
+ })
+ }
+}
+
func TestChunkedTableRoundTrip(t *testing.T) {
mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
defer mem.AssertSize(t, 0)
