diff --git a/internal/backend/kernel/batchiterator_test.go b/internal/backend/kernel/batchiterator_test.go index 4eb22c8f..3bbffe70 100644 --- a/internal/backend/kernel/batchiterator_test.go +++ b/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") } } diff --git a/internal/backend/kernel/rows.go b/internal/backend/kernel/rows.go index d55f77ea..5fab0ee1 100644 --- a/internal/backend/kernel/rows.go +++ b/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 +}