diff --git a/schema.go b/schema.go index 24b0ff4b1..4330103dc 100644 --- a/schema.go +++ b/schema.go @@ -346,8 +346,7 @@ func (s *Schema) MarshalJSON() ([]byte, error) { type Alias Schema - aliasCopy := *(*Alias)(s) - aliasCopy.IdentifierFieldIDs = ids + aliasCopy := Alias{ID: s.ID, IdentifierFieldIDs: ids} return json.Marshal(struct { Type string `json:"type"` diff --git a/schema_test.go b/schema_test.go index 7fa16e345..6d794ab71 100644 --- a/schema_test.go +++ b/schema_test.go @@ -24,6 +24,7 @@ import ( "path/filepath" "runtime" "strings" + "sync" "testing" "github.com/apache/iceberg-go" @@ -2293,3 +2294,48 @@ func TestVisitGeoSchemaWithSchemaVisitorPerPrimitiveType(t *testing.T) { assert.Equal(t, 1, v.geometryCalls) assert.Equal(t, 1, v.geographyCalls) } + +func TestSchemaMarshalJSONConcurrentLazyLookups(t *testing.T) { + for range 32 { + schema := iceberg.NewSchemaWithIdentifiers(17, nil, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64, Required: true}, + iceberg.NestedField{ID: 2, Name: "data", Type: iceberg.PrimitiveTypes.String}, + ) + start := make(chan struct{}) + var wg sync.WaitGroup + for range 8 { + wg.Go(func() { + <-start + for range 8 { + _, err := json.Marshal(schema) + assert.NoError(t, err) + } + }) + wg.Go(func() { + <-start + _, found := schema.FindFieldByID(1) + assert.True(t, found) + _, found = schema.FindFieldByName("data") + assert.True(t, found) + _, found = schema.FindFieldByNameCaseInsensitive("DATA") + assert.True(t, found) + name, found := schema.FindColumnName(2) + assert.True(t, found) + assert.Equal(t, "data", name) + }) + } + close(start) + wg.Wait() + + data, err := json.Marshal(schema) + require.NoError(t, err) + assert.JSONEq(t, `{ + "type": "struct", "schema-id": 17, "identifier-field-ids": [], + "fields": [ + {"id": 1, "name": "id", "type": "long", "required": true}, + {"id": 2, "name": "data", "type": "string", "required": false} + ] + }`, string(data)) + assert.Nil(t, schema.IdentifierFieldIDs) + } +} diff --git a/table/arrow_scanner.go b/table/arrow_scanner.go index 73a3cf8ff..160e649c9 100644 --- a/table/arrow_scanner.go +++ b/table/arrow_scanner.go @@ -869,7 +869,10 @@ func (as *arrowScan) projectedFieldIDs(rowFilter iceberg.BooleanExpression, equa } func (as *arrowScan) scanInvariants(tableProperties iceberg.Properties) (*arrowScanInvariants, error) { - projectedIDs, err := as.projectedFieldIDs(as.boundRowFilter, nil) + // Filter columns are added per task below. A local task may have a + // partition-elided residual that no longer references the original filter + // fields. + projectedIDs, err := as.projectedFieldIDs(nil, nil) if err != nil { return nil, err } @@ -883,6 +886,18 @@ func (as *arrowScan) scanInvariants(tableProperties iceberg.Properties) (*arrowS } func (as *arrowScan) addTaskProjectedFieldIDs(invariants *arrowScanInvariants, tasks []FileScanTask) error { + // Tasks without a residual still use the scan's original filter. Add those + // fields once, then add only the actual residual fields for other tasks. + for _, task := range tasks { + if task.Residual == nil { + if err := addFilterFieldIDs(invariants.projectedIDs, as.boundRowFilter); err != nil { + return err + } + + break + } + } + for _, task := range tasks { if task.Residual == nil { continue diff --git a/table/arrow_scanner_bench_test.go b/table/arrow_scanner_bench_test.go index 659c0ee7b..36720f96e 100644 --- a/table/arrow_scanner_bench_test.go +++ b/table/arrow_scanner_bench_test.go @@ -298,3 +298,158 @@ func BenchmarkArrowScanAddTaskProjectedFieldIDs(b *testing.B) { }) } } + +var benchmarkTaskResidualRows int64 + +func BenchmarkArrowScanTaskResidual(b *testing.B) { + const rowCount = 32_768 + + schema := iceberg.NewSchema(1, + iceberg.NestedField{ID: 1, Name: "tenant_id", Type: iceberg.PrimitiveTypes.String}, + iceberg.NestedField{ID: 2, Name: "amount", Type: iceberg.PrimitiveTypes.Int64}, + iceberg.NestedField{ID: 3, Name: "payload", Type: iceberg.PrimitiveTypes.String}, + ) + projectedSchema, err := schema.Select(true, "payload") + if err != nil { + b.Fatal(err) + } + metadata, err := NewMetadata( + schema, iceberg.UnpartitionedSpec, UnsortedSortOrder, "mem://benchmark/identity-residual", nil, + ) + if err != nil { + b.Fatal(err) + } + + arrowSchema, err := SchemaToArrowSchema(schema, nil, true, false) + if err != nil { + b.Fatal(err) + } + record := mustLoadRecordBatchFromJSON(arrowSchema, benchmarkTaskResidualJSON(rowCount)) + defer record.Release() + arrowTable := array.NewTableFromRecords(arrowSchema, []arrow.RecordBatch{record}) + defer arrowTable.Release() + + dataPath := "mem://benchmark/identity-residual/data.parquet" + fs := iceio.NewMemFS() + writer, err := fs.Create(dataPath) + if err != nil { + b.Fatal(err) + } + if err := pqarrow.WriteTable(arrowTable, writer, record.NumRows(), + parquet.NewWriterProperties(parquet.WithStats(true)), pqarrow.DefaultWriterProps()); err != nil { + b.Fatal(err) + } + if err := writer.Close(); err != nil { + b.Fatal(err) + } + file, err := fs.Open(dataPath) + if err != nil { + b.Fatal(err) + } + fileInfo, err := file.Stat() + if closeErr := file.Close(); err != nil { + b.Fatal(err) + } else if closeErr != nil { + b.Fatal(closeErr) + } + + dataFileBuilder, err := iceberg.NewDataFileBuilder( + *iceberg.UnpartitionedSpec, + iceberg.EntryContentData, + dataPath, + iceberg.ParquetFile, + nil, + nil, + nil, + rowCount, + fileInfo.Size(), + ) + if err != nil { + b.Fatal(err) + } + task := FileScanTask{ + File: dataFileBuilder.Build(), + Start: 0, + Length: fileInfo.Size(), + } + tasks := []FileScanTask{task} + + identityFilter := iceberg.EqualTo(iceberg.Reference("tenant_id"), "acme") + mixedFilter := iceberg.NewAnd( + identityFilter, + iceberg.GreaterThan(iceberg.Reference("amount"), int64(100)), + ) + mixedResidual := iceberg.GreaterThan(iceberg.Reference("amount"), int64(100)) + for _, tc := range []struct { + name string + filter iceberg.BooleanExpression + residual iceberg.BooleanExpression + expectedRows int64 + }{ + {name: "identity_only/original_filter", filter: identityFilter, expectedRows: rowCount}, + {name: "identity_only/always_true_residual", filter: identityFilter, residual: iceberg.AlwaysTrue{}, expectedRows: rowCount}, + {name: "mixed/original_filter", filter: mixedFilter, expectedRows: rowCount / 2}, + {name: "mixed/amount_residual", filter: mixedFilter, residual: mixedResidual, expectedRows: rowCount / 2}, + } { + boundFilter, err := iceberg.BindExpr(schema, tc.filter, true) + if err != nil { + b.Fatal(err) + } + + b.Run(tc.name, func(b *testing.B) { + tasks[0].Residual = tc.residual + b.ReportAllocs() + b.ReportMetric(float64(rowCount), "rows/op") + b.ResetTimer() + for b.Loop() { + scan := &arrowScan{ + metadata: metadata, + fs: fs, + scanSchema: schema, + projectedSchema: projectedSchema, + boundRowFilter: boundFilter, + caseSensitive: true, + rowLimit: -1, + concurrency: 1, + } + _, records, err := scan.GetRecords(b.Context(), tasks) + if err != nil { + b.Fatal(err) + } + + var rows int64 + for record, err := range records { + if err != nil { + b.Fatal(err) + } + rows += int64(record.NumRows()) + record.Release() + } + if rows != tc.expectedRows { + b.Fatalf("unexpected row count: got %d, want %d", rows, tc.expectedRows) + } + benchmarkTaskResidualRows = rows + } + }) + } +} + +func benchmarkTaskResidualJSON(rowCount int) string { + var result strings.Builder + result.Grow(rowCount * 64) + result.WriteByte('[') + for i := range rowCount { + if i > 0 { + result.WriteByte(',') + } + amount := 50 + if i%2 != 0 { + amount = 150 + } + fmt.Fprintf(&result, `{"tenant_id":"acme","amount":%d,"payload":"payload-%d"}`, + amount, i) + } + result.WriteByte(']') + + return result.String() +} diff --git a/table/arrow_scanner_test.go b/table/arrow_scanner_test.go index 8e9cd511d..0955791b4 100644 --- a/table/arrow_scanner_test.go +++ b/table/arrow_scanner_test.go @@ -149,7 +149,7 @@ func TestArrowScanSnapshotsInvariants(t *testing.T) { assert.Equal(t, 1, metadata.nameMappingCalls) } -func TestArrowScanAddTaskProjectedFieldIDsSkipsNilResiduals(t *testing.T) { +func TestArrowScanAddTaskProjectedFieldIDsUsesTaskResiduals(t *testing.T) { schema := iceberg.NewSchema(1, iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64}, iceberg.NestedField{ID: 2, Name: "data", Type: iceberg.PrimitiveTypes.String}, @@ -162,9 +162,9 @@ func TestArrowScanAddTaskProjectedFieldIDsSkipsNilResiduals(t *testing.T) { scanSchema: schema, boundRowFilter: boundFilter, } - invariants := &arrowScanInvariants{projectedIDs: set[int]{1: {}}} + invariants := &arrowScanInvariants{projectedIDs: set[int]{}} tasks := []FileScanTask{ - {}, + {Residual: iceberg.AlwaysTrue{}}, {Residual: iceberg.EqualTo(iceberg.Reference("data"), "value")}, {}, } @@ -174,6 +174,36 @@ func TestArrowScanAddTaskProjectedFieldIDsSkipsNilResiduals(t *testing.T) { assert.Equal(t, set[int]{1: {}, 2: {}}, invariants.projectedIDs) } +func TestArrowScanInvariantsStartWithRequestedFields(t *testing.T) { + schema := iceberg.NewSchema(1, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64}, + iceberg.NestedField{ID: 2, Name: "data", Type: iceberg.PrimitiveTypes.String}, + ) + projected, err := schema.Select(true, "data") + require.NoError(t, err) + metadata, err := NewMetadata( + schema, iceberg.UnpartitionedSpec, UnsortedSortOrder, "mem://test/table", iceberg.Properties{}, + ) + require.NoError(t, err) + boundFilter, err := iceberg.BindExpr(schema, + iceberg.EqualTo(iceberg.Reference("id"), int64(1)), true) + require.NoError(t, err) + + scanner := &arrowScan{ + metadata: metadata, + scanSchema: schema, + projectedSchema: projected, + boundRowFilter: boundFilter, + } + invariants, err := scanner.scanInvariants(metadata.Properties()) + require.NoError(t, err) + assert.Equal(t, set[int]{2: {}}, invariants.projectedIDs) + + err = scanner.addTaskProjectedFieldIDs(invariants, []FileScanTask{{Residual: iceberg.AlwaysTrue{}}}) + require.NoError(t, err) + assert.Equal(t, set[int]{2: {}}, invariants.projectedIDs) +} + func TestEnrichRecordsWithPosDeleteFields(t *testing.T) { testSchema := arrow.NewSchema([]arrow.Field{ {Name: "first_name", Type: &arrow.StringType{}, Nullable: false}, diff --git a/table/partition_residual.go b/table/partition_residual.go new file mode 100644 index 000000000..b97590df0 --- /dev/null +++ b/table/partition_residual.go @@ -0,0 +1,467 @@ +// 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 ( + "math" + "slices" + + "github.com/apache/iceberg-go" +) + +// partitionResidualPlan partially evaluates a bound row filter using the +// values of identity partition fields for one partition spec. Predicates that +// cannot be proven from the partition remain in the residual. +type partitionResidualPlan struct { + schema *iceberg.Schema + root *partitionResidualNode + sources map[int]partitionResidualSource +} + +type partitionResidualSource struct { + partitionFieldIDs []int + path []int +} + +type partitionResidualNodeKind uint8 + +const ( + partitionResidualOpaque partitionResidualNodeKind = iota + partitionResidualTrue + partitionResidualFalse + partitionResidualPredicateKind + partitionResidualNot + partitionResidualAnd + partitionResidualOr +) + +type partitionResidualNode struct { + kind partitionResidualNodeKind + expr iceberg.BooleanExpression + predicate *partitionResidualPredicate + left *partitionResidualNode + right *partitionResidualNode + child *partitionResidualNode +} + +type partitionResidualPredicate struct { + sourceID int + op iceberg.Operation + evaluate func(iceberg.StructLike) (bool, error) +} + +type partitionResidualPlanBuilder struct { + schema *iceberg.Schema + caseSensitive bool + identitySourceIDs map[int][]int + sources map[int]partitionResidualSource + identityPredicates int +} + +func newPartitionResidualPlan( + schema *iceberg.Schema, + spec *iceberg.PartitionSpec, + filter iceberg.BooleanExpression, + caseSensitive bool, +) *partitionResidualPlan { + if filter == nil || spec == nil { + return nil + } + + identitySourceIDs := make(map[int][]int) + for _, field := range spec.Fields() { + if !isIdentityPartitionField(field) { + continue + } + + sourceID := field.SourceID() + if sourceID <= 0 { + continue + } + + identitySourceIDs[sourceID] = append(identitySourceIDs[sourceID], field.FieldID) + } + if len(identitySourceIDs) == 0 { + return nil + } + + builder := &partitionResidualPlanBuilder{ + schema: schema, + caseSensitive: caseSensitive, + identitySourceIDs: identitySourceIDs, + sources: make(map[int]partitionResidualSource), + } + root, err := iceberg.VisitExpr(filter, builder) + if err != nil || builder.identityPredicates == 0 { + return nil + } + + return &partitionResidualPlan{ + schema: schema, + root: root, + sources: builder.sources, + } +} + +func isIdentityPartitionField(field iceberg.PartitionField) bool { + if len(field.SourceIDs) != 1 { + return false + } + + switch transform := field.Transform.(type) { + case iceberg.IdentityTransform: + return true + case *iceberg.IdentityTransform: + return transform != nil + default: + return false + } +} + +func (b *partitionResidualPlanBuilder) VisitTrue() *partitionResidualNode { + return &partitionResidualNode{kind: partitionResidualTrue, expr: iceberg.AlwaysTrue{}} +} + +func (b *partitionResidualPlanBuilder) VisitFalse() *partitionResidualNode { + return &partitionResidualNode{kind: partitionResidualFalse, expr: iceberg.AlwaysFalse{}} +} + +func (b *partitionResidualPlanBuilder) VisitNot(child *partitionResidualNode) *partitionResidualNode { + return &partitionResidualNode{ + kind: partitionResidualNot, + expr: iceberg.NewNot(child.expr), + child: child, + } +} + +func (b *partitionResidualPlanBuilder) VisitAnd(left, right *partitionResidualNode) *partitionResidualNode { + return &partitionResidualNode{ + kind: partitionResidualAnd, + expr: iceberg.NewAnd(left.expr, right.expr), + left: left, + right: right, + } +} + +func (b *partitionResidualPlanBuilder) VisitOr(left, right *partitionResidualNode) *partitionResidualNode { + return &partitionResidualNode{ + kind: partitionResidualOr, + expr: iceberg.NewOr(left.expr, right.expr), + left: left, + right: right, + } +} + +func (b *partitionResidualPlanBuilder) VisitUnbound(pred iceberg.UnboundPredicate) *partitionResidualNode { + return &partitionResidualNode{kind: partitionResidualOpaque, expr: pred} +} + +func (b *partitionResidualPlanBuilder) VisitBound(pred iceberg.BoundPredicate) *partitionResidualNode { + if literal, ok := pred.(iceberg.BoundLiteralPredicate); ok && partitionResidualValueIsNaN(literal.Literal()) { + return &partitionResidualNode{kind: partitionResidualOpaque, expr: pred} + } + + ref, ok := pred.Term().(iceberg.BoundReference) + if !ok { + return &partitionResidualNode{kind: partitionResidualOpaque, expr: pred} + } + + partitionFieldIDs, ok := b.identitySourceIDs[ref.Field().ID] + if !ok || !partitionResidualPredicateSupported(pred) { + return &partitionResidualNode{kind: partitionResidualOpaque, expr: pred} + } + + unbound, err := iceberg.TranslateColumnNames(pred, b.schema) + if err != nil { + return &partitionResidualNode{kind: partitionResidualOpaque, expr: pred} + } + + evaluate, err := iceberg.ExpressionEvaluator(b.schema, unbound, b.caseSensitive) + if err != nil { + return &partitionResidualNode{kind: partitionResidualOpaque, expr: pred} + } + + path := ref.PosPath() + if !partitionResidualPathSupported(b.schema, path) { + return &partitionResidualNode{kind: partitionResidualOpaque, expr: pred} + } + + source, exists := b.sources[ref.Field().ID] + if !exists { + source = partitionResidualSource{path: slices.Clone(path)} + } + if !slices.Equal(source.path, path) { + return &partitionResidualNode{kind: partitionResidualOpaque, expr: pred} + } + for _, partitionFieldID := range partitionFieldIDs { + if !slices.Contains(source.partitionFieldIDs, partitionFieldID) { + source.partitionFieldIDs = append(source.partitionFieldIDs, partitionFieldID) + } + } + b.sources[ref.Field().ID] = source + b.identityPredicates++ + + return &partitionResidualNode{ + kind: partitionResidualPredicateKind, + expr: pred, + predicate: &partitionResidualPredicate{ + sourceID: ref.Field().ID, + op: pred.Op(), + evaluate: evaluate, + }, + } +} + +func partitionResidualPredicateSupported(pred iceberg.BoundPredicate) bool { + switch pred.(type) { + case iceberg.BoundUnaryPredicate, iceberg.BoundLiteralPredicate, iceberg.BoundSetPredicate: + return true + default: + return false + } +} + +func partitionResidualPathSupported(schema *iceberg.Schema, path []int) bool { + if len(path) == 0 { + return false + } + + fields := schema.Fields() + for i, pos := range path { + if pos < 0 || pos >= len(fields) { + return false + } + if i == len(path)-1 { + return true + } + + nested, ok := fields[pos].Type.(*iceberg.StructType) + if !ok || nested == nil { + return false + } + fields = nested.FieldList + } + + return false +} + +// residual returns the task residual and whether at least one partition value +// was used. A false changed result means the caller should keep a nil task +// residual and use the scan filter as the conservative fallback. +func (p *partitionResidualPlan) residual(partition map[int]any) (iceberg.BooleanExpression, bool) { + record := make(partitionSourceRecord, p.schema.NumFields()) + knownSources := make(map[int]bool, len(p.sources)) + for sourceID, source := range p.sources { + value, ok := partitionValue(partition, source.partitionFieldIDs) + if !ok || !setPartitionSourceValue(record, p.schema, source.path, value) { + continue + } + + knownSources[sourceID] = value == nil || partitionResidualValueIsNaN(value) + } + + residual, changed, _, _ := p.root.residual(knownSources, record) + if !changed { + return nil, false + } + + return residual, true +} + +func partitionValue(partition map[int]any, fieldIDs []int) (any, bool) { + for _, fieldID := range fieldIDs { + value, ok := partition[fieldID] + if !ok { + continue + } + if _, unknown := value.(iceberg.AboveMaxLiteral); unknown { + return nil, false + } + if _, unknown := value.(iceberg.BelowMinLiteral); unknown { + return nil, false + } + if literal, isLiteral := value.(iceberg.Literal); isLiteral { + return literal.Any(), true + } + + return value, true + } + + return nil, false +} + +type partitionSourceRecord []any + +func (r partitionSourceRecord) Size() int { return len(r) } +func (r partitionSourceRecord) Get(pos int) any { return r[pos] } +func (r partitionSourceRecord) Set(pos int, value any) { r[pos] = value } + +func setPartitionSourceValue( + record partitionSourceRecord, + schema *iceberg.Schema, + path []int, + value any, +) bool { + if len(path) == 0 { + return false + } + + if len(path) == 1 { + pos := path[0] + if pos < 0 || pos >= schema.NumFields() { + return false + } + record[pos] = value + + return true + } + + fields := schema.Fields() + current := record + for i, pos := range path { + if pos < 0 || pos >= len(fields) { + return false + } + if i == len(path)-1 { + current[pos] = value + + return true + } + + nested, ok := fields[pos].Type.(*iceberg.StructType) + if !ok || nested == nil { + return false + } + + child, ok := current[pos].(partitionSourceRecord) + if !ok { + child = make(partitionSourceRecord, len(nested.FieldList)) + current[pos] = child + } + current = child + fields = nested.FieldList + } + + return false +} + +func (n *partitionResidualNode) residual( + knownSources map[int]bool, + record iceberg.StructLike, +) (iceberg.BooleanExpression, bool, bool, bool) { + switch n.kind { + case partitionResidualOpaque: + return n.expr, false, false, false + case partitionResidualTrue: + return iceberg.AlwaysTrue{}, false, true, true + case partitionResidualFalse: + return iceberg.AlwaysFalse{}, false, true, false + case partitionResidualPredicateKind: + requiresRowEvaluation, known := knownSources[n.predicate.sourceID] + if !known { + return n.expr, false, false, false + } + // Scalar comparisons order nulls and NaNs, while Arrow's row + // filters propagate nulls and use IEEE floating-point comparisons. + // Keep those predicates intact, including beneath NOT. + if requiresRowEvaluation && n.predicate.op != iceberg.OpIsNull && n.predicate.op != iceberg.OpNotNull { + return n.expr, false, false, false + } + + value, err := n.predicate.evaluate(record) + if err != nil { + return n.expr, false, false, false + } + if value { + return iceberg.AlwaysTrue{}, true, true, true + } + + return iceberg.AlwaysFalse{}, true, true, false + case partitionResidualNot: + child, changed, exact, value := n.child.residual(knownSources, record) + if !exact { + return iceberg.NewNot(child), changed, false, false + } + + return boolExpressionForValue(!value), changed, true, !value + case partitionResidualAnd: + left, leftChanged, leftExact, leftValue := n.left.residual(knownSources, record) + if leftExact && !leftValue { + return iceberg.AlwaysFalse{}, leftChanged, true, false + } + + right, rightChanged, rightExact, rightValue := n.right.residual(knownSources, record) + changed := leftChanged || rightChanged + switch { + case leftExact && leftValue: + return right, changed, rightExact, rightValue + case rightExact && !rightValue: + return iceberg.AlwaysFalse{}, changed, true, false + case rightExact && rightValue: + return left, changed, leftExact, leftValue + case leftExact && rightExact: + return iceberg.AlwaysTrue{}, changed, true, true + default: + return iceberg.NewAnd(left, right), changed, false, false + } + case partitionResidualOr: + left, leftChanged, leftExact, leftValue := n.left.residual(knownSources, record) + if leftExact && leftValue { + return iceberg.AlwaysTrue{}, leftChanged, true, true + } + + right, rightChanged, rightExact, rightValue := n.right.residual(knownSources, record) + changed := leftChanged || rightChanged + switch { + case leftExact && !leftValue: + return right, changed, rightExact, rightValue + case rightExact && rightValue: + return iceberg.AlwaysTrue{}, changed, true, true + case rightExact && !rightValue: + return left, changed, leftExact, leftValue + case leftExact && rightExact: + return iceberg.AlwaysFalse{}, changed, true, false + default: + return iceberg.NewOr(left, right), changed, false, false + } + } + + return n.expr, false, false, false +} + +func partitionResidualValueIsNaN(value any) bool { + if literal, ok := value.(iceberg.Literal); ok { + value = literal.Any() + } + switch value := value.(type) { + case float32: + return math.IsNaN(float64(value)) + case float64: + return math.IsNaN(value) + default: + return false + } +} + +func boolExpressionForValue(value bool) iceberg.BooleanExpression { + if value { + return iceberg.AlwaysTrue{} + } + + return iceberg.AlwaysFalse{} +} diff --git a/table/partition_residual_null_test.go b/table/partition_residual_null_test.go new file mode 100644 index 000000000..7392d061a --- /dev/null +++ b/table/partition_residual_null_test.go @@ -0,0 +1,127 @@ +// 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 ( + "fmt" + "math" + "testing" + + "github.com/apache/iceberg-go" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestPartitionResidualPreservesArrowNullFiltering(t *testing.T) { + checkPartitionResidualArrowFiltering(t, iceberg.PrimitiveTypes.Int32, nil, "null", + iceberg.Int32Literal(5), []iceberg.Transform{iceberg.IdentityTransform{}}) +} + +func TestPartitionResidualPreservesArrowNaNFiltering(t *testing.T) { + for _, tt := range []struct { + name string + value any + json string + literal iceberg.Literal + }{ + {"null", nil, "null", iceberg.Float64Literal(5)}, + {"nan partition", math.NaN(), `"NaN"`, iceberg.Float64Literal(5)}, + {"nan literal", float64(5), "5", iceberg.Float64Literal(math.NaN())}, + {"nan partition and literal", math.NaN(), `"NaN"`, iceberg.Float64Literal(math.NaN())}, + } { + t.Run(tt.name, func(t *testing.T) { + checkPartitionResidualArrowFiltering(t, iceberg.PrimitiveTypes.Float64, tt.value, tt.json, + tt.literal, []iceberg.Transform{iceberg.IdentityTransform{}}) + }) + } +} + +func checkPartitionResidualArrowFiltering(t *testing.T, typ iceberg.Type, value any, jsonValue string, + literal iceberg.Literal, transforms []iceberg.Transform, +) { + t.Helper() + schema := iceberg.NewSchema(1, + iceberg.NestedField{ID: 1, Name: "id", Type: typ}, + iceberg.NestedField{ID: 2, Name: "flag", Type: iceberg.PrimitiveTypes.Bool, Required: true}, + ) + arrowSchema, err := SchemaToArrowSchema(schema, nil, true, false) + require.NoError(t, err) + countRows := func(t *testing.T, filter iceberg.BooleanExpression) int64 { + t.Helper() + plan, err := compileFileFilterPlan(schema, filter, true, true, false) + require.NoError(t, err) + if plan.dropFile { + return 0 + } + record := mustLoadRecordBatchFromJSON(arrowSchema, + fmt.Sprintf(`[{"id":%s,"flag":true},{"id":%s,"flag":false}]`, jsonValue, jsonValue)) + if process := plan.recordProcessor(t.Context()); process != nil { + record, err = process(record) + require.NoError(t, err) + } + defer record.Release() + + return record.NumRows() + } + ref := iceberg.Reference("id") + otherLiteral, err := iceberg.Int32Literal(1).To(typ) + require.NoError(t, err) + predicates := []iceberg.BooleanExpression{ + iceberg.IsNull(ref), iceberg.NotNull(ref), + iceberg.SetPredicate(iceberg.OpIn, ref, []iceberg.Literal{literal, otherLiteral}), + iceberg.SetPredicate(iceberg.OpNotIn, ref, []iceberg.Literal{literal, otherLiteral}), + } + if typ.Equals(iceberg.PrimitiveTypes.Float64) { + predicates = append(predicates, iceberg.IsNaN(ref), iceberg.NotNaN(ref)) + } + for _, op := range []iceberg.Operation{ + iceberg.OpEQ, iceberg.OpNEQ, iceberg.OpLT, iceberg.OpLTEQ, iceberg.OpGT, iceberg.OpGTEQ, + } { + predicates = append(predicates, iceberg.LiteralPredicate(op, ref, literal)) + } + flag := iceberg.EqualTo(iceberg.Reference("flag"), true) + for _, transform := range transforms { + t.Run(transform.String(), func(t *testing.T) { + spec := iceberg.NewPartitionSpec(iceberg.PartitionField{ + SourceIDs: []int{1}, FieldID: 1000, Name: "partition", Transform: transform, + }) + for _, predicate := range predicates { + for _, filter := range []iceberg.BooleanExpression{ + predicate, iceberg.NewNot(predicate), + iceberg.NewAnd(predicate, flag), iceberg.NewOr(predicate, flag), + iceberg.NewNot(iceberg.NewAnd(predicate, flag)), + iceberg.NewNot(iceberg.NewOr(predicate, flag)), + } { + t.Run(filter.String(), func(t *testing.T) { + bound, err := iceberg.BindExpr(schema, filter, true) + require.NoError(t, err) + residualPlan := newPartitionResidualPlan(schema, &spec, bound, true) + residual := bound + if residualPlan != nil { + candidate, changed := residualPlan.residual(map[int]any{1000: value}) + if changed { + residual = candidate + } + } + assert.Equal(t, countRows(t, bound), countRows(t, residual), "residual=%s", residual) + }) + } + } + }) + } +} diff --git a/table/partition_residual_test.go b/table/partition_residual_test.go new file mode 100644 index 000000000..80e1c54a1 --- /dev/null +++ b/table/partition_residual_test.go @@ -0,0 +1,383 @@ +// 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 ( + "bytes" + "context" + "testing" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/decimal128" + "github.com/apache/arrow-go/v18/parquet" + "github.com/apache/arrow-go/v18/parquet/pqarrow" + "github.com/apache/iceberg-go" + iceio "github.com/apache/iceberg-go/io" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func partitionResidualTestSchema() *iceberg.Schema { + return iceberg.NewSchema(1, + iceberg.NestedField{ + ID: 1, Name: "tenant_id", Type: iceberg.PrimitiveTypes.String, + }, + iceberg.NestedField{ + ID: 2, Name: "amount", Type: iceberg.PrimitiveTypes.Int64, + }, + iceberg.NestedField{ + ID: 3, Name: "payload", Type: iceberg.PrimitiveTypes.String, + }, + ) +} + +func partitionResidualTestSpec(transform iceberg.Transform) iceberg.PartitionSpec { + return iceberg.NewPartitionSpecID(0, iceberg.PartitionField{ + SourceIDs: []int{1}, + FieldID: 1000, + Name: "tenant_id", + Transform: transform, + }) +} + +func boundPartitionResidualPlan( + t *testing.T, + filter iceberg.BooleanExpression, + transform iceberg.Transform, +) *partitionResidualPlan { + t.Helper() + + schema := partitionResidualTestSchema() + bound, err := iceberg.BindExpr(schema, filter, true) + require.NoError(t, err) + spec := partitionResidualTestSpec(transform) + + return newPartitionResidualPlan(schema, &spec, bound, true) +} + +func TestPartitionResidualPlanElidesSatisfiedIdentityPredicate(t *testing.T) { + filter := iceberg.NewAnd( + iceberg.EqualTo(iceberg.Reference("tenant_id"), "acme"), + iceberg.GreaterThan(iceberg.Reference("amount"), int64(100)), + ) + plan := boundPartitionResidualPlan(t, filter, iceberg.IdentityTransform{}) + require.NotNil(t, plan) + + residual, changed := plan.residual(map[int]any{1000: "acme"}) + require.True(t, changed) + + want, err := iceberg.BindExpr(partitionResidualTestSchema(), + iceberg.GreaterThan(iceberg.Reference("amount"), int64(100)), true) + require.NoError(t, err) + assert.True(t, residual.Equals(want), "expected %s, got %s", want, residual) +} + +func TestPartitionResidualPlanElidesIdentityPredicateToAlwaysTrue(t *testing.T) { + plan := boundPartitionResidualPlan(t, + iceberg.EqualTo(iceberg.Reference("tenant_id"), "acme"), + iceberg.IdentityTransform{}) + require.NotNil(t, plan) + + residual, changed := plan.residual(map[int]any{1000: "acme"}) + require.True(t, changed) + assert.Equal(t, iceberg.AlwaysTrue{}, residual) +} + +func TestPartitionResidualPlanUsesNullIdentityValues(t *testing.T) { + schema := partitionResidualTestSchema() + filter := iceberg.IsNull(iceberg.Reference("tenant_id")) + bound, err := iceberg.BindExpr(schema, filter, true) + require.NoError(t, err) + spec := partitionResidualTestSpec(iceberg.IdentityTransform{}) + plan := newPartitionResidualPlan(schema, &spec, bound, true) + require.NotNil(t, plan) + + residual, changed := plan.residual(map[int]any{1000: nil}) + require.True(t, changed) + assert.Equal(t, iceberg.AlwaysTrue{}, residual) +} + +func TestPartitionResidualPlanNormalizesDecodedLiteralValues(t *testing.T) { + schema := iceberg.NewSchema(1, iceberg.NestedField{ + ID: 1, Name: "price", Type: iceberg.DecimalTypeOf(10, 2), + }) + value := iceberg.Decimal{Val: decimal128.FromI64(123), Scale: 2} + filter := iceberg.EqualTo(iceberg.Reference("price"), value) + bound, err := iceberg.BindExpr(schema, filter, true) + require.NoError(t, err) + spec := iceberg.NewPartitionSpecID(0, iceberg.PartitionField{ + SourceIDs: []int{1}, FieldID: 1000, Name: "price", Transform: iceberg.IdentityTransform{}, + }) + plan := newPartitionResidualPlan(schema, &spec, bound, true) + require.NotNil(t, plan) + + residual, changed := plan.residual(map[int]any{1000: iceberg.DecimalLiteral(value)}) + require.True(t, changed) + assert.Equal(t, iceberg.AlwaysTrue{}, residual) +} + +func TestPartitionResidualPlanPreservesUnknownAndNonIdentityPredicates(t *testing.T) { + filter := iceberg.NewAnd( + iceberg.EqualTo(iceberg.Reference("tenant_id"), "acme"), + iceberg.GreaterThan(iceberg.Reference("amount"), int64(100)), + ) + + identityPlan := boundPartitionResidualPlan(t, filter, iceberg.IdentityTransform{}) + residual, changed := identityPlan.residual(nil) + assert.False(t, changed) + assert.Nil(t, residual) + + nonIdentityPlan := boundPartitionResidualPlan(t, filter, iceberg.BucketTransform{NumBuckets: 16}) + assert.Nil(t, nonIdentityPlan) +} + +func TestPartitionResidualPlanDoesNotEvaluateTransformedPredicates(t *testing.T) { + filter := iceberg.EqualTo( + iceberg.NewUnboundTransform(iceberg.BucketTransform{NumBuckets: 16}, iceberg.Reference("tenant_id")), + int32(1), + ) + plan := boundPartitionResidualPlan(t, filter, iceberg.IdentityTransform{}) + assert.Nil(t, plan) +} + +func TestPartitionResidualPlanHandlesOrResiduals(t *testing.T) { + filter := iceberg.NewOr( + iceberg.EqualTo(iceberg.Reference("tenant_id"), "acme"), + iceberg.GreaterThan(iceberg.Reference("amount"), int64(100)), + ) + plan := boundPartitionResidualPlan(t, filter, iceberg.IdentityTransform{}) + require.NotNil(t, plan) + + residual, changed := plan.residual(map[int]any{1000: "acme"}) + require.True(t, changed) + assert.Equal(t, iceberg.AlwaysTrue{}, residual) + + residual, changed = plan.residual(map[int]any{1000: "other"}) + require.True(t, changed) + want, err := iceberg.BindExpr(partitionResidualTestSchema(), + iceberg.GreaterThan(iceberg.Reference("amount"), int64(100)), true) + require.NoError(t, err) + assert.True(t, residual.Equals(want), "expected %s, got %s", want, residual) +} + +func TestPartitionResidualPlanHandlesNestedIdentityFields(t *testing.T) { + nested := &iceberg.StructType{FieldList: []iceberg.NestedField{ + {ID: 2, Name: "tenant_id", Type: iceberg.PrimitiveTypes.String}, + {ID: 3, Name: "amount", Type: iceberg.PrimitiveTypes.Int64}, + }} + schema := iceberg.NewSchema(1, iceberg.NestedField{ + ID: 1, Name: "details", Type: nested, + }) + spec := iceberg.NewPartitionSpecID(0, iceberg.PartitionField{ + SourceIDs: []int{2}, FieldID: 1000, Name: "tenant_id", Transform: iceberg.IdentityTransform{}, + }) + filter := iceberg.NewAnd( + iceberg.EqualTo(iceberg.Reference("details.tenant_id"), "acme"), + iceberg.GreaterThan(iceberg.Reference("details.amount"), int64(100)), + ) + bound, err := iceberg.BindExpr(schema, filter, true) + require.NoError(t, err) + plan := newPartitionResidualPlan(schema, &spec, bound, true) + require.NotNil(t, plan) + + residual, changed := plan.residual(map[int]any{1000: "acme"}) + require.True(t, changed) + want, err := iceberg.BindExpr(schema, + iceberg.GreaterThan(iceberg.Reference("details.amount"), int64(100)), true) + require.NoError(t, err) + assert.True(t, residual.Equals(want), "expected %s, got %s", want, residual) +} + +func TestPlanFilesLocalSetsIdentityPartitionResidual(t *testing.T) { + ctx := context.Background() + const tableLocation = "mem://identity-residual" + + schema := partitionResidualTestSchema() + spec := partitionResidualTestSpec(iceberg.IdentityTransform{}) + metadata, err := NewMetadata(schema, &spec, UnsortedSortOrder, tableLocation, nil) + require.NoError(t, err) + fs := iceio.NewMemFS() + dataPath := tableLocation + "/data/file.parquet" + dataSize := writePartitionResidualParquet(t, fs, dataPath, schema, `[ + {"tenant_id":"acme","amount":50,"payload":"low"}, + {"tenant_id":"acme","amount":150,"payload":"keep"}, + {"tenant_id":"acme","amount":250,"payload":"high"} + ]`) + + dataFileBuilder, err := iceberg.NewDataFileBuilder( + spec, + iceberg.EntryContentData, + dataPath, + iceberg.ParquetFile, + map[int]any{1000: "acme"}, + nil, + nil, + 3, + dataSize, + ) + require.NoError(t, err) + + snapshotID := int64(1) + entry := iceberg.NewManifestEntryBuilder( + iceberg.EntryStatusADDED, + &snapshotID, + dataFileBuilder.Build(), + ).SequenceNum(1).Build() + manifestPath := tableLocation + "/metadata/manifest.avro" + var manifestBuffer bytes.Buffer + manifest, err := iceberg.WriteManifest( + manifestPath, &manifestBuffer, 2, spec, schema, snapshotID, []iceberg.ManifestEntry{entry}, + ) + require.NoError(t, err) + + manifestListPath := tableLocation + "/metadata/snap-1.avro" + var manifestListBuffer bytes.Buffer + sequenceNumber := int64(1) + require.NoError(t, iceberg.WriteManifestList( + 2, &manifestListBuffer, snapshotID, nil, &sequenceNumber, 0, + []iceberg.ManifestFile{manifest}, + )) + + require.NoError(t, fs.WriteFile(manifestPath, manifestBuffer.Bytes())) + require.NoError(t, fs.WriteFile(manifestListPath, manifestListBuffer.Bytes())) + + metadataBuilder, err := MetadataBuilderFromBase(metadata, "") + require.NoError(t, err) + schemaID := metadata.CurrentSchema().ID + require.NoError(t, metadataBuilder.AddSnapshot(&Snapshot{ + SnapshotID: snapshotID, + SequenceNumber: 1, + TimestampMs: metadata.LastUpdatedMillis() + 1, + ManifestList: manifestListPath, + Summary: &Summary{Operation: OpAppend}, + SchemaID: &schemaID, + })) + require.NoError(t, metadataBuilder.SetSnapshotRef(MainBranch, snapshotID, BranchRef)) + built, err := metadataBuilder.Build() + require.NoError(t, err) + + tbl := New( + Identifier{"db", "identity-residual"}, + built, + tableLocation+"/metadata/metadata.json", + func(context.Context) (iceio.IO, error) { return fs, nil }, + nil, + ) + filter := iceberg.NewAnd( + iceberg.EqualTo(iceberg.Reference("tenant_id"), "acme"), + iceberg.GreaterThan(iceberg.Reference("amount"), int64(100)), + ) + scan := tbl.Scan(WithRowFilter(filter), WithSelectedFields("payload")) + tasks, err := scan.PlanFiles(ctx) + require.NoError(t, err) + require.Len(t, tasks, 1) + require.NotNil(t, tasks[0].Residual) + + want, err := iceberg.BindExpr(schema, + iceberg.GreaterThan(iceberg.Reference("amount"), int64(100)), true) + require.NoError(t, err) + assert.True(t, tasks[0].Residual.Equals(want), + "expected %s, got %s", want, tasks[0].Residual) + + resultSchema, records, err := scan.ReadTasks(ctx, tasks) + require.NoError(t, err) + require.Equal(t, 1, resultSchema.NumFields()) + require.Equal(t, "payload", resultSchema.Field(0).Name) + + var payloads []string + for record, err := range records { + require.NoError(t, err) + values := record.Column(0).(*array.String) + for i := range values.Len() { + payloads = append(payloads, values.Value(i)) + } + record.Release() + } + assert.Equal(t, []string{"keep", "high"}, payloads) +} + +func writePartitionResidualParquet( + t *testing.T, + fs *iceio.MemFS, + path string, + schema *iceberg.Schema, + jsonData string, +) int64 { + t.Helper() + + arrowSchema, err := SchemaToArrowSchema(schema, nil, true, false) + require.NoError(t, err) + record := mustLoadRecordBatchFromJSON(arrowSchema, jsonData) + defer record.Release() + + tbl := array.NewTableFromRecords(arrowSchema, []arrow.RecordBatch{record}) + defer tbl.Release() + + writer, err := fs.Create(path) + require.NoError(t, err) + require.NoError(t, pqarrow.WriteTable(tbl, writer, record.NumRows(), + parquet.NewWriterProperties(parquet.WithStats(true)), pqarrow.DefaultWriterProps())) + require.NoError(t, writer.Close()) + + file, err := fs.Open(path) + require.NoError(t, err) + defer file.Close() + info, err := file.Stat() + require.NoError(t, err) + + return info.Size() +} + +var partitionResidualBenchmarkSink iceberg.BooleanExpression + +func BenchmarkPartitionResidualPlanning(b *testing.B) { + schema := partitionResidualTestSchema() + spec := partitionResidualTestSpec(iceberg.IdentityTransform{}) + filter := iceberg.NewAnd( + iceberg.EqualTo(iceberg.Reference("tenant_id"), "acme"), + iceberg.GreaterThan(iceberg.Reference("amount"), int64(100)), + ) + bound, err := iceberg.BindExpr(schema, filter, true) + if err != nil { + b.Fatal(err) + } + plan := newPartitionResidualPlan(schema, &spec, bound, true) + if plan == nil { + b.Fatal("expected an identity partition residual plan") + } + + partitions := make([]map[int]any, 4096) + for i := range partitions { + value := "acme" + if i%2 == 0 { + value = "other" + } + partitions[i] = map[int]any{1000: value} + } + + b.ReportAllocs() + b.ResetTimer() + for b.Loop() { + for _, partition := range partitions { + partitionResidualBenchmarkSink, _ = plan.residual(partition) + } + } + b.StopTimer() + b.ReportMetric(float64(len(partitions)), "files/op") +} diff --git a/table/scanner.go b/table/scanner.go index f5599d932..276ed84af 100644 --- a/table/scanner.go +++ b/table/scanner.go @@ -1117,6 +1117,18 @@ func (scan *Scan) planFilesLocal(ctx context.Context, acc *scanMetricsAccumulato return nil, err } + var boundRowFilter iceberg.BooleanExpression + if scan.rowFilter != nil { + boundRowFilter, err = iceberg.BindExpr(schema, scan.rowFilter, scan.caseSensitive) + if err != nil { + return nil, err + } + } + var residualPlans map[int]*partitionResidualPlan + if boundRowFilter != nil { + residualPlans = make(map[int]*partitionResidualPlan) + } + // Step 3: Index positional deletes and match them to data files. posDeleteIndex, err := buildPositionalDeleteIndex(entries.positionalDeleteEntries) if err != nil { @@ -1161,6 +1173,20 @@ func (scan *Scan) planFilesLocal(ctx context.Context, acc *scanMetricsAccumulato Start: 0, Length: e.DataFile().FileSizeBytes(), } + if boundRowFilter != nil { + specID := int(e.DataFile().SpecID()) + residualPlan, found := residualPlans[specID] + if !found { + residualPlan = newPartitionResidualPlan( + schema, scan.metadata.PartitionSpecByID(specID), boundRowFilter, scan.caseSensitive) + residualPlans[specID] = residualPlan + } + if residualPlan != nil { + if residual, simplified := residualPlan.residual(dataFilePartition(e.DataFile())); simplified { + task.Residual = residual + } + } + } // Row lineage constants: readers use these to synthesize _row_id and // _last_updated_sequence_number when requested. Per spec the // synthesized _last_updated_sequence_number is the manifest entry's @@ -1435,7 +1461,7 @@ type FileScanTask struct { DeletionVectorFiles []iceberg.DataFile // deletion vectors (puffin files) Start, Length int64 // Residual is the portion of the scan filter that must still be evaluated - // for this task. Remote planners may simplify the original filter using + // for this task. Local or remote planners may simplify the original filter using // file metadata; nil means the caller did not provide a task residual. // ReadTasks applies the scan's original row filter and each task residual. Residual iceberg.BooleanExpression