diff --git a/catalog/rest/fetch_scan_tasks_validation.go b/catalog/rest/fetch_scan_tasks_validation.go index 45803eb8e..ebb5ea185 100644 --- a/catalog/rest/fetch_scan_tasks_validation.go +++ b/catalog/rest/fetch_scan_tasks_validation.go @@ -33,12 +33,16 @@ func validatePlanningTaskEnvelope(data []byte, status PlanStatus, tasks ScanTask return err } - if status != PlanStatusCompleted { - for _, name := range []string{"plan-tasks", "file-scan-tasks", "delete-files"} { - raw, ok := fields[name] - if ok && !isJSONNull(raw) { - return fmt.Errorf("%w: %s response includes %s for status %q", ErrRESTError, endpoint, name, status) - } + for _, name := range []string{"plan-tasks", "file-scan-tasks", "delete-files"} { + raw, ok := fields[name] + if !ok { + continue + } + if isJSONNull(raw) { + return fmt.Errorf("%w: %s response field %s must not be null", ErrRESTError, endpoint, name) + } + if status != PlanStatusCompleted { + return fmt.Errorf("%w: %s response includes %s for status %q", ErrRESTError, endpoint, name, status) } } diff --git a/catalog/rest/scan_planning.go b/catalog/rest/scan_planning.go index b44399221..93ddfb6d1 100644 --- a/catalog/rest/scan_planning.go +++ b/catalog/rest/scan_planning.go @@ -30,6 +30,7 @@ import ( "math/rand/v2" "net/http" "net/url" + "runtime" "slices" "strconv" "strings" @@ -41,6 +42,7 @@ import ( iceio "github.com/apache/iceberg-go/io" "github.com/apache/iceberg-go/table" "github.com/google/uuid" + "golang.org/x/sync/errgroup" "golang.org/x/sync/semaphore" ) @@ -247,7 +249,8 @@ func (r *Catalog) PlanFiles(ctx context.Context, req table.ScanPlanningRequest) "%w: unexpected plan status %q from planTableScan", ErrRESTError, resp.Status) } - envelopes, err := r.collectScanTasks(ctx, req.Identifier, completed.ScanTasks) + envelopes, err := r.collectScanTasksWithConcurrency( + ctx, req.Identifier, completed.ScanTasks, req.MaxConcurrency) if err != nil { cleanup() @@ -293,35 +296,153 @@ func (r *Catalog) planIOBaseProps(req table.ScanPlanningRequest) iceberg.Propert return props } -// collectScanTasks expands plan-task handles into their task envelopes, walking -// the fanout: a fetchScanTasks response can itself return more plan-tasks. A -// handle is fetched at most once; a server that re-issues one would otherwise -// loop forever. The envelope boundaries are retained because delete-file -// references are local to each response. -func (r *Catalog) collectScanTasks(ctx context.Context, ident table.Identifier, tasks ScanTasks) ([]ScanTasks, error) { +// collectScanTasksWithConcurrency expands plan-task handles into their task +// envelopes, walking the fanout: a response can itself return more plan-tasks. Each +// frontier is fetched concurrently, but its responses are appended in handle +// order so completion timing cannot change the result order. A handle is +// fetched at most once; a server that re-issues one would otherwise loop +// forever. The envelope boundaries are retained because delete-file references +// are local to each response. +// +// maxConcurrency bounds the number of concurrent fetches. A non-positive limit +// uses runtime.GOMAXPROCS. On a fetch error, at most the configured number of +// requests can already be in flight when cancellation reaches their contexts. +func (r *Catalog) collectScanTasksWithConcurrency( + ctx context.Context, + ident table.Identifier, + tasks ScanTasks, + maxConcurrency int, +) ([]ScanTasks, error) { + if maxConcurrency <= 0 { + maxConcurrency = runtime.GOMAXPROCS(0) + } + envelopes := []ScanTasks{tasks} - queue := append([]string(nil), tasks.PlanTasks...) - seen := make(map[string]bool, len(queue)) - for len(queue) > 0 { - handle := queue[0] - queue = queue[1:] - if seen[handle] { - continue + frontier := append([]string(nil), tasks.PlanTasks...) + seen := make(map[string]bool, len(frontier)) + for len(frontier) > 0 { + handles := make([]string, 0, len(frontier)) + for _, handle := range frontier { + if seen[handle] { + continue + } + seen[handle] = true + handles = append(handles, handle) + } + if len(handles) == 0 { + break } - seen[handle] = true - resp, err := r.FetchScanTasks(ctx, ident, FetchScanTasksRequest{PlanTask: handle}) + responses, err := r.fetchScanTaskFrontier(ctx, ident, handles, maxConcurrency) if err != nil { return nil, err } - envelopes = append(envelopes, resp.ScanTasks) - queue = append(queue, resp.PlanTasks...) + + var nextFrontier []string + for _, response := range responses { + envelopes = append(envelopes, response.ScanTasks) + nextFrontier = append(nextFrontier, response.PlanTasks...) + } + frontier = nextFrontier } return envelopes, nil } +// fetchScanTaskFrontier fetches one breadth-first frontier concurrently while +// placing responses back into handle order. When a handle fails, unfinished +// later handles are canceled, but earlier handles are allowed to finish so the +// first error in the old serial handle order remains deterministic. +func (r *Catalog) fetchScanTaskFrontier( + ctx context.Context, + ident table.Identifier, + handles []string, + maxConcurrency int, +) ([]FetchScanTasksResponse, error) { + responses := make([]FetchScanTasksResponse, len(handles)) + errs := make([]error, len(handles)) + if len(handles) == 0 { + return responses, nil + } + if maxConcurrency <= 0 { + maxConcurrency = runtime.GOMAXPROCS(0) + } + maxConcurrency = min(maxConcurrency, len(handles)) + + requestCtxs := make([]context.Context, len(handles)) + requestCancels := make([]context.CancelFunc, len(handles)) + for i := range handles { + requestCtxs[i], requestCancels[i] = context.WithCancel(ctx) + } + defer func() { + for _, cancel := range requestCancels { + cancel() + } + }() + + var ( + mu sync.Mutex + completed = make([]bool, len(handles)) + siblingCanceled = make([]bool, len(handles)) + lowestFailureIdx = -1 + ) + + var group errgroup.Group + group.SetLimit(maxConcurrency) + + for i, handle := range handles { + group.Go(func() error { + response, fetchErr := r.FetchScanTasks(requestCtxs[i], ident, FetchScanTasksRequest{PlanTask: handle}) + + mu.Lock() + defer mu.Unlock() + completed[i] = true + if fetchErr == nil { + responses[i] = response + + return nil + } + + errs[i] = fmt.Errorf("fetching scan tasks for handle %q: %w", handle, fetchErr) + if siblingCanceled[i] || (lowestFailureIdx >= 0 && i > lowestFailureIdx) { + return fetchErr + } + + if lowestFailureIdx < 0 || i < lowestFailureIdx { + lowestFailureIdx = i + for j := i + 1; j < len(handles); j++ { + if completed[j] || siblingCanceled[j] { + continue + } + siblingCanceled[j] = true + requestCancels[j]() + } + } + + return fetchErr + }) + } + + if waitErr := group.Wait(); waitErr != nil { + if err := ctx.Err(); err != nil { + return nil, err + } + + mu.Lock() + defer mu.Unlock() + for i, err := range errs { + if err != nil && !siblingCanceled[i] { + return nil, err + } + } + + return nil, waitErr + } + + return responses, nil +} + // remoteScanTasks decodes each server task envelope into domain FileScanTasks. // The decoder must run before envelopes are combined because delete-file // references are indexes into the envelope-local delete-files array. diff --git a/catalog/rest/scan_planning_bench_test.go b/catalog/rest/scan_planning_bench_test.go new file mode 100644 index 000000000..531bd2a12 --- /dev/null +++ b/catalog/rest/scan_planning_bench_test.go @@ -0,0 +1,73 @@ +// 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 rest + +import ( + "context" + "fmt" + "net/http" + "net/http/httptest" + "net/url" + "testing" + "time" + + "github.com/apache/iceberg-go/table" +) + +func BenchmarkCollectScanTasks64Handles(b *testing.B) { + const ( + handleCount = 64 + latency = 10 * time.Millisecond + ) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { + time.Sleep(latency) + _, _ = w.Write([]byte(`{"file-scan-tasks":[]}`)) + })) + b.Cleanup(server.Close) + + serverURL, err := url.Parse(server.URL) + if err != nil { + b.Fatal(err) + } + + catalog := &Catalog{ + baseURI: serverURL.JoinPath("v1"), + cl: server.Client(), + endpoints: newEndpointSet([]endpoint{endpointFetchScanTasks}), + } + handles := make([]string, handleCount) + for i := range handles { + handles[i] = fmt.Sprintf("task-%d", i) + } + tasks := ScanTasks{PlanTasks: handles} + ident := table.Identifier{"db", "tbl"} + + for _, maxConcurrency := range []int{1, 2, 4, 8, 16, 32, 64} { + b.Run(fmt.Sprintf("concurrency-%d", maxConcurrency), func(b *testing.B) { + b.ResetTimer() + for range b.N { + _, err := catalog.collectScanTasksWithConcurrency( + context.Background(), ident, tasks, maxConcurrency) + if err != nil { + b.Fatal(err) + } + } + }) + } +} diff --git a/catalog/rest/scan_planning_test.go b/catalog/rest/scan_planning_test.go index a45972f6a..093eca2cd 100644 --- a/catalog/rest/scan_planning_test.go +++ b/catalog/rest/scan_planning_test.go @@ -27,6 +27,7 @@ import ( "net/url" "strconv" "strings" + "sync" "sync/atomic" "testing" "time" @@ -213,9 +214,15 @@ func TestPlanTableScanResponseRejectsInvalidStatusEnvelope(t *testing.T) { for i, payload := range []string{ `{"status":"submitted","plan-id":"abc","file-scan-tasks":[]}`, + `{"status":"submitted","plan-id":"abc","file-scan-tasks":null}`, `{"status":"submitted","plan-id":"abc","delete-files":[]}`, + `{"status":"submitted","plan-id":"abc","delete-files":null}`, `{"status":"failed","plan-tasks":[]}`, + `{"status":"failed","plan-tasks":null}`, `{"status":"completed","plan-id":"abc","delete-files":[{}]}`, + `{"status":"completed","plan-id":"abc","plan-tasks":null}`, + `{"status":"completed","plan-id":"abc","file-scan-tasks":null}`, + `{"status":"completed","plan-id":"abc","delete-files":null}`, `{"status":"failed","plan-id":"abc"}`, } { t.Run(fmt.Sprintf("payload-%d", i), func(t *testing.T) { @@ -315,9 +322,15 @@ func TestFetchPlanningResultResponseValidation(t *testing.T) { for i, payload := range []string{ `{"status":"submitted","plan-tasks":[]}`, + `{"status":"submitted","plan-tasks":null}`, `{"status":"submitted","delete-files":[]}`, + `{"status":"submitted","delete-files":null}`, `{"status":"cancelled","file-scan-tasks":[]}`, + `{"status":"cancelled","file-scan-tasks":null}`, `{"status":"completed","delete-files":[{}]}`, + `{"status":"completed","plan-tasks":null}`, + `{"status":"completed","file-scan-tasks":null}`, + `{"status":"completed","delete-files":null}`, } { t.Run(fmt.Sprintf("payload-%d", i), func(t *testing.T) { t.Parallel() @@ -1218,6 +1231,7 @@ func TestPlanFilesPollsSubmittedPlan(t *testing.T) { func TestPlanFilesExpandsPlanTasks(t *testing.T) { t.Parallel() + var mu sync.Mutex var fetched []string cat := newScanPlanningTestCatalog(t, []endpoint{endpointPlanTableScan, endpointFetchScanTasks}, func(mux *http.ServeMux) { mux.HandleFunc("/v1/namespaces/db/tables/tbl/plan", func(w http.ResponseWriter, req *http.Request) { @@ -1227,7 +1241,9 @@ func TestPlanFilesExpandsPlanTasks(t *testing.T) { mux.HandleFunc("/v1/namespaces/db/tables/tbl/tasks", func(w http.ResponseWriter, req *http.Request) { var body FetchScanTasksRequest require.NoError(t, json.NewDecoder(req.Body).Decode(&body)) + mu.Lock() fetched = append(fetched, body.PlanTask) + mu.Unlock() switch body.PlanTask { case "h1": _, err := w.Write([]byte(`{"plan-tasks":["h2"]}`)) @@ -1242,9 +1258,212 @@ func TestPlanFilesExpandsPlanTasks(t *testing.T) { result, err := cat.PlanFiles(context.Background(), planFilesReq()) require.NoError(t, err) assert.Empty(t, result.Tasks) + mu.Lock() + defer mu.Unlock() assert.Equal(t, []string{"h1", "h2"}, fetched) } +func TestCollectScanTasksFetchesFrontierConcurrentlyInOrder(t *testing.T) { + t.Parallel() + + h1Started := make(chan struct{}) + h2Started := make(chan struct{}) + releaseH1 := make(chan struct{}) + var releaseOnce sync.Once + release := func() { releaseOnce.Do(func() { close(releaseH1) }) } + + cat := newScanPlanningTestCatalog(t, []endpoint{endpointFetchScanTasks}, func(mux *http.ServeMux) { + mux.HandleFunc("/v1/namespaces/db/tables/tbl/tasks", func(w http.ResponseWriter, req *http.Request) { + var body FetchScanTasksRequest + if err := json.NewDecoder(req.Body).Decode(&body); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + + return + } + + switch body.PlanTask { + case "h1": + close(h1Started) + <-releaseH1 + _, _ = w.Write([]byte(`{"file-scan-tasks":[{"data-file":{"file-path":"h1"}}]}`)) + case "h2": + close(h2Started) + _, _ = w.Write([]byte(`{"file-scan-tasks":[{"data-file":{"file-path":"h2"}}]}`)) + default: + http.Error(w, "unexpected plan task", http.StatusBadRequest) + } + }) + }) + t.Cleanup(release) + + type outcome struct { + envelopes []ScanTasks + err error + } + done := make(chan outcome, 1) + go func() { + envelopes, err := cat.collectScanTasksWithConcurrency(t.Context(), table.Identifier{"db", "tbl"}, ScanTasks{ + PlanTasks: []string{"h1", "h2"}, + }, 2) + done <- outcome{envelopes: envelopes, err: err} + }() + + for name, started := range map[string]<-chan struct{}{ + "h1": h1Started, + "h2": h2Started, + } { + select { + case <-started: + case <-time.After(time.Second): + t.Fatalf("timed out waiting for %s to start", name) + } + } + + release() + select { + case result := <-done: + require.NoError(t, result.err) + require.Len(t, result.envelopes, 3) + require.Len(t, result.envelopes[1].FileScanTasks, 1) + require.Len(t, result.envelopes[2].FileScanTasks, 1) + require.NotNil(t, result.envelopes[1].FileScanTasks[0].DataFile) + require.NotNil(t, result.envelopes[2].FileScanTasks[0].DataFile) + assert.Equal(t, "h1", result.envelopes[1].FileScanTasks[0].DataFile.FilePath) + assert.Equal(t, "h2", result.envelopes[2].FileScanTasks[0].DataFile.FilePath) + case <-time.After(time.Second): + t.Fatal("timed out waiting for frontier fetches") + } +} + +type orderedScanTaskErrorTransport struct { + planTaskReturned chan struct{} +} + +func (t *orderedScanTaskErrorTransport) RoundTrip(req *http.Request) (*http.Response, error) { + var body FetchScanTasksRequest + if err := json.NewDecoder(req.Body).Decode(&body); err != nil { + return nil, err + } + + var errType string + switch body.PlanTask { + case "table": + select { + case <-t.planTaskReturned: + case <-req.Context().Done(): + return nil, req.Context().Err() + } + errType = errTypeNoSuchTable + case "plan-task": + close(t.planTaskReturned) + errType = errTypeNoSuchPlanTask + default: + return nil, fmt.Errorf("unexpected plan task %q", body.PlanTask) + } + + data := fmt.Sprintf(`{"error":{"message":%q,"type":%q,"code":404}}`, errType, errType) + + return &http.Response{ + StatusCode: http.StatusNotFound, + Header: http.Header{"Content-Type": {"application/json"}}, + Body: io.NopCloser(strings.NewReader(data)), + ContentLength: int64(len(data)), + Request: req, + }, nil +} + +func TestCollectScanTasksReturnsFirstErrorInHandleOrder(t *testing.T) { + t.Parallel() + + cat := newScanPlanningTestCatalog(t, []endpoint{endpointFetchScanTasks}, nil) + cat.cl = &http.Client{Transport: &orderedScanTaskErrorTransport{ + planTaskReturned: make(chan struct{}), + }} + + envelopes, err := cat.collectScanTasksWithConcurrency(t.Context(), table.Identifier{"db", "tbl"}, ScanTasks{ + PlanTasks: []string{"table", "plan-task"}, + }, 2) + require.ErrorIs(t, err, catalog.ErrNoSuchTable) + assert.Contains(t, err.Error(), `handle "table"`) + assert.NotErrorIs(t, err, ErrNoSuchPlanTask) + assert.Nil(t, envelopes) +} + +func TestCollectScanTasksBoundsFrontierConcurrency(t *testing.T) { + t.Parallel() + + const maxConcurrency = 8 + const handleCount = maxConcurrency + 1 + started := make(chan string, handleCount) + release := make(chan struct{}) + var releaseOnce sync.Once + finish := func() { releaseOnce.Do(func() { close(release) }) } + + cat := newScanPlanningTestCatalog(t, []endpoint{endpointFetchScanTasks}, func(mux *http.ServeMux) { + mux.HandleFunc("/v1/namespaces/db/tables/tbl/tasks", func(w http.ResponseWriter, req *http.Request) { + var body FetchScanTasksRequest + if err := json.NewDecoder(req.Body).Decode(&body); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + + return + } + + started <- body.PlanTask + <-release + _, _ = w.Write([]byte(`{"file-scan-tasks":[]}`)) + }) + }) + t.Cleanup(finish) + + handles := make([]string, handleCount) + for i := range handles { + handles[i] = fmt.Sprintf("h%d", i) + } + done := make(chan error, 1) + go func() { + _, err := cat.collectScanTasksWithConcurrency(t.Context(), table.Identifier{"db", "tbl"}, ScanTasks{ + PlanTasks: handles, + }, maxConcurrency) + done <- err + }() + + for range maxConcurrency { + select { + case <-started: + case <-time.After(time.Second): + t.Fatal("timed out waiting for bounded frontier workers") + } + } + + select { + case handle := <-started: + t.Fatalf("frontier exceeded concurrency limit with %s", handle) + case <-time.After(50 * time.Millisecond): + } + + finishOne := func() { + select { + case release <- struct{}{}: + case <-time.After(time.Second): + t.Fatal("timed out releasing a frontier worker") + } + } + finishOne() + select { + case <-started: + case <-time.After(time.Second): + t.Fatal("timed out starting the queued frontier handle") + } + finish() + + select { + case err := <-done: + require.NoError(t, err) + case <-time.After(time.Second): + t.Fatal("timed out waiting for bounded frontier fetches") + } +} + // TestPlanFilesFanoutCycleTerminates guards the seen-set: a server that re-issues // a handle it already returned must not loop forever. func TestPlanFilesFanoutCycleTerminates(t *testing.T) { @@ -1685,3 +1904,94 @@ func TestPlanFilesPropagatesFailure(t *testing.T) { _, err := cat.PlanFiles(context.Background(), planFilesReq()) require.ErrorIs(t, err, ErrPlanFailed) } + +func TestCollectScanTasksDeduplicatesAcrossFrontiers(t *testing.T) { + t.Parallel() + + children := map[string][]string{ + "a": {"c", "d", "a"}, + "b": {"d", "e"}, + "c": {"b"}, + } + var mu sync.Mutex + calls := make(map[string]int) + cat := newScanPlanningTestCatalog(t, []endpoint{endpointFetchScanTasks}, func(mux *http.ServeMux) { + mux.HandleFunc("/v1/namespaces/db/tables/tbl/tasks", func(w http.ResponseWriter, req *http.Request) { + var body FetchScanTasksRequest + if err := json.NewDecoder(req.Body).Decode(&body); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + + return + } + mu.Lock() + calls[body.PlanTask]++ + mu.Unlock() + response := ScanTasks{ + PlanTasks: children[body.PlanTask], + FileScanTasks: []RESTFileScanTask{{DataFile: &RESTDataFile{ + RESTContentFile: RESTContentFile{FilePath: body.PlanTask}, + }}}, + } + _ = json.NewEncoder(w).Encode(response) + }) + }) + + envelopes, err := cat.collectScanTasksWithConcurrency(t.Context(), table.Identifier{"db", "tbl"}, ScanTasks{ + PlanTasks: []string{"a", "a", "b"}, + }, 2) + require.NoError(t, err) + require.Len(t, envelopes, 6) + for i, name := range []string{"a", "b", "c", "d", "e"} { + require.Len(t, envelopes[i+1].FileScanTasks, 1) + assert.Equal(t, name, envelopes[i+1].FileScanTasks[0].DataFile.FilePath) + } + mu.Lock() + defer mu.Unlock() + assert.Equal(t, map[string]int{"a": 1, "b": 1, "c": 1, "d": 1, "e": 1}, calls) +} + +func TestCollectScanTasksCancelsSiblingRequestsOnFailure(t *testing.T) { + t.Parallel() + + started := make(chan struct{}) + cancelled := make(chan struct{}) + ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second) + defer cancel() + cat := newScanPlanningTestCatalog(t, []endpoint{endpointFetchScanTasks}, func(mux *http.ServeMux) { + mux.HandleFunc("/v1/namespaces/db/tables/tbl/tasks", func(w http.ResponseWriter, req *http.Request) { + var body FetchScanTasksRequest + if err := json.NewDecoder(req.Body).Decode(&body); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + + return + } + switch body.PlanTask { + case "slow": + close(started) + <-req.Context().Done() + close(cancelled) + case "failed": + select { + case <-started: + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusBadRequest) + _, _ = w.Write([]byte(`{"error":{"message":"invalid handle","type":"BadRequestException","code":400}}`)) + case <-req.Context().Done(): + } + default: + http.Error(w, "unexpected handle", http.StatusBadRequest) + } + }) + }) + + envelopes, err := cat.collectScanTasksWithConcurrency(ctx, table.Identifier{"db", "tbl"}, ScanTasks{ + PlanTasks: []string{"failed", "slow"}, + }, 2) + require.ErrorIs(t, err, ErrBadRequest) + assert.Nil(t, envelopes) + select { + case <-cancelled: + case <-ctx.Done(): + t.Fatal("sibling request was not cancelled after a fetch failure") + } +} diff --git a/schema.go b/schema.go index 996577a12..42501789e 100644 --- a/schema.go +++ b/schema.go @@ -352,6 +352,8 @@ func (s *Schema) MarshalJSON() ([]byte, error) { type Alias Schema + // Keep this literal in sync with every JSON-marshaled Schema field. Copying + // the whole Schema would also copy its lazy atomic caches after they are used. aliasCopy := Alias{ID: s.ID, IdentifierFieldIDs: ids} return json.Marshal(struct { diff --git a/schema_test.go b/schema_test.go index 5661453e6..f0773a188 100644 --- a/schema_test.go +++ b/schema_test.go @@ -22,6 +22,7 @@ import ( "fmt" "os" "path/filepath" + "reflect" "runtime" "strings" "sync" @@ -679,6 +680,42 @@ func TestSerializeSchema(t *testing.T) { }`, string(data)) } +func TestMarshalSchemaIncludesExportedFields(t *testing.T) { + schema := iceberg.NewSchemaWithIdentifiers(17, []int{1}, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64, Required: true}, + ) + data, err := json.Marshal(schema) + require.NoError(t, err) + var marshaled map[string]json.RawMessage + require.NoError(t, json.Unmarshal(data, &marshaled)) + + require.Contains(t, marshaled, "type") + require.Contains(t, marshaled, "fields") + wantKeys := 2 + value := reflect.ValueOf(schema).Elem() + for i := range value.NumField() { + field := value.Type().Field(i) + tag, tagged := field.Tag.Lookup("json") + name, _, _ := strings.Cut(tag, ",") + if !field.IsExported() || !tagged || name == "-" { + continue + } + if name == "" { + name = field.Name + } + wantKeys++ + + // A nonzero fixture catches omitted assignments even when the alias + // still emits the field's key with a zero value. Extend it for new fields. + require.False(t, value.Field(i).IsZero(), "populate Schema.%s in the fixture", field.Name) + want, err := json.Marshal(value.Field(i).Interface()) + require.NoError(t, err) + require.Contains(t, marshaled, name) + assert.JSONEq(t, string(want), string(marshaled[name]), "Schema.%s", field.Name) + } + assert.Len(t, marshaled, wantKeys) +} + func TestUnmarshalSchema(t *testing.T) { var schema iceberg.Schema require.NoError(t, json.Unmarshal([]byte(`{ @@ -2302,6 +2339,8 @@ func TestVisitGeoSchemaWithSchemaVisitorPerPrimitiveType(t *testing.T) { assert.Equal(t, 1, v.geographyCalls) } +// This test is intended to be run with -race; without the detector, the old +// MarshalJSON implementation also passes these assertions. func TestSchemaMarshalJSONConcurrentLazyLookups(t *testing.T) { for range 32 { schema := iceberg.NewSchemaWithIdentifiers(17, nil, diff --git a/table/scan_planning.go b/table/scan_planning.go index 24ce690ae..c41fca5ad 100644 --- a/table/scan_planning.go +++ b/table/scan_planning.go @@ -105,6 +105,9 @@ type ScanPlanningRequest struct { RowFilter iceberg.BooleanExpression MinRowsRequested *int64 StatsFields []string + // MaxConcurrency is the maximum number of concurrent operations a planner + // may use for this scan. Zero means use the planner's default. + MaxConcurrency int // CaseSensitive must carry the Scan's value (which defaults to true), not // Go's false zero value, or the wire request would flip the spec default. // Nil means use the scan default. diff --git a/table/scan_planning_test.go b/table/scan_planning_test.go index b9489bb65..1363abdb0 100644 --- a/table/scan_planning_test.go +++ b/table/scan_planning_test.go @@ -399,12 +399,14 @@ func TestScanPlanningRemoteResolvesDefaultProjectionAndSchema(t *testing.T) { planner := &fakeScanPlanner{supports: true} scan := (&Table{metadata: metadata}).Scan( WithScanPlanningMode(ScanPlanningRemote), + WithMaxConcurrency(3), ) scan.planner = planner _, err = scan.PlanFiles(context.Background()) require.NoError(t, err) assert.Equal(t, []string{"id"}, planner.receivedRequest.SelectedFields) + assert.Equal(t, 3, planner.receivedRequest.MaxConcurrency) assert.True(t, planner.receivedRequest.Schema.Equals(metadata.CurrentSchema())) require.NotNil(t, planner.receivedRequest.UseSnapshotSchema) assert.False(t, *planner.receivedRequest.UseSnapshotSchema) diff --git a/table/scanner.go b/table/scanner.go index c30c76ed8..40b596224 100644 --- a/table/scanner.go +++ b/table/scanner.go @@ -1689,6 +1689,7 @@ func (scan *Scan) planFilesRemote(ctx context.Context) ([]FileScanTask, error) { SelectedFields: selectedFields, RowFilter: scan.rowFilter, MinRowsRequested: minRowsRequested, + MaxConcurrency: scan.concurrency, CaseSensitive: &caseSensitive, UseSnapshotSchema: &useSnapshotSchema, })