feat: expose Payload.WritePayload to allow serializing into IPC format

arrow/ipc/file_writer.go

@@ -203,6 +203,12 @@ func (p *Payload) SerializeBody(w io.Writer) error {
 	return nil
 }
 
+// WritePayload serializes the payload in IPC format
+// into the provided writer.
+func (p *Payload) WritePayload(w io.Writer) (int, error) {
+	return writeIPCPayload(w, *p)
+}
+
 func (p *Payload) Release() {
 	if p.meta != nil {
 		p.meta.Release()

arrow/ipc/ipc.go

@@ -71,6 +71,7 @@ type config struct {
 	ensureNativeEndian bool
 	noAutoSchema       bool
 	emitDictDeltas     bool
+	skipEmittingSchema bool
 	minSpaceSavings    *float64
 }
 

arrow/ipc/writer.go

@@ -91,8 +91,9 @@ type Writer struct {
 
 	// map of the last written dictionaries by id
 	// so we can avoid writing the same dictionary over and over
-	lastWrittenDicts map[int64]arrow.Array
-	emitDictDeltas   bool
+	lastWrittenDicts   map[int64]arrow.Array
+	emitDictDeltas     bool
+	skipEmittingSchema bool
 }
 
 // NewWriterWithPayloadWriter constructs a writer with the provided payload writer
@@ -116,14 +117,15 @@ func NewWriterWithPayloadWriter(pw PayloadWriter, opts ...Option) *Writer {
 func NewWriter(w io.Writer, opts ...Option) *Writer {
 	cfg := newConfig(opts...)
 	return &Writer{
-		w:              w,
-		mem:            cfg.alloc,
-		pw:             &streamWriter{w: w},
-		schema:         cfg.schema,
-		codec:          cfg.codec,
-		emitDictDeltas: cfg.emitDictDeltas,
-		compressNP:     cfg.compressNP,
-		compressors:    make([]compressor, cfg.compressNP),
+		w:                  w,
+		mem:                cfg.alloc,
+		pw:                 &streamWriter{w: w},
+		schema:             cfg.schema,
+		codec:              cfg.codec,
+		emitDictDeltas:     cfg.emitDictDeltas,
+		skipEmittingSchema: cfg.skipEmittingSchema,
+		compressNP:         cfg.compressNP,
+		compressors:        make([]compressor, cfg.compressNP),
 	}
 }
 
@@ -276,6 +278,11 @@ func (w *Writer) start() error {
 	w.mapper.ImportSchema(w.schema)
 	w.lastWrittenDicts = make(map[int64]arrow.Array)
 
+	// skip writing schema payloads
+	if w.skipEmittingSchema {
+		return nil
+	}
+
 	// write out schema payloads
 	ps := payloadFromSchema(w.schema, w.mem, &w.mapper)
 	defer ps.Release()

arrow/ipc/writer_test.go

@@ -332,3 +332,28 @@ func TestGetPayloads(t *testing.T) {
 
 	assert.Truef(t, array.RecordEqual(rec, got), "expected: %s\ngot: %s", rec, got)
 }
+
+func TestWritePayload(t *testing.T) {
+	mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+	defer mem.AssertSize(t, 0)
+
+	bldr := array.NewRecordBuilder(mem, arrow.NewSchema([]arrow.Field{{Name: "col", Type: arrow.PrimitiveTypes.Int8}}, nil))
+	bldr.Field(0).(*array.Int8Builder).AppendValues([]int8{1, 2, 3, 4, 5}, nil)
+	rec := bldr.NewRecord()
+	defer rec.Release()
+
+	var buf bytes.Buffer
+	p, err := GetRecordBatchPayload(rec, WithAllocator(mem))
+	defer p.Release()
+	require.NoError(t, err)
+
+	_, err = p.WritePayload(&buf)
+	require.NoError(t, err)
+
+	r := NewMessageReader(&buf, WithAllocator(mem))
+	defer r.Release()
+
+	msg, err := r.Message()
+	require.NoError(t, err)
+	require.True(t, msg.Type() == MessageRecordBatch)
+}