-
Notifications
You must be signed in to change notification settings - Fork 54
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
feat(partitioning): Add partition specs/fields and basic transform interface #2
Merged
Merged
Changes from 1 commit
Commits
Show all changes
9 commits
Select commit
Hold shift + click to select a range
bb70919
feat(partitioning): Add partition specs/fields and basic transform in…
zeroshade 8b796d3
update rat-exclude for rat-results
zeroshade 6430a94
consistency of comment
zeroshade 7ee1981
create transforms.go, case insensitive parse
zeroshade 5583e0e
update version check
zeroshade 38a7484
update from review comments
zeroshade e8c3e38
Update transforms.go
zeroshade 080a71c
update go.sum
zeroshade beeaed2
add test for spec with void AND non-void transform
zeroshade File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Next
Next commit
feat(partitioning): Add partition specs/fields and basic transform in…
…terface
- Loading branch information
commit bb709193aba1452d19c20160a1bf3f3d432d71fd
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,367 @@ | ||
// 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 ( | ||
"encoding" | ||
"encoding/json" | ||
"fmt" | ||
"strconv" | ||
"strings" | ||
|
||
"golang.org/x/exp/slices" | ||
) | ||
|
||
const ( | ||
partitionDataIDStart = 1000 | ||
nastra marked this conversation as resolved.
Show resolved
Hide resolved
|
||
InitialPartitionSpecID = 0 | ||
) | ||
|
||
// UnpartitionedPartitionSpec is the default unpartitioned spec which can | ||
// be used for comparisons or to just provide a convenience for referencing | ||
// the same unpartitioned spec object. | ||
var UnpartitionedPartitionSpec = &PartitionSpec{id: 0} | ||
zeroshade marked this conversation as resolved.
Show resolved
Hide resolved
|
||
|
||
// PartitionField represents how one partition value is derived from the | ||
// source column by transformation. | ||
type PartitionField struct { | ||
// SourceID is the source column ID of the table's schema | ||
zeroshade marked this conversation as resolved.
Show resolved
Hide resolved
|
||
SourceID int `json:"source-id"` | ||
// FieldID is the partition field id across all the table partition specs | ||
FieldID int `json:"field-id"` | ||
// Name is the name of the partition field itself | ||
Name string `json:"name"` | ||
// Transform is the transform used to produce the partition value | ||
Transform Transform `json:"transform"` | ||
} | ||
|
||
func (p *PartitionField) String() string { | ||
return fmt.Sprintf("%d: %s: %s(%d)", p.FieldID, p.Name, p.Transform, p.SourceID) | ||
} | ||
|
||
func (p *PartitionField) UnmarshalJSON(b []byte) error { | ||
type Alias PartitionField | ||
aux := struct { | ||
TransformString string `json:"transform"` | ||
*Alias | ||
}{ | ||
Alias: (*Alias)(p), | ||
} | ||
|
||
err := json.Unmarshal(b, &aux) | ||
if err != nil { | ||
return err | ||
} | ||
|
||
if p.Transform, err = ParseTransform(aux.TransformString); err != nil { | ||
return err | ||
} | ||
|
||
return nil | ||
} | ||
|
||
// PartitionSpec captures the transformation from table data to partition values | ||
type PartitionSpec struct { | ||
// any change to a PartitionSpec will produce a new spec id | ||
id int | ||
fields []PartitionField | ||
|
||
// this is populated by initialize after creation | ||
sourceIdToFields map[int][]PartitionField | ||
} | ||
|
||
func NewPartitionSpec(fields ...PartitionField) PartitionSpec { | ||
return NewPartitionSpecID(InitialPartitionSpecID, fields...) | ||
} | ||
|
||
func NewPartitionSpecID(id int, fields ...PartitionField) PartitionSpec { | ||
ret := PartitionSpec{id: id, fields: fields} | ||
ret.initialize() | ||
return ret | ||
} | ||
|
||
// CompatibleWith returns true if this partition spec is considered | ||
// compatible with the passed in partition spec. This means that the two | ||
// specs have equivalent field lists regardless of the spec id. | ||
func (ps *PartitionSpec) CompatibleWith(other *PartitionSpec) bool { | ||
if ps == other { | ||
return true | ||
} | ||
|
||
if len(ps.fields) != len(other.fields) { | ||
return false | ||
} | ||
|
||
return slices.EqualFunc(ps.fields, other.fields, func(left, right PartitionField) bool { | ||
zeroshade marked this conversation as resolved.
Show resolved
Hide resolved
|
||
return left.SourceID == right.SourceID && left.Name == right.Name && | ||
left.Transform == right.Transform | ||
}) | ||
} | ||
|
||
// Equals returns true iff the field lists are the same AND the spec id | ||
// is the same between this partition spec and the provided one. | ||
func (ps *PartitionSpec) Equals(other PartitionSpec) bool { | ||
return ps.id == other.id && slices.Equal(ps.fields, other.fields) | ||
} | ||
|
||
func (ps PartitionSpec) MarshalJSON() ([]byte, error) { | ||
if ps.fields == nil { | ||
ps.fields = []PartitionField{} | ||
} | ||
return json.Marshal(struct { | ||
ID int `json:"spec-id"` | ||
Fields []PartitionField `json:"fields"` | ||
}{ps.id, ps.fields}) | ||
} | ||
|
||
func (ps *PartitionSpec) UnmarshalJSON(b []byte) error { | ||
aux := struct { | ||
ID int `json:"spec-id"` | ||
Fields []PartitionField `json:"fields"` | ||
}{ID: ps.id, Fields: ps.fields} | ||
|
||
if err := json.Unmarshal(b, &aux); err != nil { | ||
return err | ||
} | ||
|
||
ps.id, ps.fields = aux.ID, aux.Fields | ||
ps.initialize() | ||
return nil | ||
} | ||
|
||
func (ps *PartitionSpec) initialize() { | ||
ps.sourceIdToFields = make(map[int][]PartitionField) | ||
for _, f := range ps.fields { | ||
ps.sourceIdToFields[f.SourceID] = | ||
append(ps.sourceIdToFields[f.SourceID], f) | ||
} | ||
} | ||
|
||
func (ps *PartitionSpec) ID() int { return ps.id } | ||
func (ps *PartitionSpec) NumFields() int { return len(ps.fields) } | ||
func (ps *PartitionSpec) Field(i int) PartitionField { return ps.fields[i] } | ||
func (ps *PartitionSpec) IsUnpartitioned() bool { return len(ps.fields) == 0 } | ||
zeroshade marked this conversation as resolved.
Show resolved
Hide resolved
|
||
func (ps *PartitionSpec) FieldsBySourceID(fieldID int) []PartitionField { | ||
return slices.Clone(ps.sourceIdToFields[fieldID]) | ||
} | ||
|
||
func (ps PartitionSpec) String() string { | ||
var b strings.Builder | ||
b.WriteByte('[') | ||
for i, f := range ps.fields { | ||
if i == 0 { | ||
b.WriteString("\n") | ||
} | ||
b.WriteString("\t") | ||
b.WriteString(f.String()) | ||
b.WriteString("\n") | ||
} | ||
b.WriteByte(']') | ||
|
||
return b.String() | ||
} | ||
|
||
func (ps *PartitionSpec) LastAssignedFieldID() int { | ||
if len(ps.fields) == 0 { | ||
return partitionDataIDStart | ||
} | ||
|
||
id := ps.fields[0].FieldID | ||
for _, f := range ps.fields[1:] { | ||
if f.FieldID > id { | ||
id = f.FieldID | ||
} | ||
} | ||
return id | ||
} | ||
|
||
// PartitionType produces a struct of the partition spec. | ||
// | ||
// The partition fields should be optional: | ||
// - All partition transforms are required to produce null if the input value | ||
// is null. This can happen when the source column is optional. | ||
// - Partition fields may be added later, in which case not all files would | ||
// have the result field and it may be null. | ||
// | ||
// There is a case where we can guarantee that a partition field in the first | ||
// and only parittion spec that uses a required source column will never be | ||
// null, but it doesn't seem worth tracking this case. | ||
func (ps *PartitionSpec) PartitionType(schema *Schema) *StructType { | ||
nestedFields := []NestedField{} | ||
for _, field := range ps.fields { | ||
sourceType, ok := schema.FindTypeByID(field.SourceID) | ||
if !ok { | ||
continue | ||
} | ||
resultType := field.Transform.ResultType(sourceType) | ||
nestedFields = append(nestedFields, NestedField{ | ||
ID: field.FieldID, | ||
Name: field.Name, | ||
Type: resultType, | ||
Required: false, | ||
}) | ||
} | ||
return &StructType{FieldList: nestedFields} | ||
} | ||
|
||
// ParseTransform takes the string representation of a transform as | ||
// defined in the iceberg spec, and produces the appropriate Transform | ||
// object or an error if the string is not a valid transform string. | ||
func ParseTransform(s string) (Transform, error) { | ||
switch { | ||
case strings.HasPrefix(s, "bucket"): | ||
matches := regexFromBrackets.FindStringSubmatch(s) | ||
if len(matches) != 2 { | ||
break | ||
} | ||
|
||
n, _ := strconv.Atoi(matches[1]) | ||
return BucketTransform{N: n}, nil | ||
case strings.HasPrefix(s, "truncate"): | ||
matches := regexFromBrackets.FindStringSubmatch(s) | ||
if len(matches) != 2 { | ||
break | ||
} | ||
|
||
n, _ := strconv.Atoi(matches[1]) | ||
return TruncateTransform{W: n}, nil | ||
default: | ||
switch s { | ||
case "identity": | ||
return IdentityTransform{}, nil | ||
case "void": | ||
return VoidTransform{}, nil | ||
case "year": | ||
return YearTransform{}, nil | ||
case "month": | ||
return MonthTransform{}, nil | ||
case "day": | ||
return DayTransform{}, nil | ||
case "hour": | ||
return HourTransform{}, nil | ||
} | ||
} | ||
|
||
return nil, fmt.Errorf("%w: %s", ErrInvalidTransform, s) | ||
} | ||
|
||
// Transform is an interface for the various Transformation types | ||
// in partition specs. Currently they do not yet provide actual | ||
// transformation functions or implementation. That will come later as | ||
// data reading gets implemented. | ||
type Transform interface { | ||
fmt.Stringer | ||
encoding.TextMarshaler | ||
ResultType(t Type) Type | ||
} | ||
|
||
// IdentityTransform uses the identity function, performing no transformation | ||
// but instead partitioning on the value itself. | ||
type IdentityTransform struct{} | ||
zeroshade marked this conversation as resolved.
Show resolved
Hide resolved
|
||
|
||
func (t IdentityTransform) MarshalText() ([]byte, error) { | ||
return []byte(t.String()), nil | ||
} | ||
|
||
func (IdentityTransform) String() string { return "identity" } | ||
|
||
func (IdentityTransform) ResultType(t Type) Type { return t } | ||
|
||
// VoidTransform is a transformation that always returns nil. | ||
type VoidTransform struct{} | ||
|
||
func (t VoidTransform) MarshalText() ([]byte, error) { | ||
return []byte(t.String()), nil | ||
} | ||
|
||
func (VoidTransform) String() string { return "void" } | ||
|
||
func (VoidTransform) ResultType(t Type) Type { return t } | ||
|
||
// BucketTransform transforms values into a bucket partition value. It is | ||
// parameterized by a number of buckets. Bucket partition transforms use | ||
// a 32-bit hash of the source value to produce a positive value by mod | ||
// the bucket number. | ||
type BucketTransform struct { | ||
N int | ||
} | ||
|
||
func (t BucketTransform) MarshalText() ([]byte, error) { | ||
return []byte(t.String()), nil | ||
} | ||
|
||
func (t BucketTransform) String() string { return fmt.Sprintf("bucket[%d]", t.N) } | ||
|
||
func (BucketTransform) ResultType(Type) Type { return PrimitiveTypes.Int32 } | ||
|
||
// TruncateTransform is a transformation for truncating a value to a specified width. | ||
type TruncateTransform struct { | ||
W int | ||
} | ||
|
||
func (t TruncateTransform) MarshalText() ([]byte, error) { | ||
return []byte(t.String()), nil | ||
} | ||
|
||
func (t TruncateTransform) String() string { return fmt.Sprintf("truncate[%d]", t.W) } | ||
|
||
func (TruncateTransform) ResultType(t Type) Type { return t } | ||
|
||
// YearTransform transforms a datetime value into a year value. | ||
type YearTransform struct{} | ||
|
||
func (t YearTransform) MarshalText() ([]byte, error) { | ||
return []byte(t.String()), nil | ||
} | ||
|
||
func (YearTransform) String() string { return "year" } | ||
|
||
func (YearTransform) ResultType(Type) Type { return PrimitiveTypes.Int32 } | ||
|
||
// MonthTransform transforms a datetime value into a month value. | ||
type MonthTransform struct{} | ||
|
||
func (t MonthTransform) MarshalText() ([]byte, error) { | ||
return []byte(t.String()), nil | ||
} | ||
|
||
func (MonthTransform) String() string { return "month" } | ||
|
||
func (MonthTransform) ResultType(Type) Type { return PrimitiveTypes.Int32 } | ||
|
||
// DayTransform transforms a datetime value into a date value. | ||
type DayTransform struct{} | ||
|
||
func (t DayTransform) MarshalText() ([]byte, error) { | ||
return []byte(t.String()), nil | ||
} | ||
|
||
func (DayTransform) String() string { return "day" } | ||
|
||
func (DayTransform) ResultType(Type) Type { return PrimitiveTypes.Date } | ||
|
||
// HourTransform transforms a datetime value into an hour value. | ||
type HourTransform struct{} | ||
|
||
func (t HourTransform) MarshalText() ([]byte, error) { | ||
return []byte(t.String()), nil | ||
} | ||
|
||
func (HourTransform) String() string { return "hour" } | ||
|
||
func (HourTransform) ResultType(Type) Type { return PrimitiveTypes.Int32 } |
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Generally I am really against using the experimental packages for anything that will be used likely in production. There recently has been a lot of people ignoring this and creating havoc in the release cycles of Go as seen in the following issues:
golang/go#51317
golang/go#61374
I think it would be better at this time to avoid using anything in the exp branch.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
@zeroshade thoughts on this one?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Sorry I forgot to respond here on the PR because we ended up discussing this in the slack channel.
Tl;DR: the standard for Go in general is to target N-1 versions (ie. right now the latest is Go 1.21, so we want to target Go 1.20 + Go 1.21). This slices package was only added to the Go stdlib in Go 1.21, so we still need to point at
golang.org/x/exp/slices
until/unless we drop support for Go 1.20, at which point we can update this to just use the stdlib version of the package.We're still about 6 months or so before Go 1.22 would get released, but this should be fine for now. Once Go 1.22 is released I'll replace these with the stdlib package, letting us maintain the N-1 compatibility.