laskoviymishka commented on code in PR #2149: URL: https://github.com/apache/iceberg-go/pull/2149#discussion_r4337002562
########## table/scan_task_finalization_test.go: ########## @@ -0,0 +1,645 @@ +// 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" + "testing" + "time" + + "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, signalOnce: new(sync.Once), + } + } + 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{} + signalOnce *sync.Once +} + +func (f finalizationOrderFile) Partition() map[int]any { + if f.signal != nil { + f.signalOnce.Do(func() { close(f.signal) }) + } + if f.wait != nil { + select { + case <-f.wait: + case <-time.After(5 * time.Second): + panic("timed out waiting for later worker to reach its error") + } + } + + return f.DataFile.Partition() +} + +func TestFinalizePlannedTasksUnsplitReusesInput(t *testing.T) { + // Exercise the concurrent worker path even on single-core CI runners. + setFinalizationTestProcs(t, 4) + schema := simpleSchema() + metadata, err := NewMetadata( + schema, + iceberg.UnpartitionedSpec, + UnsortedSortOrder, + "mem://table", + nil, + ) + require.NoError(t, err) + file := finalizationTestDataFile( + t, *iceberg.UnpartitionedSpec, "mem://table/data/file.parquet", nil, 100, nil) + tasks := make([]FileScanTask, 512) + for i := range tasks { + tasks[i] = FileScanTask{File: file, Start: 0, Length: file.FileSizeBytes()} + } + scan := &Scan{metadata: metadata, rowFilter: iceberg.AlwaysTrue{}, caseSensitive: true, concurrency: 8} + + got, err := scan.finalizePlannedTasks(tasks, schema, &scanMetricsAccumulator{}) + require.NoError(t, err) + require.Len(t, got, len(tasks)) + assert.Same(t, &tasks[0], &got[0]) + assert.Same(t, &tasks[len(tasks)-1], &got[len(got)-1]) + assert.Nil(t, got[0].Residual) +} + +func TestFinalizePlannedTasksManyWorkersCoverEveryTask(t *testing.T) { + // A rounded-up chunk size can put the last worker past the input when + // there are more workers than tasks per worker and a short remainder. + const workers = 300 + previous := runtime.GOMAXPROCS(workers) + defer runtime.GOMAXPROCS(previous) + + schema := simpleSchema() + metadata, err := NewMetadata(schema, iceberg.UnpartitionedSpec, UnsortedSortOrder, + "mem://table", iceberg.Properties{ReadSplitTargetSizeKey: "25"}) + require.NoError(t, err) + file := finalizationTestDataFile(t, *iceberg.UnpartitionedSpec, + "mem://table/data/unsplit.parquet", nil, 100, nil) + splitFile := finalizationTestDataFile(t, *iceberg.UnpartitionedSpec, + "mem://table/data/split.parquet", nil, 100, []int64{8, 25, 50, 75}) + lastFile := finalizationTestDataFile(t, *iceberg.UnpartitionedSpec, + "mem://table/data/last.parquet", nil, 100, nil) + tasks := make([]FileScanTask, workers*scanTaskFinalizationMinTasksPerWorker+1) + firstRowIDs := make([]int64, len(tasks)) + for i := range tasks { + firstRowIDs[i] = int64(i * 100) + tasks[i] = FileScanTask{File: file, Length: 100, FirstRowID: &firstRowIDs[i]} + } + tasks[0].File = splitFile + tasks[len(tasks)-1].File = lastFile + scan := &Scan{metadata: metadata, rowFilter: iceberg.AlwaysTrue{}, concurrency: workers} + + var acc scanMetricsAccumulator + got, err := scan.finalizePlannedTasks(tasks, schema, &acc) + require.NoError(t, err) + require.Len(t, got, len(tasks)+3) + for i := range 4 { + assert.Equal(t, splitFile.FilePath(), got[i].File.FilePath()) + assert.Equal(t, int64(i*25), got[i].Start) + assert.Equal(t, int64(25), got[i].Length) + assert.Same(t, tasks[0].FirstRowID, got[i].FirstRowID) + } + for i, task := range got[4 : len(got)-1] { + if task.File != file || task.Start != 0 || task.Length != 100 || task.FirstRowID != tasks[i+1].FirstRowID { + t.Fatalf("unsplit task %d changed or moved: %+v", i+1, task) + } + } + assert.Equal(t, lastFile.FilePath(), got[len(got)-1].File.FilePath()) + assert.Same(t, tasks[len(tasks)-1].FirstRowID, got[len(got)-1].FirstRowID) + assert.Equal(t, int64(len(tasks)), acc.resultDataFiles) + assert.Equal(t, int64(len(tasks)*100), acc.totalFileSize) +} + +// serialFinalizationReference reproduces the pre-parallelization loop. It is +// deliberately independent of finalizeTaskRange and the worker merge logic. +func serialFinalizationReference( + scan *Scan, + tasks []FileScanTask, + schema *iceberg.Schema, +) ([]FileScanTask, scanMetricsAccumulator, error) { + var acc scanMetricsAccumulator + var filter iceberg.BooleanExpression + if scan.rowFilter != nil && !scan.rowFilter.Equals(iceberg.AlwaysTrue{}) { + var err error + filter, err = iceberg.BindExpr(schema, scan.rowFilter, scan.caseSensitive) + if err != nil { + return nil, acc, err + } + } + evaluators := make(map[int]*partitionResidualEvaluator) + acc.resultDataFiles = int64(len(tasks)) + result := make([]FileScanTask, 0, len(tasks)) + target := scan.metadata.Properties().GetInt64(ReadSplitTargetSizeKey, ReadSplitTargetSizeDefault) + for _, task := range tasks { + if filter != nil { + specID := int(task.File.SpecID()) + evaluator, ok := evaluators[specID] + if !ok { + var err error + evaluator, err = newPartitionResidualEvaluator( + schema, scan.metadata.PartitionSpecByID(specID), filter, scan.caseSensitive) + if err != nil { + return nil, acc, fmt.Errorf("build partition residual evaluator for spec %d: %w", specID, err) + } + evaluators[specID] = evaluator + } + if evaluator != nil { + var err error + var simplified bool + task.Residual, simplified, err = evaluator.residual(dataFilePartition(task.File)) + if err != nil { + return nil, acc, fmt.Errorf( + "evaluate partition residual for %s: %w", task.File.FilePath(), err) + } + if !simplified { + task.Residual = nil + } + } + } + acc.addResultDeleteMetrics(task) + acc.totalFileSize += task.File.FileSizeBytes() + if splits, ok := splitParquetScanTask(task, target); ok { + result = append(result, splits...) + } else { + result = append(result, task) + } + } + + return result, acc, nil +} + +func TestFinalizePlannedTasksMixedSplitsAgainstIndependentReference(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) + + dvBuilder, err := iceberg.NewDataFileBuilder( + *iceberg.UnpartitionedSpec, iceberg.EntryContentPosDeletes, + "mem://table/deletes.dv.puffin", iceberg.PuffinFile, nil, nil, nil, 1, 2048) + require.NoError(t, err) + dv := dvBuilder.ContentSizeInBytes(17).ContentOffset(100).Build() + + for _, count := range []int{0, 1, 255, 256, 257, 512, 1024} { + t.Run(fmt.Sprintf("files=%d", count), func(t *testing.T) { + tasks := make([]FileScanTask, count) + splits := 0 + for i := range tasks { + offsets := []int64(nil) + if i%7 < 2 || i == count-1 || i == 255 || i == 256 { + offsets = []int64{8, 25, 50, 75} + splits++ + } + file := finalizationTestDataFile(t, *iceberg.UnpartitionedSpec, + fmt.Sprintf("mem://table/data/file-%d.parquet", i), nil, 100, offsets) + tasks[i] = FileScanTask{File: file, Length: 100} + if i%9 == 0 { + tasks[i].DeleteFiles = []iceberg.DataFile{file} + } + if i%11 == 0 { + tasks[i].EqualityDeleteFiles = []iceberg.DataFile{file} + } + if i%13 == 0 { + tasks[i].DeletionVectorFiles = []iceberg.DataFile{dv} + } + } + scan := &Scan{ + metadata: metadata, rowFilter: iceberg.AlwaysTrue{}, Review Comment: Could one case here use a real row filter? Every split test runs with `AlwaysTrue`, and the residual test passes `nil` offsets ([L131](https://github.com/apache/iceberg-go/blob/59ccb1e2718a8721186384457f210fdefd2190cd/table/scan_task_finalization_test.go#L131)), so nothing checks that split tasks keep the `Residual` written at [L172](https://github.com/apache/iceberg-go/blob/59ccb1e2718a8721186384457f210fdefd2190cd/table/scan_task_finalization.go#L172). As I read it, a refactor that sets residuals after splitting would still pass. A partitioned spec with `EqualTo(id, 7)` checked against `serialFinalizationReference` would pin it. -- 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]
