laskoviymishka commented on code in PR #2381: URL: https://github.com/apache/iceberg-go/pull/2381#discussion_r4508045302
########## table/dynamic_partition_overwrite_test.go: ########## @@ -0,0 +1,510 @@ +// 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" + "slices" + "testing" + + "github.com/apache/iceberg-go" + iceio "github.com/apache/iceberg-go/io" + "github.com/stretchr/testify/require" +) + +func TestDynamicPartitionMatcher(t *testing.T) { + spec := iceberg.NewPartitionSpecID(7, + iceberg.PartitionField{ + SourceIDs: []int{1}, + FieldID: 1000, + Name: "bucket_id", + Transform: iceberg.BucketTransform{NumBuckets: 16}, + }, + iceberg.PartitionField{ + SourceIDs: []int{2}, + FieldID: 1001, + Name: "region", + Transform: iceberg.IdentityTransform{}, + }, + ) + + files := []iceberg.DataFile{ + newTestDataFile(t, spec, "file://new-1.parquet", + map[int]any{1000: int32(4), 1001: "north"}), + newTestDataFile(t, spec, "file://new-2.parquet", + map[int]any{1000: int32(4), 1001: "north"}), + newTestDataFile(t, spec, "file://new-3.parquet", + map[int]any{1000: int32(5), 1001: "south"}), + } + keys, err := dynamicPartitionKeys(files) + require.NoError(t, err) + require.Len(t, keys, 2) + for _, tc := range []struct { + partition map[int]any + want bool + }{ + {map[int]any{1000: int32(4), 1001: "north"}, true}, + {map[int]any{1000: int64(5), 1001: "south"}, true}, + {map[int]any{1000: int32(4), 1001: "south"}, false}, + {map[int]any{1000: int32(7), 1001: "north"}, false}, + } { + key, err := canonicalPartitionKey(int32(spec.ID()), tc.partition) + require.NoError(t, err) + _, found := keys[key] + require.Equal(t, tc.want, found) + } +} + +func TestDynamicPartitionOverwriteDataFiles(t *testing.T) { + ctx := context.Background() + spec := partitionedSpec() + seed, fs := createTestTransactionWithMemIO(t, spec) + + original := []iceberg.DataFile{ + newTestDataFile(t, spec, "file://old-a.parquet", map[int]any{1000: int32(1)}), + newTestDataFile(t, spec, "file://old-a2.parquet", map[int]any{1000: int32(1)}), + newTestDataFile(t, spec, "file://old-b.parquet", map[int]any{1000: int32(2)}), + } + require.NoError(t, seed.AddDataFiles(ctx, original, nil)) + seedMeta, err := seed.meta.Build() + require.NoError(t, err) + tbl := New(Identifier{"db", "tbl"}, seedMeta, "metadata.json", + func(context.Context) (iceio.IO, error) { return fs, nil }, nil) + + txn := tbl.NewTransaction() + replacement := []iceberg.DataFile{ + newTestDataFile(t, spec, "file://new-a.parquet", map[int]any{1000: int32(1)}), + newTestDataFile(t, spec, "file://new-a2.parquet", map[int]any{1000: int32(1)}), + } + + require.NoError(t, txn.DynamicPartitionOverwriteDataFiles(ctx, replacement, + iceberg.Properties{"source": "dynamic-test"})) + + updated, err := txn.meta.Build() + require.NoError(t, err) + snapshot := updated.CurrentSnapshot() + require.NotNil(t, snapshot) + require.NotNil(t, snapshot.Summary) + require.Equal(t, OpOverwrite, snapshot.Summary.Operation) + require.Equal(t, "true", snapshot.Summary.Properties["replace-partitions"]) + require.Equal(t, "dynamic-test", snapshot.Summary.Properties["source"]) + + var paths []string + for df, err := range snapshot.dataFiles(ctx, fs, nil) { + require.NoError(t, err) + if df.ContentType() == iceberg.EntryContentData { + paths = append(paths, df.FilePath()) + } + } + slices.Sort(paths) + require.Equal(t, []string{"file://new-a.parquet", "file://new-a2.parquet", "file://old-b.parquet"}, paths) + require.NotContains(t, txn.pinnedRefs, MainBranch, "existing branches use partition-scoped conflict validation") +} + +func TestDynamicPartitionOverwriteRejectsAlreadyLivePath(t *testing.T) { + spec := partitionedSpec() + seed, fs := createTestTransactionWithMemIO(t, spec) + existing := newTestDataFile(t, spec, "file://existing.parquet", map[int]any{1000: int32(1)}) + require.NoError(t, seed.AddDataFiles(t.Context(), []iceberg.DataFile{existing}, nil)) + seedMeta, err := seed.meta.Build() + require.NoError(t, err) + tbl := New(Identifier{"db", "tbl"}, seedMeta, "metadata.json", + func(context.Context) (iceio.IO, error) { return fs, nil }, nil) + + txn := tbl.NewTransaction() + err = txn.DynamicPartitionOverwriteDataFiles(t.Context(), []iceberg.DataFile{existing}, nil) + require.ErrorIs(t, err, ErrInvalidOperation) + require.ErrorContains(t, err, "already exists") + require.Empty(t, txn.meta.updates) + require.Empty(t, txn.pinnedRefs) +} + +func TestDynamicPartitionOverwriteAppendsNewPartition(t *testing.T) { + ctx := context.Background() + spec := partitionedSpec() + txn, fs := createTestTransactionWithMemIO(t, spec) + newFile := newTestDataFile(t, spec, "file://first.parquet", map[int]any{1000: int32(10)}) + + require.NoError(t, txn.DynamicPartitionOverwriteDataFiles(ctx, []iceberg.DataFile{newFile}, nil)) + metadata, err := txn.meta.Build() + require.NoError(t, err) + snapshot := metadata.CurrentSnapshot() + require.NotNil(t, snapshot) + var files []string + for df, err := range snapshot.dataFiles(ctx, fs, nil) { + require.NoError(t, err) + files = append(files, df.FilePath()) + } + require.Equal(t, []string{"file://first.parquet"}, files) + require.Equal(t, OpOverwrite, snapshot.Summary.Operation) + require.Equal(t, "true", snapshot.Summary.Properties["replace-partitions"]) + require.Contains(t, txn.pinnedRefs, MainBranch) +} + +func TestDynamicPartitionOverwriteUnpartitionedReplacesAllFiles(t *testing.T) { + ctx := context.Background() + spec := iceberg.NewPartitionSpec() + seed, fs := createTestTransactionWithMemIO(t, spec) + original := []iceberg.DataFile{ + newTestDataFile(t, spec, "file://old-1.parquet", nil), + newTestDataFile(t, spec, "file://old-2.parquet", nil), + } + require.NoError(t, seed.AddDataFiles(ctx, original, nil)) + seedMeta, err := seed.meta.Build() + require.NoError(t, err) + tbl := New(Identifier{"db", "tbl"}, seedMeta, "metadata.json", + func(context.Context) (iceio.IO, error) { return fs, nil }, nil) + + txn := tbl.NewTransaction() + replacement := newTestDataFile(t, spec, "file://replacement.parquet", nil) + require.NoError(t, txn.DynamicPartitionOverwriteDataFiles(ctx, []iceberg.DataFile{replacement}, nil)) + updated, err := txn.meta.Build() + require.NoError(t, err) + snapshot := updated.CurrentSnapshot() + require.Equal(t, OpOverwrite, snapshot.Summary.Operation) + require.Equal(t, "true", snapshot.Summary.Properties["replace-partitions"]) + var paths []string + for df, err := range snapshot.dataFiles(ctx, fs, nil) { + require.NoError(t, err) + if df.ContentType() == iceberg.EntryContentData { + paths = append(paths, df.FilePath()) + } + } + require.Equal(t, []string{"file://replacement.parquet"}, paths) +} + +func TestDynamicPartitionOverwriteEmptyInputIsNoOp(t *testing.T) { + txn, _ := createTestTransactionWithMemIO(t, partitionedSpec()) + require.NoError(t, txn.DynamicPartitionOverwriteDataFiles(t.Context(), nil, nil)) + require.Empty(t, txn.meta.updates) + require.Empty(t, txn.pinnedRefs) +} + +func TestDynamicPartitionOverwriteRespectsCanceledContext(t *testing.T) { + spec := partitionedSpec() + txn, _ := createTestTransactionWithMemIO(t, spec) + ctx, cancel := context.WithCancel(t.Context()) + cancel() + + file := newTestDataFile(t, spec, "file://canceled.parquet", map[int]any{1000: int32(5)}) + err := txn.DynamicPartitionOverwriteDataFiles(ctx, []iceberg.DataFile{file}, nil) + require.ErrorIs(t, err, context.Canceled) + require.Empty(t, txn.meta.updates) + require.Empty(t, txn.pinnedRefs) +} + +func TestDynamicPartitionOverwriteRejectsUnsupportedLayouts(t *testing.T) { Review Comment: nit: this test is still named `TestDynamicPartitionOverwriteRejectsUnsupportedLayouts`, but its first subtest, `format v3 row lineage`, asserts a successful v3 overwrite (`FirstRowID` and `AddedRows` set). Looks like a leftover from when v3 was rejected. Could we move it into its own `TestDynamicPartitionOverwriteV3RowLineage`, so a v3 regression doesn't show up under a "Rejects" name? ########## table/dynamic_partition_overwrite.go: ########## @@ -0,0 +1,444 @@ +// 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" + "maps" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/iceberg-go" + "github.com/google/uuid" +) + +// dynamicPartitionKeys normalizes all partition values before comparing them. +// A manifest decoder may represent the same logical day as an int32 while the +// writer uses iceberg.Date. The shared canonical encoder handles both forms, +// decimals, bytes, UUIDs, and nulls without lossy string formatting. +func dynamicPartitionKeys(files []iceberg.DataFile) (map[string]struct{}, error) { + keys := make(map[string]struct{}, len(files)) + for _, file := range files { + key, err := canonicalPartitionKey(file.SpecID(), dataFilePartition(file)) + if err != nil { + return nil, fmt.Errorf("dynamic partition overwrite file %q: %w", file.FilePath(), err) + } + keys[key] = struct{}{} + } + + return keys, nil +} + +func (t *Transaction) dynamicOverwriteSpec() (*iceberg.PartitionSpec, error) { + meta, err := t.txnMeta() + if err != nil { + return nil, err + } + + if meta.formatVersion > 3 { + return nil, fmt.Errorf("%w: dynamic partition overwrite does not support format v%d", ErrInvalidOperation, meta.formatVersion) + } + + spec, err := meta.CurrentSpec() + if err != nil { + return nil, err + } + if spec == nil { + return nil, fmt.Errorf("%w: dynamic partition overwrite requires a current partition spec", ErrInvalidOperation) + } + if err := checkNoUnknownTransform(spec); err != nil { + return nil, err + } + + return spec, nil +} + +// validateDynamicOverwriteConflicts uses exact stored partition tuples so +// transformed and null partitions cannot become false-negative row predicates. +// It rejects concurrent additions and removals in touched partitions while +// allowing unrelated partition appends during refresh-and-replay. +func validateDynamicOverwriteConflicts(cc *conflictContext, specID int, unpartitioned bool, touched map[string]struct{}) error { + if cc == nil || len(cc.concurrent) == 0 { + return nil + } + for _, snapshot := range cc.concurrent { + manifests, err := cc.manifestsFor(snapshot) + if err != nil { + return fmt.Errorf("dynamic partition overwrite: concurrent manifests for snapshot %d: %w", snapshot.SnapshotID, err) + } + for _, manifest := range manifests { + for entry, err := range manifest.Entries(cc.ctx, cc.manifestReadIO(), false) { + if err != nil { + return fmt.Errorf("dynamic partition overwrite: read concurrent manifest %s: %w", manifest.FilePath(), err) + } + if entry.SnapshotID() != snapshot.SnapshotID || entry.Status() == iceberg.EntryStatusEXISTING { + continue + } + file := entry.DataFile() + conflict := ErrConflictingDataFiles + if file.ContentType() != iceberg.EntryContentData { + conflict = ErrConflictingDeleteFiles + } + // No partition scope exists for an unpartitioned overwrite. + // Every concurrent change matters, including older spec IDs. + if unpartitioned { + return fmt.Errorf("%w: concurrent snapshot %d changed file %q during an unpartitioned overwrite", + conflict, snapshot.SnapshotID, file.FilePath()) + } + if int(file.SpecID()) != specID { + continue + } + key, err := canonicalPartitionKey(file.SpecID(), dataFilePartition(file)) + if err != nil { + return fmt.Errorf("dynamic partition overwrite: concurrent file %q: %w", file.FilePath(), err) + } + if _, ok := touched[key]; ok { + return fmt.Errorf("%w: concurrent snapshot %d changed file %q in an overwritten partition", + conflict, snapshot.SnapshotID, file.FilePath()) + } + } + } + } + + return nil +} + +// DynamicPartitionOverwriteDataFiles atomically replaces all live data in +// the partitions represented by filesToAdd. Empty input is a no-op; an +// unpartitioned table has one partition. New partitions still use OpOverwrite. +// +// Added files must use the table's current partition spec; existing live files +// under older specs are preserved unless the current spec is unpartitioned, +// in which case all existing files are replaced. Supports v1-v3, including v3 row lineage +// and associated deletion vectors. On v3, a new data file may omit first_row_id +// and inherit it from its manifest. +// Concurrent writes to touched partitions fail while independent appends to +// other partitions can be retried unless delete-file removal prevents replay. +// A missing target branch is pinned together with main to guard its origin. +// Default partition and sort order changes also fail. +// +// All input files must already exist. Uncommitted data files remain the +// caller's responsibility after a staging or commit failure. By default this +// sets the table's name mapping; WithoutAutoNameMapping disables that step. +// WithoutDuplicateCheck is ignored to avoid adding an already live file twice. +func (t *Transaction) DynamicPartitionOverwriteDataFiles( + ctx context.Context, filesToAdd []iceberg.DataFile, snapshotProps iceberg.Properties, opts ...WriteOption, +) error { + spec, err := t.dynamicOverwriteSpec() + if err != nil { + return err + } + if len(filesToAdd) == 0 { + return nil + } + if err := ctx.Err(); err != nil { + return err + } + + var cfg dataFileCfg + for _, option := range opts { + option(&cfg) + } + if cfg.rewriteSemantics { + return fmt.Errorf("%w: dynamic partition overwrite cannot use data-preserving rewrite semantics", ErrInvalidOperation) + } + + // Dynamic overwrite adds new rows, so v3 first_row_id can be inherited + // from the manifest. Java ReplacePartitions uses the same inheritance. + addedPaths, err := t.validateDataFilesToAdd(filesToAdd, "DynamicPartitionOverwriteDataFiles", false) + if err != nil { + return err + } + for _, df := range filesToAdd { + if int(df.SpecID()) != spec.ID() { + return fmt.Errorf("%w: dynamic partition overwrite file %q uses spec %d, current spec is %d", + ErrInvalidOperation, df.FilePath(), df.SpecID(), spec.ID()) + } + } + + fs, err := t.tbl.fsF(ctx) + if err != nil { + return err + } + wfs, err := requireWriteFileIO(fs) + if err != nil { + return err + } + + keys, err := dynamicPartitionKeys(filesToAdd) + if err != nil { + return err + } + var toDelete, deleteFilesToRemove []iceberg.DataFile + unpartitioned := spec.IsUnpartitioned() + survivingDataPaths := make(map[string]struct{}) + var hasSurvivingOldSpecData bool + validationPlan := newPartitionValidationPlan(spec) + if snapshot := t.planningSnapshot(t.meta); snapshot != nil { + for df, err := range snapshot.dataFiles(ctx, fs, nil) { + if err != nil { + return err Review Comment: nit: the `snapshot.dataFiles` loop returns the manifest read error bare, while the delete-entries loop below wraps it as `dynamic partition overwrite: read delete files: %w`. Could we wrap this one the same way (`read data files`), so a manifest failure here keeps the operation context? -- 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: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
