laskoviymishka commented on code in PR #2149: URL: https://github.com/apache/iceberg-go/pull/2149#discussion_r4328220752
########## table/scan_task_finalization.go: ########## @@ -0,0 +1,211 @@ +// 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 ( + "fmt" + "runtime" + "sync" + + "github.com/apache/iceberg-go" +) + +// CPU-heavy finalization needs enough work to amortize worker and merge costs. +const scanTaskFinalizationMinTasksPerWorker = 256 + +type scanTaskFinalizationResult struct { + metrics scanMetricsAccumulator + tasks []FileScanTask + expanded bool + err error +} + +// finalizePlannedTasks applies residuals, records result metrics, and expands +// Parquet byte-range splits. The caller transfers ownership of plannedTasks: +// workers mutate residuals in that slice, even when an error is returned. +// Callers must not reuse it after this call. If no task splits, the returned +// slice aliases the original backing array. +func (scan *Scan) finalizePlannedTasks( + plannedTasks []FileScanTask, + schema *iceberg.Schema, + acc *scanMetricsAccumulator, +) ([]FileScanTask, error) { + var boundRowFilter iceberg.BooleanExpression + if scan.rowFilter != nil && !scan.rowFilter.Equals(iceberg.AlwaysTrue{}) { + var err error + boundRowFilter, err = iceberg.BindExpr(schema, scan.rowFilter, scan.caseSensitive) + if err != nil { + return nil, err + } + } + + var residualEvaluators *keyDefaultMapErr[int, *partitionResidualEvaluator] + if boundRowFilter != nil { + residualEvaluators = newKeyDefaultMapWrapErr(func(specID int) (*partitionResidualEvaluator, error) { + return newPartitionResidualEvaluator( + schema, scan.metadata.PartitionSpecByID(specID), boundRowFilter, scan.caseSensitive) + }) + } + + splitTargetSize := scan.metadata.Properties().GetInt64( + ReadSplitTargetSizeKey, ReadSplitTargetSizeDefault) + acc.resultDataFiles = int64(len(plannedTasks)) + if len(plannedTasks) == 0 { + return []FileScanTask{}, nil + } + + workerCount := min(max(scan.concurrency, 1), runtime.GOMAXPROCS(0), + max(1, len(plannedTasks)/scanTaskFinalizationMinTasksPerWorker)) + results := make([]scanTaskFinalizationResult, workerCount) + + finalizeRange := func(worker, start, end int) { + defer func() { + if recovered := recover(); recovered != nil { + results[worker].err = fmt.Errorf("finalize scan tasks [%d:%d]: panic: %v", Review Comment: This is a side effect of the recover I asked for last round. It wraps `finalizeRange`, so it also catches the inline `workerCount == 1` path, which covers every scan under 512 tasks. A panic there used to crash with a full stack. Now it comes back as a plain error with no stack, and `%v` hides it from `errors.Is`/`As` when the panic value is an error. I'd either move the recover into the `go func` so only spawned workers get it, or keep it shared and add `debug.Stack()` plus `%w` when `recovered` is an error. Either way, a one-worker panic test would pin down which behavior we want, since `WorkerPanicReturnsError` only covers the concurrent path. ########## table/scan_task_finalization.go: ########## @@ -0,0 +1,211 @@ +// 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 ( + "fmt" + "runtime" + "sync" + + "github.com/apache/iceberg-go" +) + +// CPU-heavy finalization needs enough work to amortize worker and merge costs. +const scanTaskFinalizationMinTasksPerWorker = 256 + +type scanTaskFinalizationResult struct { + metrics scanMetricsAccumulator + tasks []FileScanTask + expanded bool + err error +} + +// finalizePlannedTasks applies residuals, records result metrics, and expands +// Parquet byte-range splits. The caller transfers ownership of plannedTasks: +// workers mutate residuals in that slice, even when an error is returned. +// Callers must not reuse it after this call. If no task splits, the returned +// slice aliases the original backing array. +func (scan *Scan) finalizePlannedTasks( + plannedTasks []FileScanTask, + schema *iceberg.Schema, + acc *scanMetricsAccumulator, +) ([]FileScanTask, error) { + var boundRowFilter iceberg.BooleanExpression + if scan.rowFilter != nil && !scan.rowFilter.Equals(iceberg.AlwaysTrue{}) { + var err error + boundRowFilter, err = iceberg.BindExpr(schema, scan.rowFilter, scan.caseSensitive) + if err != nil { + return nil, err + } + } + + var residualEvaluators *keyDefaultMapErr[int, *partitionResidualEvaluator] + if boundRowFilter != nil { + residualEvaluators = newKeyDefaultMapWrapErr(func(specID int) (*partitionResidualEvaluator, error) { + return newPartitionResidualEvaluator( + schema, scan.metadata.PartitionSpecByID(specID), boundRowFilter, scan.caseSensitive) + }) + } + + splitTargetSize := scan.metadata.Properties().GetInt64( + ReadSplitTargetSizeKey, ReadSplitTargetSizeDefault) + acc.resultDataFiles = int64(len(plannedTasks)) + if len(plannedTasks) == 0 { + return []FileScanTask{}, nil + } + + workerCount := min(max(scan.concurrency, 1), runtime.GOMAXPROCS(0), + max(1, len(plannedTasks)/scanTaskFinalizationMinTasksPerWorker)) + results := make([]scanTaskFinalizationResult, workerCount) + + finalizeRange := func(worker, start, end int) { + defer func() { + if recovered := recover(); recovered != nil { + results[worker].err = fmt.Errorf("finalize scan tasks [%d:%d]: panic: %v", + start, end, recovered) + } + }() + results[worker] = scan.finalizeTaskRange( + plannedTasks, start, end, residualEvaluators, splitTargetSize) + } + + if workerCount == 1 { + finalizeRange(0, 0, len(plannedTasks)) + } else { + var wg sync.WaitGroup + chunkSize, remainder := len(plannedTasks)/workerCount, len(plannedTasks)%workerCount + start := 0 + for worker := range workerCount { + // Distribute the remainder instead of rounding every range up: + // even large worker counts must stay within the input slice. + end := start + chunkSize + if worker < remainder { + end++ + } + + wg.Add(1) + go func(worker, start, end int) { + defer wg.Done() + finalizeRange(worker, start, end) + }(worker, start, end) + start = end + } + wg.Wait() + } + + // Ranges follow task order, regardless of the order workers completed. + for i := range results { + if results[i].err != nil { + return nil, results[i].err + } + } + + totalTasks := 0 + anySplit := false + for i := range results { + mergeTaskFinalizationMetrics(acc, &results[i].metrics) + anySplit = anySplit || results[i].expanded + totalTasks += len(results[i].tasks) + } + if !anySplit { + return plannedTasks, nil + } + if workerCount == 1 { + return results[0].tasks, nil + } + + // Workers that split files emit their ranges in task order. Unsplit + // workers reuse their input ranges without temporary replacement lists. + // Concatenate once; do not hold an allocation per individual split. + finalized := make([]FileScanTask, 0, totalTasks) Review Comment: The peak-memory issue from last round is still here on the all-split path. Each worker grows `result.tasks` by append-doubling, and then this loop copies everything into `finalized`, so input, slack and output are all live at once. Pre-sizing `result.tasks` on the first split (roughly `(end-start)*2`, or an estimate from the split offsets) would cover most of it. ########## table/scan_task_finalization_planfiles_bench_test.go: ########## @@ -0,0 +1,169 @@ +// 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 ( + "bytes" + "context" + "fmt" + "testing" + + "github.com/apache/iceberg-go" +) + +const ( + planFilesFinalizationBenchmarkManifestCount = 8 + planFilesFinalizationBenchmarkFilesPerManifest = 2_500 +) + +var planFilesFinalizationBenchmarkSink int + +// BenchmarkPlanFilesTaskFinalization measures the effect of scan-task +// finalization in the context of local PlanFiles. The benchmark fixture is +// shared by the pre-change and parallel implementations, so copy this file to +// the PR base worktree when comparing the two revisions. +func BenchmarkPlanFilesTaskFinalization(b *testing.B) { + tbl := newPlanFilesFinalizationBenchmarkTable(b) + fileCount := planFilesFinalizationBenchmarkManifestCount * planFilesFinalizationBenchmarkFilesPerManifest + + for _, tc := range []struct { + name string + rowFilter iceberg.BooleanExpression + }{ + {name: "plain", rowFilter: iceberg.AlwaysTrue{}}, + {name: "identity_residual", rowFilter: iceberg.EqualTo(iceberg.Reference("id"), int32(7))}, + } { + b.Run(tc.name, func(b *testing.B) { + scan := tbl.Scan(WithMaxConcurrency(8), WithRowFilter(tc.rowFilter)) + b.ReportAllocs() + b.ReportMetric(float64(fileCount), "files/op") + for b.Loop() { + tasks, err := scan.PlanFiles(context.Background()) + if err != nil { + b.Fatal(err) + } + if len(tasks) != fileCount { + b.Fatalf("PlanFiles returned %d tasks, want %d", len(tasks), fileCount) + } + planFilesFinalizationBenchmarkSink = len(tasks) + } + }) + } +} + +func newPlanFilesFinalizationBenchmarkTable(b testing.TB) *Table { + b.Helper() + + const ( + tableLocation = "mem://plan-files-finalization-benchmark" + sequenceNum = int64(1) + ) + snapshotID := int64(1) + + fs := newTrackingIO() + schema := simpleSchema() + spec := partitionedSpec() + meta, err := NewMetadata(schema, &spec, UnsortedSortOrder, tableLocation, nil) + if err != nil { + b.Fatal(err) + } + builder, err := MetadataBuilderFromBase(meta, "") + if err != nil { + b.Fatal(err) + } + + manifests := make([]iceberg.ManifestFile, planFilesFinalizationBenchmarkManifestCount) + for manifestIndex := range manifests { + entries := make([]iceberg.ManifestEntry, planFilesFinalizationBenchmarkFilesPerManifest) + for entryIndex := range entries { + dataPath := fmt.Sprintf("%s/data-%d-%d.parquet", tableLocation, manifestIndex, entryIndex) + fileBuilder, err := iceberg.NewDataFileBuilder( + spec, + iceberg.EntryContentData, + dataPath, + iceberg.ParquetFile, + map[int]any{1000: int32(7)}, Review Comment: Every file here has spec 0, partition value 7, no split offsets and no deletes. So the bench only runs the plain path and a residual that always simplifies to AlwaysTrue, and the in-memory IO pushes finalization's share up. I'd vary the partition value (`i%4`), add a split-offsets case and a case with a positional delete manifest plus a DV, and ideally interleave a second spec. If we keep it this simple, the doc comment should say it's the no-split, no-delete lower bound. ########## table/scan_task_finalization_test.go: ########## @@ -0,0 +1,574 @@ +// 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 ( + "fmt" + "runtime" + "testing" + + "github.com/apache/iceberg-go" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func setFinalizationTestProcs(t *testing.T, count int) { + t.Helper() + previous := runtime.GOMAXPROCS(count) + t.Cleanup(func() { runtime.GOMAXPROCS(previous) }) +} + +func TestFinalizePlannedTasksParallelMatchesSerial(t *testing.T) { + // Exercise the concurrent worker path even on single-core CI runners. + setFinalizationTestProcs(t, 4) + schema := iceberg.NewSchema(1, iceberg.NestedField{ + ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64, Required: true, + }) + metadata, err := NewMetadata( + schema, + iceberg.UnpartitionedSpec, + UnsortedSortOrder, + "mem://table", + iceberg.Properties{ReadSplitTargetSizeKey: "25"}, + ) + require.NoError(t, err) + + tasks := make([]FileScanTask, 1024) + for i := range tasks { + file := finalizationTestDataFile( + t, *iceberg.UnpartitionedSpec, + fmt.Sprintf("mem://table/data/file-%d.parquet", i), nil, + 100, []int64{8, 25, 50, 75}) + tasks[i] = FileScanTask{File: file, Start: 0, Length: file.FileSizeBytes()} + if i%3 == 0 { + tasks[i].DeleteFiles = []iceberg.DataFile{file} + } + if i%5 == 0 { + tasks[i].EqualityDeleteFiles = []iceberg.DataFile{file} + } + } + + scan := &Scan{ + metadata: metadata, rowFilter: iceberg.AlwaysTrue{}, caseSensitive: true, concurrency: 8, + } + serial, serialMetrics, err := serialFinalizationReference( + scan, append([]FileScanTask(nil), tasks...), schema) + require.NoError(t, err) + var parallelMetrics scanMetricsAccumulator + parallel, err := scan.finalizePlannedTasks( + append([]FileScanTask(nil), tasks...), schema, ¶llelMetrics) + require.NoError(t, err) + + assert.Equal(t, serial, parallel) + assert.Equal(t, serialMetrics, parallelMetrics) + require.Len(t, parallel, len(tasks)*4) + for i := range tasks { + for split, start := range []int64{0, 25, 50, 75} { + task := parallel[i*4+split] + assert.Equal(t, fmt.Sprintf("mem://table/data/file-%d.parquet", i), task.File.FilePath()) + assert.Equal(t, start, task.Start) + assert.Equal(t, int64(25), task.Length) + } + } + assert.Equal(t, int64(len(tasks)), parallelMetrics.resultDataFiles) + assert.Equal(t, int64(len(tasks)*100), parallelMetrics.totalFileSize) + assert.Equal(t, int64(342), parallelMetrics.positionalDeleteFiles) + assert.Equal(t, int64(205), parallelMetrics.equalityDeleteFiles) + assert.Equal(t, int64(547), parallelMetrics.resultDeleteFiles) + assert.Equal(t, int64(54_700), parallelMetrics.totalDeleteFileSize) +} + +func TestFinalizePlannedTasksParallelResidualsMatchSerial(t *testing.T) { + // Exercise the concurrent worker path even on single-core CI runners. + setFinalizationTestProcs(t, 4) + schema := simpleSchema() + spec := partitionedSpec() + unpartitioned := iceberg.NewPartitionSpecID(1) + metadata, err := NewMetadata( + schema, + &spec, + UnsortedSortOrder, + "mem://table", + nil, + ) + require.NoError(t, err) + builder, err := MetadataBuilderFromBase(metadata, "") + require.NoError(t, err) + require.NoError(t, builder.AddPartitionSpec(&unpartitioned, false)) + metadata, err = builder.Build() + require.NoError(t, err) + + tasks := make([]FileScanTask, 1024) + for i := range tasks { + partitionValue := int32(7) + if i%2 != 0 { + partitionValue = 8 + } + fileSpec := spec + if i%4 < 2 { + fileSpec = unpartitioned + } + file := finalizationTestDataFile( + t, fileSpec, + fmt.Sprintf("mem://table/data/file-%d.parquet", i), + map[int]any{1000: partitionValue}, 100, nil) + tasks[i] = FileScanTask{File: file, Start: 0, Length: file.FileSizeBytes()} + } + + filter := iceberg.EqualTo(iceberg.Reference("id"), int32(7)) + scan := &Scan{metadata: metadata, rowFilter: filter, caseSensitive: true, concurrency: 8} + serial, serialMetrics, err := serialFinalizationReference( + scan, append([]FileScanTask(nil), tasks...), schema) + require.NoError(t, err) + var parallelMetrics scanMetricsAccumulator + parallel, err := scan.finalizePlannedTasks( + append([]FileScanTask(nil), tasks...), schema, ¶llelMetrics) + require.NoError(t, err) + + assert.Equal(t, serial, parallel) + assert.Equal(t, serialMetrics, parallelMetrics) + for i, task := range parallel { + if i%4 < 2 { + assert.Nil(t, task.Residual, "task %d has no partition projection", i) + } else if i%2 == 0 { + _, ok := task.Residual.(iceberg.AlwaysTrue) + assert.True(t, ok, "task %d should simplify to always true", i) + } else { + _, ok := task.Residual.(iceberg.AlwaysFalse) + assert.True(t, ok, "task %d should simplify to always false", i) + } + } +} + +func TestFinalizePlannedTasksReturnsFirstErrorInTaskOrder(t *testing.T) { + schema := simpleSchema() + spec := partitionedSpec() + metadata, err := NewMetadata(schema, &spec, UnsortedSortOrder, "mem://table", nil) + require.NoError(t, err) + + tasks := make([]FileScanTask, 1024) + for i := range tasks { + var partitionValue any = int32(7) + if i == 255 || i == 256 { + // The earlier failing task is at the end of worker range 0; + // worker 1 can fail first at the start of its own range. + partitionValue = "invalid integer partition" + } + file := finalizationTestDataFile(t, spec, + fmt.Sprintf("mem://table/data/file-%d.parquet", i), + map[int]any{1000: partitionValue}, 100, nil) + tasks[i] = FileScanTask{File: file, Length: file.FileSizeBytes()} + } + + for _, concurrency := range []int{1, 8} { + t.Run(fmt.Sprintf("concurrency=%d", concurrency), func(t *testing.T) { + scan := &Scan{ + metadata: metadata, concurrency: concurrency, caseSensitive: true, + rowFilter: iceberg.EqualTo(iceberg.Reference("id"), int32(7)), + } + input := append([]FileScanTask(nil), tasks...) + if concurrency > 1 { + // With four CPU workers and 1024 tasks, worker 0 owns + // [0:256] and worker 1 owns [256:512]. + previous := runtime.GOMAXPROCS(4) + defer runtime.GOMAXPROCS(previous) + + laterErrorReached := make(chan struct{}) + input[255].File = finalizationOrderFile{ + DataFile: input[255].File, wait: laterErrorReached, + } + input[256].File = finalizationOrderFile{ + DataFile: input[256].File, signal: laterErrorReached, + } + } + result, planErr := scan.finalizePlannedTasks(input, schema, &scanMetricsAccumulator{}) + require.ErrorContains(t, planErr, "evaluate partition residual for mem://table/data/file-255.parquet") + assert.Nil(t, result) + }) + } +} + +// Controls which worker reaches its partition error first, without relying on +// scheduler timing. The wrapper intentionally does not forward the built-in +// borrowed-partition interface, so residual evaluation calls Partition. +type finalizationOrderFile struct { + iceberg.DataFile + wait <-chan struct{} + signal chan struct{} +} + +func (f finalizationOrderFile) Partition() map[int]any { + if f.signal != nil { + close(f.signal) Review Comment: `close(f.signal)` has no `sync.Once`. If `Partition()` runs twice on task 256, the second close panics and the recover turns it into `results[1].err`. Worker 0's error still wins, so the test passes without proving the ordering. Separately, `<-f.wait` has no timeout, so if the worker split ever shifts, this hangs until the 10-minute test timeout. I'd add a `*sync.Once` field (the wrapper is a value type) and `select` on `time.After(5 * time.Second)` with a panic, which the recover turns into a failed `ErrorContains`. ########## table/scan_task_finalization.go: ########## @@ -0,0 +1,211 @@ +// 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 ( + "fmt" + "runtime" + "sync" + + "github.com/apache/iceberg-go" +) + +// CPU-heavy finalization needs enough work to amortize worker and merge costs. +const scanTaskFinalizationMinTasksPerWorker = 256 Review Comment: Nothing measures 256 yet. Both benches run at 20K to 100K tasks where it never binds, so the first parallel point at 512 tasks has nothing behind it. A quick sweep in the micro-bench (512, 1K, 2K, 4K) would show where break-even sits for `plain` vs `identity_residual`. Otherwise, call it a heuristic in the comment. Do we have a rough sense of where it crosses over? -- 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]
