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 1b50ce30f perf(table): reuse flat primitive Arrow columns during scans
(#2236)
1b50ce30f is described below
commit 1b50ce30fa71b28f1bfcf7433d0af1b616cd6d18
Author: Minh Vu <[email protected]>
AuthorDate: Tue Oct 20 19:30:42 2026 +0200
perf(table): reuse flat primitive Arrow columns during scans (#2236)
* perf(table): reuse flat primitive Arrow columns during scans
Signed-off-by: Minh Vu <[email protected]>
* test(table): guard primitive projection fast path against mismatched
columns
---------
Signed-off-by: Minh Vu <[email protected]>
---
table/arrow_scanner.go | 62 +++++++-
table/arrow_scanner_projection_fastpath_test.go | 199 ++++++++++++++++++++++++
2 files changed, 257 insertions(+), 4 deletions(-)
diff --git a/table/arrow_scanner.go b/table/arrow_scanner.go
index 89c7dfdb2..3a75df002 100644
--- a/table/arrow_scanner.go
+++ b/table/arrow_scanner.go
@@ -1550,10 +1550,11 @@ func (c *preparedFileReadPlanCache) entry(path string)
*preparedFileReadEntry {
// a scan. The snapshots returned by Metadata are owned by the scan and are
// safe to share read-only across file workers.
type arrowScanInvariants struct {
- tableSchema *iceberg.Schema
- tableProperties iceberg.Properties
- projectedIDs set[int]
- nameMapping iceberg.NameMapping
+ tableSchema *iceberg.Schema
+ tableProperties iceberg.Properties
+ projectedIDs set[int]
+ nameMapping iceberg.NameMapping
+ outputArrowSchema *arrow.Schema
}
// collectLeafIDs recursively collects leaf field IDs from a type
@@ -2460,6 +2461,45 @@ func pruningFilterHasMissingInitialDefault(
return false, nil
}
+// canReuseFlatProjection detects the common case of flat, already-typed
+// batches. Full schema projection still handles nested types, schema
evolution,
+// promotions, defaults, and changed column ordering. A matching Arrow type
alone
+// is not enough: the file field ID and Iceberg type must also agree.
+func canReuseFlatProjection(requested, fileSchema *iceberg.Schema, input,
output *arrow.Schema) bool {
+ if requested == nil || fileSchema == nil || input == nil || output ==
nil {
+ return false
+ }
+
+ fields := requested.FieldsRef(iceinternal.SchemaRef{})
+ if len(fields) != input.NumFields() || len(fields) !=
output.NumFields() {
+ return false
+ }
+
+ for i, field := range fields {
+ if _, ok := field.Type.(iceberg.PrimitiveType); !ok {
+ return false
+ }
+ fileField, ok := fileSchema.FindFieldByIDRef(field.ID,
iceinternal.SchemaRef{})
+ if !ok || fileField.Name != field.Name ||
!fileField.Type.Equals(field.Type) {
+ return false
+ }
+
+ source, target := input.Field(i), output.Field(i)
+ if source.Name != field.Name || target.Name != field.Name ||
+ !arrow.TypeEqual(source.Type, target.Type) ||
len(input.FieldIndices(field.Name)) != 1 {
+ return false
+ }
+ if id, ok := source.Metadata.GetValue(ArrowParquetFieldIDKey);
ok {
+ parsed, err := strconv.Atoi(id)
+ if err != nil || parsed != field.ID {
+ return false
+ }
+ }
+ }
+
+ return true
+}
+
func (as *arrowScan) recordsFromTask(ctx context.Context, task
tblutils.Enumerated[FileScanTask], sink *recordSink, positionalDeleteBitmap
*dv.RoaringPositionBitmap, dvBitmap *dv.RoaringPositionBitmap, eqDeleteSets
[]*equalityDeleteSet, invariants *arrowScanInvariants) (err error) {
defer func() {
if err != nil {
@@ -2588,9 +2628,22 @@ func (as *arrowScan) recordsFromTask(ctx
context.Context, task tblutils.Enumerat
pipeline = append(pipeline, filterFunc)
}
+ var lastInputSchema *arrow.Schema
+ var reuseFlatProjection bool
pipeline = append(pipeline, func(r arrow.RecordBatch)
(arrow.RecordBatch, error) {
defer r.Release()
+ if r.Schema() != lastInputSchema {
+ lastInputSchema = r.Schema()
+ reuseFlatProjection = canReuseFlatProjection(
+ as.projectedSchema, readSchema,
lastInputSchema, invariants.outputArrowSchema)
+ }
+ if reuseFlatProjection {
+ // Rewrap existing columns with the projected field
metadata. Do not
+ // rebuild the arrays or walk the Iceberg schema for
every batch.
+ return
array.NewRecordBatch(invariants.outputArrowSchema, r.Columns(), r.NumRows()),
nil
+ }
+
return ToRequestedSchema(ctx, as.projectedSchema, readSchema,
r, SchemaOptions{
UseLargeTypes: as.useLargeTypes,
TableProperties: invariants.tableProperties,
@@ -2952,6 +3005,7 @@ func (as *arrowScan) GetRecords(ctx context.Context,
tasks []FileScanTask) (*arr
if err != nil {
return nil, nil, err
}
+ invariants.outputArrowSchema = resultSchema
if err := as.addTaskProjectedFieldIDs(invariants, tasks); err != nil {
return nil, nil, err
}
diff --git a/table/arrow_scanner_projection_fastpath_test.go
b/table/arrow_scanner_projection_fastpath_test.go
new file mode 100644
index 000000000..b51f1a5e1
--- /dev/null
+++ b/table/arrow_scanner_projection_fastpath_test.go
@@ -0,0 +1,199 @@
+// 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 table
+
+import (
+ "context"
+ "fmt"
+ "strconv"
+ "testing"
+
+ "github.com/apache/arrow-go/v18/arrow"
+ "github.com/apache/arrow-go/v18/arrow/array"
+ "github.com/apache/arrow-go/v18/arrow/memory"
+ "github.com/apache/iceberg-go"
+ "github.com/stretchr/testify/require"
+)
+
+func projectionPrimitiveSchemas(t testing.TB) (*iceberg.Schema, *arrow.Schema,
*arrow.Schema) {
+ t.Helper()
+ requested := iceberg.NewSchema(1,
+ iceberg.NestedField{ID: 1, Name: "id", Type:
iceberg.PrimitiveTypes.Int64, Required: true},
+ iceberg.NestedField{ID: 2, Name: "city", Type:
iceberg.PrimitiveTypes.String, Doc: "city label"},
+ )
+ output, err := SchemaToArrowSchemaWithOptions(requested,
ArrowSchemaOptions{})
+ require.NoError(t, err)
+
+ fields := output.Fields()
+ for i := range fields {
+ fields[i].Metadata = arrow.NewMetadata(
+ []string{ArrowParquetFieldIDKey},
[]string{strconv.Itoa(i + 1)})
+ }
+ meta := arrow.NewMetadata([]string{"parquet-source"},
[]string{"source-metadata"})
+ input := arrow.NewSchema(fields, &meta)
+
+ return requested, input, output
+}
+
+func TestCanReuseFlatProjection(t *testing.T) {
+ requested, input, output := projectionPrimitiveSchemas(t)
+ require.True(t, canReuseFlatProjection(requested, requested, input,
output))
+
+ fields := input.Fields()
+ fields[0], fields[1] = fields[1], fields[0]
+ require.False(t, canReuseFlatProjection(requested, requested,
arrow.NewSchema(fields, nil), output),
+ "reordered columns must be projected")
+
+ fields = input.Fields()
+ fields[0].Metadata =
arrow.NewMetadata([]string{ArrowParquetFieldIDKey}, []string{"12"})
+ require.False(t, canReuseFlatProjection(requested, requested,
arrow.NewSchema(fields, nil), output),
+ "incorrect Arrow field IDs must not bypass projection")
+
+ fields = input.Fields()
+ fields[0].Name = "previous_name"
+ require.False(t, canReuseFlatProjection(requested, requested,
arrow.NewSchema(fields, nil), output),
+ "renamed Arrow fields must be projected")
+
+ fields = input.Fields()
+ fields[0].Type = arrow.PrimitiveTypes.Int32
+ require.False(t, canReuseFlatProjection(requested, requested,
arrow.NewSchema(fields, nil), output),
+ "mismatched physical Arrow types must not be rewrapped")
+
+ fields = input.Fields()
+ fields[1].Name = fields[0].Name
+ require.False(t, canReuseFlatProjection(requested, requested,
arrow.NewSchema(fields, nil), output),
+ "duplicate Arrow field names must not bypass projection")
+
+ fields = append(input.Fields(), arrow.Field{Name: "unexpected", Type:
arrow.PrimitiveTypes.Int32})
+ require.False(t, canReuseFlatProjection(requested, requested,
arrow.NewSchema(fields, nil), output),
+ "extra input columns must still follow schema projection")
+
+ oldName := iceberg.NewSchema(0,
+ iceberg.NestedField{ID: 1, Name: "previous_name", Type:
iceberg.PrimitiveTypes.Int64},
+ iceberg.NestedField{ID: 2, Name: "city", Type:
iceberg.PrimitiveTypes.String},
+ )
+ require.False(t, canReuseFlatProjection(requested, oldName, input,
output),
+ "renamed Iceberg fields must be projected")
+
+ oldType := iceberg.NewSchema(0,
+ iceberg.NestedField{ID: 1, Name: "id", Type:
iceberg.PrimitiveTypes.Int32},
+ iceberg.NestedField{ID: 2, Name: "city", Type:
iceberg.PrimitiveTypes.String},
+ )
+ require.False(t, canReuseFlatProjection(requested, oldType, input,
output),
+ "type promotions must be projected")
+
+ missing := iceberg.NewSchema(0,
+ iceberg.NestedField{ID: 1, Name: "id", Type:
iceberg.PrimitiveTypes.Int64},
+ )
+ require.False(t, canReuseFlatProjection(requested, missing, input,
output),
+ "missing fields may need defaults")
+
+ nested := iceberg.NewSchema(0, iceberg.NestedField{
+ ID: 1, Name: "struct", Type: &iceberg.StructType{FieldList:
[]iceberg.NestedField{
+ {ID: 2, Name: "value", Type:
iceberg.PrimitiveTypes.Int64},
+ }},
+ })
+ nestedArrow, err := SchemaToArrowSchema(nested, nil, false, false)
+ require.NoError(t, err)
+ require.False(t, canReuseFlatProjection(nested, nested, nestedArrow,
nestedArrow),
+ "nested fields require the full projection visitor")
+}
+
+func TestReuseFlatProjectionPreservesRowsAndMetadata(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+ defer mem.AssertSize(t, 0)
+ requested, inputSchema, outputSchema := projectionPrimitiveSchemas(t)
+ builder := array.NewRecordBuilder(mem, inputSchema)
+ defer builder.Release()
+
+ builder.Field(0).(*array.Int64Builder).AppendValues([]int64{7, 8, 9},
nil)
+ cities := builder.Field(1).(*array.StringBuilder)
+ cities.Append("Berlin")
+ cities.AppendNull()
+ cities.Append("Vienna")
+ batch := builder.NewRecordBatch()
+ defer batch.Release()
+
+ require.True(t, canReuseFlatProjection(requested, requested,
batch.Schema(), outputSchema))
+ expected, err := ToRequestedSchema(context.Background(), requested,
requested, batch, SchemaOptions{})
+ require.NoError(t, err)
+ defer expected.Release()
+ got := array.NewRecordBatch(outputSchema, batch.Columns(),
batch.NumRows())
+ defer got.Release()
+
+ require.True(t, array.RecordEqual(expected, got))
+ require.True(t, arrowSchemaEqual(expected.Schema(), got.Schema()))
+ require.Same(t, batch.Column(0), got.Column(0), "input arrays must be
reused")
+ require.Same(t, batch.Column(1), got.Column(1), "input arrays must be
reused")
+ require.Equal(t, 1, got.Column(1).NullN())
+}
+
+func BenchmarkFlatPrimitiveProjection(b *testing.B) {
+ for _, width := range []int{4, 32} {
+ b.Run(fmt.Sprintf("fields=%d", width), func(b *testing.B) {
+ fields := make([]iceberg.NestedField, width)
+ for i := range fields {
+ fields[i] = iceberg.NestedField{
+ ID: i + 1, Name:
fmt.Sprintf("field_%d", i), Type: iceberg.PrimitiveTypes.Int64,
+ }
+ }
+ requested := iceberg.NewSchema(1, fields...)
+ output, err := SchemaToArrowSchema(requested, nil,
false, false)
+ if err != nil {
+ b.Fatal(err)
+ }
+ inputFields := output.Fields()
+ for i := range inputFields {
+ inputFields[i].Metadata = arrow.NewMetadata(
+ []string{ArrowParquetFieldIDKey},
[]string{strconv.Itoa(i + 1)})
+ }
+ builder :=
array.NewRecordBuilder(memory.DefaultAllocator, arrow.NewSchema(inputFields,
nil))
+ defer builder.Release()
+ values := make([]int64, 1024)
+ for i := range values {
+ values[i] = int64(i)
+ }
+ for i := range inputFields {
+
builder.Field(i).(*array.Int64Builder).AppendValues(values, nil)
+ }
+ batch := builder.NewRecordBatch()
+ defer batch.Release()
+ if !canReuseFlatProjection(requested, requested,
batch.Schema(), output) {
+ b.Fatal("expected zero-copy projection")
+ }
+
+ b.Run("schema_visitor", func(b *testing.B) {
+ b.ReportAllocs()
+ for b.Loop() {
+ out, err :=
ToRequestedSchema(context.Background(), requested, requested, batch,
SchemaOptions{})
+ if err != nil {
+ b.Fatal(err)
+ }
+ out.Release()
+ }
+ })
+ b.Run("zero_copy", func(b *testing.B) {
+ b.ReportAllocs()
+ for b.Loop() {
+ out := array.NewRecordBatch(output,
batch.Columns(), batch.NumRows())
+ out.Release()
+ }
+ })
+ })
+ }
+}
