This is an automated email from the ASF dual-hosted git repository.
zeroshade pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-go.git
The following commit(s) were added to refs/heads/main by this push:
new 4818a62c feat(arrow/ipc): add Reader.ContinueFrom for streams split
across blobs (#1346)
4818a62c is described below
commit 4818a62cdce5fca5939401636edaf40e1bea8f85
Author: Matt Topol <[email protected]>
AuthorDate: Mon Oct 5 16:52:54 2026 -0400
feat(arrow/ipc): add Reader.ContinueFrom for streams split across blobs
(#1346)
### Rationale for this change
Fixes #1325. Some callers decode one logical IPC stream that arrives as
a sequence of blobs, where only the first blob carries the schema.
Through v18.7.0 they could reuse one `ipc.Reader` by swapping the
underlying `io.Reader`. Since v18.8.0, `done` stays set once a blob
reaches EOF and can't be reset, so every blob after the first decodes as
zero rows with a nil error.
### What changes are included in this PR?
- `(*Reader).ContinueFrom(io.Reader) error` and
`(*Reader).ContinueFromMessageReader(MessageReader) error` resume a
reader on a new source. The schema and dictionaries already read are
kept. Both methods also:
- discard the current record batch, release the previous message reader,
and clear `done`;
- refuse to continue a reader that has already failed, so an earlier
error can't be hidden;
- refuse to continue a reader that hasn't read its schema yet.
- As requested in
https://github.com/apache/arrow-go/issues/1325#issuecomment-5878340460,
a schema message where a record batch is expected now fails with
`arrow/ipc: unexpected schema message: …`. That marks the reader done,
and the error is sticky. Before, this gave the generic
invalid-message-type error. The other invalid-type errors now also have
their closing parenthesis back.
### Are these changes tested?
Yes. `TestReaderContinueFrom` splits a single writer's output, with an
int64 column and a dictionary column, into three blobs and checks:
- all batches round-trip through both `Next` and `Read`, and
dictionaries carry across blobs;
- a continuation that starts with a schema message fails loudly and
stays failed;
- `ContinueFrom` on a reader whose schema hasn't been read yet returns
an error and leaves the reader usable.
The test uses a checked allocator. Removing the `done` reset makes it
fail.
### Are there any user-facing changes?
Yes: two new methods on `ipc.Reader`, and a clearer error when a schema
message appears mid-stream.
### Message size limits (follow-up commit)
`NewReader` built its `messageReader` without the configured limits.
Those limits were added in #955, which never touched `reader.go`, so a
stream `Reader` ignored `WithMetadataSizeLimit`/`WithBodySizeLimit` and
enforced no limit at all. `NewReader` now creates its message reader
with the options it was given, and `ContinueFrom` reuses the same
limits.
**Behavior change:** stream readers created with `NewReader` now enforce
the documented default limits, 64 MiB for metadata and 256 MiB for
bodies. A stream with larger messages needs
`WithBodySizeLimit`/`WithMetadataSizeLimit`; passing a value ≤ 0
disables the limit. `TestReaderSizeLimits` covers both `NewReader` and
`ContinueFrom`.
---
arrow/ipc/reader.go | 73 ++++++++++++++++++++--
arrow/ipc/reader_test.go | 155 +++++++++++++++++++++++++++++++++++++++++++++++
2 files changed, 224 insertions(+), 4 deletions(-)
diff --git a/arrow/ipc/reader.go b/arrow/ipc/reader.go
index e15994e0..edd50fc8 100644
--- a/arrow/ipc/reader.go
+++ b/arrow/ipc/reader.go
@@ -52,6 +52,9 @@ type Reader struct {
expectedSchema *arrow.Schema
mem memory.Allocator
+ // size limits for the message readers this Reader creates itself
+ maxMetadataSize int64
+ maxBodySize int64
}
// NewReaderFromMessageReader allows constructing a new reader object with the
@@ -80,6 +83,8 @@ func NewReaderFromMessageReader(r MessageReader, opts
...Option) (reader *Reader
mem: cfg.alloc,
ensureNativeEndian: cfg.ensureNativeEndian,
expectedSchema: cfg.schema,
+ maxMetadataSize: cfg.maxMetadataSize,
+ maxBodySize: cfg.maxBodySize,
}
reader.refCount.Add(1)
@@ -105,8 +110,7 @@ func NewReader(r io.Reader, opts ...Option) (rr *Reader,
err error) {
}
}()
cfg := newConfig(opts...)
- mr := &messageReader{r: r, mem: cfg.alloc}
- mr.refCount.Add(1)
+ mr := NewMessageReader(r, opts...)
rr = &Reader{
r: mr,
refCount: atomic.Int64{},
@@ -115,6 +119,8 @@ func NewReader(r io.Reader, opts ...Option) (rr *Reader,
err error) {
mem: cfg.alloc,
ensureNativeEndian: cfg.ensureNativeEndian,
expectedSchema: cfg.schema,
+ maxMetadataSize: cfg.maxMetadataSize,
+ maxBodySize: cfg.maxBodySize,
}
rr.refCount.Add(1)
@@ -285,8 +291,17 @@ func (r *Reader) next() bool {
return false
}
- if got, want := msg.Type(), MessageRecordBatch; got != want {
- r.err = fmt.Errorf("arrow/ipc: invalid message type (got=%v,
want=%v", got, want)
+ switch msg.Type() {
+ case MessageRecordBatch:
+ case MessageSchema:
+ // a schema message is only valid at the very start of a
stream; a
+ // continuation source that begins with one is a separate
stream, and
+ // skipping it would silently drop or misread its batches
+ r.err = errors.New("arrow/ipc: unexpected schema message: the
stream's schema has already been read")
+ r.done = true
+ return false
+ default:
+ r.err = fmt.Errorf("arrow/ipc: invalid message type (got=%v,
want=%v)", msg.Type(), MessageRecordBatch)
return false
}
@@ -328,4 +343,54 @@ func (r *Reader) Read() (arrow.RecordBatch, error) {
return r.rec, nil
}
+// ContinueFrom resumes reading from src as a continuation of the same logical
+// stream, keeping the schema and dictionaries already read. It is
+// [Reader.ContinueFromMessageReader] for a plain byte stream, read with the
+// allocator and message size limits the Reader was created with.
+func (r *Reader) ContinueFrom(src io.Reader) error {
+ mr := NewMessageReader(src, WithAllocator(r.mem),
+ WithMetadataSizeLimit(r.maxMetadataSize),
WithBodySizeLimit(r.maxBodySize))
+ if err := r.ContinueFromMessageReader(mr); err != nil {
+ mr.Release()
+ return err
+ }
+ return nil
+}
+
+// ContinueFromMessageReader resumes reading from mr as a continuation of the
+// same logical stream, keeping the schema and dictionaries already read. This
+// supports a stream that arrives as a sequence of independent blobs where only
+// the first one carries the schema.
+//
+// mr must start with record batch or dictionary batch messages. If it starts
+// with a schema message, the next call to Next or Read fails with an error
+// rather than reading the batches that follow.
+//
+// The schema must already have been read, and the reader must not have
+// failed: continuing after an error could hide corruption, so the error is
+// returned instead. Any messages left unread in the previous source are
+// discarded, as is the current record batch. On success the reader takes
+// ownership of mr and releases it, as NewReaderFromMessageReader does; on
+// error mr is left to the caller. ContinueFromMessageReader must not be
+// called concurrently with Next or Read.
+func (r *Reader) ContinueFromMessageReader(mr MessageReader) error {
+ if r.err != nil {
+ return fmt.Errorf("arrow/ipc: cannot continue a reader that
failed: %w", r.err)
+ }
+ if r.schema == nil {
+ return errors.New("arrow/ipc: cannot continue a reader that has
not read its schema")
+ }
+
+ if r.rec != nil {
+ r.rec.Release()
+ r.rec = nil
+ }
+ if r.r != nil {
+ r.r.Release()
+ }
+ r.r = mr
+ r.done = false
+ return nil
+}
+
var _ array.RecordReader = (*Reader)(nil)
diff --git a/arrow/ipc/reader_test.go b/arrow/ipc/reader_test.go
index 29db5985..0f3dcfe3 100644
--- a/arrow/ipc/reader_test.go
+++ b/arrow/ipc/reader_test.go
@@ -257,3 +257,158 @@ func BenchmarkIPC(b *testing.B) {
})
}
}
+
+// writeSplitStream writes recs with a single Writer and returns the bytes it
+// produced for each record as a separate blob: the first blob carries the
+// schema and dictionaries, and the last one also carries the end-of-stream
+// marker.
+func writeSplitStream(t *testing.T, mem memory.Allocator, schema
*arrow.Schema, recs []arrow.RecordBatch) [][]byte {
+ t.Helper()
+ var buf bytes.Buffer
+ w := NewWriter(&buf, WithSchema(schema), WithAllocator(mem))
+ blobs := make([][]byte, 0, len(recs))
+ for _, rec := range recs {
+ require.NoError(t, w.Write(rec))
+ blobs = append(blobs, bytes.Clone(buf.Bytes()))
+ buf.Reset()
+ }
+ require.NoError(t, w.Close())
+ blobs[len(blobs)-1] = append(blobs[len(blobs)-1], buf.Bytes()...)
+ return blobs
+}
+
+func TestReaderContinueFrom(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+ defer mem.AssertSize(t, 0)
+
+ schema := arrow.NewSchema([]arrow.Field{
+ {Name: "i", Type: arrow.PrimitiveTypes.Int64},
+ {Name: "d", Type: &arrow.DictionaryType{IndexType:
arrow.PrimitiveTypes.Int32, ValueType: arrow.BinaryTypes.String}},
+ }, nil)
+
+ recs := make([]arrow.RecordBatch, 3)
+ for n := range recs {
+ b := array.NewRecordBuilder(mem, schema)
+ b.Field(0).(*array.Int64Builder).AppendValues([]int64{int64(2 *
n), int64(2*n + 1)}, nil)
+ require.NoError(t,
b.Field(1).(*array.BinaryDictionaryBuilder).AppendString("foo"))
+ require.NoError(t,
b.Field(1).(*array.BinaryDictionaryBuilder).AppendString("bar"))
+ recs[n] = b.NewRecordBatch()
+ b.Release()
+ defer recs[n].Release()
+ }
+ blobs := writeSplitStream(t, mem, schema, recs)
+
+ readAll := func(t *testing.T, useRead bool) {
+ rdr, err := NewReader(bytes.NewReader(blobs[0]),
WithAllocator(mem))
+ require.NoError(t, err)
+ defer rdr.Release()
+
+ var got []arrow.RecordBatch
+ for i, blob := range blobs {
+ if i > 0 {
+ require.NoError(t,
rdr.ContinueFrom(bytes.NewReader(blob)))
+ }
+ for {
+ var rec arrow.RecordBatch
+ if useRead {
+ rec, err = rdr.Read()
+ if err == io.EOF {
+ break
+ }
+ require.NoError(t, err)
+ } else {
+ if !rdr.Next() {
+ require.NoError(t, rdr.Err())
+ break
+ }
+ rec = rdr.RecordBatch()
+ }
+ rec.Retain()
+ got = append(got, rec)
+ }
+ }
+
+ require.Len(t, got, len(recs))
+ for i, rec := range got {
+ assert.Truef(t, array.RecordEqual(recs[i], rec), "batch
%d: got %v, want %v", i, rec, recs[i])
+ rec.Release()
+ }
+ }
+
+ t.Run("Next", func(t *testing.T) { readAll(t, false) })
+ t.Run("Read", func(t *testing.T) { readAll(t, true) })
+
+ t.Run("SchemaMessageInContinuation", func(t *testing.T) {
+ rdr, err := NewReader(bytes.NewReader(blobs[0]),
WithAllocator(mem))
+ require.NoError(t, err)
+ defer rdr.Release()
+ for rdr.Next() {
+ }
+ require.NoError(t, rdr.Err())
+
+ // a whole new stream, schema message included, is not a
continuation
+ require.NoError(t,
rdr.ContinueFrom(bytes.NewReader(bytes.Join(blobs, nil))))
+ assert.False(t, rdr.Next())
+ assert.ErrorContains(t, rdr.Err(), "unexpected schema message")
+
+ _, err = rdr.Read()
+ assert.ErrorContains(t, err, "unexpected schema message")
+
+ // the failure is sticky: continuing again must not hide it
+ assert.ErrorContains(t,
rdr.ContinueFrom(bytes.NewReader(blobs[1])), "unexpected schema message")
+ })
+
+ t.Run("SchemaNotRead", func(t *testing.T) {
+ rdr, err := NewReader(bytes.NewReader(blobs[0]),
WithAllocator(mem), WithDelayReadSchema(true))
+ require.NoError(t, err)
+ defer rdr.Release()
+
+ assert.ErrorContains(t,
rdr.ContinueFrom(bytes.NewReader(blobs[1])), "has not read its schema")
+ // the reader is untouched and still reads its original source
+ assert.True(t, rdr.Next())
+ assert.EqualValues(t, 2, rdr.RecordBatch().NumRows())
+ })
+}
+
+func TestReaderSizeLimits(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+ defer mem.AssertSize(t, 0)
+
+ schema := arrow.NewSchema([]arrow.Field{{Name: "i", Type:
arrow.PrimitiveTypes.Int64}}, nil)
+ // a small batch whose body fits under the limit, then one whose body
+ // (1000 int64 values) does not
+ recs := make([]arrow.RecordBatch, 2)
+ for n, rows := range []int{2, 1000} {
+ b := array.NewRecordBuilder(mem, schema)
+ b.Field(0).(*array.Int64Builder).AppendValues(make([]int64,
rows), nil)
+ recs[n] = b.NewRecordBatch()
+ b.Release()
+ defer recs[n].Release()
+ }
+ blobs := writeSplitStream(t, mem, schema, recs)
+ const bodyLimit = 1024
+
+ t.Run("NewReader", func(t *testing.T) {
+ rdr, err := NewReader(bytes.NewReader(bytes.Join(blobs, nil)),
WithAllocator(mem), WithBodySizeLimit(bodyLimit))
+ require.NoError(t, err)
+ defer rdr.Release()
+
+ require.True(t, rdr.Next())
+ assert.False(t, rdr.Next())
+ assert.ErrorContains(t, rdr.Err(), "exceeds limit 1024")
+ })
+
+ t.Run("ContinueFrom", func(t *testing.T) {
+ rdr, err := NewReader(bytes.NewReader(blobs[0]),
WithAllocator(mem), WithBodySizeLimit(bodyLimit))
+ require.NoError(t, err)
+ defer rdr.Release()
+
+ require.True(t, rdr.Next())
+ require.False(t, rdr.Next())
+ require.NoError(t, rdr.Err())
+
+ require.NoError(t, rdr.ContinueFrom(bytes.NewReader(blobs[1])))
+ assert.False(t, rdr.Next())
+ assert.ErrorContains(t, rdr.Err(), "exceeds limit 1024")
+ })
+}
