(arrow) branch main updated: GH-43443: [Go] [IPC] Infer schema from first record if not specified (#43484)

[jo...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

joellubi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow.git


The following commit(s) were added to refs/heads/main by this push:
     new e9f6667291 GH-43443: [Go] [IPC] Infer schema from first record if not specified (#43484)
e9f6667291 is described below

commit e9f6667291b68ef5d82b4a193fdd84c8ef06a2cf
Author: Joel Lubinitsky <33...@users.noreply.github.com>
AuthorDate: Wed Jul 31 15:07:15 2024 -0500

    GH-43443: [Go] [IPC] Infer schema from first record if not specified (#43484)
    
    
    
    ### Rationale for this change
    
    Fixes: #43443
    
    Makes usage of the IPC writer and any writers that use it such the flight writer simpler.
    
    ### What changes are included in this PR?
    
    - Infer schema from first record if schema is not specified
    - IPC and Flight tests
    
    ### Are these changes tested?
    
    Yes
    
    ### Are there any user-facing changes?
    
    Any `ipc.Writer` that does not specify the optional argument `ipc.WithSchema` will no longer return an error as long as the incoming stream of records has a consistent schema.
    
    * GitHub Issue: #43443
    
    Authored-by: Joel Lubinitsky <jo...@gmail.com>
    Signed-off-by: Joel Lubinitsky <jo...@gmail.com>
---
 go/arrow/flight/flight_test.go | 35 +++++++++++++++++++++++++++++++++++
 go/arrow/ipc/writer.go         |  8 ++++++--
 go/arrow/ipc/writer_test.go    | 19 +++++++++++++++++++
 3 files changed, 60 insertions(+), 2 deletions(-)

diff --git a/go/arrow/flight/flight_test.go b/go/arrow/flight/flight_test.go
index fe896f39a2..a03d839e94 100755
--- a/go/arrow/flight/flight_test.go
+++ b/go/arrow/flight/flight_test.go
@@ -23,11 +23,13 @@ import (
 	"io"
 	"testing"
 
+	"github.com/apache/arrow/go/v18/arrow"
 	"github.com/apache/arrow/go/v18/arrow/array"
 	"github.com/apache/arrow/go/v18/arrow/flight"
 	"github.com/apache/arrow/go/v18/arrow/internal/arrdata"
 	"github.com/apache/arrow/go/v18/arrow/ipc"
 	"github.com/apache/arrow/go/v18/arrow/memory"
+	"github.com/stretchr/testify/require"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/credentials/insecure"
@@ -449,3 +451,36 @@ func TestReaderError(t *testing.T) {
 		t.Fatal("should have errored")
 	}
 }
+
+func TestWriterInferSchema(t *testing.T) {
+	recs, ok := arrdata.Records["primitives"]
+	require.True(t, ok)
+
+	fs := flightStreamWriter{}
+	w := flight.NewRecordWriter(&fs)
+
+	for _, rec := range recs {
+		require.NoError(t, w.Write(rec))
+	}
+
+	require.NoError(t, w.Close())
+}
+
+func TestWriterInconsistentSchema(t *testing.T) {
+	recs, ok := arrdata.Records["primitives"]
+	require.True(t, ok)
+
+	schema := arrow.NewSchema([]arrow.Field{{Name: "unknown", Type: arrow.PrimitiveTypes.Int8}}, nil)
+	fs := flightStreamWriter{}
+	w := flight.NewRecordWriter(&fs, ipc.WithSchema(schema))
+
+	require.ErrorContains(t, w.Write(recs[0]), "arrow/ipc: tried to write record batch with different schema")
+	require.NoError(t, w.Close())
+}
+
+type flightStreamWriter struct{}
+
+// Send implements flight.DataStreamWriter.
+func (f *flightStreamWriter) Send(data *flight.FlightData) error { return nil }
+
+var _ flight.DataStreamWriter = (*flightStreamWriter)(nil)
diff --git a/go/arrow/ipc/writer.go b/go/arrow/ipc/writer.go
index ca4f77d35e..02c67635bb 100644
--- a/go/arrow/ipc/writer.go
+++ b/go/arrow/ipc/writer.go
@@ -159,15 +159,19 @@ func (w *Writer) Write(rec arrow.Record) (err error) {
 		}
 	}()
 
+	incomingSchema := rec.Schema()
+
 	if !w.started {
+		if w.schema == nil {
+			w.schema = incomingSchema
+		}
 		err := w.start()
 		if err != nil {
 			return err
 		}
 	}
 
-	schema := rec.Schema()
-	if schema == nil || !schema.Equal(w.schema) {
+	if incomingSchema == nil || !incomingSchema.Equal(w.schema) {
 		return errInconsistentSchema
 	}
 
diff --git a/go/arrow/ipc/writer_test.go b/go/arrow/ipc/writer_test.go
index e5683243e4..60d811e68e 100644
--- a/go/arrow/ipc/writer_test.go
+++ b/go/arrow/ipc/writer_test.go
@@ -235,3 +235,22 @@ func TestWriteWithCompressionAndMinSavings(t *testing.T) {
 		}
 	}
 }
+
+func TestWriterInferSchema(t *testing.T) {
+	bldr := array.NewRecordBuilder(memory.DefaultAllocator, 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
+	w := NewWriter(&buf)
+
+	require.NoError(t, w.Write(rec))
+	require.NoError(t, w.Close())
+
+	r, err := NewReader(&buf)
+	require.NoError(t, err)
+	defer r.Release()
+
+	require.True(t, r.Schema().Equal(rec.Schema()))
+}