From a5baa1f8918c848fb93a39d47d52b43e57eaed4d Mon Sep 17 00:00:00 2001 From: Neelesh Salian Date: Fri, 24 Jul 2026 12:03:32 -0700 Subject: [PATCH] feat(table): variant predicate pushdown and reject variant partition source --- partitions.go | 3 + partitions_test.go | 12 ++ table/arrow_scanner.go | 32 ++- table/evaluators.go | 177 +++++++---------- table/evaluators_test.go | 66 +++++++ table/internal/variant_bounds.go | 89 ++++----- table/internal/variant_bounds_test.go | 73 +++++-- table/variant_residual.go | 193 ++++++++++++++++++ table/variant_residual_test.go | 69 +++++++ table/variant_shredded_write_test.go | 259 +++++++++++++++++++++++++ variant_cast.go | 269 ++++++++++++++++++++++++++ variant_cast_test.go | 77 ++++++++ variant_extract.go | 248 ++++++++++++++++++++++++ variant_extract_test.go | 61 ++++++ variant_path.go | 132 +++++++++++++ variant_path_test.go | 85 ++++++++ visitors.go | 100 ++++++++-- 17 files changed, 1731 insertions(+), 214 deletions(-) create mode 100644 table/variant_residual.go create mode 100644 table/variant_residual_test.go create mode 100644 variant_cast.go create mode 100644 variant_cast_test.go create mode 100644 variant_extract.go create mode 100644 variant_extract_test.go create mode 100644 variant_path.go create mode 100644 variant_path_test.go diff --git a/partitions.go b/partitions.go index 6d2364865..7d09ebbbc 100644 --- a/partitions.go +++ b/partitions.go @@ -260,6 +260,9 @@ func (p *PartitionSpec) addSpecFieldInternal(targetName string, field NestedFiel if err := validateTransform(transform); err != nil { return err } + if _, ok := field.Type.(VariantType); ok { + return fmt.Errorf("%w: cannot partition by %s source field: %s", ErrInvalidArgument, field.Type, targetName) + } for _, existingField := range p.fields { if existingField.Name == targetName { return errors.New("duplicate partition name: " + targetName) diff --git a/partitions_test.go b/partitions_test.go index f0abdd588..47ef7cf96 100644 --- a/partitions_test.go +++ b/partitions_test.go @@ -176,6 +176,18 @@ func TestPartitionSpecRejectsInvalidBucketTransform(t *testing.T) { require.ErrorContains(t, err, "numBuckets > 0") } +func TestPartitionSpecRejectsVariantSource(t *testing.T) { + schema := iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "v", Type: iceberg.VariantType{}}, + ) + + _, err := iceberg.NewPartitionSpecOpts( + iceberg.AddPartitionFieldBySourceID(1, "v_p", iceberg.IdentityTransform{}, schema, nil), + ) + require.ErrorIs(t, err, iceberg.ErrInvalidArgument) + require.ErrorContains(t, err, "cannot partition by") +} + func TestPartitionSpec_MarshalTextRejectsInvalidBucketTransform(t *testing.T) { spec := iceberg.NewPartitionSpecID(3, iceberg.PartitionField{ diff --git a/table/arrow_scanner.go b/table/arrow_scanner.go index 02da7fa1d..f12f42a0b 100644 --- a/table/arrow_scanner.go +++ b/table/arrow_scanner.go @@ -727,7 +727,7 @@ func (as *arrowScan) getRecordFilter(ctx context.Context, fileSchema *iceberg.Sc return nil, false, nil } - translatedFilter, err := iceberg.TranslateColumnNames(as.boundRowFilter, fileSchema) + translatedFilter, extracts, err := iceberg.TranslateColumnNamesForScan(as.boundRowFilter, fileSchema) if err != nil { return nil, false, err } @@ -736,23 +736,35 @@ func (as *arrowScan) getRecordFilter(ctx context.Context, fileSchema *iceberg.Sc return nil, true, nil } - translatedFilter, err = iceberg.BindExpr(fileSchema, translatedFilter, as.caseSensitive) + filterSchema := fileSchema + if len(extracts) > 0 { + filterSchema, err = augmentSchemaWithExtracts(fileSchema, extracts) + if err != nil { + return nil, false, err + } + } + + translatedFilter, err = iceberg.BindExpr(filterSchema, translatedFilter, as.caseSensitive) if err != nil { return nil, false, err } - if !translatedFilter.Equals(iceberg.AlwaysTrue{}) { - extSet, recordFilter, err := substrait.ConvertExpr(fileSchema, translatedFilter, as.caseSensitive) - if err != nil { - return nil, false, err - } + if translatedFilter.Equals(iceberg.AlwaysTrue{}) { + return nil, false, nil + } - ctx = exprs.WithExtensionIDSet(ctx, exprs.NewExtensionSetDefault(*extSet)) + extSet, recordFilter, err := substrait.ConvertExpr(filterSchema, translatedFilter, as.caseSensitive) + if err != nil { + return nil, false, err + } - return filterRecords(ctx, recordFilter), false, nil + ctx = exprs.WithExtensionIDSet(ctx, exprs.NewExtensionSetDefault(*extSet)) + base := filterRecords(ctx, recordFilter) + if len(extracts) == 0 { + return base, false, nil } - return nil, false, nil + return as.extractResidualFilter(ctx, extracts, base), false, nil } // fieldIndexByID returns the index of the field carrying fieldID in its Arrow diff --git a/table/evaluators.go b/table/evaluators.go index bebbcdbcd..6d278ec6e 100644 --- a/table/evaluators.go +++ b/table/evaluators.go @@ -828,6 +828,10 @@ func (m *inclusiveMetricsEval) VisitBound(pred iceberg.BoundPredicate) bool { } func (m *inclusiveMetricsEval) VisitIsNull(t iceberg.BoundTerm) bool { + if _, ok := t.(iceberg.BoundExtract); ok { + return rowsMightMatch + } + fieldID := t.Ref().Field().ID if cnt, exists := m.nullCounts[fieldID]; exists && cnt == 0 { return rowsCannotMatch @@ -848,6 +852,10 @@ func (m *inclusiveMetricsEval) VisitNotNull(t iceberg.BoundTerm) bool { } func (m *inclusiveMetricsEval) VisitIsNan(t iceberg.BoundTerm) bool { + if _, ok := t.(iceberg.BoundExtract); ok { + return rowsMightMatch + } + fieldID := t.Ref().Field().ID if cnt, exists := m.nanCounts[fieldID]; exists && cnt == 0 { return rowsCannotMatch @@ -862,6 +870,10 @@ func (m *inclusiveMetricsEval) VisitIsNan(t iceberg.BoundTerm) bool { } func (m *inclusiveMetricsEval) VisitNotNan(t iceberg.BoundTerm) bool { + if _, ok := t.(iceberg.BoundExtract); ok { + return rowsMightMatch + } + fieldID := t.Ref().Field().ID if m.containsNansOnly(fieldID) { @@ -871,25 +883,44 @@ func (m *inclusiveMetricsEval) VisitNotNan(t iceberg.BoundTerm) bool { return rowsMightMatch } -func (m *inclusiveMetricsEval) VisitLess(t iceberg.BoundTerm, lit iceberg.Literal) bool { - field := t.Ref().Field() - fieldID := field.ID +// boundFor decodes the file bound for term t from raw: a scalar for a reference, or the +// variant sub-path value for an extract; ok is false when raw is nil or not castable. +func (m *inclusiveMetricsEval) boundFor(t iceberg.BoundTerm, raw []byte) (iceberg.Literal, bool) { + if raw == nil { + return nil, false + } - if m.containsNullsOnly(fieldID) || m.containsNansOnly(fieldID) { - return rowsCannotMatch + if ext, ok := t.(iceberg.BoundExtract); ok { + lit, found, err := internal.VariantBoundLiteral(raw, ext.Path(), ext.Type().(iceberg.PrimitiveType)) + if err != nil { + panic(err) + } + + return lit, found } + field := t.Ref().Field() if _, ok := field.Type.(iceberg.PrimitiveType); !ok { panic(fmt.Errorf("%w: expected iceberg.PrimitiveType, got %s", iceberg.ErrInvalidTypeString, field.Type)) } - if lowerBoundBytes := m.lowerBounds[fieldID]; lowerBoundBytes != nil { - lowerBound, err := iceberg.LiteralFromBytes(field.Type, lowerBoundBytes) - if err != nil { - panic(err) - } + lit, err := iceberg.LiteralFromBytes(field.Type, raw) + if err != nil { + panic(err) + } + + return lit, true +} + +func (m *inclusiveMetricsEval) VisitLess(t iceberg.BoundTerm, lit iceberg.Literal) bool { + fieldID := t.Ref().Field().ID + + if m.containsNullsOnly(fieldID) || m.containsNansOnly(fieldID) { + return rowsCannotMatch + } + if lowerBound, ok := m.boundFor(t, m.lowerBounds[fieldID]); ok { if m.isNan(lowerBound) { // nan indicates unreliable bounds return rowsMightMatch @@ -904,24 +935,13 @@ func (m *inclusiveMetricsEval) VisitLess(t iceberg.BoundTerm, lit iceberg.Litera } func (m *inclusiveMetricsEval) VisitLessEqual(t iceberg.BoundTerm, lit iceberg.Literal) bool { - field := t.Ref().Field() - fieldID := field.ID + fieldID := t.Ref().Field().ID if m.containsNullsOnly(fieldID) || m.containsNansOnly(fieldID) { return rowsCannotMatch } - if _, ok := field.Type.(iceberg.PrimitiveType); !ok { - panic(fmt.Errorf("%w: expected iceberg.PrimitiveType, got %s", - iceberg.ErrInvalidTypeString, field.Type)) - } - - if lowerBoundBytes := m.lowerBounds[fieldID]; lowerBoundBytes != nil { - lowerBound, err := iceberg.LiteralFromBytes(field.Type, lowerBoundBytes) - if err != nil { - panic(err) - } - + if lowerBound, ok := m.boundFor(t, m.lowerBounds[fieldID]); ok { if m.isNan(lowerBound) { // nan indicates unreliable bounds return rowsMightMatch @@ -936,24 +956,13 @@ func (m *inclusiveMetricsEval) VisitLessEqual(t iceberg.BoundTerm, lit iceberg.L } func (m *inclusiveMetricsEval) VisitGreater(t iceberg.BoundTerm, lit iceberg.Literal) bool { - field := t.Ref().Field() - fieldID := field.ID + fieldID := t.Ref().Field().ID if m.containsNullsOnly(fieldID) || m.containsNansOnly(fieldID) { return rowsCannotMatch } - if _, ok := field.Type.(iceberg.PrimitiveType); !ok { - panic(fmt.Errorf("%w: expected iceberg.PrimitiveType, got %s", - iceberg.ErrInvalidTypeString, field.Type)) - } - - if upperBoundBytes := m.upperBounds[fieldID]; upperBoundBytes != nil { - upperBound, err := iceberg.LiteralFromBytes(field.Type, upperBoundBytes) - if err != nil { - panic(err) - } - + if upperBound, ok := m.boundFor(t, m.upperBounds[fieldID]); ok { if getCmpLiteral(upperBound)(upperBound, lit) <= 0 { if m.isNan(upperBound) { return rowsMightMatch @@ -967,24 +976,13 @@ func (m *inclusiveMetricsEval) VisitGreater(t iceberg.BoundTerm, lit iceberg.Lit } func (m *inclusiveMetricsEval) VisitGreaterEqual(t iceberg.BoundTerm, lit iceberg.Literal) bool { - field := t.Ref().Field() - fieldID := field.ID + fieldID := t.Ref().Field().ID if m.containsNullsOnly(fieldID) || m.containsNansOnly(fieldID) { return rowsCannotMatch } - if _, ok := field.Type.(iceberg.PrimitiveType); !ok { - panic(fmt.Errorf("%w: expected iceberg.PrimitiveType, got %s", - iceberg.ErrInvalidTypeString, field.Type)) - } - - if upperBoundBytes := m.upperBounds[fieldID]; upperBoundBytes != nil { - upperBound, err := iceberg.LiteralFromBytes(field.Type, upperBoundBytes) - if err != nil { - panic(err) - } - + if upperBound, ok := m.boundFor(t, m.upperBounds[fieldID]); ok { if getCmpLiteral(upperBound)(upperBound, lit) < 0 { if m.isNan(upperBound) { return rowsMightMatch @@ -998,47 +996,28 @@ func (m *inclusiveMetricsEval) VisitGreaterEqual(t iceberg.BoundTerm, lit iceber } func (m *inclusiveMetricsEval) VisitEqual(t iceberg.BoundTerm, lit iceberg.Literal) bool { - field := t.Ref().Field() - fieldID := field.ID + fieldID := t.Ref().Field().ID if m.containsNullsOnly(fieldID) || m.containsNansOnly(fieldID) { return rowsCannotMatch } - if _, ok := field.Type.(iceberg.PrimitiveType); !ok { - panic(fmt.Errorf("%w: expected iceberg.PrimitiveType, got %s", - iceberg.ErrInvalidTypeString, field.Type)) - } - - var cmp func(iceberg.Literal, iceberg.Literal) int - if lowerBoundBytes := m.lowerBounds[fieldID]; lowerBoundBytes != nil { - lowerBound, err := iceberg.LiteralFromBytes(field.Type, lowerBoundBytes) - if err != nil { - panic(err) - } - + if lowerBound, ok := m.boundFor(t, m.lowerBounds[fieldID]); ok { if m.isNan(lowerBound) { return rowsMightMatch } - cmp = getCmpLiteral(lowerBound) - if cmp(lowerBound, lit) == 1 { + if getCmpLiteral(lowerBound)(lowerBound, lit) == 1 { return rowsCannotMatch } } - if upperBoundBytes := m.upperBounds[fieldID]; upperBoundBytes != nil { - upperBound, err := iceberg.LiteralFromBytes(field.Type, upperBoundBytes) - if err != nil { - panic(err) - } - + if upperBound, ok := m.boundFor(t, m.upperBounds[fieldID]); ok { if m.isNan(upperBound) { return rowsMightMatch } - cmp = getCmpLiteral(upperBound) - if cmp(upperBound, lit) == -1 { + if getCmpLiteral(upperBound)(upperBound, lit) == -1 { return rowsCannotMatch } } @@ -1051,8 +1030,7 @@ func (m *inclusiveMetricsEval) VisitNotEqual(iceberg.BoundTerm, iceberg.Literal) } func (m *inclusiveMetricsEval) VisitIn(t iceberg.BoundTerm, s iceberg.Set[iceberg.Literal]) bool { - field := t.Ref().Field() - fieldID := field.ID + fieldID := t.Ref().Field().ID if m.containsNullsOnly(fieldID) || m.containsNansOnly(fieldID) { return rowsCannotMatch @@ -1063,18 +1041,8 @@ func (m *inclusiveMetricsEval) VisitIn(t iceberg.BoundTerm, s iceberg.Set[iceber return rowsMightMatch } - if _, ok := field.Type.(iceberg.PrimitiveType); !ok { - panic(fmt.Errorf("%w: expected iceberg.PrimitiveType, got %s", - iceberg.ErrInvalidTypeString, field.Type)) - } - values := s.Members() - if lowerBoundBytes := m.lowerBounds[fieldID]; lowerBoundBytes != nil { - lowerBound, err := iceberg.LiteralFromBytes(field.Type, lowerBoundBytes) - if err != nil { - panic(lowerBound) - } - + if lowerBound, ok := m.boundFor(t, m.lowerBounds[fieldID]); ok { if m.isNan(lowerBound) { return rowsMightMatch } @@ -1085,12 +1053,7 @@ func (m *inclusiveMetricsEval) VisitIn(t iceberg.BoundTerm, s iceberg.Set[iceber } } - if upperBoundBytes := m.upperBounds[fieldID]; upperBoundBytes != nil { - upperBound, err := iceberg.LiteralFromBytes(field.Type, upperBoundBytes) - if err != nil { - panic(err) - } - + if upperBound, ok := m.boundFor(t, m.upperBounds[fieldID]); ok { if m.isNan(upperBound) { return rowsMightMatch } @@ -1112,18 +1075,12 @@ func (m *inclusiveMetricsEval) VisitNotIn(iceberg.BoundTerm, iceberg.Set[iceberg } func (m *inclusiveMetricsEval) VisitStartsWith(t iceberg.BoundTerm, lit iceberg.Literal) bool { - field := t.Ref().Field() - fieldID := field.ID + fieldID := t.Ref().Field().ID if m.containsNullsOnly(fieldID) { return rowsCannotMatch } - if _, ok := field.Type.(iceberg.PrimitiveType); !ok { - panic(fmt.Errorf("%w: expected iceberg.PrimitiveType, got %s", - iceberg.ErrInvalidTypeString, field.Type)) - } - var prefix string if val, ok := lit.(iceberg.TypedLiteral[string]); ok { prefix = val.Value() @@ -1133,12 +1090,7 @@ func (m *inclusiveMetricsEval) VisitStartsWith(t iceberg.BoundTerm, lit iceberg. lenPrefix := len(prefix) - if lowerBoundBytes := m.lowerBounds[fieldID]; lowerBoundBytes != nil { - lowerBound, err := iceberg.LiteralFromBytes(field.Type, lowerBoundBytes) - if err != nil { - panic(err) - } - + if lowerBound, ok := m.boundFor(t, m.lowerBounds[fieldID]); ok { var v string switch l := lowerBound.(type) { case iceberg.TypedLiteral[string]: @@ -1156,12 +1108,7 @@ func (m *inclusiveMetricsEval) VisitStartsWith(t iceberg.BoundTerm, lit iceberg. } } - if upperBoundBytes := m.upperBounds[fieldID]; upperBoundBytes != nil { - upperBound, err := iceberg.LiteralFromBytes(field.Type, upperBoundBytes) - if err != nil { - panic(err) - } - + if upperBound, ok := m.boundFor(t, m.upperBounds[fieldID]); ok { var v string switch u := upperBound.(type) { case iceberg.TypedLiteral[string]: @@ -1183,6 +1130,10 @@ func (m *inclusiveMetricsEval) VisitStartsWith(t iceberg.BoundTerm, lit iceberg. } func (m *inclusiveMetricsEval) VisitNotStartsWith(t iceberg.BoundTerm, lit iceberg.Literal) bool { + if _, ok := t.(iceberg.BoundExtract); ok { + return rowsMightMatch + } + field := t.Ref().Field() fieldID := field.ID @@ -1289,6 +1240,10 @@ func (m *strictMetricsEval) VisitUnbound(iceberg.UnboundPredicate) bool { } func (m *strictMetricsEval) VisitBound(pred iceberg.BoundPredicate) bool { + if _, ok := pred.Term().(iceberg.BoundExtract); ok { + return rowsMightNotMatch + } + return iceberg.VisitBoundPredicate(pred, m) } diff --git a/table/evaluators_test.go b/table/evaluators_test.go index ef0c82d55..7761f3c3b 100644 --- a/table/evaluators_test.go +++ b/table/evaluators_test.go @@ -22,6 +22,7 @@ import ( "math" "testing" + "github.com/apache/arrow-go/v18/parquet/variant" "github.com/apache/iceberg-go" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -3120,3 +3121,68 @@ func TestBloomPredicateCollector(t *testing.T) { assert.Empty(t, preds) }) } + +// TestVariantExtractPruningGuards confirms extract terms are never pruned by null/nan-count guards. +func TestVariantExtractPruningGuards(t *testing.T) { + sc := iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64, Required: true}, + iceberg.NestedField{ID: 2, Name: "payload", Type: iceberg.VariantType{}}, + ) + extI := iceberg.Extract("payload", "$.a", iceberg.PrimitiveTypes.Int64) + extF := iceberg.Extract("payload", "$.a", iceberg.PrimitiveTypes.Float64) + extS := iceberg.Extract("payload", "$.b", iceberg.PrimitiveTypes.String) + + for _, tt := range []struct { + name string + expr iceberg.BooleanExpression + file *mockDataFile + }{ + {"is_null", iceberg.IsNull(extI), &mockDataFile{count: 10, valueCounts: map[int]int64{2: 10}, nullCounts: map[int]int64{2: 0}}}, + {"is_nan", iceberg.IsNaN(extF), &mockDataFile{count: 10, valueCounts: map[int]int64{2: 10}, nanCounts: map[int]int64{2: 0}}}, + {"not_nan", iceberg.NotNaN(extF), &mockDataFile{count: 10, valueCounts: map[int]int64{2: 10}, nanCounts: map[int]int64{2: 10}}}, + {"not_starts_with", iceberg.NotStartsWith(extS, "r"), &mockDataFile{count: 10, valueCounts: map[int]int64{2: 10}, nullCounts: map[int]int64{99: 0}}}, + } { + t.Run(tt.name, func(t *testing.T) { + eval, err := newInclusiveMetricsEvaluator(sc, tt.expr, true, true) + require.NoError(t, err) + res, err := eval(tt.file) + require.NoError(t, err) + assert.Equal(t, rowsMightMatch, res, "variant extract guard must never prune") + }) + } +} + +// TestStrictMetricsExtractNeverMatches confirms the strict evaluator treats extract terms +// conservatively (rowsMightNotMatch) instead of decoding the variant bytes and erroring. +func TestStrictMetricsExtractNeverMatches(t *testing.T) { + sc := iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64, Required: true}, + iceberg.NestedField{ID: 2, Name: "payload", Type: iceberg.VariantType{}}, + ) + + var b variant.Builder + start := b.Offset() + entries := []variant.FieldEntry{b.NextField(start, "$['a']")} + require.NoError(t, b.AppendInt(5)) + require.NoError(t, b.FinishObject(start, entries)) + v, err := b.Build() + require.NoError(t, err) + bound := append(append([]byte{}, v.Metadata().Bytes()...), v.Bytes()...) + + // A file whose variant column carries bounds is what triggered the LiteralFromBytes(variant) error. + file := &mockDataFile{ + count: 10, + valueCounts: map[int]int64{2: 10}, + nullCounts: map[int]int64{2: 0}, + lowerBounds: map[int][]byte{2: bound}, + upperBounds: map[int][]byte{2: bound}, + } + + expr := iceberg.EqualTo(iceberg.Extract("payload", "$.a", iceberg.PrimitiveTypes.Int64), int64(5)) + eval, err := newStrictMetricsEvaluator(sc, expr, true, true) + require.NoError(t, err) + + res, err := eval(file) + require.NoError(t, err) + assert.Equal(t, rowsMightNotMatch, res) +} diff --git a/table/internal/variant_bounds.go b/table/internal/variant_bounds.go index 241f7b25c..c21802c36 100644 --- a/table/internal/variant_bounds.go +++ b/table/internal/variant_bounds.go @@ -90,7 +90,7 @@ func walkVariantTyped(typedPath []string, valuePaths []string, fields []string, return } *out = append(*out, variantLeaf{ - jsonPath: normalizedVariantPath(fields), + jsonPath: iceberg.NormalizeVariantPath(fields), typedPath: joinPath(typedPath), valuePaths: cloneStrs(valuePaths), icebergType: it, @@ -151,60 +151,6 @@ func cloneStrs(s []string) []string { return append([]string(nil), s...) } func joinPath(s []string) string { return strings.Join(s, ".") } -// normalizedVariantPath builds the spec's RFC-9535 normalized JSON path. -func normalizedVariantPath(fields []string) string { - if len(fields) == 0 { - return "$" - } - - var b strings.Builder - b.WriteByte('$') - for _, f := range fields { - b.WriteString("['") - b.WriteString(rfc9535Escape(f)) - b.WriteString("']") - } - - return b.String() -} - -func rfc9535Escape(name string) string { - if strings.IndexFunc(name, func(r rune) bool { - return r < 0x20 || r == '\'' || r == '\\' - }) < 0 { - return name - } - - var b strings.Builder - b.Grow(len(name) + 4) - for _, r := range name { - switch r { - case '\b': - b.WriteString(`\b`) - case '\t': - b.WriteString(`\t`) - case '\f': - b.WriteString(`\f`) - case '\n': - b.WriteString(`\n`) - case '\r': - b.WriteString(`\r`) - case '\'': - b.WriteString(`\'`) - case '\\': - b.WriteString(`\\`) - default: - if r < 0x20 { - fmt.Fprintf(&b, `\u%04x`, r) - } else { - b.WriteRune(r) - } - } - } - - return b.String() -} - // variantFieldBound is a shredded field's lower bound and upper bound; a nil upper is omitted. type variantFieldBound struct { jsonPath string @@ -367,3 +313,36 @@ func appendDecimalToVariant(b *variant.Builder, t iceberg.DecimalType, v any) er return fmt.Errorf("variant bounds: unsupported decimal value %T", v) } + +// VariantBoundLiteral decodes the stored variant bound object at the normalized path and casts it to typ. +func VariantBoundLiteral(raw []byte, path string, typ iceberg.PrimitiveType) (lit iceberg.Literal, ok bool, err error) { + defer func() { + if r := recover(); r != nil { + lit, ok, err = nil, false, fmt.Errorf("variant bound decode: %v", r) + } + }() + + meta, err := variant.NewMetadata(raw) + if err != nil { + return nil, false, err + } + + v, err := variant.New(raw[:meta.SizeBytes()], raw[meta.SizeBytes():]) + if err != nil { + return nil, false, err + } + + obj, isObj := v.Value().(variant.ObjectValue) + if !isObj { + return nil, false, nil + } + + field, err := obj.ValueByKey(path) + if err != nil { + return nil, false, nil + } + + lit, ok = iceberg.CastVariantLiteral(field.Value, typ) + + return lit, ok, nil +} diff --git a/table/internal/variant_bounds_test.go b/table/internal/variant_bounds_test.go index 794c7b592..12c37e40d 100644 --- a/table/internal/variant_bounds_test.go +++ b/table/internal/variant_bounds_test.go @@ -35,27 +35,6 @@ import ( "github.com/stretchr/testify/require" ) -func TestNormalizedVariantPathEscaping(t *testing.T) { - for _, tt := range []struct { - name string - fields []string - want string - }{ - {"root", nil, "$"}, - {"plain", []string{"event_type"}, "$['event_type']"}, - {"dotted name kept literal", []string{"user.name"}, "$['user.name']"}, - {"nested", []string{"location", "latitude"}, "$['location']['latitude']"}, - {"single quote escaped", []string{"o'brien"}, `$['o\'brien']`}, - {"backslash escaped", []string{`a\b`}, `$['a\\b']`}, - {"newline escaped", []string{"a\nb"}, `$['a\nb']`}, - {"other control char hex-escaped", []string{"a\x01b"}, `$['a\u0001b']`}, - } { - t.Run(tt.name, func(t *testing.T) { - assert.Equal(t, tt.want, normalizedVariantPath(tt.fields)) - }) - } -} - func TestEnumerateVariantLeavesNestedObject(t *testing.T) { // Object {a:int64, n:int16, location:{latitude:float64}, tags:[]string} inner := arrow.StructOf(arrow.Field{Name: "latitude", Type: arrow.PrimitiveTypes.Float64}) @@ -478,3 +457,55 @@ func TestSerializeVariantBoundsDropsNilUpper(t *testing.T) { assert.Equal(t, wantUpper, upper, "dropped upper must be omitted from the upper object") assert.NotEqual(t, lower, upper, "lower still carries $['b']") } + +func buildBoundObject(t *testing.T) []byte { + t.Helper() + + var b variant.Builder + start := b.Offset() + entries := []variant.FieldEntry{b.NextField(start, "$['a']")} + require.NoError(t, b.AppendInt(42)) + entries = append(entries, b.NextField(start, "$['b']")) + require.NoError(t, b.AppendString("hello")) + require.NoError(t, b.FinishObject(start, entries)) + v, err := b.Build() + require.NoError(t, err) + + return append(append([]byte{}, v.Metadata().Bytes()...), v.Bytes()...) +} + +func TestVariantBoundLiteral(t *testing.T) { + raw := buildBoundObject(t) + + lit, ok, err := VariantBoundLiteral(raw, "$['a']", iceberg.PrimitiveTypes.Int64) + require.NoError(t, err) + require.True(t, ok) + assert.Equal(t, int64(42), lit.Any()) + + lit, ok, err = VariantBoundLiteral(raw, "$['b']", iceberg.PrimitiveTypes.String) + require.NoError(t, err) + require.True(t, ok) + assert.Equal(t, "hello", lit.Any()) +} + +func TestVariantBoundLiteralMissingKey(t *testing.T) { + _, ok, err := VariantBoundLiteral(buildBoundObject(t), "$['missing']", iceberg.PrimitiveTypes.Int64) + require.NoError(t, err) + assert.False(t, ok) +} + +func TestVariantBoundLiteralTypeMismatch(t *testing.T) { + // $['a'] is an integer; requesting a string bound must not match. + _, ok, err := VariantBoundLiteral(buildBoundObject(t), "$['a']", iceberg.PrimitiveTypes.String) + require.NoError(t, err) + assert.False(t, ok) +} + +func TestVariantBoundLiteralMalformed(t *testing.T) { + // Malformed / truncated bytes must return an error, never panic (arrow-go A2 hardening). + for _, raw := range [][]byte{{}, {0x01}, {0xff, 0xff, 0xff}, {0x01, 0x00, 0x00, 0x7f, 0x7f}} { + require.NotPanics(t, func() { + _, _, _ = VariantBoundLiteral(raw, "$['a']", iceberg.PrimitiveTypes.Int64) + }) + } +} diff --git a/table/variant_residual.go b/table/variant_residual.go new file mode 100644 index 000000000..5336182a3 --- /dev/null +++ b/table/variant_residual.go @@ -0,0 +1,193 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +package table + +import ( + "context" + "fmt" + "strconv" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/compute" + "github.com/apache/arrow-go/v18/arrow/extensions" + "github.com/apache/arrow-go/v18/arrow/memory" + "github.com/apache/iceberg-go" + "github.com/google/uuid" +) + +// augmentSchemaWithExtracts returns fileSchema plus one primitive column per variant extract term. +func augmentSchemaWithExtracts(fileSchema *iceberg.Schema, cols []iceberg.VariantExtractColumn) (*iceberg.Schema, error) { + fields := fileSchema.Fields() + for _, c := range cols { + fields = append(fields, iceberg.NestedField{ + ID: c.FieldID, + Name: c.Name, + Type: c.Term.Type().(iceberg.PrimitiveType), + }) + } + + return iceberg.NewSchema(fileSchema.ID, fields...), nil +} + +// buildExtractColumn materializes one variant extract term into a typed Arrow array over rec. +func buildExtractColumn(col iceberg.VariantExtractColumn, rec arrow.RecordBatch, mem memory.Allocator) (arrow.Array, arrow.Field, error) { + typ := col.Term.Type().(iceberg.PrimitiveType) + dt, err := TypeToArrowType(typ, false, false) + if err != nil { + return nil, arrow.Field{}, err + } + + bldr := array.NewBuilder(mem, dt) + defer bldr.Release() + + n := int(rec.NumRows()) + varIdx := fieldIndexByID(rec.Schema(), col.Term.Ref().Field().ID) + varr, _ := columnAt(rec, varIdx).(*extensions.VariantArray) + + for i := 0; i < n; i++ { + if varr == nil || varr.IsNull(i) { + bldr.AppendNull() + + continue + } + + v, verr := varr.Value(i) + if verr != nil { + bldr.AppendNull() + + continue + } + + lit, ok := col.Term.ExtractValue(v) + if !ok { + bldr.AppendNull() + + continue + } + + if aerr := appendExtractLiteral(bldr, lit); aerr != nil { + return nil, arrow.Field{}, aerr + } + } + + field := arrow.Field{ + Name: col.Name, + Type: dt, + Nullable: true, + Metadata: arrow.NewMetadata([]string{ArrowParquetFieldIDKey}, []string{strconv.Itoa(col.FieldID)}), + } + + return bldr.NewArray(), field, nil +} + +func columnAt(rec arrow.RecordBatch, idx int) arrow.Array { + if idx < 0 { + return nil + } + + return rec.Column(idx) +} + +// appendExtractLiteral appends a decoded extract literal to its typed builder. +func appendExtractLiteral(bldr array.Builder, lit iceberg.Literal) error { + switch b := bldr.(type) { + case *array.BooleanBuilder: + b.Append(lit.Any().(bool)) + case *array.Int32Builder: + b.Append(lit.Any().(int32)) + case *array.Int64Builder: + b.Append(lit.Any().(int64)) + case *array.Float32Builder: + b.Append(lit.Any().(float32)) + case *array.Float64Builder: + b.Append(lit.Any().(float64)) + case *array.StringBuilder: + b.Append(lit.Any().(string)) + case *array.BinaryBuilder: + b.Append(lit.Any().([]byte)) + case *array.FixedSizeBinaryBuilder: + b.Append(lit.Any().([]byte)) + case *extensions.UUIDBuilder: + b.Append(lit.Any().(uuid.UUID)) + case *array.Date32Builder: + b.Append(arrow.Date32(lit.Any().(iceberg.Date))) + case *array.Time64Builder: + b.Append(arrow.Time64(lit.Any().(iceberg.Time))) + case *array.TimestampBuilder: + switch v := lit.Any().(type) { + case iceberg.Timestamp: + b.Append(arrow.Timestamp(v)) + case iceberg.TimestampNano: + b.Append(arrow.Timestamp(v)) + default: + return fmt.Errorf("%w: variant extract timestamp value %T", iceberg.ErrNotImplemented, v) + } + case *array.Decimal128Builder: + b.Append(lit.Any().(iceberg.Decimal).Val) + default: + return fmt.Errorf("%w: variant extract target builder %T", iceberg.ErrNotImplemented, bldr) + } + + return nil +} + +// extractResidualFilter appends derived extract columns to each batch, runs base, then strips them. +func (as *arrowScan) extractResidualFilter(ctx context.Context, cols []iceberg.VariantExtractColumn, base recProcessFn) recProcessFn { + mem := compute.GetAllocator(ctx) + + return func(rec arrow.RecordBatch) (arrow.RecordBatch, error) { + origSchema := rec.Schema() + origN := int(rec.NumCols()) + + derived := make([]arrow.Array, 0, len(cols)) + fields := make([]arrow.Field, 0, len(cols)) + for _, c := range cols { + arr, field, err := buildExtractColumn(c, rec, mem) + if err != nil { + for _, a := range derived { + a.Release() + } + rec.Release() + + return nil, err + } + derived = append(derived, arr) + fields = append(fields, field) + } + + augFields := append(append([]arrow.Field{}, origSchema.Fields()...), fields...) + md := origSchema.Metadata() + augSchema := arrow.NewSchema(augFields, &md) + augRec := array.NewRecordBatch(augSchema, append(rec.Columns(), derived...), rec.NumRows()) + rec.Release() + for _, a := range derived { + a.Release() + } + + filtered, err := base(augRec) + if err != nil { + return nil, err + } + + out := array.NewRecordBatch(origSchema, filtered.Columns()[:origN], filtered.NumRows()) + filtered.Release() + + return out, nil + } +} diff --git a/table/variant_residual_test.go b/table/variant_residual_test.go new file mode 100644 index 000000000..b655b8cda --- /dev/null +++ b/table/variant_residual_test.go @@ -0,0 +1,69 @@ +// 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/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/memory" + "github.com/apache/iceberg-go" + "github.com/google/uuid" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestAppendExtractLiteral covers every extract target builder, derived through the same +// TypeToArrowType path buildExtractColumn uses, including uuid and fixed which previously +// fell through to the default ErrNotImplemented and aborted the scan. +func TestAppendExtractLiteral(t *testing.T) { + mem := memory.DefaultAllocator + + for _, tt := range []struct { + name string + typ iceberg.PrimitiveType + lit iceberg.Literal + }{ + {"int64", iceberg.PrimitiveTypes.Int64, iceberg.NewLiteral(int64(5))}, + {"string", iceberg.PrimitiveTypes.String, iceberg.NewLiteral("hi")}, + {"binary", iceberg.PrimitiveTypes.Binary, iceberg.NewLiteral([]byte{1, 2})}, + {"uuid", iceberg.PrimitiveTypes.UUID, iceberg.NewLiteral(uuid.UUID{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16})}, + {"fixed", iceberg.FixedTypeOf(16), iceberg.NewLiteral(make([]byte, 16))}, + } { + t.Run(tt.name, func(t *testing.T) { + dt, err := TypeToArrowType(tt.typ, false, false) + require.NoError(t, err) + + bldr := array.NewBuilder(mem, dt) + defer bldr.Release() + + require.NoError(t, appendExtractLiteral(bldr, tt.lit)) + arr := bldr.NewArray() + defer arr.Release() + assert.Equal(t, 1, arr.Len()) + }) + } +} + +func TestAppendExtractLiteralUnsupported(t *testing.T) { + bldr := array.NewBuilder(memory.DefaultAllocator, arrow.ListOf(arrow.PrimitiveTypes.Int64)) + defer bldr.Release() + + require.ErrorIs(t, appendExtractLiteral(bldr, iceberg.NewLiteral(int64(1))), iceberg.ErrNotImplemented) +} diff --git a/table/variant_shredded_write_test.go b/table/variant_shredded_write_test.go index bef0b49e5..e67e8da92 100644 --- a/table/variant_shredded_write_test.go +++ b/table/variant_shredded_write_test.go @@ -1156,6 +1156,265 @@ func TestShreddedVariantWriteChildStats(t *testing.T) { assert.Equal(t, buildObj(5_000_000_007), df.UpperBoundValues()[variantField], "upper bound object") } +// TestShreddedVariantExtractPruning prunes variant_get predicates using the written bounds. +func TestShreddedVariantExtractPruning(t *testing.T) { + files := writeVariantTable(t, iceberg.Properties{ + PropertyFormatVersion: "3", + ParquetShredVariantsKey: "true", + }) + require.Len(t, files, 1) + df := files[0] + + sc := iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64, Required: true}, + iceberg.NestedField{ID: 2, Name: "payload", Type: iceberg.VariantType{}, Required: false}, + ) + + extInt := func(op iceberg.Operation, v int64) iceberg.BooleanExpression { + return iceberg.LiteralPredicate(op, + iceberg.Extract("payload", "$.a", iceberg.PrimitiveTypes.Int64), + iceberg.NewLiteral(v)) + } + + for _, tt := range []struct { + name string + expr iceberg.BooleanExpression + want bool + }{ + {"eq in range", extInt(iceberg.OpEQ, 5_000_000_003), rowsMightMatch}, + {"eq above max", extInt(iceberg.OpEQ, 6_000_000_000), rowsCannotMatch}, + {"eq below min", extInt(iceberg.OpEQ, 4_000_000_000), rowsCannotMatch}, + {"gt at max", extInt(iceberg.OpGT, 5_000_000_007), rowsCannotMatch}, + {"lt at min", extInt(iceberg.OpLT, 5_000_000_000), rowsCannotMatch}, + {"gt in range", extInt(iceberg.OpGT, 5_000_000_000), rowsMightMatch}, + {"missing path never prunes", iceberg.LiteralPredicate(iceberg.OpEQ, + iceberg.Extract("payload", "$.missing", iceberg.PrimitiveTypes.Int64), + iceberg.NewLiteral(int64(1))), rowsMightMatch}, + {"starts_with match", iceberg.LiteralPredicate(iceberg.OpStartsWith, + iceberg.Extract("payload", "$.b", iceberg.PrimitiveTypes.String), + iceberg.NewLiteral("r")), rowsMightMatch}, + {"starts_with no match", iceberg.LiteralPredicate(iceberg.OpStartsWith, + iceberg.Extract("payload", "$.b", iceberg.PrimitiveTypes.String), + iceberg.NewLiteral("z")), rowsCannotMatch}, + } { + t.Run(tt.name, func(t *testing.T) { + eval, err := newInclusiveMetricsEvaluator(sc, tt.expr, true, true) + require.NoError(t, err) + res, err := eval(df) + require.NoError(t, err) + assert.Equal(t, tt.want, res) + }) + } +} + +// TestShreddedVariantExtractResidualScan confirms residual filtering returns only matching rows. +func TestShreddedVariantExtractResidualScan(t *testing.T) { + files := writeVariantTable(t, iceberg.Properties{ + PropertyFormatVersion: "3", + ParquetShredVariantsKey: "true", + }) + require.Len(t, files, 1) + + p := strings.TrimPrefix(files[0].FilePath(), "file://") + f, err := os.Open(p) + require.NoError(t, err) + defer f.Close() + + tbl, err := pqarrow.ReadTable(context.Background(), f, nil, pqarrow.ArrowReadProperties{}, memory.DefaultAllocator) + require.NoError(t, err) + defer tbl.Release() + + fileSchema, err := ArrowSchemaToIceberg(tbl.Schema(), false, nil) + require.NoError(t, err) + + cols := make([]arrow.Array, tbl.NumCols()) + for i := range cols { + cols[i] = tbl.Column(i).Data().Chunk(0) + } + rec := array.NewRecordBatch(tbl.Schema(), cols, tbl.NumRows()) + + // payload.a == 5_000_000_003 matches exactly one of the 8 rows. + pred := iceberg.LiteralPredicate(iceberg.OpEQ, + iceberg.Extract("payload", "$.a", iceberg.PrimitiveTypes.Int64), + iceberg.NewLiteral(int64(5_000_000_003))) + bound, err := iceberg.BindExpr(fileSchema, pred, true) + require.NoError(t, err) + + as := &arrowScan{boundRowFilter: bound, caseSensitive: true} + fn, skip, err := as.getRecordFilter(context.Background(), fileSchema) + require.NoError(t, err) + require.False(t, skip) + require.NotNil(t, fn) + + out, err := fn(rec) + require.NoError(t, err) + defer out.Release() + + require.Equal(t, int64(1), out.NumRows(), "only the a==5_000_000_003 row survives residual filtering") + require.EqualValues(t, tbl.NumCols(), out.NumCols(), "synthetic extract columns must be stripped from the result") + + pv := out.Column(out.Schema().FieldIndices("payload")[0]).(*extensions.VariantArray) + v, err := pv.Value(0) + require.NoError(t, err) + j, err := v.MarshalJSON() + require.NoError(t, err) + assert.JSONEq(t, `{"a":5000000003,"b":"row"}`, string(j)) +} + +// TestNonShreddedVariantExtractResidualScan confirms residual filtering on a non-shredded variant column. +func TestNonShreddedVariantExtractResidualScan(t *testing.T) { + files := writeVariantTable(t, iceberg.Properties{ + PropertyFormatVersion: "3", + }) + require.Len(t, files, 1) + + p := strings.TrimPrefix(files[0].FilePath(), "file://") + f, err := os.Open(p) + require.NoError(t, err) + defer f.Close() + + tbl, err := pqarrow.ReadTable(context.Background(), f, nil, pqarrow.ArrowReadProperties{}, memory.DefaultAllocator) + require.NoError(t, err) + defer tbl.Release() + + require.False(t, tbl.Column(tbl.Schema().FieldIndices("payload")[0]).Data().Chunk(0).(*extensions.VariantArray).IsShredded(), + "payload must be non-shredded for this test") + + fileSchema, err := ArrowSchemaToIceberg(tbl.Schema(), false, nil) + require.NoError(t, err) + + cols := make([]arrow.Array, tbl.NumCols()) + for i := range cols { + cols[i] = tbl.Column(i).Data().Chunk(0) + } + rec := array.NewRecordBatch(tbl.Schema(), cols, tbl.NumRows()) + + pred := iceberg.LiteralPredicate(iceberg.OpEQ, + iceberg.Extract("payload", "$.a", iceberg.PrimitiveTypes.Int64), + iceberg.NewLiteral(int64(5_000_000_003))) + bound, err := iceberg.BindExpr(fileSchema, pred, true) + require.NoError(t, err) + + as := &arrowScan{boundRowFilter: bound, caseSensitive: true} + fn, skip, err := as.getRecordFilter(context.Background(), fileSchema) + require.NoError(t, err) + require.False(t, skip) + require.NotNil(t, fn) + + out, err := fn(rec) + require.NoError(t, err) + defer out.Release() + + require.Equal(t, int64(1), out.NumRows(), "only the a==5_000_000_003 row survives residual filtering") + require.EqualValues(t, tbl.NumCols(), out.NumCols(), "synthetic extract columns must be stripped from the result") +} + +// TestShreddedVariantExtractResidualNoLeak asserts the residual filter releases every Arrow buffer. +func TestShreddedVariantExtractResidualNoLeak(t *testing.T) { + files := writeVariantTable(t, iceberg.Properties{ + PropertyFormatVersion: "3", + ParquetShredVariantsKey: "true", + }) + require.Len(t, files, 1) + + p := strings.TrimPrefix(files[0].FilePath(), "file://") + f, err := os.Open(p) + require.NoError(t, err) + defer f.Close() + + tbl, err := pqarrow.ReadTable(context.Background(), f, nil, pqarrow.ArrowReadProperties{}, memory.DefaultAllocator) + require.NoError(t, err) + defer tbl.Release() + + fileSchema, err := ArrowSchemaToIceberg(tbl.Schema(), false, nil) + require.NoError(t, err) + + cols := make([]arrow.Array, tbl.NumCols()) + for i := range cols { + cols[i] = tbl.Column(i).Data().Chunk(0) + } + rec := array.NewRecordBatch(tbl.Schema(), cols, tbl.NumRows()) + + pred := iceberg.LiteralPredicate(iceberg.OpEQ, + iceberg.Extract("payload", "$.a", iceberg.PrimitiveTypes.Int64), + iceberg.NewLiteral(int64(5_000_000_003))) + bound, err := iceberg.BindExpr(fileSchema, pred, true) + require.NoError(t, err) + + checked := memory.NewCheckedAllocator(memory.DefaultAllocator) + ctx := compute.WithAllocator(context.Background(), checked) + + as := &arrowScan{boundRowFilter: bound, caseSensitive: true} + fn, _, err := as.getRecordFilter(ctx, fileSchema) + require.NoError(t, err) + require.NotNil(t, fn) + + out, err := fn(rec) + require.NoError(t, err) + out.Release() + + checked.AssertSize(t, 0) +} + +// TestVariantExtractResidualCombined exercises AND/OR filters over multiple extract terms. +func TestVariantExtractResidualCombined(t *testing.T) { + files := writeVariantTable(t, iceberg.Properties{ + PropertyFormatVersion: "3", + ParquetShredVariantsKey: "true", + }) + require.Len(t, files, 1) + + p := strings.TrimPrefix(files[0].FilePath(), "file://") + f, err := os.Open(p) + require.NoError(t, err) + defer f.Close() + + tbl, err := pqarrow.ReadTable(context.Background(), f, nil, pqarrow.ArrowReadProperties{}, memory.DefaultAllocator) + require.NoError(t, err) + defer tbl.Release() + + fileSchema, err := ArrowSchemaToIceberg(tbl.Schema(), false, nil) + require.NoError(t, err) + + eqA := func(v int64) iceberg.BooleanExpression { + return iceberg.LiteralPredicate(iceberg.OpEQ, iceberg.Extract("payload", "$.a", iceberg.PrimitiveTypes.Int64), iceberg.NewLiteral(v)) + } + + for _, tt := range []struct { + name string + expr iceberg.BooleanExpression + want int64 + }{ + {"and two distinct extracts", iceberg.NewAnd(eqA(5_000_000_003), + iceberg.LiteralPredicate(iceberg.OpEQ, iceberg.Extract("payload", "$.b", iceberg.PrimitiveTypes.String), iceberg.NewLiteral("row"))), 1}, + {"or same extract twice", iceberg.NewOr(eqA(5_000_000_003), eqA(5_000_000_005)), 2}, + {"and extract excludes all", iceberg.NewAnd(eqA(5_000_000_003), eqA(5_000_000_005)), 0}, + } { + t.Run(tt.name, func(t *testing.T) { + cols := make([]arrow.Array, tbl.NumCols()) + for i := range cols { + cols[i] = tbl.Column(i).Data().Chunk(0) + } + rec := array.NewRecordBatch(tbl.Schema(), cols, tbl.NumRows()) + + bound, err := iceberg.BindExpr(fileSchema, tt.expr, true) + require.NoError(t, err) + + as := &arrowScan{boundRowFilter: bound, caseSensitive: true} + fn, _, err := as.getRecordFilter(context.Background(), fileSchema) + require.NoError(t, err) + require.NotNil(t, fn) + + out, err := fn(rec) + require.NoError(t, err) + defer out.Release() + + require.Equal(t, tt.want, out.NumRows()) + require.EqualValues(t, tbl.NumCols(), out.NumCols(), "synthetic columns stripped") + }) + } +} + // TestShreddedVariantWriteNullFieldKeepsBound: int64+null field still gets a bound. func TestShreddedVariantWriteNullFieldKeepsBound(t *testing.T) { mem := memory.DefaultAllocator diff --git a/variant_cast.go b/variant_cast.go new file mode 100644 index 000000000..0d2bc0cb0 --- /dev/null +++ b/variant_cast.go @@ -0,0 +1,269 @@ +// 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 iceberg + +import ( + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/decimal" + "github.com/apache/arrow-go/v18/arrow/decimal128" + "github.com/apache/arrow-go/v18/parquet/variant" + "github.com/google/uuid" +) + +const ( + microsPerDay = int64(86_400_000_000) + nanosPerDay = int64(86_400_000_000_000) + nanosPerMicro = int64(1_000) +) + +// CastVariantLiteral casts a leaf variant value to typ and wraps it as a Literal. +func CastVariantLiteral(v variant.Value, typ PrimitiveType) (Literal, bool) { + result, ok := castVariantValue(v, typ) + if !ok { + return nil, false + } + + lit := literalFromCastValue(result) + if lit == nil { + return nil, false + } + if !lit.Type().Equals(typ) { + conv, err := lit.To(typ) + if err != nil { + return nil, false + } + + lit = conv + } + + return lit, true +} + +// castVariantValue casts a leaf variant value to the Go value backing typ. +func castVariantValue(v variant.Value, typ PrimitiveType) (any, bool) { + raw := v.Value() + if raw == nil { + return nil, false + } + + if r, ok := exactVariantMatch(v.Type(), raw, typ); ok { + return r, true + } + + switch t := typ.(type) { + case Int32Type: + switch n := raw.(type) { + case int8: + return int32(n), true + case int16: + return int32(n), true + } + case Int64Type: + switch n := raw.(type) { + case int8: + return int64(n), true + case int16: + return int64(n), true + case int32: + return int64(n), true + } + case Float64Type: + if f, ok := raw.(float32); ok { + return float64(f), true + } + case FixedType: + if b, ok := raw.([]byte); ok && len(b) == t.Len() { + return b, true + } + case DecimalType: + return castVariantDecimal(raw, t) + case BooleanType: + if b, ok := raw.(bool); ok { + return b, true + } + case TimestampType, TimestampTzType: + return castVariantToMicros(v.Type(), raw) + case TimestampNsType, TimestampTzNsType: + return castVariantToNanos(v.Type(), raw) + case DateType: + return castVariantToDate(v.Type(), raw) + } + + return nil, false +} + +// exactVariantMatch returns raw coerced to the Go type backing typ when the variant physical type matches typ exactly. +func exactVariantMatch(pt variant.Type, raw any, typ PrimitiveType) (any, bool) { + switch typ.(type) { + case Int32Type: + if pt == variant.Int32 { + return raw.(int32), true + } + case Int64Type: + if pt == variant.Int64 { + return raw.(int64), true + } + case Float32Type: + if pt == variant.Float { + return raw.(float32), true + } + case Float64Type: + if pt == variant.Double { + return raw.(float64), true + } + case DateType: + if pt == variant.Date { + return Date(raw.(arrow.Date32)), true + } + case TimestampType: + if pt == variant.TimestampMicrosNTZ { + return Timestamp(raw.(arrow.Timestamp)), true + } + case TimestampTzType: + if pt == variant.TimestampMicros { + return Timestamp(raw.(arrow.Timestamp)), true + } + case TimestampNsType: + if pt == variant.TimestampNanosNTZ { + return TimestampNano(raw.(arrow.Timestamp)), true + } + case TimestampTzNsType: + if pt == variant.TimestampNanos { + return TimestampNano(raw.(arrow.Timestamp)), true + } + case TimeType: + if pt == variant.Time { + return Time(raw.(arrow.Time64)), true + } + case UUIDType: + if pt == variant.UUID { + return raw.(uuid.UUID), true + } + case StringType: + if pt == variant.String { + return raw.(string), true + } + case BinaryType: + if pt == variant.Binary { + return raw.([]byte), true + } + } + + return nil, false +} + +func castVariantDecimal(raw any, typ DecimalType) (any, bool) { + switch d := raw.(type) { + case variant.DecimalValue[decimal.Decimal32]: + if int(d.Scale) != typ.Scale() { + return nil, false + } + + return Decimal{Val: decimal128.FromI64(int64(d.Value.(decimal.Decimal32))), Scale: int(d.Scale)}, true + case variant.DecimalValue[decimal.Decimal64]: + if int(d.Scale) != typ.Scale() { + return nil, false + } + + return Decimal{Val: decimal128.FromI64(int64(d.Value.(decimal.Decimal64))), Scale: int(d.Scale)}, true + case variant.DecimalValue[decimal.Decimal128]: + if int(d.Scale) != typ.Scale() { + return nil, false + } + + return Decimal{Val: d.Value.(decimal.Decimal128), Scale: int(d.Scale)}, true + } + + return nil, false +} + +func castVariantToMicros(pt variant.Type, raw any) (any, bool) { + switch pt { + case variant.TimestampNanos, variant.TimestampNanosNTZ: + return Timestamp(floorDiv(int64(raw.(arrow.Timestamp)), nanosPerMicro)), true + case variant.Date: + return Timestamp(int64(raw.(arrow.Date32)) * microsPerDay), true + } + + return nil, false +} + +func castVariantToNanos(pt variant.Type, raw any) (any, bool) { + switch pt { + case variant.TimestampMicros, variant.TimestampMicrosNTZ: + return TimestampNano(int64(raw.(arrow.Timestamp)) * nanosPerMicro), true + case variant.Date: + return TimestampNano(int64(raw.(arrow.Date32)) * nanosPerDay), true + } + + return nil, false +} + +func castVariantToDate(pt variant.Type, raw any) (any, bool) { + switch pt { + case variant.TimestampMicros, variant.TimestampMicrosNTZ: + return Date(floorDiv(int64(raw.(arrow.Timestamp)), microsPerDay)), true + case variant.TimestampNanos, variant.TimestampNanosNTZ: + return Date(floorDiv(int64(raw.(arrow.Timestamp)), nanosPerDay)), true + } + + return nil, false +} + +// floorDiv divides rounding toward negative infinity. +func floorDiv(a, b int64) int64 { + q := a / b + if (a%b != 0) && ((a < 0) != (b < 0)) { + q-- + } + + return q +} + +func literalFromCastValue(result any) Literal { + switch r := result.(type) { + case bool: + return NewLiteral(r) + case int32: + return NewLiteral(r) + case int64: + return NewLiteral(r) + case float32: + return NewLiteral(r) + case float64: + return NewLiteral(r) + case Date: + return NewLiteral(r) + case Time: + return NewLiteral(r) + case Timestamp: + return NewLiteral(r) + case TimestampNano: + return NewLiteral(r) + case string: + return NewLiteral(r) + case []byte: + return NewLiteral(r) + case uuid.UUID: + return NewLiteral(r) + case Decimal: + return NewLiteral(r) + } + + return nil +} diff --git a/variant_cast_test.go b/variant_cast_test.go new file mode 100644 index 000000000..41b10c024 --- /dev/null +++ b/variant_cast_test.go @@ -0,0 +1,77 @@ +// 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 iceberg + +import ( + "testing" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/decimal" + "github.com/apache/arrow-go/v18/arrow/decimal128" + "github.com/apache/arrow-go/v18/parquet/variant" + "github.com/google/uuid" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func scalarVariant(t *testing.T, build func(*variant.Builder) error) variant.Value { + t.Helper() + var b variant.Builder + require.NoError(t, build(&b)) + v, err := b.Build() + require.NoError(t, err) + + return v +} + +func TestCastVariantLiteral(t *testing.T) { + testUUID := uuid.UUID{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16} + + for _, tt := range []struct { + name string + build func(*variant.Builder) error + typ PrimitiveType + want any // nil means not castable + }{ + {"small int to int64", func(b *variant.Builder) error { return b.AppendInt(5) }, PrimitiveTypes.Int64, int64(5)}, + {"small int to int32", func(b *variant.Builder) error { return b.AppendInt(5) }, PrimitiveTypes.Int32, int32(5)}, + {"int64 exact", func(b *variant.Builder) error { return b.AppendInt(5_000_000_000) }, PrimitiveTypes.Int64, int64(5_000_000_000)}, + {"float32 widens to float64", func(b *variant.Builder) error { return b.AppendFloat32(1.5) }, PrimitiveTypes.Float64, float64(1.5)}, + {"boolean exact", func(b *variant.Builder) error { return b.AppendBool(true) }, PrimitiveTypes.Bool, true}, + {"string exact", func(b *variant.Builder) error { return b.AppendString("hi") }, PrimitiveTypes.String, "hi"}, + {"binary exact", func(b *variant.Builder) error { return b.AppendBinary([]byte{1, 2, 3}) }, PrimitiveTypes.Binary, []byte{1, 2, 3}}, + {"date exact", func(b *variant.Builder) error { return b.AppendDate(arrow.Date32(100)) }, PrimitiveTypes.Date, Date(100)}, + {"timestamp micros exact", func(b *variant.Builder) error { return b.AppendTimestamp(arrow.Timestamp(123), true, false) }, PrimitiveTypes.Timestamp, Timestamp(123)}, + {"uuid exact", func(b *variant.Builder) error { return b.AppendUUID(testUUID) }, PrimitiveTypes.UUID, testUUID}, + {"decimal scale match", func(b *variant.Builder) error { return b.AppendDecimal8(2, decimal.Decimal64(1234)) }, DecimalTypeOf(10, 2), Decimal{Val: decimal128.FromI64(1234), Scale: 2}}, + {"decimal scale mismatch", func(b *variant.Builder) error { return b.AppendDecimal8(2, decimal.Decimal64(1234)) }, DecimalTypeOf(10, 3), nil}, + {"string not castable to int64", func(b *variant.Builder) error { return b.AppendString("hi") }, PrimitiveTypes.Int64, nil}, + {"nanos floor to micros pre-epoch", func(b *variant.Builder) error { return b.AppendTimestamp(arrow.Timestamp(-1500), false, false) }, PrimitiveTypes.Timestamp, Timestamp(-2)}, + } { + t.Run(tt.name, func(t *testing.T) { + lit, ok := CastVariantLiteral(scalarVariant(t, tt.build), tt.typ) + if tt.want == nil { + assert.False(t, ok) + + return + } + require.True(t, ok) + assert.Equal(t, tt.want, lit.Any()) + }) + } +} diff --git a/variant_extract.go b/variant_extract.go new file mode 100644 index 000000000..4433359cd --- /dev/null +++ b/variant_extract.go @@ -0,0 +1,248 @@ +// 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 iceberg + +import ( + "fmt" + + "github.com/apache/arrow-go/v18/parquet/variant" + "github.com/google/uuid" +) + +// BoundExtract is a bound variant sub-path term used for metrics pruning and residual evaluation. +type BoundExtract interface { + BoundTerm + + Path() string + // ExtractValue navigates v to this term's path and casts the leaf to the target type. + ExtractValue(v variant.Value) (Literal, bool) +} + +// Extract creates an unbound variant sub-path term for a dotted JSONPath. +func Extract(ref Reference, path string, typ PrimitiveType) UnboundTerm { + return &unboundExtract{ref: ref, path: path, typ: typ} +} + +type unboundExtract struct { + ref Reference + path string + typ PrimitiveType +} + +func (*unboundExtract) isTerm() {} +func (u *unboundExtract) String() string { + return fmt.Sprintf("extract(%s, path=%s, type=%s)", u.ref, u.path, u.typ) +} + +func (u *unboundExtract) Ref() Reference { return u.ref } +func (u *unboundExtract) Path() string { return u.path } +func (u *unboundExtract) Type() PrimitiveType { return u.typ } + +func (u *unboundExtract) Equals(other UnboundTerm) bool { + rhs, ok := other.(*unboundExtract) + if !ok { + return false + } + + sameType := (u.typ == nil && rhs.typ == nil) || + (u.typ != nil && rhs.typ != nil && u.typ.Equals(rhs.typ)) + + return u.ref == rhs.ref && u.path == rhs.path && sameType +} + +func (u *unboundExtract) Bind(schema *Schema, caseSensitive bool) (BoundTerm, error) { + bound, err := u.ref.Bind(schema, caseSensitive) + if err != nil { + return nil, err + } + if _, ok := bound.Type().(VariantType); !ok { + return nil, fmt.Errorf("%w: cannot bind extract, not a variant: %s", ErrInvalidArgument, u.ref) + } + if u.typ == nil { + return nil, fmt.Errorf("%w: cannot bind extract, target type is required", ErrInvalidArgument) + } + if !isVariantExtractTarget(u.typ) { + return nil, fmt.Errorf("%w: cannot bind extract, unsupported target type: %s", ErrInvalidArgument, u.typ) + } + + fields, err := parseVariantPath(u.path) + if err != nil { + return nil, err + } + + acc, ok := schema.accessorForField(bound.Ref().Field().ID) + if !ok { + return nil, ErrInvalidSchema + } + + return createBoundExtract(bound.Ref(), fields, NormalizeVariantPath(fields), u.typ, acc), nil +} + +// isVariantExtractTarget reports whether typ is a supported extract target (the set createBoundExtract handles). +func isVariantExtractTarget(typ PrimitiveType) bool { + switch typ.(type) { + case BooleanType, Int32Type, Int64Type, Float32Type, Float64Type, + DateType, TimeType, TimestampType, TimestampTzType, TimestampNsType, TimestampTzNsType, + StringType, FixedType, BinaryType, DecimalType, UUIDType: + return true + } + + return false +} + +var _ BoundExtract = (*boundExtract[int32])(nil) + +type boundExtract[T LiteralType] struct { + ref BoundReference + fields []string + path string + typ PrimitiveType + acc accessor +} + +func createBoundExtract(ref BoundReference, fields []string, path string, typ PrimitiveType, acc accessor) BoundTerm { + switch typ.(type) { + case BooleanType: + return &boundExtract[bool]{ref: ref, fields: fields, path: path, typ: typ, acc: acc} + case Int32Type: + return &boundExtract[int32]{ref: ref, fields: fields, path: path, typ: typ, acc: acc} + case Int64Type: + return &boundExtract[int64]{ref: ref, fields: fields, path: path, typ: typ, acc: acc} + case Float32Type: + return &boundExtract[float32]{ref: ref, fields: fields, path: path, typ: typ, acc: acc} + case Float64Type: + return &boundExtract[float64]{ref: ref, fields: fields, path: path, typ: typ, acc: acc} + case DateType: + return &boundExtract[Date]{ref: ref, fields: fields, path: path, typ: typ, acc: acc} + case TimeType: + return &boundExtract[Time]{ref: ref, fields: fields, path: path, typ: typ, acc: acc} + case TimestampType, TimestampTzType: + return &boundExtract[Timestamp]{ref: ref, fields: fields, path: path, typ: typ, acc: acc} + case TimestampNsType, TimestampTzNsType: + return &boundExtract[TimestampNano]{ref: ref, fields: fields, path: path, typ: typ, acc: acc} + case StringType: + return &boundExtract[string]{ref: ref, fields: fields, path: path, typ: typ, acc: acc} + case FixedType, BinaryType: + return &boundExtract[[]byte]{ref: ref, fields: fields, path: path, typ: typ, acc: acc} + case DecimalType: + return &boundExtract[Decimal]{ref: ref, fields: fields, path: path, typ: typ, acc: acc} + case UUIDType: + return &boundExtract[uuid.UUID]{ref: ref, fields: fields, path: path, typ: typ, acc: acc} + } + panic("unhandled variant extract target type: " + typ.String()) +} + +func (*boundExtract[T]) isTerm() {} +func (b *boundExtract[T]) String() string { + return fmt.Sprintf("extract(%s, path=%s, type=%s)", b.ref, b.path, b.typ) +} + +func (b *boundExtract[T]) Ref() BoundReference { return b.ref } +func (b *boundExtract[T]) Type() Type { return b.typ } +func (b *boundExtract[T]) Path() string { return b.path } + +func (b *boundExtract[T]) Equals(other BoundTerm) bool { + rhs, ok := other.(*boundExtract[T]) + if !ok { + return false + } + + return b.ref.Equals(rhs.ref) && b.path == rhs.path && b.typ.Equals(rhs.typ) +} + +// navigateVariant walks the nested variant object by member names, returning the leaf value or invalid if the path is absent. +func navigateVariant(v variant.Value, fields []string) (variant.Value, bool) { + for _, name := range fields { + obj, ok := v.Value().(variant.ObjectValue) + if !ok { + return variant.Value{}, false + } + + field, err := obj.ValueByKey(name) + if err != nil { + return variant.Value{}, false + } + + v = field.Value + } + + return v, true +} + +// leafValue navigates the nested variant object from the value stored at this term's reference. +func (b *boundExtract[T]) leafValue(st StructLike) (variant.Value, bool) { + raw := b.acc.Get(st) + v, ok := raw.(variant.Value) + if !ok { + return variant.Value{}, false + } + + return navigateVariant(v, b.fields) +} + +// ExtractValue navigates v to this term's path and casts the leaf to the target type. +func (b *boundExtract[T]) ExtractValue(v variant.Value) (Literal, bool) { + leaf, ok := navigateVariant(v, b.fields) + if !ok { + return nil, false + } + + return CastVariantLiteral(leaf, b.typ) +} + +func (b *boundExtract[T]) eval(st StructLike) Optional[T] { + v, ok := b.leafValue(st) + if !ok { + return Optional[T]{} + } + + result, ok := castVariantValue(v, b.typ) + if !ok { + return Optional[T]{} + } + + val, ok := result.(T) + if !ok { + return Optional[T]{} + } + + return Optional[T]{Valid: true, Val: val} +} + +func (b *boundExtract[T]) evalToLiteral(st StructLike) Optional[Literal] { + v := b.eval(st) + if !v.Valid { + return Optional[Literal]{} + } + + lit := NewLiteral(v.Val) + if !lit.Type().Equals(b.typ) { + conv, err := lit.To(b.typ) + if err != nil { + return Optional[Literal]{} + } + + lit = conv + } + + return Optional[Literal]{Val: lit, Valid: true} +} + +func (b *boundExtract[T]) evalIsNull(st StructLike) bool { + return !b.eval(st).Valid +} diff --git a/variant_extract_test.go b/variant_extract_test.go new file mode 100644 index 000000000..5a2dceb1f --- /dev/null +++ b/variant_extract_test.go @@ -0,0 +1,61 @@ +// 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 iceberg + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func extractBindSchema() *Schema { + return NewSchema(0, + NestedField{ID: 1, Name: "payload", Type: VariantType{}}, + NestedField{ID: 2, Name: "name", Type: PrimitiveTypes.String}, + ) +} + +func TestExtractBind(t *testing.T) { + term, err := Extract("payload", "$.a.b", PrimitiveTypes.Int64).Bind(extractBindSchema(), true) + require.NoError(t, err) + + be, ok := term.(BoundExtract) + require.True(t, ok) + assert.Equal(t, "$['a']['b']", be.Path()) + assert.True(t, PrimitiveTypes.Int64.Equals(be.Type())) + assert.Equal(t, 1, be.Ref().Field().ID) +} + +func TestExtractBindRejects(t *testing.T) { + for _, tt := range []struct { + name string + term UnboundTerm + }{ + {"non-variant source", Extract("name", "$.a", PrimitiveTypes.Int64)}, + {"nil target type", Extract("payload", "$.a", nil)}, + {"unknown target type", Extract("payload", "$.a", UnknownType{})}, + {"bracket path", Extract("payload", "$['a']", PrimitiveTypes.Int64)}, + {"unknown field", Extract("missing", "$.a", PrimitiveTypes.Int64)}, + } { + t.Run(tt.name, func(t *testing.T) { + _, err := tt.term.Bind(extractBindSchema(), true) + require.Error(t, err) + }) + } +} diff --git a/variant_path.go b/variant_path.go new file mode 100644 index 000000000..0f52ad0ea --- /dev/null +++ b/variant_path.go @@ -0,0 +1,132 @@ +// 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 iceberg + +import ( + "fmt" + "strings" +) + +// NormalizeVariantPath renders member names as the spec's RFC-9535 normalized JSON path. +func NormalizeVariantPath(fields []string) string { + if len(fields) == 0 { + return "$" + } + + var b strings.Builder + b.WriteByte('$') + for _, f := range fields { + b.WriteString("['") + b.WriteString(rfc9535Escape(f)) + b.WriteString("']") + } + + return b.String() +} + +func rfc9535Escape(name string) string { + if strings.IndexFunc(name, func(r rune) bool { + return r < 0x20 || r == '\'' || r == '\\' + }) < 0 { + return name + } + + var b strings.Builder + b.Grow(len(name) + 4) + for _, r := range name { + switch r { + case '\b': + b.WriteString(`\b`) + case '\t': + b.WriteString(`\t`) + case '\f': + b.WriteString(`\f`) + case '\n': + b.WriteString(`\n`) + case '\r': + b.WriteString(`\r`) + case '\'': + b.WriteString(`\'`) + case '\\': + b.WriteString(`\\`) + default: + if r < 0x20 { + fmt.Fprintf(&b, `\u%04x`, r) + } else { + b.WriteRune(r) + } + } + } + + return b.String() +} + +// parseVariantPath parses a dot-shorthand variant path ($.a.b) into its member names. +func parseVariantPath(path string) ([]string, error) { + if strings.ContainsAny(path, "[]") { + return nil, fmt.Errorf("%w: unsupported variant path, contains bracket: %q", ErrInvalidArgument, path) + } + if strings.Contains(path, "*") { + return nil, fmt.Errorf("%w: unsupported variant path, contains wildcard: %q", ErrInvalidArgument, path) + } + if strings.Contains(path, "..") { + return nil, fmt.Errorf("%w: unsupported variant path, contains recursive descent: %q", ErrInvalidArgument, path) + } + + parts := strings.Split(path, ".") + if parts[0] != "$" { + return nil, fmt.Errorf("%w: invalid variant path, does not start with $: %q", ErrInvalidArgument, path) + } + + names := parts[1:] + for _, name := range names { + if !isRFC9535MemberName(name) { + return nil, fmt.Errorf("%w: invalid variant path %q (%q has invalid characters)", ErrInvalidArgument, path, name) + } + } + + return names, nil +} + +// isRFC9535MemberName reports whether name is a valid RFC-9535 member-name shorthand. +func isRFC9535MemberName(name string) bool { + for i, r := range name { + if i == 0 { + if !isRFC9535NameFirst(r) { + return false + } + + continue + } + if isRFC9535NameFirst(r) || (r >= '0' && r <= '9') { + continue + } + + return false + } + + return name != "" +} + +func isRFC9535NameFirst(r rune) bool { + return r == '_' || + (r >= 'A' && r <= 'Z') || + (r >= 'a' && r <= 'z') || + (r >= 0x80 && r <= 0xD7FF) || + (r >= 0xE000 && r <= 0x10FFFF) +} diff --git a/variant_path_test.go b/variant_path_test.go new file mode 100644 index 000000000..83235c941 --- /dev/null +++ b/variant_path_test.go @@ -0,0 +1,85 @@ +// 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 iceberg + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestNormalizeVariantPathEscaping(t *testing.T) { + for _, tt := range []struct { + name string + fields []string + want string + }{ + {"root", nil, "$"}, + {"plain", []string{"event_type"}, "$['event_type']"}, + {"dotted name kept literal", []string{"user.name"}, "$['user.name']"}, + {"nested", []string{"location", "latitude"}, "$['location']['latitude']"}, + {"single quote escaped", []string{"o'brien"}, `$['o\'brien']`}, + {"backslash escaped", []string{`a\b`}, `$['a\\b']`}, + {"newline escaped", []string{"a\nb"}, `$['a\nb']`}, + {"other control char hex-escaped", []string{"a\x01b"}, "$['a\\u0001b']"}, + } { + t.Run(tt.name, func(t *testing.T) { + assert.Equal(t, tt.want, NormalizeVariantPath(tt.fields)) + }) + } +} + +func TestParseVariantPath(t *testing.T) { + for _, tt := range []struct { + name string + path string + want []string + }{ + {"root", "$", []string{}}, + {"single member", "$.event_id", []string{"event_id"}}, + {"nested members", "$.location.latitude", []string{"location", "latitude"}}, + {"underscore and digits", "$._a1.b2", []string{"_a1", "b2"}}, + {"non-ascii first char", "$.naïve", []string{"naïve"}}, + } { + t.Run(tt.name, func(t *testing.T) { + got, err := parseVariantPath(tt.path) + require.NoError(t, err) + assert.Equal(t, tt.want, got) + }) + } +} + +func TestParseVariantPathRejects(t *testing.T) { + for _, tt := range []struct { + name string + path string + }{ + {"bracket notation", "$['event_id']"}, + {"wildcard", "$.*"}, + {"recursive descent", "$..event_id"}, + {"missing root", "event_id"}, + {"leading digit member", "$.1abc"}, + {"empty member", "$."}, + } { + t.Run(tt.name, func(t *testing.T) { + _, err := parseVariantPath(tt.path) + require.ErrorIs(t, err, ErrInvalidArgument) + }) + } +} diff --git a/visitors.go b/visitors.go index 3ee8382b8..71cf21959 100644 --- a/visitors.go +++ b/visitors.go @@ -471,44 +471,110 @@ func (expressionFieldIDs) VisitBound(pred BoundPredicate) map[int]struct{} { // the field IDs in the file schema. If columns don't exist they are replaced with // AlwaysFalse or AlwaysTrue depending on the operator. func TranslateColumnNames(expr BooleanExpression, fileSchema *Schema) (BooleanExpression, error) { - return VisitExpr(expr, columnNameTranslator{fileSchema: fileSchema}) + res, extracts, err := TranslateColumnNamesForScan(expr, fileSchema) + if err != nil { + return nil, err + } + if len(extracts) > 0 { + return nil, fmt.Errorf("%w: variant extract terms require TranslateColumnNamesForScan", ErrNotImplemented) + } + + return res, nil +} + +// VariantExtractColumn describes a synthetic column that materializes one variant extract term for filtering. +type VariantExtractColumn struct { + Term BoundExtract + FieldID int + Name string } -type columnNameTranslator struct { +// TranslateColumnNamesForScan translates a bound filter to the file schema, mapping variant extract terms to synthetic reference columns. +func TranslateColumnNamesForScan(expr BooleanExpression, fileSchema *Schema) (BooleanExpression, []VariantExtractColumn, error) { + tr := &scanTranslator{ + fileSchema: fileSchema, + nextID: fileSchema.HighestFieldID() + 1, + byKey: map[string]int{}, + } + + res, err := VisitExpr(expr, tr) + if err != nil { + return nil, nil, err + } + + return res, tr.columns, nil +} + +type scanTranslator struct { fileSchema *Schema + nextID int + nameSeq int + byKey map[string]int + columns []VariantExtractColumn } -func (columnNameTranslator) VisitTrue() BooleanExpression { return AlwaysTrue{} } -func (columnNameTranslator) VisitFalse() BooleanExpression { return AlwaysFalse{} } -func (columnNameTranslator) VisitNot(child BooleanExpression) BooleanExpression { +func (*scanTranslator) VisitTrue() BooleanExpression { return AlwaysTrue{} } +func (*scanTranslator) VisitFalse() BooleanExpression { return AlwaysFalse{} } +func (*scanTranslator) VisitNot(child BooleanExpression) BooleanExpression { return NewNot(child) } -func (columnNameTranslator) VisitAnd(left, right BooleanExpression) BooleanExpression { +func (*scanTranslator) VisitAnd(left, right BooleanExpression) BooleanExpression { return NewAnd(left, right) } -func (columnNameTranslator) VisitOr(left, right BooleanExpression) BooleanExpression { +func (*scanTranslator) VisitOr(left, right BooleanExpression) BooleanExpression { return NewOr(left, right) } -func (columnNameTranslator) VisitUnbound(pred UnboundPredicate) BooleanExpression { +func (*scanTranslator) VisitUnbound(pred UnboundPredicate) BooleanExpression { panic(fmt.Errorf("%w: expected bound predicate, got: %s", ErrInvalidArgument, pred.Term())) } -func (c columnNameTranslator) VisitBound(pred BoundPredicate) BooleanExpression { - fileColName, found := c.fileSchema.FindColumnName(pred.Term().Ref().Field().ID) - if !found { - // in the case of schema evolution, the column might not be present - // in the file schema when reading older data - if pred.Op() == OpIsNull { - return AlwaysTrue{} +func (t *scanTranslator) extractRef(ext BoundExtract) Reference { + key := ext.String() + idx, ok := t.byKey[key] + if !ok { + idx = len(t.columns) + t.columns = append(t.columns, VariantExtractColumn{ + Term: ext, + FieldID: t.nextID, + Name: t.syntheticName(), + }) + t.nextID++ + t.byKey[key] = idx + } + + return Reference(t.columns[idx].Name) +} + +// syntheticName returns a derived-column name absent from the file schema. +func (t *scanTranslator) syntheticName() string { + for { + name := fmt.Sprintf("_variant_extract_%d", t.nameSeq) + t.nameSeq++ + if _, found := t.fileSchema.FindFieldByName(name); !found { + return name } + } +} - return AlwaysFalse{} +func (t *scanTranslator) VisitBound(pred BoundPredicate) BooleanExpression { + var ref Reference + if ext, ok := pred.Term().(BoundExtract); ok { + ref = t.extractRef(ext) + } else { + fileColName, found := t.fileSchema.FindColumnName(pred.Term().Ref().Field().ID) + if !found { + if pred.Op() == OpIsNull { + return AlwaysTrue{} + } + + return AlwaysFalse{} + } + ref = Reference(fileColName) } - ref := Reference(fileColName) switch p := pred.(type) { case BoundUnaryPredicate: return p.AsUnbound(ref)