feat: Support setting IPC options in FlightSQL call options

arrow/flight/flightsql/client.go

@@ -247,6 +247,18 @@ func (c *Client) ExecuteSubstraitUpdate(ctx context.Context, plan SubstraitPlan,
 // The provided RecordReader will be retained for the duration of the call, but it is the caller's
 // responsibility to release the original reference.
 func (c *Client) ExecuteIngest(ctx context.Context, rdr array.RecordReader, reqOptions *ExecuteIngestOpts, opts ...grpc.CallOption) (int64, error) {
+	return c.executeIngest(ctx, rdr, reqOptions, nil, opts...)
+}
+
+// ExecuteIngestWithIPC is like ExecuteIngest, and also allows configuring IPC
+// stream writer options such as compression.
+// The provided RecordReader will be retained for the duration of the call, but it is the caller's
+// responsibility to release the original reference.
+func (c *Client) ExecuteIngestWithIPC(ctx context.Context, rdr array.RecordReader, reqOptions *ExecuteIngestOpts, ipcOpts []ipc.Option, opts ...grpc.CallOption) (int64, error) {
+	return c.executeIngest(ctx, rdr, reqOptions, ipcOpts, opts...)
+}
+
+func (c *Client) executeIngest(ctx context.Context, rdr array.RecordReader, reqOptions *ExecuteIngestOpts, ipcOpts []ipc.Option, opts ...grpc.CallOption) (int64, error) {
 	var (
 		err          error
 		desc         *flight.FlightDescriptor
@@ -274,7 +286,10 @@ func (c *Client) ExecuteIngest(ctx context.Context, rdr array.RecordReader, reqO
 		return 0, err
 	}
 
-	wr = flight.NewRecordWriter(stream, ipc.WithAllocator(c.Alloc), ipc.WithSchema(rdr.Schema()))
+	writerOpts := make([]ipc.Option, 0, 2+len(ipcOpts))
+	writerOpts = append(writerOpts, ipc.WithAllocator(c.Alloc), ipc.WithSchema(rdr.Schema()))
+	writerOpts = append(writerOpts, ipcOpts...)
+	wr = flight.NewRecordWriter(stream, writerOpts...)
 	defer wr.Close()
 
 	wr.SetFlightDescriptor(desc)

arrow/flight/flightsql/client_test.go

@@ -27,7 +27,10 @@ import (
 	"github.com/apache/arrow-go/v18/arrow/flight"
 	"github.com/apache/arrow-go/v18/arrow/flight/flightsql"
 	pb "github.com/apache/arrow-go/v18/arrow/flight/gen/flight"
+	"github.com/apache/arrow-go/v18/arrow/internal/flatbuf"
+	"github.com/apache/arrow-go/v18/arrow/ipc"
 	"github.com/apache/arrow-go/v18/arrow/memory"
+	flatbuffers "github.com/google/flatbuffers/go"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/suite"
 	"google.golang.org/grpc"
@@ -163,6 +166,7 @@ func getAction(cmd proto.Message) *flight.Action {
 func (s *FlightSqlClientSuite) SetupTest() {
 	s.mockClient = FlightServiceClientMock{}
 	s.sqlClient.Client = &s.mockClient
+	s.sqlClient.Alloc = memory.DefaultAllocator
 	s.callOpts = []grpc.CallOption{grpc.EmptyCallOption{}}
 }
 
@@ -649,6 +653,147 @@ func (s *FlightSqlClientSuite) TestExecuteUpdate() {
 	s.EqualValues(100, num)
 }
 
+func (s *FlightSqlClientSuite) TestExecuteIngestWithIPCOptions() {
+	schema := arrow.NewSchema([]arrow.Field{{Name: "id", Type: arrow.PrimitiveTypes.Int64}}, nil)
+	rec, _, err := array.RecordFromJSON(memory.DefaultAllocator, schema, strings.NewReader(`[{"id": 1}]`))
+	s.Require().NoError(err)
+	defer rec.Release()
+
+	rdr, err := array.NewRecordReader(schema, []arrow.RecordBatch{rec})
+	s.Require().NoError(err)
+	defer rdr.Release()
+
+	request := &flightsql.ExecuteIngestOpts{
+		Table:                  "target_table",
+		TableDefinitionOptions: &flightsql.TableDefinitionOptions{},
+	}
+
+	result := &pb.DoPutUpdateResult{RecordCount: 1}
+	resdata, _ := proto.Marshal(result)
+
+	mockedPut := &mockDoPutClient{}
+	defer mockedPut.AssertExpectations(s.T())
+
+	var sent []*flight.FlightData
+	mockedPut.On("Send", mock.AnythingOfType("*flight.FlightData")).Run(func(args mock.Arguments) {
+		sent = append(sent, proto.Clone(args.Get(0).(*flight.FlightData)).(*flight.FlightData))
+	}).Return(nil)
+	mockedPut.On("CloseSend").Return(nil)
+	mockedPut.On("Recv").Return(&pb.PutResult{AppMetadata: resdata}, nil).Once()
+	mockedPut.On("Recv").Return(&pb.PutResult{}, io.EOF).Once()
+
+	s.mockClient.On("DoPut", s.callOpts).Return(mockedPut, nil)
+
+	count, err := s.sqlClient.ExecuteIngestWithIPC(
+		context.Background(),
+		rdr,
+		request,
+		[]ipc.Option{ipc.WithLZ4()},
+		s.callOpts...,
+	)
+	s.Require().NoError(err)
+	s.EqualValues(1, count)
+
+	var rbCompression *flatbuf.BodyCompression
+	for _, fd := range sent {
+		if len(fd.DataHeader) == 0 {
+			continue
+		}
+
+		msg := flatbuf.GetRootAsMessage(fd.DataHeader, 0)
+		if msg.HeaderType() != flatbuf.MessageHeaderRecordBatch {
+			continue
+		}
+
+		var header flatbuffers.Table
+		if !msg.Header(&header) {
+			continue
+		}
+
+		var batch flatbuf.RecordBatch
+		batch.Init(header.Bytes, header.Pos)
+		rbCompression = batch.Compression(nil)
+		break
+	}
+
+	if s.NotNil(rbCompression, "record batch should include compression metadata") {
+		s.Equal(flatbuf.CompressionTypeLZ4_FRAME, rbCompression.Codec())
+	}
+}
+
+func (s *FlightSqlClientSuite) TestExecuteIngestWithSchemaOverrideOption() {
+	dataSchema := arrow.NewSchema([]arrow.Field{{Name: "id", Type: arrow.PrimitiveTypes.Int64}}, nil)
+	overrideSchema := arrow.NewSchema([]arrow.Field{{Name: "name", Type: arrow.BinaryTypes.String}}, nil)
+
+	rec, _, err := array.RecordFromJSON(memory.DefaultAllocator, dataSchema, strings.NewReader(`[{"id": 1}]`))
+	s.Require().NoError(err)
+	defer rec.Release()
+
+	rdr, err := array.NewRecordReader(dataSchema, []arrow.RecordBatch{rec})
+	s.Require().NoError(err)
+	defer rdr.Release()
+
+	request := &flightsql.ExecuteIngestOpts{
+		Table:                  "target_table",
+		TableDefinitionOptions: &flightsql.TableDefinitionOptions{},
+	}
+
+	mockedPut := &mockDoPutClient{}
+	defer mockedPut.AssertExpectations(s.T())
+	mockedPut.On("Send", mock.AnythingOfType("*flight.FlightData")).Return(nil)
+
+	s.mockClient.On("DoPut", s.callOpts).Return(mockedPut, nil)
+
+	_, err = s.sqlClient.ExecuteIngestWithIPC(
+		context.Background(),
+		rdr,
+		request,
+		[]ipc.Option{ipc.WithSchema(overrideSchema)},
+		s.callOpts...,
+	)
+	s.Error(err)
+	s.ErrorContains(err, "different schema")
+}
+
+func (s *FlightSqlClientSuite) TestExecuteIngestWithSliceOptions() {
+	schema := arrow.NewSchema([]arrow.Field{{Name: "id", Type: arrow.PrimitiveTypes.Int64}}, nil)
+	rec, _, err := array.RecordFromJSON(memory.DefaultAllocator, schema, strings.NewReader(`[{"id": 1}]`))
+	s.Require().NoError(err)
+	defer rec.Release()
+
+	rdr, err := array.NewRecordReader(schema, []arrow.RecordBatch{rec})
+	s.Require().NoError(err)
+	defer rdr.Release()
+
+	request := &flightsql.ExecuteIngestOpts{
+		Table:                  "target_table",
+		TableDefinitionOptions: &flightsql.TableDefinitionOptions{},
+	}
+
+	result := &pb.DoPutUpdateResult{RecordCount: 1}
+	resdata, _ := proto.Marshal(result)
+
+	mockedPut := &mockDoPutClient{}
+	defer mockedPut.AssertExpectations(s.T())
+	mockedPut.On("Send", mock.AnythingOfType("*flight.FlightData")).Return(nil)
+	mockedPut.On("CloseSend").Return(nil)
+	mockedPut.On("Recv").Return(&pb.PutResult{AppMetadata: resdata}, nil).Once()
+	mockedPut.On("Recv").Return(&pb.PutResult{}, io.EOF).Once()
+
+	s.mockClient.On("DoPut", s.callOpts).Return(mockedPut, nil)
+
+	ipcOpts := []ipc.Option{ipc.WithLZ4()}
+	count, err := s.sqlClient.ExecuteIngestWithIPC(
+		context.Background(),
+		rdr,
+		request,
+		ipcOpts,
+		s.callOpts...,
+	)
+	s.Require().NoError(err)
+	s.EqualValues(1, count)
+}
+
 func (s *FlightSqlClientSuite) TestGetSqlInfo() {
 	sqlInfo := []flightsql.SqlInfo{
 		flightsql.SqlInfoFlightSqlServerName,