zeroshade commented on code in PR #2383: URL: https://github.com/apache/iceberg-go/pull/2383#discussion_r4470714448
########## table/rewrite_by_path.go: ########## @@ -0,0 +1,293 @@ +// 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" + "math" + "runtime" + "slices" + + "github.com/apache/iceberg-go" +) + +// StageRewrittenDataFilesByPath stages a replace snapshot using Parquet files +// already written by another engine. It does not read or rewrite data rows. +// Call Transaction.Commit to persist the staged snapshot. +// +// This helper supports unpartitioned v1 and v2 tables, including historical +// unpartitioned specs, and only when no removed file has applicable deletes. +// The replacement files +// must contain the same live rows and the same total number of physical rows. +// Column value, null, and NaN counts are compared when recorded, but neither +// footer statistics nor record counts can prove row-level equivalence. The caller is responsible +// for preserving the rows. For delete-aware compaction and partitioned tables, +// use RewriteDataFiles or NewRewrite with FileToDataFile instead. +// +// Both path lists must be non-empty, unique, and disjoint. Removed paths must +// be live data files in the target branch and added paths must not already be +// referenced there. Format v3 is not supported because path metadata alone +// cannot preserve row lineage. +func (t *Transaction) StageRewrittenDataFilesByPath( + ctx context.Context, + filesToDelete, filesToAdd []string, + snapshotProps iceberg.Properties, +) error { + meta, err := t.txnMeta() + if err != nil { + return err + } + if err := ctx.Err(); err != nil { + return err + } + if meta.formatVersion >= 3 { + return fmt.Errorf("%w: path-based rewrites do not support v3 row lineage", iceberg.ErrNotImplemented) + } + + spec, err := meta.CurrentSpec() + if err != nil { + return err + } + if spec.NumFields() != 0 { + return fmt.Errorf("%w: path-based rewrites do not support partitioned tables", iceberg.ErrNotImplemented) + } + if len(filesToDelete) == 0 || len(filesToAdd) == 0 { + return fmt.Errorf("%w: a data-file rewrite requires both removed and added files", ErrInvalidOperation) + } + + toDelete := make(map[string]struct{}, len(filesToDelete)) + for _, path := range filesToDelete { + if path == "" { + return fmt.Errorf("%w: removed file paths must be non-empty", ErrInvalidOperation) + } + if _, exists := toDelete[path]; exists { + return fmt.Errorf("%w: duplicate removed file path %q", ErrInvalidOperation, path) + } + toDelete[path] = struct{}{} + } + + toAdd := make(map[string]struct{}, len(filesToAdd)) + for _, path := range filesToAdd { + if path == "" { + return fmt.Errorf("%w: added file paths must be non-empty", ErrInvalidOperation) + } + if _, exists := toAdd[path]; exists { + return fmt.Errorf("%w: duplicate added file path %q", ErrInvalidOperation, path) + } + if _, exists := toDelete[path]; exists { + return fmt.Errorf("%w: cannot remove and add the same file %q", ErrInvalidOperation, path) + } + toAdd[path] = struct{}{} + } + + snapshot := t.planningSnapshot(meta) + if snapshot == nil { + return fmt.Errorf("%w: cannot rewrite files in a table without an existing snapshot", ErrInvalidOperation) + } + + fs, err := t.tbl.fsF(ctx) + if err != nil { + return err + } + if _, err := requireWriteFileIO(fs); err != nil { + return err + } + + // The scan below only exposes delete files that apply to planned data + // files. A live delete file with no matching data file still reserves its + // path in the table, so check all delete manifest entries before staging. + for entry, err := range snapshot.entries(fs, iceberg.ManifestContentDeletes, true) { + if err != nil { + return err + } + if err := ctx.Err(); err != nil { + return err + } + + path := entry.DataFile().FilePath() + if _, exists := toAdd[path]; exists { + return fmt.Errorf("%w: added file %q is already referenced by the table as a delete file", + ErrInvalidOperation, path) + } + } + + // Plan the target snapshot to discover delete files that apply to each + // source file. Parquet splitting can yield several tasks per data file. + scan, err := t.Scan(WithSnapshotID(snapshot.SnapshotID)) + if err != nil { + return err + } + tasks, err := scan.PlanFiles(ctx) + if err != nil { + return err + } + + removed := make([]iceberg.DataFile, 0, len(filesToDelete)) + found := make(map[string]struct{}, len(filesToDelete)) + originalCounts := newPathRewriteCounts() + var originalRecordCount int64 + for _, task := range tasks { + df := task.File + path := df.FilePath() + if _, exists := toAdd[path]; exists { + return fmt.Errorf("%w: added file %q is already referenced by the table", ErrInvalidOperation, path) + } + if _, selected := toDelete[path]; !selected { + continue + } + if len(task.DeleteFiles) > 0 || len(task.EqualityDeleteFiles) > 0 || len(task.DeletionVectorFiles) > 0 { + return fmt.Errorf("%w: cannot rewrite data file %q with applicable delete files; use NewRewrite or RewriteDataFiles", ErrInvalidOperation, path) + } + if _, duplicateTask := found[path]; duplicateTask { + continue + } + if df.ContentType() != iceberg.EntryContentData { + return fmt.Errorf("%w: cannot rewrite non-data file %q", ErrInvalidOperation, path) + } + sourceSpec, err := meta.GetSpecByID(int(df.SpecID())) + if err != nil { + return fmt.Errorf("resolving partition spec for source file %q: %w", path, err) + } + if sourceSpec.NumFields() != 0 { + return fmt.Errorf("%w: source data file %q uses partitioned spec %d", + iceberg.ErrNotImplemented, path, df.SpecID()) + } + if df.Count() < 0 || df.Count() > math.MaxInt64-originalRecordCount { + return fmt.Errorf("%w: invalid total record count for removed data files", ErrInvalidOperation) + } + if err := originalCounts.add(df); err != nil { + return err + } + + originalRecordCount += df.Count() + found[path] = struct{}{} + removed = append(removed, df) + } + + if len(found) != len(toDelete) { + missing := make([]string, 0, len(toDelete)-len(found)) + for path := range toDelete { + if _, exists := found[path]; !exists { + missing = append(missing, path) + } + } + slices.Sort(missing) + + return fmt.Errorf("%w: removed data file paths are not live in target branch: %v", ErrInvalidOperation, missing) + } + + // Reading footers does not require a name mapping. Do not modify the + // transaction until the replacement metadata and counts are validated. + added, err := filesToDataFiles(ctx, fs, meta, filesToAdd, runtime.GOMAXPROCS(0)) + if err != nil { + return err + } + + replacementCounts := newPathRewriteCounts() + var replacementRecordCount int64 + for _, df := range added { + if df.Count() < 0 || df.Count() > math.MaxInt64-replacementRecordCount { + return fmt.Errorf("%w: invalid total record count for replacement data files", ErrInvalidOperation) + } + if err := replacementCounts.add(df); err != nil { + return err + } + + replacementRecordCount += df.Count() + } + if replacementRecordCount != originalRecordCount { + return fmt.Errorf("%w: path-based rewrite record count changed from %d to %d", + ErrInvalidOperation, originalRecordCount, replacementRecordCount) + } + if err := originalCounts.validateReplacement(replacementCounts); err != nil { + return err + } + + rewrite := t.NewRewrite(snapshotProps) + for _, df := range removed { + rewrite.DeleteFile(df) + } + for _, df := range added { + rewrite.AddDataFile(df) + } + + return rewrite.Commit(ctx) +} + +type pathRewriteCounts struct { + values map[int]int64 + nulls map[int]int64 + nans map[int]int64 +} + +func newPathRewriteCounts() pathRewriteCounts { + return pathRewriteCounts{ + values: make(map[int]int64), + nulls: make(map[int]int64), + nans: make(map[int]int64), + } +} + +func (c pathRewriteCounts) add(df iceberg.DataFile) error { + for _, metric := range []struct { + name string + file map[int]int64 + total map[int]int64 + }{ + {"value", df.ValueCounts(), c.values}, + {"null", df.NullValueCounts(), c.nulls}, + {"NaN", df.NaNValueCounts(), c.nans}, + } { + for fieldID, count := range metric.file { + if count < 0 || count > math.MaxInt64-metric.total[fieldID] { + return fmt.Errorf("%w: invalid column %d %s count in data file %q", ErrInvalidOperation, fieldID, metric.name, df.FilePath()) + } + metric.total[fieldID] += count + } + } + + return nil +} + +func (c pathRewriteCounts) validateReplacement(replacement pathRewriteCounts) error { + for _, metric := range []struct { + name string + original map[int]int64 + replacement map[int]int64 + }{ + {"value", c.values, replacement.values}, + {"null", c.nulls, replacement.nulls}, + {"NaN", c.nans, replacement.nans}, + } { + for fieldID, originalCount := range metric.original { + replacementCount, exists := metric.replacement[fieldID] + if !exists || replacementCount != originalCount { + return fmt.Errorf("%w: path-based rewrite changed column %d %s count from %d to %d", + ErrInvalidOperation, fieldID, metric.name, originalCount, replacementCount) + } + } + for fieldID := range metric.replacement { + if _, exists := metric.original[fieldID]; !exists { + return fmt.Errorf("%w: path-based rewrite has new %s-count metric for column %d", ErrInvalidOperation, metric.name, fieldID) Review Comment: This rejects byte-identical rewrites whenever the originals recorded statistics for fewer columns than iceberg-go computes from the replacement footers. Repro: create the table with `write.metadata.metrics.column.s=none` and append (the original records value counts for column 1 only), remove that property, copy the file byte for byte to a new path, then call `StageRewrittenDataFilesByPath(original → copy)`; it fails with `path-based rewrite has new value-count metric for column 2`. The same probably happens for Spark-written files on tables wider than 100 columns, since Java only records statistics for the first 100 by default (`write.metadata.metrics.max-inferred-column-defaults`). Checking only the columns present in the originals still catches lost or emptied columns, so this direction can be dropped. -- 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]
