This is an automated email from the ASF dual-hosted git repository.
zeroshade pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/iceberg-go.git
The following commit(s) were added to refs/heads/main by this push:
new d6831edf0 feat(table): import bucket-partitioned Parquet files (#2395)
d6831edf0 is described below
commit d6831edf0f98c7554d00fb2957ac2b51ba897822
Author: Minh Vu <[email protected]>
AuthorDate: Wed Nov 4 19:11:40 2026 +0100
feat(table): import bucket-partitioned Parquet files (#2395)
* feat(table): import bucket-partitioned Parquet into v1 tables
Signed-off-by: Minh Vu <[email protected]>
* fix(table): preserve file import API compatibility
Signed-off-by: Minh Vu <[email protected]>
* fix(table): snapshot partition values before concurrent imports
Signed-off-by: Minh Vu <[email protected]>
* fix(table): reject unknown partition transforms during file import
* test(table): prevent unknown transform bypass via supplied partition tuple
* fix(table): validate supplied partition values against Parquet bounds
* style(table): sort arrow utils imports
* fix(table): preserve addFilesOperation comparability
---------
Signed-off-by: Minh Vu <[email protected]>
---
table/arrow_utils.go | 249 ++++++++++++++++++++++++---
table/file_partition_values_test.go | 74 +++++++++
table/internal/utils.go | 6 +-
table/table_test.go | 2 +-
table/transaction.go | 40 ++++-
table/v1_partition_import_test.go | 323 ++++++++++++++++++++++++++++++++++++
6 files changed, 666 insertions(+), 28 deletions(-)
diff --git a/table/arrow_utils.go b/table/arrow_utils.go
index 1cc15095e..896909d4f 100644
--- a/table/arrow_utils.go
+++ b/table/arrow_utils.go
@@ -25,6 +25,7 @@ import (
"fmt"
"iter"
"log/slog"
+ "reflect"
"regexp"
"slices"
"strconv"
@@ -1853,7 +1854,7 @@ func computeStatsPlan(sc *iceberg.Schema, props
iceberg.Properties) (result map[
return result, nil
}
-func filesToDataFiles(ctx context.Context, fileIO iceio.IO, meta
*MetadataBuilder, filePaths []string, concurrency int) (_ []iceberg.DataFile,
err error) {
+func filesToDataFiles(ctx context.Context, fileIO iceio.IO, meta
*MetadataBuilder, filePaths []string, concurrency int, partitionValuesByPath
map[string]map[int]any) (_ []iceberg.DataFile, err error) {
partitionSpec, err := meta.CurrentSpec()
if err != nil || partitionSpec == nil {
return nil, fmt.Errorf("%w: cannot add files without a current
spec", err)
@@ -1884,7 +1885,7 @@ func filesToDataFiles(ctx context.Context, fileIO
iceio.IO, meta *MetadataBuilde
// Pass 0 so AddDataFilesByPath does not claim
sort_order_id on files
// it did not write (#1184). Callers that know the
file's
// sort layout use FileToDataFile with an explicit id.
- dataFiles[i] = fileToDataFile(ctx, fileIO, filePath,
currentSchema, currentSpec, 0, meta.props)
+ dataFiles[i] = fileToDataFile(ctx, fileIO, filePath,
currentSchema, currentSpec, 0, meta.props, partitionValuesByPath[filePath])
return nil
})
@@ -1897,6 +1898,160 @@ func filesToDataFiles(ctx context.Context, fileIO
iceio.IO, meta *MetadataBuilde
return dataFiles, nil
}
+// FileToDataFileOption configures the conversion of an existing Parquet file.
+type FileToDataFileOption func(*fileToDataFileConfig)
+
+type fileToDataFileConfig struct {
+ partitionValues map[int]any
+}
+
+// WithFilePartitionValues supplies the partition tuple of an existing data
file,
+// keyed by partition field ID (not source column ID). This is required for
+// transforms such as bucket[N], which cannot be inferred from Parquet
statistics.
+// For a spec containing a bucket transform, all non-void fields must be
supplied.
+// Each supplied value must match the transform's result type. A nil value
+// explicitly represents a null partition. Values for order-preserving
+// transforms are checked against Parquet bounds when available. Callers remain
+// responsible for ensuring that all rows agree with the partition values.
+func WithFilePartitionValues(values map[int]any) FileToDataFileOption {
+ partitionValues := cloneFilePartitionTuple(values)
+
+ return func(cfg *fileToDataFileConfig) {
+ cfg.partitionValues = partitionValues
+ }
+}
+
+// cloneFilePartitionTuple prevents the caller from changing partition values
+// while concurrently imported Parquet files are being inspected.
+func cloneFilePartitionTuple(values map[int]any) map[int]any {
+ if values == nil {
+ return nil
+ }
+ cloned := make(map[int]any, len(values))
+ for id, value := range values {
+ if raw, ok := value.([]byte); ok {
+ value = bytes.Clone(raw)
+ }
+ cloned[id] = value
+ }
+
+ return cloned
+}
+
+func validateFileBucketValue(field iceberg.PartitionField, value any,
numBuckets int) error {
+ if value == nil {
+ return nil
+ }
+
+ bucket, ok := value.(int32)
+ if !ok || bucket < 0 || int64(bucket) >= int64(numBuckets) {
+ return fmt.Errorf("%w: partition field %q (id %d) requires an
int32 bucket in [0, %d), got %v (%T)",
+ iceberg.ErrInvalidArgument, field.Name, field.FieldID,
numBuckets, value, value)
+ }
+
+ return nil
+}
+
+func validFilePartitionValueType(typ iceberg.Type, value any) bool {
+ switch t := typ.(type) {
+ case iceberg.BooleanType:
+ _, ok := value.(bool)
+
+ return ok
+ case iceberg.Int32Type:
+ _, ok := value.(int32)
+
+ return ok
+ case iceberg.Int64Type:
+ _, ok := value.(int64)
+
+ return ok
+ case iceberg.Float32Type:
+ _, ok := value.(float32)
+
+ return ok
+ case iceberg.Float64Type:
+ _, ok := value.(float64)
+
+ return ok
+ case iceberg.StringType:
+ _, ok := value.(string)
+
+ return ok
+ case iceberg.BinaryType:
+ _, ok := value.([]byte)
+
+ return ok
+ case iceberg.FixedType:
+ v, ok := value.([]byte)
+
+ return ok && len(v) == t.Len()
+ case iceberg.DateType:
+ _, ok := value.(iceberg.Date)
+
+ return ok
+ case iceberg.TimeType:
+ _, ok := value.(iceberg.Time)
+
+ return ok
+ case iceberg.TimestampType, iceberg.TimestampTzType:
+ _, ok := value.(iceberg.Timestamp)
+
+ return ok
+ case iceberg.TimestampNsType, iceberg.TimestampTzNsType:
+ _, ok := value.(iceberg.TimestampNano)
+
+ return ok
+ case iceberg.UUIDType:
+ _, ok := value.(uuid.UUID)
+
+ return ok
+ case iceberg.DecimalType:
+ v, ok := value.(iceberg.Decimal)
+
+ return ok && v.Scale == t.Scale() &&
v.Val.FitsInPrecision(int32(t.Precision()))
+ default:
+ return false
+ }
+}
+
+func validateFilePartitionValues(schema *iceberg.Schema, spec
iceberg.PartitionSpec, values map[int]any) error {
+ if err := validatePartitionValues(spec, values); err != nil {
+ return fmt.Errorf("%w: %v", iceberg.ErrInvalidArgument, err)
+ }
+
+ for _, field := range spec.Fields() {
+ value, supplied := values[field.FieldID]
+ if !supplied {
+ continue
+ }
+
+ switch transform := field.Transform.(type) {
+ case iceberg.BucketTransform:
+ if err := validateFileBucketValue(field, value,
transform.NumBuckets); err != nil {
+ return err
+ }
+ case iceberg.VoidTransform:
+ if value != nil {
+ return fmt.Errorf("%w: void partition field %q
(id %d) must be null",
+ iceberg.ErrInvalidArgument, field.Name,
field.FieldID)
+ }
+ }
+ if value != nil {
+ source, ok :=
schema.FindFieldByID(field.FirstSourceID())
+ if !ok {
+ return fmt.Errorf("%w: source field for
partition %q not found", iceberg.ErrInvalidSchema, field.Name)
+ }
+ resultType := field.Transform.ResultType(source.Type)
+ if !validFilePartitionValueType(resultType, value) {
+ return fmt.Errorf("%w: partition field %q
expects %s, got %T", iceberg.ErrInvalidArgument, field.Name, resultType, value)
+ }
+ }
+ }
+
+ return nil
+}
+
// FileToDataFile builds a DataFile for an existing parquet file at
// filePath, reading its footer to populate record count, file size,
// column sizes, value/null counts and lower/upper bounds, and to infer
@@ -1905,7 +2060,13 @@ func filesToDataFiles(ctx context.Context, fileIO
iceio.IO, meta *MetadataBuilde
// footers can convey the file's sort layout (spec data-file field
// sort_order_id); pass 0 to make no claim. Panics from the unexported
// conversion are recovered into an error.
-func FileToDataFile(ctx context.Context, fileIO iceio.IO, filePath string,
currentSchema *iceberg.Schema, currentSpec iceberg.PartitionSpec, sortOrderID
int, props iceberg.Properties) (df iceberg.DataFile, err error) {
+func FileToDataFile(ctx context.Context, fileIO iceio.IO, filePath string,
currentSchema *iceberg.Schema, currentSpec iceberg.PartitionSpec, sortOrderID
int, props iceberg.Properties) (iceberg.DataFile, error) {
+ return FileToDataFileWithOptions(ctx, fileIO, filePath, currentSchema,
currentSpec, sortOrderID, props)
+}
+
+// FileToDataFileWithOptions builds a DataFile like FileToDataFile, with
options
+// to supply partition values that cannot be inferred from Parquet statistics.
+func FileToDataFileWithOptions(ctx context.Context, fileIO iceio.IO, filePath
string, currentSchema *iceberg.Schema, currentSpec iceberg.PartitionSpec,
sortOrderID int, props iceberg.Properties, opts ...FileToDataFileOption) (df
iceberg.DataFile, err error) {
defer func() {
if r := recover(); r != nil {
switch e := r.(type) {
@@ -1917,10 +2078,24 @@ func FileToDataFile(ctx context.Context, fileIO
iceio.IO, filePath string, curre
}
}()
- return fileToDataFile(ctx, fileIO, filePath, currentSchema,
currentSpec, sortOrderID, props), nil
+ cfg := fileToDataFileConfig{}
+ for _, opt := range opts {
+ opt(&cfg)
+ }
+
+ return fileToDataFile(ctx, fileIO, filePath, currentSchema,
currentSpec, sortOrderID, props, cfg.partitionValues), nil
}
-func fileToDataFile(ctx context.Context, fileIO iceio.IO, filePath string,
currentSchema *iceberg.Schema, currentSpec iceberg.PartitionSpec, sortOrderID
int, props iceberg.Properties) iceberg.DataFile {
+func fileToDataFile(ctx context.Context, fileIO iceio.IO, filePath string,
currentSchema *iceberg.Schema, currentSpec iceberg.PartitionSpec, sortOrderID
int, props iceberg.Properties, supplied map[int]any) iceberg.DataFile {
+ // Explicit partition tuples must not bypass the Iceberg requirement
that
+ // writers reject unknown partition transforms.
+ if err := checkNoUnknownTransform(¤tSpec); err != nil {
+ panic(err)
+ }
+ if err := validateFilePartitionValues(currentSchema, currentSpec,
supplied); err != nil {
+ panic(err)
+ }
+
format := tblutils.FormatFromFileName(filePath)
rdr := must(format.Open(ctx, fileIO, filePath))
defer rdr.Close()
@@ -1945,29 +2120,59 @@ func fileToDataFile(ctx context.Context, fileIO
iceio.IO, filePath string, curre
arrSchema,
)
- partitionValues := make(map[int]any)
- if !currentSpec.Equals(*iceberg.UnpartitionedSpec) {
- for _, field := range currentSpec.Fields() {
- if !field.Transform.PreservesOrder() {
- panic(fmt.Errorf("cannot infer partition value
from parquet metadata for a non-linear partition field: %s with transform %s",
field.Name, field.Transform))
+ // For bucket specs, require the complete tuple instead of combining
caller
+ // values with footer inference, which cannot prove null partition
values.
+ hasBucket := false
+ for _, field := range currentSpec.Fields() {
+ switch field.Transform.(type) {
+ case iceberg.BucketTransform:
+ hasBucket = true
+ }
+ }
+ partitionValues := make(map[int]any, currentSpec.NumFields())
+ for _, field := range currentSpec.Fields() {
+ if value, ok := supplied[field.FieldID]; ok {
+ if field.Transform.PreservesOrder() {
+ // Keep the same lower/upper bound checks as
the inference path.
+ inferred := statistics.PartitionValue(field,
currentSchema)
+ if inferred != nil && !reflect.DeepEqual(value,
inferred) {
+ panic(fmt.Errorf("%w: supplied
partition value for field %q (id %d) does not match Parquet bounds: got %v,
expected %v",
+ iceberg.ErrInvalidArgument,
field.Name, field.FieldID, value, inferred))
+ }
}
+ partitionValues[field.FieldID] = value
- partitionVal := statistics.PartitionValue(field,
currentSchema)
- if partitionVal != nil {
- partitionValues[field.FieldID] = partitionVal
- }
+ continue
+ }
+ switch field.Transform.(type) {
+ case iceberg.VoidTransform:
+ // Iceberg v1 uses void transforms to drop a partition
field
+ // without changing its field ID or position.
+ partitionValues[field.FieldID] = nil
+
+ continue
+ }
+ if hasBucket || !field.Transform.PreservesOrder() {
+ panic(fmt.Errorf("%w: cannot infer partition field %q
(id %d) with transform %s; supply WithFilePartitionValues",
+ iceberg.ErrInvalidArgument, field.Name,
field.FieldID, field.Transform))
+ }
+
+ partitionVal := statistics.PartitionValue(field, currentSchema)
+ if partitionVal != nil {
+ partitionValues[field.FieldID] = partitionVal
}
}
return statistics.ToDataFile(tblutils.DataFileOpts{
- Schema: currentSchema,
- Spec: currentSpec,
- Path: filePath,
- Format: iceberg.ParquetFile,
- Content: iceberg.EntryContentData,
- FileSize: rdr.SourceFileSize(),
- PartitionValues: partitionValues,
- SortOrderID: sortOrderID,
+ Schema: currentSchema,
+ Spec: currentSpec,
+ Path: filePath,
+ Format: iceberg.ParquetFile,
+ Content: iceberg.EntryContentData,
+ FileSize: rdr.SourceFileSize(),
+ PartitionValues: partitionValues,
+ PreserveNullPartitionValues: true,
+ SortOrderID: sortOrderID,
})
}
diff --git a/table/file_partition_values_test.go
b/table/file_partition_values_test.go
new file mode 100644
index 000000000..a6134c686
--- /dev/null
+++ b/table/file_partition_values_test.go
@@ -0,0 +1,74 @@
+// 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"
+ "testing"
+
+ "github.com/apache/arrow-go/v18/arrow/decimal128"
+ "github.com/apache/iceberg-go"
+ "github.com/google/uuid"
+ "github.com/stretchr/testify/require"
+)
+
+func TestValidFilePartitionValueType(t *testing.T) {
+ fixed, err := iceberg.FixedTypeOf(2)
+ require.NoError(t, err)
+ dec, err := iceberg.DecimalTypeOf(3, 1)
+ require.NoError(t, err)
+ tests := []struct {
+ typ iceberg.Type
+ valid, invalid any
+ }{
+ {iceberg.PrimitiveTypes.Bool, true, "true"},
+ {iceberg.PrimitiveTypes.Int32, int32(1), int64(1)},
+ {iceberg.PrimitiveTypes.Int64, int64(1), int32(1)},
+ {iceberg.PrimitiveTypes.Float32, float32(1), float64(1)},
+ {iceberg.PrimitiveTypes.Float64, float64(1), float32(1)},
+ {iceberg.PrimitiveTypes.String, "x", []byte("x")},
+ {iceberg.PrimitiveTypes.Binary, []byte("x"), "x"},
+ {fixed, []byte("ab"), []byte("a")},
+ {iceberg.PrimitiveTypes.Date, iceberg.Date(1), int32(1)},
+ {iceberg.PrimitiveTypes.Time, iceberg.Time(1), int64(1)},
+ {iceberg.PrimitiveTypes.Timestamp, iceberg.Timestamp(1),
int64(1)},
+ {iceberg.PrimitiveTypes.TimestampTz, iceberg.Timestamp(1),
int64(1)},
+ {iceberg.PrimitiveTypes.UUID, uuid.UUID{}, []byte("x")},
+ {dec, iceberg.Decimal{Val: decimal128.FromI64(123), Scale: 1},
iceberg.Decimal{Val: decimal128.FromI64(1000), Scale: 1}},
+ }
+ for _, tc := range tests {
+ t.Run(tc.typ.String(), func(t *testing.T) {
+ require.True(t, validFilePartitionValueType(tc.typ,
tc.valid))
+ require.False(t, validFilePartitionValueType(tc.typ,
tc.invalid))
+ })
+ }
+}
+
+func TestFileToDataFileRejectsUnknownTransformWithExplicitPartition(t
*testing.T) {
+ unknown, err := iceberg.ParseTransform("opaque-user-transform")
+ require.NoError(t, err)
+ spec := iceberg.NewPartitionSpecID(1, iceberg.PartitionField{
+ SourceIDs: []int{1}, FieldID: 1000, Name: "opaque", Transform:
unknown,
+ })
+
+ // No storage is needed: the invalid write is rejected before opening
the file.
+ _, err = FileToDataFileWithOptions(context.Background(), nil,
"unused.parquet",
+ simpleSchema(), spec, 0, nil,
+ WithFilePartitionValues(map[int]any{1000: "opaque-value"}))
+ require.ErrorIs(t, err, iceberg.ErrInvalidTransform)
+}
diff --git a/table/internal/utils.go b/table/internal/utils.go
index ca52adf4e..2bbbc2c4c 100644
--- a/table/internal/utils.go
+++ b/table/internal/utils.go
@@ -292,6 +292,8 @@ type DataFileOpts struct {
Content iceberg.ManifestEntryContent
FileSize int64
PartitionValues map[int]any
+ // PreserveNullPartitionValues distinguishes an explicit null from an
omitted value.
+ PreserveNullPartitionValues bool
// SortOrderID claims the file's rows are fully sorted by that order;
zero
// makes no claim and leaves the field absent.
SortOrderID int
@@ -313,9 +315,9 @@ func (d *DataFileStatistics) ToDataFile(opts DataFileOpts)
iceberg.DataFile {
for _, field := range opts.Spec.Fields() {
sourceField, sourceFieldFound :=
opts.Schema.FindFieldByIDRef(
field.FirstSourceID(),
iceberginternal.SchemaRef{})
- partitionVal := opts.PartitionValues[field.FieldID]
+ partitionVal, supplied :=
opts.PartitionValues[field.FieldID]
switch {
- case partitionVal != nil:
+ case partitionVal != nil || (supplied &&
opts.PreserveNullPartitionValues):
// prioritizing caller-supplied value.
fieldIDToPartitionData[field.FieldID] =
partitionVal
case field.Transform.PreservesOrder() &&
sourceFieldFound:
diff --git a/table/table_test.go b/table/table_test.go
index a20272dfd..aa99a39ad 100644
--- a/table/table_test.go
+++ b/table/table_test.go
@@ -834,7 +834,7 @@ func (t *TableWritingTestSuite)
TestAddFilesToBucketPartitionedTableFails() {
tx := tbl.NewTransaction()
err := tx.AddDataFilesByPath(t.ctx, files, nil, false)
t.Error(err)
- t.ErrorContains(err, "cannot infer partition value from parquet
metadata for a non-linear partition field: baz_bucket_3 with transform
bucket[3]")
+ t.ErrorContains(err, "cannot infer partition field \"baz_bucket_3\"")
}
func (t *TableWritingTestSuite)
TestAddFilesToPartitionedTableFailsLowerAndUpperMismatch() {
diff --git a/table/transaction.go b/table/transaction.go
index 67ffbb08c..f3b182952 100644
--- a/table/transaction.go
+++ b/table/transaction.go
@@ -1082,7 +1082,7 @@ func (t *Transaction) ReplaceDataFilesByPath(ctx
context.Context, filesToDelete,
if err != nil {
return err
}
- dataFiles, err := filesToDataFiles(ctx, fs, meta, filesToAdd, 1)
+ dataFiles, err := filesToDataFiles(ctx, fs, meta, filesToAdd, 1, nil)
if err != nil {
return err
}
@@ -2246,7 +2246,12 @@ func (t *Transaction) replaceFiles(ctx context.Context,
dataFilesToDelete, dataF
type AddDataFilesByPathOption func(addFilesOp *addFilesOperation)
type addFilesOperation struct {
- concurrency int
+ concurrency int
+ partitionValues *addFilesPartitionValues
+}
+
+type addFilesPartitionValues struct {
+ byPath map[string]map[int]any
}
// WithAddDataFilesByPathConcurrency overrides the concurrency used by
AddDataFilesByPath.
@@ -2259,8 +2264,28 @@ func WithAddDataFilesByPathConcurrency(concurrency int)
AddDataFilesByPathOption
}
}
+// WithAddDataFilesByPathPartitionValues supplies the partition tuple for each
+// existing data file, indexed first by file path and then by partition field
ID.
+// Bucket partitions cannot be inferred from Parquet min/max statistics, so the
+// caller must provide all non-void partition fields for bucket specs. For
other
+// specs, unlisted fields are inferred when their transforms preserve order.
+// Void fields are set to null.
+// All supplied values must be correct for every row of their file.
+func WithAddDataFilesByPathPartitionValues(values map[string]map[int]any)
AddDataFilesByPathOption {
+ partitionValues := make(map[string]map[int]any, len(values))
+ for path, tuple := range values {
+ partitionValues[path] = cloneFilePartitionTuple(tuple)
+ }
+
+ return func(op *addFilesOperation) {
+ op.partitionValues = &addFilesPartitionValues{byPath:
partitionValues}
+ }
+}
+
// AddDataFilesByPath adds data files at the provided paths to the table.
// Unlike AddDataFiles, it opens each file to read its schema and statistics.
+// Use WithAddDataFilesByPathPartitionValues for partition transforms whose
+// values cannot be inferred from column statistics (for example, bucket[N]).
func (t *Transaction) AddDataFilesByPath(ctx context.Context, filePaths
[]string, snapshotProps iceberg.Properties, ignoreDuplicates bool, opts
...AddDataFilesByPathOption) error {
meta, err := t.txnMeta()
if err != nil {
@@ -2280,6 +2305,15 @@ func (t *Transaction) AddDataFilesByPath(ctx
context.Context, filePaths []string
}
set[filePath] = struct{}{}
}
+ var partitionValues map[string]map[int]any
+ if addFilesOp.partitionValues != nil {
+ partitionValues = addFilesOp.partitionValues.byPath
+ }
+ for filePath := range partitionValues {
+ if _, ok := set[filePath]; !ok {
+ return fmt.Errorf("%w: partition values supplied for
path %q not present in AddDataFilesByPath", iceberg.ErrInvalidArgument,
filePath)
+ }
+ }
fs, err := t.tbl.fsF(ctx)
if err != nil {
@@ -2329,7 +2363,7 @@ func (t *Transaction) AddDataFilesByPath(ctx
context.Context, filePaths []string
}
updater := t.appendSnapshotProducer(meta, wfs, snapshotProps)
- dataFiles, err := filesToDataFiles(ctx, fs, meta, filePaths,
addFilesOp.concurrency)
+ dataFiles, err := filesToDataFiles(ctx, fs, meta, filePaths,
addFilesOp.concurrency, partitionValues)
if err != nil {
return err
}
diff --git a/table/v1_partition_import_test.go
b/table/v1_partition_import_test.go
new file mode 100644
index 000000000..f18d4d5f7
--- /dev/null
+++ b/table/v1_partition_import_test.go
@@ -0,0 +1,323 @@
+// 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_test
+
+import (
+ "context"
+ "os"
+ "path/filepath"
+ "testing"
+
+ "github.com/apache/iceberg-go"
+ iceio "github.com/apache/iceberg-go/io"
+ "github.com/apache/iceberg-go/table"
+ "github.com/stretchr/testify/require"
+)
+
+// A v1 table with an externally written, bucket-partitioned Parquet file.
+// Its rows share one bucket, which is required by the Iceberg specification.
+func makeV1BucketImportFixture(t *testing.T) (*table.Table, iceio.LocalFS,
string, int, int32) {
+ t.Helper()
+
+ dir := t.TempDir()
+ schema := iceberg.NewSchema(1,
+ iceberg.NestedField{ID: 1, Name: "id", Type:
iceberg.PrimitiveTypes.Int32, Required: true},
+ )
+ bucket := iceberg.BucketTransform{NumBuckets: 8}
+ spec := iceberg.NewPartitionSpecID(0, iceberg.PartitionField{
+ SourceIDs: []int{1}, FieldID: 1000, Name: "id_bucket",
Transform: bucket,
+ })
+ meta, err := table.NewMetadata(schema, &spec, table.UnsortedSortOrder,
dir, iceberg.Properties{
+ table.PropertyFormatVersion: "1",
+ })
+ require.NoError(t, err)
+ require.Equal(t, 1, meta.FormatVersion())
+
+ parquetBytes, _ := parquetMetaFromSchema(t, schema,
"[{\"id\":7},{\"id\":7}]")
+ path := filepath.Join(dir, "data.parquet")
+ require.NoError(t, os.WriteFile(path, parquetBytes, 0o600))
+
+ result := bucket.Apply(iceberg.Optional[iceberg.Literal]{
+ Valid: true,
+ Val: iceberg.Int32Literal(7),
+ })
+ require.True(t, result.Valid)
+ bucketID := result.Val.(iceberg.Int32Literal).Value()
+ currentSpec := meta.PartitionSpec()
+ fieldID := currentSpec.Field(0).FieldID
+
+ fs := iceio.LocalFS{}
+ cat := &mergeCatalog{meta: meta}
+ tbl := table.New(table.Identifier{"default", "bucket_import"}, meta,
filepath.Join(dir, "metadata", "00000.json"),
+ func(context.Context) (iceio.IO, error) { return fs, nil }, cat)
+
+ return tbl, fs, path, fieldID, bucketID
+}
+
+func TestV1FileToDataFileExplicitBucketPartition(t *testing.T) {
+ tbl, fs, path, fieldID, bucketID := makeV1BucketImportFixture(t)
+ spec := tbl.Metadata().PartitionSpec()
+
+ _, err := table.FileToDataFileWithOptions(t.Context(), fs, path,
tbl.Metadata().CurrentSchema(), spec, 0, nil)
+ require.ErrorContains(t, err, "cannot infer partition field")
+
+ file, err := table.FileToDataFileWithOptions(t.Context(), fs, path,
tbl.Metadata().CurrentSchema(), spec, 0, nil,
+ table.WithFilePartitionValues(map[int]any{fieldID: bucketID}))
+ require.NoError(t, err)
+ require.Equal(t, bucketID, file.Partition()[fieldID])
+ require.EqualValues(t, 2, file.Count())
+ require.Positive(t, file.FileSizeBytes())
+
+ t.Run("invalid bucket type", func(t *testing.T) {
+ _, err := table.FileToDataFileWithOptions(t.Context(), fs,
path, tbl.Metadata().CurrentSchema(), spec, 0, nil,
+ table.WithFilePartitionValues(map[int]any{fieldID:
int(1)}))
+ require.ErrorIs(t, err, iceberg.ErrInvalidArgument)
+ })
+
+ t.Run("out of range bucket", func(t *testing.T) {
+ _, err := table.FileToDataFileWithOptions(t.Context(), fs,
path, tbl.Metadata().CurrentSchema(), spec, 0, nil,
+ table.WithFilePartitionValues(map[int]any{fieldID:
int32(8)}))
+ require.ErrorIs(t, err, iceberg.ErrInvalidArgument)
+ })
+
+ t.Run("negative bucket", func(t *testing.T) {
+ _, err := table.FileToDataFileWithOptions(t.Context(), fs,
path, tbl.Metadata().CurrentSchema(), spec, 0, nil,
+ table.WithFilePartitionValues(map[int]any{fieldID:
int32(-1)}))
+ require.ErrorIs(t, err, iceberg.ErrInvalidArgument)
+ })
+
+ t.Run("unknown partition field", func(t *testing.T) {
+ _, err := table.FileToDataFileWithOptions(t.Context(), fs,
path, tbl.Metadata().CurrentSchema(), spec, 0, nil,
+ table.WithFilePartitionValues(map[int]any{fieldID +
1000: bucketID}))
+ require.ErrorIs(t, err, iceberg.ErrInvalidArgument)
+ })
+
+ t.Run("null bucket", func(t *testing.T) {
+ file, err := table.FileToDataFileWithOptions(t.Context(), fs,
path, tbl.Metadata().CurrentSchema(), spec, 0, nil,
+ table.WithFilePartitionValues(map[int]any{fieldID:
nil}))
+ require.NoError(t, err)
+ value, exists := file.Partition()[fieldID]
+ require.True(t, exists)
+ require.Nil(t, value)
+ })
+}
+
+func TestV1ImportVoidPartitionWithoutExplicitValue(t *testing.T) {
+ tbl, fs, path, _, _ := makeV1BucketImportFixture(t)
+ spec := iceberg.NewPartitionSpecID(1, iceberg.PartitionField{
+ SourceIDs: []int{1}, FieldID: 1000, Name: "old_id_bucket",
Transform: iceberg.VoidTransform{},
+ })
+
+ file, err := table.FileToDataFileWithOptions(t.Context(), fs, path,
tbl.Metadata().CurrentSchema(), spec, 0, nil)
+ require.NoError(t, err)
+ value, ok := file.Partition()[1000]
+ require.True(t, ok)
+ require.Nil(t, value)
+
+ _, err = table.FileToDataFileWithOptions(t.Context(), fs, path,
tbl.Metadata().CurrentSchema(), spec, 0, nil,
+ table.WithFilePartitionValues(map[int]any{1000: int32(0)}))
+ require.ErrorIs(t, err, iceberg.ErrInvalidArgument)
+}
+
+func TestV1AddDataFilesByPathBucketPartitionCommit(t *testing.T) {
+ tbl, fs, path, fieldID, bucketID := makeV1BucketImportFixture(t)
+ tx := tbl.NewTransaction()
+
+ err := tx.AddDataFilesByPath(t.Context(), []string{path}, nil, false,
+
table.WithAddDataFilesByPathPartitionValues(map[string]map[int]any{
+ path: {fieldID: bucketID},
+ }))
+ require.NoError(t, err)
+
+ committed, err := tx.Commit(t.Context())
+ require.NoError(t, err)
+ require.Equal(t, 1, committed.Metadata().FormatVersion())
+
+ snapshot := committed.CurrentSnapshot()
+ require.NotNil(t, snapshot)
+ manifests, err := snapshot.Manifests(fs)
+ require.NoError(t, err)
+ require.NotEmpty(t, manifests)
+
+ var imported int
+ for _, manifest := range manifests {
+ for entry, err := range manifest.Entries(fs, true) {
+ require.NoError(t, err)
+ if entry.DataFile().FilePath() == path {
+ imported++
+ require.Equal(t, bucketID,
entry.DataFile().Partition()[fieldID])
+ }
+ }
+ }
+ require.Equal(t, 1, imported)
+}
+
+func TestV1AddDataFilesByPathValidatesPartitionPaths(t *testing.T) {
+ tbl, _, path, fieldID, bucketID := makeV1BucketImportFixture(t)
+ tx := tbl.NewTransaction()
+
+ err := tx.AddDataFilesByPath(t.Context(), []string{path}, nil, false,
+
table.WithAddDataFilesByPathPartitionValues(map[string]map[int]any{
+ "other.parquet": {fieldID: bucketID},
+ }))
+ require.ErrorIs(t, err, iceberg.ErrInvalidArgument)
+
+ // A rejected option must not poison the transaction. Import without the
+ // bucket value still fails rather than writing a mispartitioned file.
+ err = tx.AddDataFilesByPath(t.Context(), []string{path}, nil, false)
+ require.ErrorContains(t, err, "cannot infer partition field")
+}
+
+func TestV1FileToDataFileExplicitPartitionValues(t *testing.T) {
+ tbl, fs, path, _, bucketID := makeV1BucketImportFixture(t)
+ spec := iceberg.NewPartitionSpecID(1,
+ iceberg.PartitionField{SourceIDs: []int{1}, FieldID: 1000,
Name: "bucket", Transform: iceberg.BucketTransform{NumBuckets: 8}},
+ iceberg.PartitionField{SourceIDs: []int{1}, FieldID: 1001,
Name: "id", Transform: iceberg.IdentityTransform{}},
+ iceberg.PartitionField{SourceIDs: []int{1}, FieldID: 1002,
Name: "old", Transform: iceberg.VoidTransform{}},
+ )
+ convert := func(values map[int]any) (iceberg.DataFile, error) {
+ return table.FileToDataFileWithOptions(t.Context(), fs, path,
tbl.Metadata().CurrentSchema(), spec, 0, nil,
+ table.WithFilePartitionValues(values))
+ }
+
+ _, err := convert(map[int]any{1000: bucketID})
+ require.ErrorContains(t, err, "cannot infer partition field")
+ _, err = convert(map[int]any{1000: bucketID, 1001: "7"})
+ require.ErrorIs(t, err, iceberg.ErrInvalidArgument)
+ file, err := convert(map[int]any{1000: bucketID, 1001: int32(7)})
+ require.NoError(t, err)
+ require.Equal(t, int32(7), file.Partition()[1001])
+ require.Nil(t, file.Partition()[1002])
+
+ // A supplied value cannot override the footer for identity partitions.
+ _, err = convert(map[int]any{1000: bucketID, 1001: int32(99)})
+ require.ErrorIs(t, err, iceberg.ErrInvalidArgument)
+ _, err = convert(map[int]any{1000: bucketID, 1001: nil})
+ require.ErrorIs(t, err, iceberg.ErrInvalidArgument)
+}
+
+func TestV1FileToDataFileChecksSuppliedOrderPreservingBounds(t *testing.T) {
+ tbl, fs, path, _, _ := makeV1BucketImportFixture(t)
+ spec := iceberg.NewPartitionSpecID(1,
+ iceberg.PartitionField{SourceIDs: []int{1}, FieldID: 1000,
Name: "id", Transform: iceberg.IdentityTransform{}},
+ iceberg.PartitionField{SourceIDs: []int{1}, FieldID: 1001,
Name: "id_truncated", Transform: iceberg.TruncateTransform{Width: 10}},
+ )
+ convert := func(values map[int]any) (iceberg.DataFile, error) {
+ return table.FileToDataFileWithOptions(t.Context(), fs, path,
tbl.Metadata().CurrentSchema(), spec, 0, nil,
+ table.WithFilePartitionValues(values))
+ }
+
+ file, err := convert(map[int]any{1000: int32(7), 1001: int32(0)})
+ require.NoError(t, err)
+ require.Equal(t, int32(7), file.Partition()[1000])
+ require.Equal(t, int32(0), file.Partition()[1001])
+
+ inferred, err := convert(nil)
+ require.NoError(t, err)
+ require.Equal(t, file.Partition(), inferred.Partition())
+
+ for _, tc := range []struct {
+ name string
+ values map[int]any
+ }{
+ {"incorrect identity", map[int]any{1000: int32(99), 1001:
int32(0)}},
+ {"incorrect truncate", map[int]any{1000: int32(7), 1001:
int32(10)}},
+ {"null identity over non-null rows", map[int]any{1000: nil,
1001: int32(0)}},
+ {"null truncate over non-null rows", map[int]any{1000:
int32(7), 1001: nil}},
+ } {
+ t.Run(tc.name, func(t *testing.T) {
+ _, err := convert(tc.values)
+ require.ErrorIs(t, err, iceberg.ErrInvalidArgument)
+ require.ErrorContains(t, err, "does not match Parquet
bounds")
+ })
+ }
+}
+
+func TestV1AddDataFilesByPathRejectsSuppliedIdentityMismatch(t *testing.T) {
+ _, fs, path, _, _ := makeV1BucketImportFixture(t)
+ dir := filepath.Dir(path)
+ schema := iceberg.NewSchema(1,
+ iceberg.NestedField{ID: 1, Name: "id", Type:
iceberg.PrimitiveTypes.Int32, Required: true},
+ )
+ spec := iceberg.NewPartitionSpecID(0, iceberg.PartitionField{
+ SourceIDs: []int{1}, FieldID: 1000, Name: "id", Transform:
iceberg.IdentityTransform{},
+ })
+ meta, err := table.NewMetadata(schema, &spec, table.UnsortedSortOrder,
dir, iceberg.Properties{
+ table.PropertyFormatVersion: "1",
+ })
+ require.NoError(t, err)
+ tbl := table.New(table.Identifier{"default", "identity_import"}, meta,
filepath.Join(dir, "metadata", "00000.json"),
+ func(context.Context) (iceio.IO, error) { return fs, nil },
&mergeCatalog{meta: meta})
+
+ for _, tc := range []struct {
+ name string
+ value any
+ }{
+ {"incorrect identity", int32(99)},
+ {"null identity", nil},
+ } {
+ t.Run(tc.name, func(t *testing.T) {
+ tx := tbl.NewTransaction()
+ err := tx.AddDataFilesByPath(t.Context(),
[]string{path}, nil, false,
+
table.WithAddDataFilesByPathPartitionValues(map[string]map[int]any{
+ path: {1000: tc.value},
+ }))
+ require.ErrorIs(t, err, iceberg.ErrInvalidArgument)
+ require.ErrorContains(t, err, "does not match Parquet
bounds")
+ })
+ }
+
+ tx := tbl.NewTransaction()
+ require.NoError(t, tx.AddDataFilesByPath(t.Context(), []string{path},
nil, false,
+
table.WithAddDataFilesByPathPartitionValues(map[string]map[int]any{
+ path: {1000: int32(7)},
+ })))
+}
+
+func TestV1ImportPartitionOptionsSnapshotCallerValues(t *testing.T) {
+ tbl, fs, path, fieldID, bucketID := makeV1BucketImportFixture(t)
+ values := map[int]any{fieldID: bucketID}
+ option := table.WithFilePartitionValues(values)
+ values[fieldID] = int32(-1)
+ file, err := table.FileToDataFileWithOptions(t.Context(), fs, path,
+ tbl.Metadata().CurrentSchema(), tbl.Metadata().PartitionSpec(),
0, nil, option)
+ require.NoError(t, err)
+ require.Equal(t, bucketID, file.Partition()[fieldID])
+
+ byPath := map[string]map[int]any{path: {fieldID: bucketID}}
+ addOption := table.WithAddDataFilesByPathPartitionValues(byPath)
+ byPath[path][fieldID] = int32(-1)
+ tx := tbl.NewTransaction()
+ require.NoError(t, tx.AddDataFilesByPath(t.Context(), []string{path},
nil, false, addOption))
+ committed, err := tx.Commit(t.Context())
+ require.NoError(t, err)
+ manifests, err := committed.CurrentSnapshot().Manifests(fs)
+ require.NoError(t, err)
+ require.NotEmpty(t, manifests)
+ var found bool
+ for _, manifest := range manifests {
+ for entry, entryErr := range manifest.Entries(fs, true) {
+ require.NoError(t, entryErr)
+ if entry.DataFile().FilePath() == path {
+ found = true
+ require.Equal(t, bucketID,
entry.DataFile().Partition()[fieldID])
+ }
+ }
+ }
+ require.True(t, found)
+}
