This is an automated email from the ASF dual-hosted git repository.
zeroshade pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-adbc.git
The following commit(s) were added to refs/heads/main by this push:
new 0a844d867f feat(go/adbc/drivermgr): support ExecutePartitions and
ReadPartition (#4963)
0a844d867f is described below
commit 0a844d867f5591d90bc70e2eaba16cf062f93130
Author: Matt Topol <[email protected]>
AuthorDate: Thu Oct 29 16:23:52 2026 +0000
feat(go/adbc/drivermgr): support ExecutePartitions and ReadPartition (#4963)
The Go driver manager wrapper stubbed `ExecutePartitions` and
`ReadPartition` as NotImplemented, even when the loaded driver supports
partitioned execution. This PR implements both.
---
go/adbc/drivermgr/wrapper.go | 50 ++++++++++++++++++++++++++++++--
go/adbc/drivermgr/wrapper_sqlite_test.go | 23 +++++++++++++++
2 files changed, 71 insertions(+), 2 deletions(-)
diff --git a/go/adbc/drivermgr/wrapper.go b/go/adbc/drivermgr/wrapper.go
index fc95d2529b..58ded6350e 100644
--- a/go/adbc/drivermgr/wrapper.go
+++ b/go/adbc/drivermgr/wrapper.go
@@ -54,6 +54,11 @@ package drivermgr
// err->release = NULL;
// }
// }
+// static void releaseAdbcPartitions(struct AdbcPartitions* partitions) {
+// if (partitions->release != NULL) {
+// partitions->release(partitions);
+// }
+// }
// struct ArrowArray* allocArr() {
// struct ArrowArray* array = (struct ArrowArray*)malloc(sizeof(struct
ArrowArray));
// memset(array, 0, sizeof(struct ArrowArray));
@@ -68,6 +73,7 @@ package drivermgr
//
import "C"
import (
+ "bytes"
"context"
"strconv"
"sync"
@@ -685,7 +691,19 @@ func (c *cnxn) Close() error {
}
func (c *cnxn) ReadPartition(_ context.Context, serializedPartition []byte)
(array.RecordReader, error) {
- return nil, &adbc.Error{Code: adbc.StatusNotImplemented}
+ var (
+ out C.struct_ArrowArrayStream
+ err C.struct_AdbcError
+ ptr *C.uint8_t
+ )
+ if len(serializedPartition) > 0 {
+ ptr = (*C.uint8_t)(unsafe.Pointer(&serializedPartition[0]))
+ }
+
+ if code := adbc.Status(C.AdbcConnectionReadPartition(c.conn, ptr,
C.size_t(len(serializedPartition)), &out, &err)); code != adbc.StatusOK {
+ return nil, toAdbcError(code, &err)
+ }
+ return getRdr(&out)
}
func (c *cnxn) SetOption(key, value string) error {
@@ -883,5 +901,33 @@ func (s *stmt) GetParameterSchema() (*arrow.Schema, error)
{
}
func (s *stmt) ExecutePartitions(context.Context) (*arrow.Schema,
adbc.Partitions, int64, error) {
- return nil, adbc.Partitions{}, 0, &adbc.Error{Code:
adbc.StatusNotImplemented}
+ var (
+ schema C.struct_ArrowSchema
+ partitions C.struct_AdbcPartitions
+ affected C.int64_t
+ err C.struct_AdbcError
+ )
+ defer
cdata.ReleaseCArrowSchema((*cdata.CArrowSchema)(unsafe.Pointer(&schema)))
+ defer C.releaseAdbcPartitions(&partitions)
+
+ if code := adbc.Status(C.AdbcStatementExecutePartitions(s.st, &schema,
&partitions, &affected, &err)); code != adbc.StatusOK {
+ return nil, adbc.Partitions{}, 0, toAdbcError(code, &err)
+ }
+
+ sc, goerr := getSchema(&schema)
+ if goerr != nil {
+ return nil, adbc.Partitions{}, int64(affected), goerr
+ }
+
+ result := adbc.Partitions{
+ NumPartitions: uint64(partitions.num_partitions),
+ PartitionIDs: make([][]byte, int(partitions.num_partitions)),
+ }
+ ids := unsafe.Slice(partitions.partitions, partitions.num_partitions)
+ lengths := unsafe.Slice(partitions.partition_lengths,
partitions.num_partitions)
+ for i, ptr := range ids {
+ result.PartitionIDs[i] =
bytes.Clone(unsafe.Slice((*byte)(unsafe.Pointer(ptr)), lengths[i]))
+ }
+
+ return sc, result, int64(affected), nil
}
diff --git a/go/adbc/drivermgr/wrapper_sqlite_test.go
b/go/adbc/drivermgr/wrapper_sqlite_test.go
index 463e0ce52f..6e7573014d 100644
--- a/go/adbc/drivermgr/wrapper_sqlite_test.go
+++ b/go/adbc/drivermgr/wrapper_sqlite_test.go
@@ -601,6 +601,29 @@ func (dm *DriverMgrSuite)
TestSqlExecuteSchemaNotImplemented() {
dm.Equal(adbc.StatusNotImplemented, adbcErr.Code)
}
+func (dm *DriverMgrSuite) TestSqlExecutePartitionsUnsupported() {
+ st, err := dm.conn.NewStatement()
+ dm.Require().NoError(err)
+ defer validation.CheckedClose(dm.T(), st)
+ dm.Require().NoError(st.SetSqlQuery("SELECT 1"))
+
+ schema, partitions, affected, err := st.ExecutePartitions(dm.ctx)
+ var adbcErr adbc.Error
+ dm.Require().ErrorAs(err, &adbcErr)
+ dm.Equal(adbc.StatusNotImplemented, adbcErr.Code)
+ dm.Nil(schema)
+ dm.Equal(adbc.Partitions{}, partitions)
+ dm.Zero(affected)
+}
+
+func (dm *DriverMgrSuite) TestReadPartitionUnsupported() {
+ rdr, err := dm.conn.ReadPartition(dm.ctx, []byte("partition"))
+ var adbcErr adbc.Error
+ dm.Require().ErrorAs(err, &adbcErr)
+ dm.Equal(adbc.StatusNotImplemented, adbcErr.Code)
+ dm.Nil(rdr)
+}
+
func (dm *DriverMgrSuite) TestSqlExecuteInvalid() {
query := "INVALID"
st, err := dm.conn.NewStatement()
