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..8be724a59 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.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 + } + b.ReportMetric(float64(rowCount), "rows/op") + }) + } +} + +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..70cf5fcb4 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,37 @@ 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..e47beddeb --- /dev/null +++ b/table/partition_residual.go @@ -0,0 +1,352 @@ +// 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" + + "github.com/apache/iceberg-go" +) + +// partitionResidualEvaluator partially evaluates a bound row filter using a +// file's partition values. It is built once per partition spec and can be +// reused for every data file in that spec. +type partitionResidualEvaluator struct { + root *partitionResidualNode + partitionType *iceberg.StructType +} + +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 { + original iceberg.BoundPredicate + projections []partitionResidualProjection +} + +type partitionResidualProjection struct { + fieldID int + + // strictComplement evaluates the inclusive projection of the negated + // predicate. If it is false, the original predicate is true for every + // value in this partition. + strictComplement func(_ iceberg.StructLike) (bool, error) + inclusive func(_ iceberg.StructLike) (bool, error) +} + +type partitionResidualBuilder struct { + spec iceberg.PartitionSpec + partitionSchema *iceberg.Schema + caseSensitive bool + hasProjection bool + err error +} + +// newPartitionResidualEvaluator creates a Java-style residual evaluator for +// one partition spec. The filter must already be bound to the scan schema. +// Unsupported transforms and predicates simply keep the original predicate. +func newPartitionResidualEvaluator( + schema *iceberg.Schema, + spec *iceberg.PartitionSpec, + filter iceberg.BooleanExpression, + caseSensitive bool, +) (*partitionResidualEvaluator, error) { + if schema == nil || spec == nil || filter == nil { + return nil, nil + } + + partitionType := spec.PartitionType(schema) + builder := &partitionResidualBuilder{ + spec: *spec, + partitionSchema: iceberg.NewSchema(0, partitionType.FieldList...), + caseSensitive: caseSensitive, + } + root, err := iceberg.VisitExpr(filter, builder) + if err != nil { + return nil, err + } + if builder.err != nil { + return nil, builder.err + } + if !builder.hasProjection { + return nil, nil + } + + return &partitionResidualEvaluator{ + root: root, + partitionType: partitionType, + }, nil +} + +func (b *partitionResidualBuilder) VisitTrue() *partitionResidualNode { + return &partitionResidualNode{kind: partitionResidualTrue, expr: iceberg.AlwaysTrue{}} +} + +func (b *partitionResidualBuilder) VisitFalse() *partitionResidualNode { + return &partitionResidualNode{kind: partitionResidualFalse, expr: iceberg.AlwaysFalse{}} +} + +func (b *partitionResidualBuilder) VisitNot(child *partitionResidualNode) *partitionResidualNode { + return &partitionResidualNode{ + kind: partitionResidualNot, + expr: iceberg.NewNot(child.expr), + child: child, + } +} + +func (b *partitionResidualBuilder) VisitAnd(left, right *partitionResidualNode) *partitionResidualNode { + return &partitionResidualNode{ + kind: partitionResidualAnd, + expr: iceberg.NewAnd(left.expr, right.expr), + left: left, + right: right, + } +} + +func (b *partitionResidualBuilder) VisitOr(left, right *partitionResidualNode) *partitionResidualNode { + return &partitionResidualNode{ + kind: partitionResidualOr, + expr: iceberg.NewOr(left.expr, right.expr), + left: left, + right: right, + } +} + +func (b *partitionResidualBuilder) VisitUnbound(pred iceberg.UnboundPredicate) *partitionResidualNode { + return &partitionResidualNode{kind: partitionResidualOpaque, expr: pred} +} + +func (b *partitionResidualBuilder) VisitBound(pred iceberg.BoundPredicate) *partitionResidualNode { + if literal, ok := pred.(iceberg.BoundLiteralPredicate); ok && partitionResidualValueIsNaN(literal.Literal()) { + return &partitionResidualNode{kind: partitionResidualOpaque, expr: pred} + } + + parts := b.spec.FieldsBySourceID(pred.Ref().Field().ID) + if len(parts) == 0 { + return &partitionResidualNode{kind: partitionResidualOpaque, expr: pred} + } + + projections := make([]partitionResidualProjection, 0, len(parts)) + for _, part := range parts { + projection, err := newPartitionResidualProjection( + b.partitionSchema, part, pred, b.caseSensitive) + if err != nil { + if b.err == nil { + b.err = fmt.Errorf("build partition residual for %s: %w", pred, err) + } + + return &partitionResidualNode{kind: partitionResidualOpaque, expr: pred} + } + if projection.strictComplement == nil && projection.inclusive == nil { + continue + } + + projections = append(projections, projection) + } + if len(projections) == 0 { + return &partitionResidualNode{kind: partitionResidualOpaque, expr: pred} + } + + b.hasProjection = true + + return &partitionResidualNode{ + kind: partitionResidualPredicateKind, + expr: pred, + predicate: &partitionResidualPredicate{ + original: pred, + projections: projections, + }, + } +} + +func newPartitionResidualProjection( + partitionSchema *iceberg.Schema, + part iceberg.PartitionField, + pred iceberg.BoundPredicate, + caseSensitive bool, +) (partitionResidualProjection, error) { + projection := partitionResidualProjection{fieldID: part.FieldID} + + negated, ok := pred.Negate().(iceberg.BoundPredicate) + if ok { + strictComplement, err := bindPartitionProjection( + partitionSchema, part, negated, caseSensitive) + if err != nil { + return projection, err + } + projection.strictComplement = strictComplement + } + + inclusive, err := bindPartitionProjection(partitionSchema, part, pred, caseSensitive) + if err != nil { + return projection, err + } + projection.inclusive = inclusive + + return projection, nil +} + +func bindPartitionProjection( + partitionSchema *iceberg.Schema, + part iceberg.PartitionField, + pred iceberg.BoundPredicate, + caseSensitive bool, +) (func(_ iceberg.StructLike) (bool, error), error) { + projected, err := part.Transform.Project(part.Name, pred) + if err != nil || projected == nil { + return nil, err + } + + return iceberg.ExpressionEvaluator(partitionSchema, projected, caseSensitive) +} + +// residual returns the portion of the original filter that still needs to be +// evaluated for a file. The changed result is false when no partition value +// simplified the filter, allowing callers to retain the nil-residual fallback. +func (p *partitionResidualEvaluator) residual( + partition map[int]any, +) (iceberg.BooleanExpression, bool, error) { + residual, changed, err := p.root.residual( + borrowedPartitionRecord{partition: partition, partitionType: p.partitionType}, + ) + if !changed { + return nil, false, err + } + + return residual, true, err +} + +func (n *partitionResidualNode) residual( + partition borrowedPartitionRecord, +) (iceberg.BooleanExpression, bool, error) { + switch n.kind { + case partitionResidualOpaque: + return n.expr, false, nil + case partitionResidualTrue, partitionResidualFalse: + return n.expr, false, nil + case partitionResidualPredicateKind: + for _, projection := range n.predicate.projections { + value, known := partition.partition[projection.fieldID] + if !known { + continue + } + // 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. + op := n.predicate.original.Op() + if op != iceberg.OpIsNull && op != iceberg.OpNotNull && + (value == nil || partitionResidualValueIsNaN(value)) { + continue + } + + if projection.strictComplement != nil { + matches, err := projection.strictComplement(partition) + if err != nil { + return nil, false, err + } + if !matches { + return iceberg.AlwaysTrue{}, true, nil + } + } + + if projection.inclusive != nil { + matches, err := projection.inclusive(partition) + if err != nil { + return nil, false, err + } + if !matches { + return iceberg.AlwaysFalse{}, true, nil + } + } + } + + return n.predicate.original, false, nil + case partitionResidualNot: + child, changed, err := n.child.residual(partition) + if err != nil { + return nil, false, err + } + + return iceberg.NewNot(child), changed, nil + case partitionResidualAnd: + left, leftChanged, err := n.left.residual(partition) + if err != nil { + return nil, false, err + } + if _, alwaysFalse := left.(iceberg.AlwaysFalse); alwaysFalse { + return left, leftChanged, nil + } + + right, rightChanged, err := n.right.residual(partition) + if err != nil { + return nil, false, err + } + + return iceberg.NewAnd(left, right), leftChanged || rightChanged, nil + case partitionResidualOr: + left, leftChanged, err := n.left.residual(partition) + if err != nil { + return nil, false, err + } + if _, alwaysTrue := left.(iceberg.AlwaysTrue); alwaysTrue { + return left, leftChanged, nil + } + + right, rightChanged, err := n.right.residual(partition) + if err != nil { + return nil, false, err + } + + return iceberg.NewOr(left, right), leftChanged || rightChanged, nil + } + + return n.expr, false, nil +} + +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 + } +} diff --git a/table/partition_residual_bench_test.go b/table/partition_residual_bench_test.go new file mode 100644 index 000000000..f7ae51bad --- /dev/null +++ b/table/partition_residual_bench_test.go @@ -0,0 +1,106 @@ +// 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 ( + "testing" + + "github.com/apache/iceberg-go" +) + +var partitionResidualBenchmarkSink iceberg.BooleanExpression + +func BenchmarkPartitionResidualPlanning(b *testing.B) { + for _, tc := range []struct { + name string + transform iceberg.Transform + filter iceberg.BooleanExpression + partition func(int) any + }{ + { + name: "identity", + transform: iceberg.IdentityTransform{}, + filter: iceberg.NewAnd( + iceberg.EqualTo(iceberg.Reference("tenant_id"), "acme"), + iceberg.GreaterThan(iceberg.Reference("amount"), int64(100)), + ), + partition: func(i int) any { + if i%2 == 0 { + return "acme" + } + + return "other" + }, + }, + { + name: "day-range", + transform: iceberg.DayTransform{}, + filter: iceberg.NewAnd( + iceberg.GreaterThanEqual(iceberg.Reference("event_ts"), "2022-11-27T10:00:00"), + iceberg.LessThan(iceberg.Reference("event_ts"), "2022-11-30T10:00:00"), + ), + partition: func(i int) any { + return iceberg.Date(19323 + i%5) + }, + }, + } { + b.Run(tc.name+"/files=4096", func(b *testing.B) { + var schema *iceberg.Schema + if tc.name == "identity" { + schema = partitionResidualTestSchema() + } else { + schema = iceberg.NewSchema(1, iceberg.NestedField{ + ID: 1, Name: "event_ts", Type: iceberg.PrimitiveTypes.Timestamp, + }) + } + spec := iceberg.NewPartitionSpec(iceberg.PartitionField{ + SourceIDs: []int{1}, FieldID: 1000, Name: "partition", Transform: tc.transform, + }) + bound, err := iceberg.BindExpr(schema, tc.filter, true) + if err != nil { + b.Fatal(err) + } + evaluator, err := newPartitionResidualEvaluator(schema, &spec, bound, true) + if err != nil { + b.Fatal(err) + } + if evaluator == nil { + b.Fatal("expected partition residual evaluator") + } + + partitions := make([]map[int]any, 4096) + for i := range partitions { + partitions[i] = map[int]any{1000: tc.partition(i)} + } + + b.ReportAllocs() + b.ResetTimer() + for b.Loop() { + for _, partition := range partitions { + residual, _, err := evaluator.residual(partition) + if err != nil { + b.Fatal(err) + } + partitionResidualBenchmarkSink = residual + } + } + b.StopTimer() + b.ReportMetric(float64(len(partitions)), "files/op") + }) + } +} diff --git a/table/partition_residual_null_test.go b/table/partition_residual_null_test.go new file mode 100644 index 000000000..42003c9c7 --- /dev/null +++ b/table/partition_residual_null_test.go @@ -0,0 +1,151 @@ +// 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{}, iceberg.TruncateTransform{Width: 10}, iceberg.BucketTransform{NumBuckets: 4}}) +} + +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) + residualEvaluator, err := newPartitionResidualEvaluator(schema, &spec, bound, true) + require.NoError(t, err) + residual := bound + if residualEvaluator != nil { + candidate, changed, err := residualEvaluator.residual(map[int]any{1000: value}) + require.NoError(t, err) + if changed { + residual = candidate + } + } + assert.Equal(t, countRows(t, bound), countRows(t, residual), "residual=%s", residual) + }) + } + } + }) + } +} + +func TestPartitionResidualEvaluatorKeepsMissingPartitionValues(t *testing.T) { + schema := iceberg.NewSchema(1, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int32}, + ) + spec := iceberg.NewPartitionSpec(iceberg.PartitionField{ + SourceIDs: []int{1}, FieldID: 1000, Name: "partition", Transform: iceberg.IdentityTransform{}, + }) + for _, filter := range []iceberg.BooleanExpression{ + iceberg.IsNull(iceberg.Reference("id")), + iceberg.NotNull(iceberg.Reference("id")), + iceberg.EqualTo(iceberg.Reference("id"), int32(5)), + } { + t.Run(filter.String(), func(t *testing.T) { + evaluator := boundPartitionResidualEvaluator(t, schema, spec, filter) + residual, changed, err := evaluator.residual(nil) + require.NoError(t, err) + assert.False(t, changed) + assert.Nil(t, residual) + }) + } +} diff --git a/table/partition_residual_read_test.go b/table/partition_residual_read_test.go new file mode 100644 index 000000000..f9bfc6021 --- /dev/null +++ b/table/partition_residual_read_test.go @@ -0,0 +1,198 @@ +// 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/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 TestReadTasksUsesIdentityPartitionResidual(t *testing.T) { + ctx := context.Background() + const tableLocation = "mem://identity-residual" + + 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}, + ) + spec := iceberg.NewPartitionSpec(iceberg.PartitionField{ + SourceIDs: []int{1}, FieldID: 1000, Name: "tenant_id", Transform: 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, + ) + identityFilter := iceberg.EqualTo(iceberg.Reference("tenant_id"), "acme") + amountFilter := iceberg.GreaterThan(iceberg.Reference("amount"), int64(100)) + for _, tt := range []struct { + name string + filter iceberg.BooleanExpression + residual iceberg.BooleanExpression + payloads []string + }{ + { + name: "identity only", + filter: identityFilter, + residual: iceberg.AlwaysTrue{}, + payloads: []string{"low", "keep", "high"}, + }, + { + name: "mixed", + filter: iceberg.NewAnd(identityFilter, amountFilter), + residual: amountFilter, + payloads: []string{"keep", "high"}, + }, + } { + t.Run(tt.name, func(t *testing.T) { + scan := tbl.Scan(WithRowFilter(tt.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, tt.residual, 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, tt.payloads, 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() +} diff --git a/table/partition_residual_test.go b/table/partition_residual_test.go new file mode 100644 index 000000000..58eb8a4da --- /dev/null +++ b/table/partition_residual_test.go @@ -0,0 +1,471 @@ +// 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/decimal128" + "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}, + ) +} + +func boundPartitionResidualEvaluator( + t *testing.T, + schema *iceberg.Schema, + spec iceberg.PartitionSpec, + filter iceberg.BooleanExpression, +) *partitionResidualEvaluator { + t.Helper() + + bound, err := iceberg.BindExpr(schema, filter, true) + require.NoError(t, err) + + evaluator, err := newPartitionResidualEvaluator(schema, &spec, bound, true) + require.NoError(t, err) + require.NotNil(t, evaluator) + + return evaluator +} + +func TestPartitionResidualEvaluatorElidesSatisfiedIdentityPredicate(t *testing.T) { + schema := partitionResidualTestSchema() + spec := iceberg.NewPartitionSpec(iceberg.PartitionField{ + SourceIDs: []int{1}, FieldID: 1000, Name: "tenant_id", Transform: iceberg.IdentityTransform{}, + }) + filter := iceberg.NewAnd( + iceberg.EqualTo(iceberg.Reference("tenant_id"), "acme"), + iceberg.GreaterThan(iceberg.Reference("amount"), int64(100)), + ) + evaluator := boundPartitionResidualEvaluator(t, schema, spec, filter) + + residual, simplified, err := evaluator.residual(map[int]any{1000: "acme"}) + require.NoError(t, err) + require.True(t, simplified) + + want, err := iceberg.BindExpr(schema, + iceberg.GreaterThan(iceberg.Reference("amount"), int64(100)), true) + require.NoError(t, err) + assert.True(t, residual.Equals(want), "expected %s, got %s", want, residual) + + residual, simplified, err = evaluator.residual(map[int]any{1000: "other"}) + require.NoError(t, err) + require.True(t, simplified) + assert.Equal(t, iceberg.AlwaysFalse{}, residual) +} + +func TestPartitionResidualEvaluatorNormalizesDecodedLiteralValues(t *testing.T) { + schema := iceberg.NewSchema(1, iceberg.NestedField{ + ID: 1, Name: "price", Type: iceberg.DecimalTypeOf(10, 2), + }) + spec := iceberg.NewPartitionSpec(iceberg.PartitionField{ + SourceIDs: []int{1}, FieldID: 1000, Name: "price", Transform: iceberg.IdentityTransform{}, + }) + value := iceberg.Decimal{Val: decimal128.FromI64(123), Scale: 2} + evaluator := boundPartitionResidualEvaluator(t, schema, spec, + iceberg.EqualTo(iceberg.Reference("price"), value)) + + residual, simplified, err := evaluator.residual(map[int]any{1000: iceberg.DecimalLiteral(value)}) + require.NoError(t, err) + require.True(t, simplified) + assert.Equal(t, iceberg.AlwaysTrue{}, residual) +} + +func TestPartitionResidualEvaluatorHandlesNestedIdentityFields(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.NewPartitionSpec(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)), + ) + evaluator := boundPartitionResidualEvaluator(t, schema, spec, filter) + + residual, simplified, err := evaluator.residual(map[int]any{1000: "acme"}) + require.NoError(t, err) + require.True(t, simplified) + 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 TestPartitionResidualEvaluatorLeavesMismatchedTransformsUnchanged(t *testing.T) { + schema := partitionResidualTestSchema() + spec := iceberg.NewPartitionSpec(iceberg.PartitionField{ + SourceIDs: []int{1}, FieldID: 1000, Name: "tenant_id", Transform: iceberg.IdentityTransform{}, + }) + filter := iceberg.EqualTo( + iceberg.NewUnboundTransform(iceberg.BucketTransform{NumBuckets: 16}, iceberg.Reference("tenant_id")), + int32(1), + ) + bound, err := iceberg.BindExpr(schema, filter, true) + require.NoError(t, err) + evaluator, err := newPartitionResidualEvaluator(schema, &spec, bound, true) + require.NoError(t, err) + assert.Nil(t, evaluator) +} + +func TestPartitionResidualEvaluatorSimplifiesBooleanCombinations(t *testing.T) { + schema := partitionResidualTestSchema() + spec := iceberg.NewPartitionSpec(iceberg.PartitionField{ + SourceIDs: []int{1}, FieldID: 1000, Name: "tenant_id", Transform: iceberg.IdentityTransform{}, + }) + + t.Run("or", func(t *testing.T) { + filter := iceberg.NewOr( + iceberg.EqualTo(iceberg.Reference("tenant_id"), "acme"), + iceberg.GreaterThan(iceberg.Reference("amount"), int64(100)), + ) + evaluator := boundPartitionResidualEvaluator(t, schema, spec, filter) + + residual, simplified, err := evaluator.residual(map[int]any{1000: "acme"}) + require.NoError(t, err) + require.True(t, simplified) + assert.Equal(t, iceberg.AlwaysTrue{}, residual) + + residual, simplified, err = evaluator.residual(map[int]any{1000: "other"}) + require.NoError(t, err) + require.True(t, simplified) + want, err := iceberg.BindExpr(schema, + iceberg.GreaterThan(iceberg.Reference("amount"), int64(100)), true) + require.NoError(t, err) + assert.True(t, residual.Equals(want)) + }) + + t.Run("not", func(t *testing.T) { + filter := iceberg.NewNot(iceberg.EqualTo(iceberg.Reference("tenant_id"), "acme")) + evaluator := boundPartitionResidualEvaluator(t, schema, spec, filter) + + residual, simplified, err := evaluator.residual(map[int]any{1000: "acme"}) + require.NoError(t, err) + require.True(t, simplified) + assert.Equal(t, iceberg.AlwaysFalse{}, residual) + + residual, simplified, err = evaluator.residual(map[int]any{1000: "other"}) + require.NoError(t, err) + require.True(t, simplified) + assert.Equal(t, iceberg.AlwaysTrue{}, residual) + }) +} + +func TestPartitionResidualEvaluatorHandlesNullAndSetPredicates(t *testing.T) { + schema := partitionResidualTestSchema() + spec := iceberg.NewPartitionSpec(iceberg.PartitionField{ + SourceIDs: []int{1}, FieldID: 1000, Name: "tenant_id", Transform: iceberg.IdentityTransform{}, + }) + + t.Run("is null", func(t *testing.T) { + evaluator := boundPartitionResidualEvaluator(t, schema, spec, + iceberg.IsNull(iceberg.Reference("tenant_id"))) + + residual, simplified, err := evaluator.residual(map[int]any{1000: nil}) + require.NoError(t, err) + require.True(t, simplified) + assert.Equal(t, iceberg.AlwaysTrue{}, residual) + + residual, simplified, err = evaluator.residual(map[int]any{1000: "acme"}) + require.NoError(t, err) + require.True(t, simplified) + assert.Equal(t, iceberg.AlwaysFalse{}, residual) + }) + + t.Run("in", func(t *testing.T) { + evaluator := boundPartitionResidualEvaluator(t, schema, spec, + iceberg.IsIn(iceberg.Reference("tenant_id"), "acme", "iceberg")) + + residual, simplified, err := evaluator.residual(map[int]any{1000: "acme"}) + require.NoError(t, err) + require.True(t, simplified) + assert.Equal(t, iceberg.AlwaysTrue{}, residual) + + residual, simplified, err = evaluator.residual(map[int]any{1000: "other"}) + require.NoError(t, err) + require.True(t, simplified) + assert.Equal(t, iceberg.AlwaysFalse{}, residual) + }) +} + +func TestPartitionResidualEvaluatorComputesDayBoundaries(t *testing.T) { + schema := iceberg.NewSchema(1, + iceberg.NestedField{ID: 1, Name: "event_ts", Type: iceberg.PrimitiveTypes.Timestamp}, + ) + spec := iceberg.NewPartitionSpec(iceberg.PartitionField{ + SourceIDs: []int{1}, FieldID: 1000, Name: "event_day", Transform: iceberg.DayTransform{}, + }) + filter := iceberg.NewAnd( + iceberg.GreaterThanEqual(iceberg.Reference("event_ts"), "2022-11-27T10:00:00"), + iceberg.LessThan(iceberg.Reference("event_ts"), "2022-11-30T10:00:00"), + ) + evaluator := boundPartitionResidualEvaluator(t, schema, spec, filter) + + lower, err := iceberg.BindExpr(schema, + iceberg.GreaterThanEqual(iceberg.Reference("event_ts"), "2022-11-27T10:00:00"), true) + require.NoError(t, err) + upper, err := iceberg.BindExpr(schema, + iceberg.LessThan(iceberg.Reference("event_ts"), "2022-11-30T10:00:00"), true) + require.NoError(t, err) + + tests := []struct { + name string + partition iceberg.Date + want iceberg.BooleanExpression + }{ + {name: "lower boundary", partition: iceberg.Date(19323), want: lower}, + {name: "interior", partition: iceberg.Date(19324), want: iceberg.AlwaysTrue{}}, + {name: "upper boundary", partition: iceberg.Date(19326), want: upper}, + {name: "outside", partition: iceberg.Date(19327), want: iceberg.AlwaysFalse{}}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + residual, simplified, err := evaluator.residual(map[int]any{1000: tt.partition}) + require.NoError(t, err) + require.True(t, simplified) + assert.True(t, residual.Equals(tt.want), "expected %s, got %s", tt.want, residual) + }) + } +} + +func TestPartitionResidualEvaluatorComputesTruncateBoundaries(t *testing.T) { + schema := iceberg.NewSchema(1, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int32}, + ) + spec := iceberg.NewPartitionSpec(iceberg.PartitionField{ + SourceIDs: []int{1}, FieldID: 1000, Name: "id_truncated", Transform: iceberg.TruncateTransform{Width: 10}, + }) + filter := iceberg.NewAnd( + iceberg.GreaterThanEqual(iceberg.Reference("id"), int32(25)), + iceberg.LessThan(iceberg.Reference("id"), int32(65)), + ) + evaluator := boundPartitionResidualEvaluator(t, schema, spec, filter) + + lower, err := iceberg.BindExpr(schema, + iceberg.GreaterThanEqual(iceberg.Reference("id"), int32(25)), true) + require.NoError(t, err) + upper, err := iceberg.BindExpr(schema, + iceberg.LessThan(iceberg.Reference("id"), int32(65)), true) + require.NoError(t, err) + + tests := []struct { + name string + partition int32 + want iceberg.BooleanExpression + }{ + {name: "before range", partition: 10, want: iceberg.AlwaysFalse{}}, + {name: "lower boundary", partition: 20, want: lower}, + {name: "interior", partition: 40, want: iceberg.AlwaysTrue{}}, + {name: "upper boundary", partition: 60, want: upper}, + {name: "after range", partition: 70, want: iceberg.AlwaysFalse{}}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + residual, simplified, err := evaluator.residual(map[int]any{1000: tt.partition}) + require.NoError(t, err) + require.True(t, simplified) + assert.True(t, residual.Equals(tt.want), "expected %s, got %s", tt.want, residual) + }) + } +} + +func TestPartitionResidualEvaluatorHandlesTransformedFilterTerms(t *testing.T) { + schema := iceberg.NewSchema(1, + iceberg.NestedField{ID: 1, Name: "event_ts", Type: iceberg.PrimitiveTypes.Timestamp}, + ) + spec := iceberg.NewPartitionSpec(iceberg.PartitionField{ + SourceIDs: []int{1}, FieldID: 1000, Name: "event_day", Transform: iceberg.DayTransform{}, + }) + filter := iceberg.EqualTo( + iceberg.NewUnboundTransform(iceberg.DayTransform{}, iceberg.Reference("event_ts")), + iceberg.Date(19323), + ) + evaluator := boundPartitionResidualEvaluator(t, schema, spec, filter) + + residual, simplified, err := evaluator.residual(map[int]any{1000: iceberg.Date(19323)}) + require.NoError(t, err) + require.True(t, simplified) + assert.Equal(t, iceberg.AlwaysTrue{}, residual) + + residual, simplified, err = evaluator.residual(map[int]any{1000: iceberg.Date(19324)}) + require.NoError(t, err) + require.True(t, simplified) + assert.Equal(t, iceberg.AlwaysFalse{}, residual) +} + +func TestPartitionResidualEvaluatorKeepsBucketPredicateConservative(t *testing.T) { + schema := partitionResidualTestSchema() + transform := iceberg.BucketTransform{NumBuckets: 16} + spec := iceberg.NewPartitionSpec(iceberg.PartitionField{ + SourceIDs: []int{1}, FieldID: 1000, Name: "tenant_bucket", Transform: transform, + }) + filter := iceberg.EqualTo(iceberg.Reference("tenant_id"), "acme") + evaluator := boundPartitionResidualEvaluator(t, schema, spec, filter) + + partition := transform.Apply(iceberg.Optional[iceberg.Literal]{ + Valid: true, Val: iceberg.StringLiteral("acme"), + }) + require.True(t, partition.Valid) + residual, simplified, err := evaluator.residual(map[int]any{1000: partition.Val.Any()}) + require.NoError(t, err) + assert.False(t, simplified, "a bucket match does not prove equality because buckets can collide") + assert.Nil(t, residual) +} + +func TestPartitionResidualEvaluatorLeavesUnpartitionedFiltersUnchanged(t *testing.T) { + schema := partitionResidualTestSchema() + spec := iceberg.NewPartitionSpec(iceberg.PartitionField{ + SourceIDs: []int{1}, FieldID: 1000, Name: "tenant_bucket", + Transform: iceberg.BucketTransform{NumBuckets: 16}, + }) + filter := iceberg.GreaterThan(iceberg.Reference("amount"), int64(100)) + + bound, err := iceberg.BindExpr(schema, filter, true) + require.NoError(t, err) + evaluator, err := newPartitionResidualEvaluator(schema, &spec, bound, true) + require.NoError(t, err) + assert.Nil(t, evaluator) +} + +func TestPlanFilesLocalPopulatesPartitionResidual(t *testing.T) { + const ( + manifestPath = "mem://default/table/metadata/manifest.avro" + manifestListPath = "mem://default/table/metadata/snap-7.avro" + dataPath = "mem://default/table/data/file.parquet" + ) + + spec := partitionedSpec() + fs := iceio.NewMemFS() + scan, schema, snapshotID := newSchemaEvolutionScanWithSnapshot( + t, &spec, fs, manifestListPath, nil) + dataFile := newTestDataFile(t, spec, dataPath, map[int]any{1000: int32(5)}) + entry := iceberg.NewManifestEntryBuilder( + iceberg.EntryStatusADDED, &snapshotID, dataFile, + ).SequenceNum(1).Build() + + var manifestBuffer bytes.Buffer + manifest, err := iceberg.WriteManifest( + manifestPath, &manifestBuffer, 2, spec, schema, snapshotID, + []iceberg.ManifestEntry{entry}, + ) + require.NoError(t, err) + require.NoError(t, fs.WriteFile(manifestPath, manifestBuffer.Bytes())) + + var listBuffer bytes.Buffer + sequenceNumber := int64(1) + require.NoError(t, iceberg.WriteManifestList( + 2, &listBuffer, snapshotID, nil, &sequenceNumber, 0, + []iceberg.ManifestFile{manifest}, + )) + require.NoError(t, fs.WriteFile(manifestListPath, listBuffer.Bytes())) + + tasks, err := scan.PlanFiles(context.Background()) + require.NoError(t, err) + require.Len(t, tasks, 1) + assert.Equal(t, iceberg.AlwaysTrue{}, tasks[0].Residual, + "the identity partition already proves id == 5") +} + +func TestPartitionResidualEvaluatorPreservesRowsAtTransformBoundaries(t *testing.T) { + schema := iceberg.NewSchema(1, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int32}, + ) + ref := iceberg.Reference("id") + filters := []iceberg.BooleanExpression{iceberg.IsNull(ref), iceberg.NotNull(ref)} + for _, boundary := range []int32{-10, -1, 0, 1, 10} { + for _, op := range []iceberg.Operation{ + iceberg.OpEQ, iceberg.OpNEQ, iceberg.OpLT, + iceberg.OpLTEQ, iceberg.OpGT, iceberg.OpGTEQ, + } { + filters = append(filters, iceberg.LiteralPredicate(op, ref, iceberg.Int32Literal(boundary))) + } + } + filters = append(filters, + iceberg.IsIn(ref, int32(-10), int32(0), int32(10)), + iceberg.NotIn(ref, int32(-10), int32(0), int32(10)), + ) + values := []any{ + nil, int32(-21), int32(-20), int32(-11), int32(-10), int32(-1), + int32(0), int32(1), int32(9), int32(10), int32(11), int32(20), int32(21), + } + for _, transform := range []iceberg.Transform{ + iceberg.IdentityTransform{}, + iceberg.TruncateTransform{Width: 10}, + iceberg.BucketTransform{NumBuckets: 4}, + } { + t.Run(transform.String(), func(t *testing.T) { + spec := iceberg.NewPartitionSpec(iceberg.PartitionField{ + SourceIDs: []int{1}, FieldID: 1000, Name: "partition", Transform: transform, + }) + for _, filter := range filters { + t.Run(filter.String(), func(t *testing.T) { + bound, err := iceberg.BindExpr(schema, filter, true) + require.NoError(t, err) + residualEvaluator, err := newPartitionResidualEvaluator(schema, &spec, bound, true) + require.NoError(t, err) + if residualEvaluator == nil { + return + } + evaluate, err := iceberg.ExpressionEvaluator(schema, filter, true) + require.NoError(t, err) + for _, value := range values { + var literal iceberg.Optional[iceberg.Literal] + if value != nil { + literal = iceberg.Optional[iceberg.Literal]{Valid: true, Val: iceberg.Int32Literal(value.(int32))} + } + partitionLiteral := transform.Apply(literal) + var partition any + if partitionLiteral.Valid { + partition = partitionLiteral.Val.Any() + } + residual, changed, err := residualEvaluator.residual(map[int]any{1000: partition}) + require.NoError(t, err) + if !changed { + continue + } + unbound, err := iceberg.TranslateColumnNames(residual, schema) + require.NoError(t, err) + evaluateResidual, err := iceberg.ExpressionEvaluator(schema, unbound, true) + require.NoError(t, err) + want, err := evaluate(partitionRecord{value}) + require.NoError(t, err) + got, err := evaluateResidual(partitionRecord{value}) + require.NoError(t, err) + assert.Equal(t, want, got, "value=%v partition=%v residual=%s", value, partition, residual) + } + }) + } + }) + } +} diff --git a/table/scanner.go b/table/scanner.go index f5599d932..f9a1f186b 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 && !scan.rowFilter.Equals(iceberg.AlwaysTrue{}) { + boundRowFilter, err = iceberg.BindExpr(schema, scan.rowFilter, scan.caseSensitive) + if err != nil { + return nil, err + } + } + var residualEvaluators map[int]*partitionResidualEvaluator + if boundRowFilter != nil { + residualEvaluators = make(map[int]*partitionResidualEvaluator) + } + // Step 3: Index positional deletes and match them to data files. posDeleteIndex, err := buildPositionalDeleteIndex(entries.positionalDeleteEntries) if err != nil { @@ -1161,6 +1173,28 @@ func (scan *Scan) planFilesLocal(ctx context.Context, acc *scanMetricsAccumulato Start: 0, Length: e.DataFile().FileSizeBytes(), } + if boundRowFilter != nil { + specID := int(e.DataFile().SpecID()) + residualEvaluator, found := residualEvaluators[specID] + if !found { + residualEvaluator, err = newPartitionResidualEvaluator( + schema, scan.metadata.PartitionSpecByID(specID), boundRowFilter, scan.caseSensitive) + if err != nil { + return nil, fmt.Errorf("build partition residual evaluator for spec %d: %w", specID, err) + } + residualEvaluators[specID] = residualEvaluator + } + if residualEvaluator != nil { + var simplified bool + task.Residual, simplified, err = residualEvaluator.residual(dataFilePartition(e.DataFile())) + if err != nil { + return nil, fmt.Errorf("evaluate partition residual for %s: %w", e.DataFile().FilePath(), err) + } + if !simplified { + task.Residual = nil + } + } + } // 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 +1469,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 and 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