From 45a1ab2d55c53de66cba00017e051159e1b64023 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Tue, 15 Sep 2026 06:27:49 -0700 Subject: [PATCH] fix(go): serialize record batch reader access --- bindings/go/table_read.go | 13 +- .../go/tests/table_read_concurrency_test.go | 142 ++++++++++++++++++ 2 files changed, 153 insertions(+), 2 deletions(-) create mode 100644 bindings/go/tests/table_read_concurrency_test.go diff --git a/bindings/go/table_read.go b/bindings/go/table_read.go index 111eaec7b..883b1696f 100644 --- a/bindings/go/table_read.go +++ b/bindings/go/table_read.go @@ -111,6 +111,7 @@ func (tr *TableRead) NewRecordBatchReader(splits []DataSplit) (*RecordBatchReade type RecordBatchReader struct { ctx context.Context lib *libRef + mu sync.Mutex readers []*paimonRecordBatchReader current int closeOnce sync.Once @@ -119,8 +120,12 @@ type RecordBatchReader struct { // NextRecord returns the next Arrow record, or io.EOF when iteration is // complete. The underlying C batch is imported via the Arrow C Data Interface // and released automatically — the caller only needs to call Release on the -// returned arrow.Record when done. +// returned arrow.Record when done. Calls on one reader are serialized; use +// separate readers for parallel reads. func (r *RecordBatchReader) NextRecord() (arrow.Record, error) { + r.mu.Lock() + defer r.mu.Unlock() + if r.readers == nil { return nil, ErrClosed } @@ -158,8 +163,12 @@ func (r *RecordBatchReader) next() (*arrowBatch, error) { return nil, io.EOF } -// Close releases the underlying C record batch readers. Safe to call multiple times. +// Close releases the underlying C record batch readers. Safe to call multiple times +// and concurrently with NextRecord. func (r *RecordBatchReader) Close() { + r.mu.Lock() + defer r.mu.Unlock() + r.closeOnce.Do(func() { freeFn := ffiRecordBatchReaderFree.symbol(r.ctx) for _, rd := range r.readers { diff --git a/bindings/go/tests/table_read_concurrency_test.go b/bindings/go/tests/table_read_concurrency_test.go new file mode 100644 index 000000000..a87099b9b --- /dev/null +++ b/bindings/go/tests/table_read_concurrency_test.go @@ -0,0 +1,142 @@ +/* + * 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 paimon_test + +import ( + "errors" + "io" + "path/filepath" + "sync" + "testing" + + "github.com/apache/arrow-go/v18/arrow" + paimon "github.com/apache/paimon-rust/bindings/go" +) + +type nextResult struct { + record arrow.Record + err error +} + +func newRecordBatchReader(t *testing.T) *paimon.RecordBatchReader { + t.Helper() + + warehouse := t.TempDir() + source := filepath.Join("testdata", "map_blob_table") + if err := copyDirectory(source, filepath.Join(warehouse, "default.db", "map_blob_table")); err != nil { + t.Fatal(err) + } + table := openTableAt(t, warehouse, "map_blob_table") + builder, err := table.NewReadBuilderWithOptions(map[string]string{"blob-as-descriptor": "true"}) + if err != nil { + t.Fatal(err) + } + t.Cleanup(builder.Close) + scan, err := builder.NewScan() + if err != nil { + t.Fatal(err) + } + t.Cleanup(scan.Close) + plan, err := scan.Plan() + if err != nil { + t.Fatal(err) + } + t.Cleanup(plan.Close) + read, err := builder.NewRead() + if err != nil { + t.Fatal(err) + } + t.Cleanup(read.Close) + reader, err := read.NewRecordBatchReader(plan.Splits()) + if err != nil { + t.Fatal(err) + } + t.Cleanup(reader.Close) + return reader +} + +func TestRecordBatchReaderConcurrentNext(t *testing.T) { + reader := newRecordBatchReader(t) + const workers = 16 + + start := make(chan struct{}) + results := make(chan nextResult, workers) + for range workers { + go func() { + <-start + record, err := reader.NextRecord() + results <- nextResult{record, err} + }() + } + close(start) + + rows := int64(0) + for range workers { + result := <-results + if result.record != nil { + rows += result.record.NumRows() + result.record.Release() + } + if result.err != nil && !errors.Is(result.err, io.EOF) { + t.Fatalf("NextRecord returned %v", result.err) + } + } + if rows != 3 { + t.Fatalf("read %d rows, want 3", rows) + } +} + +func TestRecordBatchReaderConcurrentNextAndClose(t *testing.T) { + reader := newRecordBatchReader(t) + const workers = 16 + + start := make(chan struct{}) + results := make(chan nextResult, workers) + var wg sync.WaitGroup + wg.Add(workers + 1) + for range workers { + go func() { + defer wg.Done() + <-start + record, err := reader.NextRecord() + results <- nextResult{record, err} + }() + } + go func() { + defer wg.Done() + <-start + reader.Close() + }() + close(start) + wg.Wait() + close(results) + + for result := range results { + if result.record != nil { + result.record.Release() + } + if result.err != nil && !errors.Is(result.err, io.EOF) && !errors.Is(result.err, paimon.ErrClosed) { + t.Fatalf("NextRecord returned %v", result.err) + } + } + if _, err := reader.NextRecord(); !errors.Is(err, paimon.ErrClosed) { + t.Fatalf("NextRecord after Close returned %v", err) + } +}