feat(kernel): expose processed results as Arrow IPC streams

internal/backend/kernel/batchiterator_test.go

@@ -3,27 +3,85 @@
 package kernel
 
 import (
+	"bytes"
 	"context"
 	"errors"
 	"io"
 	"testing"
 
 	"github.com/apache/arrow/go/v12/arrow"
-	dbsqlerr "github.com/databricks/databricks-sql-go/errors"
+	"github.com/apache/arrow/go/v12/arrow/array"
+	"github.com/apache/arrow/go/v12/arrow/ipc"
+	"github.com/apache/arrow/go/v12/arrow/memory"
 	dbsqlrows "github.com/databricks/databricks-sql-go/internal/rows"
 	dbrows "github.com/databricks/databricks-sql-go/rows"
 )
 
-// TestGetArrowIPCStreamsRejected pins the discoverable sentinel so callers can
-// branch on it: the kernel C ABI exports Arrow C Data, not IPC bytes.
-func TestGetArrowIPCStreamsRejected(t *testing.T) {
-	r := &kernelRows{}
-	it, err := r.GetArrowIPCStreams(context.Background())
-	if it != nil {
-		t.Errorf("GetArrowIPCStreams = %v iterator, want nil", it)
-	}
-	if !errors.Is(err, dbsqlerr.ErrNotSupportedByKernel) {
-		t.Errorf("GetArrowIPCStreams err = %v, want wrapping ErrNotSupportedByKernel", err)
+func TestGetArrowIPCStreamsRequiresSchema(t *testing.T) {
+	it, err := (&kernelRows{}).GetArrowIPCStreams(context.Background())
+	if it != nil || err == nil {
+		t.Errorf("GetArrowIPCStreams without schema = (%v, %v), want (nil, error)", it, err)
+	}
+}
+
+func TestKernelIPCStreamIterator(t *testing.T) {
+	schema := arrow.NewSchema([]arrow.Field{{Name: "value", Type: arrow.PrimitiveTypes.Int64}}, nil)
+	builder := array.NewInt64Builder(memory.DefaultAllocator)
+	builder.AppendValues([]int64{1, 2}, nil)
+	values := builder.NewArray()
+	builder.Release()
+	record := array.NewRecord(schema, []arrow.Array{values}, 2)
+	values.Release()
+
+	rows := &kernelRows{schema: schema, eof: true}
+	batches := &kernelBatchIterator{r: rows, pending: record}
+	schemaBytes, err := arrowSchemaIPCBytes(schema)
+	if err != nil {
+		t.Fatalf("arrowSchemaIPCBytes: %v", err)
+	}
+	it := &kernelIPCStreamIterator{batches: batches, schemaBytes: schemaBytes}
+
+	if !it.HasNext() {
+		t.Fatal("HasNext before Next = false, want true")
+	}
+	stream, err := it.Next()
+	if err != nil {
+		t.Fatalf("Next: %v", err)
+	}
+	reader, err := ipc.NewReader(stream)
+	if err != nil {
+		t.Fatalf("ipc.NewReader: %v", err)
+	}
+	defer reader.Release()
+	if !reader.Next() {
+		t.Fatalf("IPC reader did not yield a record: %v", reader.Err())
+	}
+	got := reader.Record()
+	if got.NumRows() != 2 || got.NumCols() != 1 {
+		t.Errorf("decoded record shape = %dx%d, want 2x1", got.NumRows(), got.NumCols())
+	}
+	gotValues := got.Column(0).(*array.Int64)
+	if gotValues.Value(0) != 1 || gotValues.Value(1) != 2 {
+		t.Errorf("decoded values = [%d, %d], want [1, 2]", gotValues.Value(0), gotValues.Value(1))
+	}
+	if it.HasNext() {
+		t.Error("HasNext after the only batch = true, want false")
+	}
+	if next, err := it.Next(); next != nil || err != io.EOF {
+		t.Errorf("Next after drain = (%v, %v), want (nil, io.EOF)", next, err)
+	}
+
+	returnedSchema, err := it.SchemaBytes()
+	if err != nil {
+		t.Fatalf("SchemaBytes: %v", err)
+	}
+	if !bytes.Equal(returnedSchema, schemaBytes) {
+		t.Error("SchemaBytes did not return the encoded result schema")
+	}
+	returnedSchema[0] ^= 0xff
+	secondSchema, _ := it.SchemaBytes()
+	if bytes.Equal(returnedSchema, secondSchema) {
+		t.Error("SchemaBytes returned mutable internal storage")
 	}
 }
 

internal/backend/kernel/rows.go

@@ -13,6 +13,7 @@ struct ArrowArray;
 import "C"
 
 import (
+	"bytes"
 	"context"
 	"database/sql/driver"
 	"fmt"
@@ -23,7 +24,7 @@ import (
 
 	"github.com/apache/arrow/go/v12/arrow"
 	"github.com/apache/arrow/go/v12/arrow/cdata"
-	dbsqlerr "github.com/databricks/databricks-sql-go/errors"
+	"github.com/apache/arrow/go/v12/arrow/ipc"
 	"github.com/databricks/databricks-sql-go/internal/arrowscan"
 	context2 "github.com/databricks/databricks-sql-go/internal/compat/context"
 	dbsqlrows "github.com/databricks/databricks-sql-go/internal/rows"
@@ -40,8 +41,7 @@ var (
 	_ driver.RowsColumnTypeDatabaseTypeName = (*kernelRows)(nil)
 	_ driver.RowsColumnTypeNullable         = (*kernelRows)(nil)
 	_ driver.RowsColumnTypeLength           = (*kernelRows)(nil)
-	// Public batch API, matching the Thrift path (rows.rows). GetArrowBatches is
-	// real; GetArrowIPCStreams rejects (kernel C ABI exports C Data, not IPC bytes).
+	// Public batch and IPC APIs, matching the Thrift path (rows.rows).
 	_ dbrows.Rows = (*kernelRows)(nil)
 )
 
@@ -312,12 +312,34 @@ func (r *kernelRows) GetArrowBatches(context.Context) (dbrows.ArrowBatchIterator
 	return &kernelBatchIterator{r: r}, nil
 }
 
-// GetArrowIPCStreams is rejected on the kernel path: the C ABI exports Arrow via
-// the C Data Interface (next_batch), not IPC bytes. Use GetArrowBatches, or the
-// Thrift backend for IPC streams.
+// GetArrowIPCStreams serializes the kernel's post-processed batches to IPC. This
+// provides a version-neutral boundary for consumers such as ADBC that use a
+// different Arrow Go major version.
 func (r *kernelRows) GetArrowIPCStreams(context.Context) (dbrows.ArrowIPCStreamIterator, error) {
-	return nil, fmt.Errorf("databricks: GetArrowIPCStreams is %w (kernel exports Arrow C Data, not IPC bytes); use GetArrowBatches",
-		dbsqlerr.ErrNotSupportedByKernel)
+	batches := &kernelBatchIterator{r: r}
+	schemaBytes, err := arrowSchemaIPCBytes(r.schema)
+	if err != nil {
+		return nil, err
+	}
+	return &kernelIPCStreamIterator{batches: batches, schemaBytes: schemaBytes}, nil
+}
+
+// arrowSchemaIPCBytes matches the Thrift path's SchemaBytes contract: an IPC
+// stream containing the schema message without the end-of-stream marker.
+func arrowSchemaIPCBytes(schema *arrow.Schema) ([]byte, error) {
+	if schema == nil {
+		return nil, fmt.Errorf("kernel: no schema available")
+	}
+	var output bytes.Buffer
+	writer := ipc.NewWriter(&output, ipc.WithSchema(schema))
+	if err := writer.Close(); err != nil {
+		return nil, fmt.Errorf("kernel: encode IPC schema: %w", err)
+	}
+	encoded := output.Bytes()
+	if len(encoded) < 8 {
+		return nil, fmt.Errorf("kernel: encoded IPC schema is truncated")
+	}
+	return append([]byte(nil), encoded[:len(encoded)-8]...), nil
 }
 
 // kernelBatchIterator adapts the kernelRows pull loop to the public
@@ -402,3 +424,43 @@ func (it *kernelBatchIterator) Schema() (*arrow.Schema, error) {
 	}
 	return it.r.schema, nil
 }
+
+// kernelIPCStreamIterator adapts processed kernel batches to the existing IPC
+// iterator API. Each returned reader owns an independent, complete IPC stream.
+type kernelIPCStreamIterator struct {
+	batches     *kernelBatchIterator
+	schemaBytes []byte
+}
+
+var _ dbrows.ArrowIPCStreamIterator = (*kernelIPCStreamIterator)(nil)
+
+func (it *kernelIPCStreamIterator) Next() (io.Reader, error) {
+	record, err := it.batches.Next()
+	if err != nil {
+		return nil, err
+	}
+	defer record.Release()
+
+	var output bytes.Buffer
+	writer := ipc.NewWriter(&output, ipc.WithSchema(record.Schema()))
+	if err := writer.Write(record); err != nil {
+		_ = writer.Close()
+		return nil, fmt.Errorf("kernel: encode IPC batch: %w", err)
+	}
+	if err := writer.Close(); err != nil {
+		return nil, fmt.Errorf("kernel: finish IPC batch: %w", err)
+	}
+	return bytes.NewReader(output.Bytes()), nil
+}
+
+func (it *kernelIPCStreamIterator) HasNext() bool {
+	return it.batches.HasNext()
+}
+
+func (it *kernelIPCStreamIterator) Close() {
+	it.batches.Close()
+}
+
+func (it *kernelIPCStreamIterator) SchemaBytes() ([]byte, error) {
+	return append([]byte(nil), it.schemaBytes...), nil
+}