zeroshade commented on code in PR #185: URL: https://github.com/apache/iceberg-go/pull/185#discussion_r1827805391
########## table/substrait/substrait.go: ########## @@ -0,0 +1,379 @@ +// 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 substrait + +import ( + _ "embed" + "fmt" + "strings" + + "github.com/apache/arrow-go/v18/arrow/compute/exprs" + "github.com/apache/iceberg-go" + "github.com/substrait-io/substrait-go/expr" + "github.com/substrait-io/substrait-go/extensions" + "github.com/substrait-io/substrait-go/types" +) + +//go:embed functions_set.yaml +var funcsetYAML string + +var ( + collection = extensions.DefaultCollection + funcSetURI = "https://github.com/apache/iceberg-go/blob/main/table/substrait/functions_set.yaml"; +) + +func init() { + if !collection.URILoaded(funcSetURI) { + if err := collection.Load(funcSetURI, strings.NewReader(funcsetYAML)); err != nil { + panic(err) + } + } +} + +func NewExtensionSet() exprs.ExtensionIDSet { + return exprs.NewExtensionSetDefault(expr.NewEmptyExtensionRegistry(&collection)) +} + +// ConvertExpr binds the provided expression to the given schema and converts it to a +// substrait expression so that it can be utilized for computation. +func ConvertExpr(schema *iceberg.Schema, e iceberg.BooleanExpression) (*expr.ExtensionRegistry, expr.Expression, error) { + base, err := ConvertSchema(schema) + if err != nil { + return nil, nil, err + } + + reg := expr.NewEmptyExtensionRegistry(&extensions.DefaultCollection) + + bldr := expr.ExprBuilder{Reg: reg, BaseSchema: &base.Struct} + b, err := iceberg.VisitExpr(e, &toSubstraitExpr{bldr: bldr, schema: schema}) + if err != nil { + return nil, nil, err + } + + out, err := b.BuildExpr() + return ®, out, err +} + +// ConvertSchema converts an Iceberg schema to a substrait NamedStruct using +// the appropriate types and column names. +func ConvertSchema(schema *iceberg.Schema) (res types.NamedStruct, err error) { + var typ types.Type + + typ, err = iceberg.Visit(schema, convertToSubstrait{}) + if err != nil { + return + } + + val := typ.(*types.StructType) + res.Struct = *val + + res.Names = make([]string, schema.NumFields()) + for i, f := range schema.Fields() { + res.Names[i] = f.Name + } + + return +} + +type convertToSubstrait struct{} + +func (convertToSubstrait) Schema(_ *iceberg.Schema, result types.Type) types.Type { + return result.WithNullability(types.NullabilityNullable) +} + +func (convertToSubstrait) Struct(_ iceberg.StructType, results []types.Type) types.Type { + return &types.StructType{ + Nullability: types.NullabilityUnspecified, + Types: results, + } +} + +func getNullability(required bool) types.Nullability { + if required { + return types.NullabilityRequired + } + return types.NullabilityNullable +} + +func (convertToSubstrait) Field(field iceberg.NestedField, result types.Type) types.Type { + return result.WithNullability(getNullability(field.Required)) +} + +func (c convertToSubstrait) List(list iceberg.ListType, elemResult types.Type) types.Type { + return &types.ListType{ + Nullability: types.NullabilityUnspecified, + Type: c.Field(list.ElementField(), elemResult), + } +} + +func (c convertToSubstrait) Map(m iceberg.MapType, keyResult, valResult types.Type) types.Type { + return &types.MapType{ + Nullability: types.NullabilityUnspecified, + Key: c.Field(m.KeyField(), keyResult), + Value: c.Field(m.ValueField(), valResult), + } +} + +func (convertToSubstrait) Primitive(iceberg.PrimitiveType) types.Type { panic("should not be called") } + +func (convertToSubstrait) VisitFixed(f iceberg.FixedType) types.Type { + return &types.FixedBinaryType{Length: int32(f.Len())} +} + +func (convertToSubstrait) VisitDecimal(d iceberg.DecimalType) types.Type { + return &types.DecimalType{Precision: int32(d.Precision()), Scale: int32(d.Scale())} +} + +func (convertToSubstrait) VisitBoolean() types.Type { return &types.BooleanType{} } +func (convertToSubstrait) VisitInt32() types.Type { return &types.Int32Type{} } +func (convertToSubstrait) VisitInt64() types.Type { return &types.Int64Type{} } +func (convertToSubstrait) VisitFloat32() types.Type { return &types.Float32Type{} } +func (convertToSubstrait) VisitFloat64() types.Type { return &types.Float64Type{} } +func (convertToSubstrait) VisitDate() types.Type { return &types.DateType{} } +func (convertToSubstrait) VisitTime() types.Type { return &types.TimeType{} } +func (convertToSubstrait) VisitTimestamp() types.Type { return &types.TimestampType{} } +func (convertToSubstrait) VisitTimestampTz() types.Type { return &types.TimestampTzType{} } +func (convertToSubstrait) VisitString() types.Type { return &types.StringType{} } +func (convertToSubstrait) VisitBinary() types.Type { return &types.BinaryType{} } +func (convertToSubstrait) VisitUUID() types.Type { return &types.UUIDType{} } + +var ( + _ iceberg.SchemaVisitorPerPrimitiveType[types.Type] = (*convertToSubstrait)(nil) +) + +var ( + boolURI = extensions.SubstraitDefaultURIPrefix + "functions_boolean.yaml" + compareURI = extensions.SubstraitDefaultURIPrefix + "functions_comparison.yaml" + stringURI = extensions.SubstraitDefaultURIPrefix + "functions_string.yaml" + + notID = extensions.ID{URI: boolURI, Name: "not"} + andID = extensions.ID{URI: boolURI, Name: "and"} + orID = extensions.ID{URI: boolURI, Name: "or"} + isNaNID = extensions.ID{URI: compareURI, Name: "is_nan"} + isNullID = extensions.ID{URI: compareURI, Name: "is_null"} + isNotNullID = extensions.ID{URI: compareURI, Name: "is_not_null"} + equalID = extensions.ID{URI: compareURI, Name: "equal"} + notEqualID = extensions.ID{URI: compareURI, Name: "not_equal"} + greaterEqualID = extensions.ID{URI: compareURI, Name: "gte"} + greaterID = extensions.ID{URI: compareURI, Name: "gt"} + lessEqualID = extensions.ID{URI: compareURI, Name: "lte"} + lessID = extensions.ID{URI: compareURI, Name: "lt"} + startsWithID = extensions.ID{URI: stringURI, Name: "starts_with"} + isInID = extensions.ID{URI: funcSetURI, Name: "is_in"} +) + +type toSubstraitExpr struct { + schema *iceberg.Schema + bldr expr.ExprBuilder +} + +func (t *toSubstraitExpr) VisitTrue() expr.Builder { + return t.bldr.Wrap(expr.NewLiteral(true, false)) +} + +func (t *toSubstraitExpr) VisitFalse() expr.Builder { + return t.bldr.Wrap(expr.NewLiteral(false, false)) +} + +func (t *toSubstraitExpr) VisitNot(child expr.Builder) expr.Builder { + return t.bldr.ScalarFunc(notID).Args(child.(expr.FuncArgBuilder)) +} + +func (t *toSubstraitExpr) VisitAnd(left, right expr.Builder) expr.Builder { + return t.bldr.ScalarFunc(andID).Args(left.(expr.FuncArgBuilder), + right.(expr.FuncArgBuilder)) +} + +func (t *toSubstraitExpr) VisitOr(left, right expr.Builder) expr.Builder { + return t.bldr.ScalarFunc(orID).Args(left.(expr.FuncArgBuilder), + right.(expr.FuncArgBuilder)) +} + +func (t *toSubstraitExpr) VisitUnbound(iceberg.UnboundPredicate) expr.Builder { + panic("can only convert bound expressions to substrait") +} + +func (t *toSubstraitExpr) VisitBound(pred iceberg.BoundPredicate) expr.Builder { + return iceberg.VisitBoundPredicate(pred, t) +} + +type substraitPrimitiveLiteralTypes interface { + bool | ~int32 | ~int64 | float32 | float64 | string +} + +func toPrimitiveSubstraitLiteral[T substraitPrimitiveLiteralTypes](v T) expr.Literal { + return expr.NewPrimitiveLiteral(v, false) +} + +func toByteSliceSubstraitLiteral[T []byte | types.UUID](v T) expr.Literal { + return expr.NewByteSliceLiteral(v, false) +} + +func toDecimalLiteral(v iceberg.DecimalLiteral) expr.Literal { + byts, _ := v.MarshalBinary() + result, _ := expr.NewLiteral(&types.Decimal{ + Scale: int32(v.Scale), + Value: byts, + Precision: int32(v.Type().(*iceberg.DecimalType).Precision()), + }, false) + return result +} + +func toFixedLiteral(v iceberg.FixedLiteral) expr.Literal { + return expr.NewFixedBinaryLiteral(types.FixedBinary(v), false) +} + +func toSubstraitLiteral(typ iceberg.Type, lit iceberg.Literal) expr.Literal { + switch lit := lit.(type) { + case iceberg.BoolLiteral: + return toPrimitiveSubstraitLiteral(bool(lit)) + case iceberg.Int32Literal: + return toPrimitiveSubstraitLiteral(int32(lit)) + case iceberg.Int64Literal: + return toPrimitiveSubstraitLiteral(int64(lit)) + case iceberg.Float32Literal: + return toPrimitiveSubstraitLiteral(float32(lit)) + case iceberg.Float64Literal: + return toPrimitiveSubstraitLiteral(float64(lit)) + case iceberg.StringLiteral: + return toPrimitiveSubstraitLiteral(string(lit)) + case iceberg.TimestampLiteral: + if typ.Equals(iceberg.PrimitiveTypes.TimestampTz) { + return toPrimitiveSubstraitLiteral(types.TimestampTz(lit)) + } + return toPrimitiveSubstraitLiteral(types.Timestamp(lit)) + case iceberg.DateLiteral: + return toPrimitiveSubstraitLiteral(types.Date(lit)) + case iceberg.TimeLiteral: + return toPrimitiveSubstraitLiteral(types.Time(lit)) + case iceberg.BinaryLiteral: + return toByteSliceSubstraitLiteral([]byte(lit)) + case iceberg.FixedLiteral: + return toFixedLiteral(lit) + case iceberg.UUIDLiteral: + return toByteSliceSubstraitLiteral(types.UUID(lit[:])) + case iceberg.DecimalLiteral: + return toDecimalLiteral(lit) + } + panic(fmt.Errorf("invalid literal type: %s", lit.Type())) +} + +func toSubstraitLiteralSet(typ iceberg.Type, lits []iceberg.Literal) expr.ListLiteralValue { + if len(lits) == 0 { + return nil + } + + out := make([]expr.Literal, len(lits)) + for i, l := range lits { + out[i] = toSubstraitLiteral(typ, l) + } + return out +} + +func (t *toSubstraitExpr) getRef(ref iceberg.BoundReference) expr.Reference { + updatedRef, err := iceberg.Reference(ref.Field().Name).Bind(t.schema, true) Review Comment: good point, i've updated everything to pass the caseSensitive down from the ScanOption -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: issues-unsubscr...@iceberg.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org --------------------------------------------------------------------- To unsubscribe, e-mail: issues-unsubscr...@iceberg.apache.org For additional commands, e-mail: issues-h...@iceberg.apache.org
