From ab7567d0388ea455b4366946c286de46702fdb5d Mon Sep 17 00:00:00 2001 From: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> Date: Wed, 30 Sep 2026 06:18:27 +0000 Subject: [PATCH] feat: support geospatial results on the kernel path Signed-off-by: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> --- CHANGELOG.md | 3 + CONNECTION_PARAMETERS.md | 1 + KERNEL_REV | 2 +- README.md | 1 + connector.go | 25 ++++++ connector_test.go | 13 +++ internal/arrowscan/arrowscan.go | 65 ++++++++++++-- internal/arrowscan/arrowscan_test.go | 84 +++++++++++++++++++ internal/arrowscan/coltype.go | 49 +++++++++++ internal/arrowscan/coltype_test.go | 37 ++++++++ internal/backend/kernel/backend.go | 36 ++++++++ internal/backend/kernel/config.go | 5 ++ .../kernel/include/databricks_kernel.h | 71 +++++++++++++--- internal/backend/kernel/kernel_test.go | 19 +++++ internal/backend/kernel/rows.go | 10 ++- internal/config/config.go | 27 +++++- internal/config/config_test.go | 33 ++++++++ kernel_config.go | 4 + kernel_config_test.go | 21 +++++ kernel_experimental_test.go | 12 +++ 20 files changed, 498 insertions(+), 20 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index cbf0dfc1..599ea4f3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,8 @@ # Release History +## Unreleased +- Add kernel-backed GEOMETRY / GEOGRAPHY results in EWKT string or binary WKB mode, selected with `WithKernelGeospatialAsString` or the `geospatialAsString` DSN parameter. Logical column metadata is preserved, and Arrow batch access retains the canonical `struct` value in binary mode. + ## v1.16.0 (2026-09-24) - Upgrade the kernel bindings to v1.1.0; the kernel dependency is now stable. - Expand the kernel backend with client query timeouts, configurable idle HTTP connections, and per-statement query tags. diff --git a/CONNECTION_PARAMETERS.md b/CONNECTION_PARAMETERS.md index ce1fdb5b..a40b6f3f 100644 --- a/CONNECTION_PARAMETERS.md +++ b/CONNECTION_PARAMETERS.md @@ -157,6 +157,7 @@ ignored with a warning and still stripped from the SEA wire. |---|---|:---:|:---:|---|---| | `useArrowNativeDecimal` | `WithArrowNativeDecimal` | ✅ | ⚠️ | `false` | Thrift: return DECIMAL as native Arrow `decimal128` (lossless string when scanned via `database/sql`). The kernel path already renders DECIMAL as the exact string regardless, so the flag is inert there. | | | `WithKernelDecimalAsFloat(b)` | ❌ | ✅ | `false` | Scan top-level DECIMAL as lossy `float64` instead of the exact string. | +| `geospatialAsString` | `WithKernelGeospatialAsString(b)` | ❌ | ✅ | kernel default (`true`) | `true`: return GEOMETRY / GEOGRAPHY as EWKT strings. `false`: return top-level values as raw WKB `[]byte`; `GetArrowBatches` retains the full `{srid,wkb}` struct. | With the default client-result settings, results otherwise render **byte-for-byte identically** on both backends (scalars, diff --git a/KERNEL_REV b/KERNEL_REV index f22f2168..f6a0e268 100644 --- a/KERNEL_REV +++ b/KERNEL_REV @@ -1 +1 @@ -80f2aee7d884994d7b0af9a9ea6078872859a9cd +ad3bc6993bca95b810839feade77ccd0ab98ece5 diff --git a/README.md b/README.md index 3953d4bb..82f47401 100644 --- a/README.md +++ b/README.md @@ -270,6 +270,7 @@ session parameter on both backends. |---|---|---|---|---| | `useArrowNativeDecimal` | `WithArrowNativeDecimal` | Thrift only (inert on kernel) | `false` | Thrift: return DECIMAL as native Arrow `decimal128` (lossless string when scanned via `database/sql`). The kernel path already renders DECIMAL as the exact string regardless. | | | `WithKernelDecimalAsFloat(b)` | SEA only | `false` | Scan top-level DECIMAL as lossy `float64` instead of the exact string. | +| `geospatialAsString` | `WithKernelGeospatialAsString(b)` | SEA only | kernel default (`true`) | `true`: return GEOMETRY / GEOGRAPHY as EWKT strings. `false`: return top-level values as raw WKB `[]byte`; `GetArrowBatches` retains the full `{srid,wkb}` struct. | See [Cloud Fetch](#cloud-fetch), [TLS](#tls), and [Proxy](#proxy) for the remaining groups. Telemetry parameters are covered under [Telemetry](#telemetry). diff --git a/connector.go b/connector.go index 99ab3646..fc756396 100644 --- a/connector.go +++ b/connector.go @@ -310,6 +310,10 @@ func withUserConfig(ucfg config.UserConfig) ConnOption { if ucfg.TokenCacheEnabledDSN { kernelExperimental(c).TokenCacheEnabled = true } + if ucfg.GeospatialAsStringDSN != nil { + asString := *ucfg.GeospatialAsStringDSN + kernelExperimental(c).GeospatialAsString = &asString + } } } @@ -665,6 +669,27 @@ func WithKernelDecimalAsFloat(asFloat bool) ConnOption { } } +// WithKernelGeospatialAsString selects how the kernel path exposes top-level +// GEOMETRY and GEOGRAPHY result values. True returns EWKT strings. False returns +// raw WKB bytes ([]byte); the logical GEOMETRY / GEOGRAPHY column type remains +// available through database/sql metadata, and the per-row SRID remains present +// in the Arrow struct returned by the driver's GetArrowBatches API. Omitting the +// option keeps the kernel default (currently EWKT strings). +// +// The conversion is entirely client-side and is never forwarded to SEA. +// Nested geospatial values follow the driver's existing complex-value contract: +// string mode renders EWKT leaves in JSON, while binary mode preserves the +// {"srid":...,"wkb":...} shape (with JSON base64 for WKB). +// +// EXPERIMENTAL, kernel-only: the default (Thrift) backend rejects this option at +// connect time. +func WithKernelGeospatialAsString(asString bool) ConnOption { + return func(c *config.Config) { + value := asString + kernelExperimental(c).GeospatialAsString = &value + } +} + // WithKernelTrustedCerts adds a PEM CA-certificate bundle to the kernel's TLS // trust store on top of the system roots — for a corporate re-signing proxy or an // on-prem CA. Required (rather than relying on SSL_CERT_FILE) because the kernel's diff --git a/connector_test.go b/connector_test.go index c2b50079..9f1648bd 100644 --- a/connector_test.go +++ b/connector_test.go @@ -450,6 +450,19 @@ func TestNewConnector(t *testing.T) { assert.Nil(t, coni.cfg.KernelExperimental) }) + t.Run("Connector test geospatialAsString=false DSN param preserves explicit false", func(t *testing.T) { + ucfg, err := config.ParseDSN("token:supersecret@databricks-host:443/sql/1.0/endpoints/abc?geospatialAsString=false") + require.NoError(t, err) + con, err := NewConnector(withUserConfig(ucfg)) + require.NoError(t, err) + + coni, ok := con.(*connector) + require.True(t, ok) + require.NotNil(t, coni.cfg.KernelExperimental) + require.NotNil(t, coni.cfg.KernelExperimental.GeospatialAsString) + assert.False(t, *coni.cfg.KernelExperimental.GeospatialAsString) + }) + t.Run("Connector test WithTransport sets HTTPClient in CloudFetchConfig", func(t *testing.T) { host := "databricks-host" accessToken := "token" diff --git a/internal/arrowscan/arrowscan.go b/internal/arrowscan/arrowscan.go index 8703addc..a6f920b3 100644 --- a/internal/arrowscan/arrowscan.go +++ b/internal/arrowscan/arrowscan.go @@ -36,8 +36,10 @@ import ( // date, timestamp, and top-level decimal (as an exact fixed-point string, // matching the Thrift path — a float64 would lose precision beyond ~17 digits). // Nested types (List/Map/Struct, and VARIANT which -// arrives nested) render to a JSON string byte-identical to the Thrift path; -// GEOMETRY arrives as a WKB/WKT string and is handled by the string arm. INTERVAL +// arrives nested) render to a JSON string byte-identical to the Thrift path. +// Top-level binary GEOMETRY / GEOGRAPHY is selected explicitly through +// ScanOptions and exposes its WKB child as []byte; string mode is handled by the +// regular string arm. INTERVAL // day-time/year-month arrive as native arrow duration/month-interval and format to // the same string the Thrift path receives pre-formatted from the server. NULLs // map to nil. A genuinely unhandled type returns an error rather than a silently @@ -106,20 +108,39 @@ func (c *StructKeyCache) keyPrefixes(st *arrow.StructType) []string { // StructKeyCache) so struct field-name keys are escaped once per result set // rather than once per row. Pass nil for the un-memoized one-shot behavior. func ScanCellCached(col arrow.Array, row int, loc *time.Location, keys *StructKeyCache) (driver.Value, error) { - return scanCell(col, row, loc, keys, false) + return scanCell(col, row, loc, keys, ScanOptions{}) } // ScanCellCachedDecimalFloat is ScanCellCached that, when decimalAsFloat is true, // scans a TOP-LEVEL Decimal128 to a lossy float64 instead of the exact string // (nested decimals still render exactly). Opt in via WithKernelDecimalAsFloat. func ScanCellCachedDecimalFloat(col arrow.Array, row int, loc *time.Location, keys *StructKeyCache, decimalAsFloat bool) (driver.Value, error) { - return scanCell(col, row, loc, keys, decimalAsFloat) + return scanCell(col, row, loc, keys, ScanOptions{DecimalAsFloat: decimalAsFloat}) } -func scanCell(col arrow.Array, row int, loc *time.Location, keys *StructKeyCache, decimalAsFloat bool) (driver.Value, error) { +// ScanOptions controls top-level result adaptations that cannot be inferred from +// an Arrow array alone. Logical GEOMETRY / GEOGRAPHY identity lives on the +// enclosing schema Field, so kernelRows sets GeospatialBinary only for a field +// carrying that metadata and a physical struct value. +type ScanOptions struct { + DecimalAsFloat bool + GeospatialBinary bool +} + +// ScanCellCachedWithOptions is the full scanner entry point used by the kernel +// rows implementation. Callers that do not need result-specific adaptations +// should continue to use ScanCellCached. +func ScanCellCachedWithOptions(col arrow.Array, row int, loc *time.Location, keys *StructKeyCache, options ScanOptions) (driver.Value, error) { + return scanCell(col, row, loc, keys, options) +} + +func scanCell(col arrow.Array, row int, loc *time.Location, keys *StructKeyCache, options ScanOptions) (driver.Value, error) { if col.IsNull(row) { return nil, nil } + if options.GeospatialBinary { + return scanGeospatialWKB(col, row) + } switch c := col.(type) { case *array.Null: return nil, nil @@ -195,7 +216,7 @@ func scanCell(col arrow.Array, row int, loc *time.Location, keys *StructKeyCache return inLocation(timestampToTime(int64(c.Value(row)), dt.Unit), loc), nil case *array.Decimal128: dt := col.DataType().(*arrow.Decimal128Type) - if decimalAsFloat { + if options.DecimalAsFloat { // Lossy fast path: float64, no per-cell string. Opt-in only. return c.Value(row).ToFloat64(dt.Scale), nil } @@ -228,6 +249,38 @@ func scanCell(col arrow.Array, row int, loc *time.Location, keys *StructKeyCache } } +// scanGeospatialWKB validates and unwraps Reyden's canonical top-level +// struct value. The outer struct null was handled by +// scanCell before this function is called. Child nulls are malformed because +// both canonical children are non-nullable; reject them rather than returning a +// byte slice whose SRID/value pairing is incomplete. +func scanGeospatialWKB(col arrow.Array, row int) (driver.Value, error) { + values, ok := col.(*array.Struct) + if !ok { + return nil, fmt.Errorf("binary geospatial column has arrow type %s, want struct", col.DataType()) + } + fields := values.DataType().(*arrow.StructType).Fields() + if len(fields) != 2 || fields[0].Name != "srid" || fields[1].Name != "wkb" || + fields[0].Type.ID() != arrow.INT32 || fields[1].Type.ID() != arrow.BINARY { + return nil, fmt.Errorf("binary geospatial column has arrow type %s, want struct", col.DataType()) + } + srids, ok := values.Field(0).(*array.Int32) + if !ok { + return nil, fmt.Errorf("binary geospatial srid child has arrow type %s, want int32", values.Field(0).DataType()) + } + wkbs, ok := values.Field(1).(*array.Binary) + if !ok { + return nil, fmt.Errorf("binary geospatial wkb child has arrow type %s, want binary", values.Field(1).DataType()) + } + if srids.IsNull(row) { + return nil, fmt.Errorf("binary geospatial value at row %d has a null srid child", row) + } + if wkbs.IsNull(row) { + return nil, fmt.Errorf("binary geospatial value at row %d has a null wkb child", row) + } + return wkbs.Value(row), nil +} + // formatDayTimeInterval renders an arrow duration (in the given time unit) as the // Thrift path's "D HH:MM:SS.nnnnnnnnn" — days, then zero-padded hours:minutes:seconds // with 9 fractional digits, negated with a leading '-'. diff --git a/internal/arrowscan/arrowscan_test.go b/internal/arrowscan/arrowscan_test.go index 31a8774b..7b5dd821 100644 --- a/internal/arrowscan/arrowscan_test.go +++ b/internal/arrowscan/arrowscan_test.go @@ -2,6 +2,7 @@ package arrowscan import ( "fmt" + "strings" "testing" "time" @@ -684,3 +685,86 @@ func TestScanCellDecimalAsFloat(t *testing.T) { t.Errorf("decimalAsFloat arm = %v, want ~99.98", fv) } } + +func TestScanCellBinaryGeospatial(t *testing.T) { + pool := memory.NewGoAllocator() + wkb := []byte{0x01, 0x01, 0x00, 0x00, 0x00} + geoType := arrow.StructOf( + arrow.Field{Name: "srid", Type: arrow.PrimitiveTypes.Int32, Nullable: false}, + arrow.Field{Name: "wkb", Type: arrow.BinaryTypes.Binary, Nullable: false}, + ) + + build := func(outerValid, sridValid, wkbValid bool) *array.Struct { + b := array.NewStructBuilder(pool, geoType) + defer b.Release() + b.Append(outerValid) + if sridValid { + b.FieldBuilder(0).(*array.Int32Builder).Append(4326) + } else { + b.FieldBuilder(0).(*array.Int32Builder).AppendNull() + } + if wkbValid { + b.FieldBuilder(1).(*array.BinaryBuilder).Append(wkb) + } else { + b.FieldBuilder(1).(*array.BinaryBuilder).AppendNull() + } + return b.NewStructArray() + } + + t.Run("unwraps WKB", func(t *testing.T) { + arr := build(true, true, true) + defer arr.Release() + got, err := ScanCellCachedWithOptions(arr, 0, nil, nil, ScanOptions{GeospatialBinary: true}) + if err != nil { + t.Fatal(err) + } + if string(got.([]byte)) != string(wkb) { + t.Errorf("WKB = %x, want %x", got, wkb) + } + }) + + t.Run("outer null wins", func(t *testing.T) { + arr := build(false, false, false) + defer arr.Release() + got, err := ScanCellCachedWithOptions(arr, 0, nil, nil, ScanOptions{GeospatialBinary: true}) + if err != nil || got != nil { + t.Errorf("outer-null geo = (%v, %v), want (nil, nil)", got, err) + } + }) + + t.Run("rejects null srid child", func(t *testing.T) { + arr := build(true, false, true) + defer arr.Release() + _, err := ScanCellCachedWithOptions(arr, 0, nil, nil, ScanOptions{GeospatialBinary: true}) + if err == nil || !strings.Contains(err.Error(), "null srid") { + t.Errorf("error = %v, want null-srid validation error", err) + } + }) + + t.Run("rejects null wkb child", func(t *testing.T) { + arr := build(true, true, false) + defer arr.Release() + _, err := ScanCellCachedWithOptions(arr, 0, nil, nil, ScanOptions{GeospatialBinary: true}) + if err == nil || !strings.Contains(err.Error(), "null wkb") { + t.Errorf("error = %v, want null-wkb validation error", err) + } + }) + + t.Run("rejects a non-canonical struct", func(t *testing.T) { + badType := arrow.StructOf( + arrow.Field{Name: "srid", Type: arrow.PrimitiveTypes.Int64}, + arrow.Field{Name: "bytes", Type: arrow.BinaryTypes.Binary}, + ) + b := array.NewStructBuilder(pool, badType) + defer b.Release() + b.Append(true) + b.FieldBuilder(0).(*array.Int64Builder).Append(4326) + b.FieldBuilder(1).(*array.BinaryBuilder).Append(wkb) + arr := b.NewStructArray() + defer arr.Release() + _, err := ScanCellCachedWithOptions(arr, 0, nil, nil, ScanOptions{GeospatialBinary: true}) + if err == nil || !strings.Contains(err.Error(), "want struct") { + t.Errorf("error = %v, want canonical-shape validation error", err) + } + }) +} diff --git a/internal/arrowscan/coltype.go b/internal/arrowscan/coltype.go index b4be2a1c..ff30ab46 100644 --- a/internal/arrowscan/coltype.go +++ b/internal/arrowscan/coltype.go @@ -4,11 +4,14 @@ import ( "database/sql" "math" "reflect" + "strings" "time" "github.com/apache/arrow/go/v12/arrow" ) +const databricksTypeNameMetadataKey = "databricks.type_name" + // ColumnTypeInfo is the per-column metadata database/sql surfaces through // sql.ColumnType. The kernel derives it from the result's Arrow schema via // ColumnTypeInfoFor — the mapping the value scanner and type reporter share, kept @@ -110,6 +113,52 @@ func ColumnTypeInfoFor(dt arrow.DataType) ColumnTypeInfo { } } +// LogicalGeospatialType returns the top-level Databricks geospatial family +// recorded by the kernel on an Arrow field. The physical value is UTF-8 in +// string mode and a struct in binary mode, so DataType alone cannot preserve +// this database type name. +func LogicalGeospatialType(field arrow.Field) (string, bool) { + name, ok := field.Metadata.GetValue(databricksTypeNameMetadataKey) + if !ok { + return "", false + } + switch strings.ToUpper(strings.TrimSpace(name)) { + case "GEOMETRY": + return "GEOMETRY", true + case "GEOGRAPHY": + return "GEOGRAPHY", true + default: + return "", false + } +} + +// IsBinaryGeospatialField reports whether field is a logical top-level +// GEOMETRY / GEOGRAPHY using the kernel's struct representation. The scanner +// performs the detailed child-shape validation when it reads a non-null value. +func IsBinaryGeospatialField(field arrow.Field) bool { + _, geospatial := LogicalGeospatialType(field) + return geospatial && field.Type.ID() == arrow.STRUCT +} + +// ColumnTypeInfoForField preserves logical GEOMETRY / GEOGRAPHY metadata while +// selecting a scan type that matches the requested physical representation. +// Other fields retain the Arrow-only mapping. +func ColumnTypeInfoForField(field arrow.Field) ColumnTypeInfo { + name, geospatial := LogicalGeospatialType(field) + if !geospatial { + return ColumnTypeInfoFor(field.Type) + } + if field.Type.ID() == arrow.STRUCT { + return varLen(name, scanTypeRawBytes) + } + if field.Type.ID() == arrow.STRING || field.Type.ID() == arrow.LARGE_STRING { + return varLen(name, scanTypeString) + } + info := ColumnTypeInfoFor(field.Type) + info.DatabaseTypeName = name + return info +} + // varLen builds a ColumnTypeInfo for a variable-length type, reporting the // unbounded length (math.MaxInt64) the Thrift path uses for such columns. func varLen(name string, scan reflect.Type) ColumnTypeInfo { diff --git a/internal/arrowscan/coltype_test.go b/internal/arrowscan/coltype_test.go index 0dbd2f8a..1ee7f581 100644 --- a/internal/arrowscan/coltype_test.go +++ b/internal/arrowscan/coltype_test.go @@ -100,3 +100,40 @@ func TestColumnTypeInfoScanTypeCoversScanner(t *testing.T) { } } } + +func TestColumnTypeInfoForFieldPreservesGeospatialType(t *testing.T) { + raw := reflect.TypeOf(sql.RawBytes{}) + str := reflect.TypeOf("") + metadata := arrow.MetadataFrom(map[string]string{ + "databricks.type_name": "geometry", + "databricks.type_text": "GEOMETRY(ANY)", + }) + geoStruct := arrow.StructOf( + arrow.Field{Name: "srid", Type: arrow.PrimitiveTypes.Int32, Nullable: false}, + arrow.Field{Name: "wkb", Type: arrow.BinaryTypes.Binary, Nullable: false}, + ) + + stringField := arrow.Field{Name: "g", Type: arrow.BinaryTypes.String, Nullable: true, Metadata: metadata} + stringInfo := ColumnTypeInfoForField(stringField) + if stringInfo.DatabaseTypeName != "GEOMETRY" || stringInfo.ScanType != str { + t.Errorf("string geo metadata = (%q, %v), want (GEOMETRY, string)", stringInfo.DatabaseTypeName, stringInfo.ScanType) + } + if IsBinaryGeospatialField(stringField) { + t.Error("UTF-8 geospatial field classified as binary") + } + + binaryField := arrow.Field{Name: "g", Type: geoStruct, Nullable: true, Metadata: metadata} + binaryInfo := ColumnTypeInfoForField(binaryField) + if binaryInfo.DatabaseTypeName != "GEOMETRY" || binaryInfo.ScanType != raw { + t.Errorf("binary geo metadata = (%q, %v), want (GEOMETRY, RawBytes)", binaryInfo.DatabaseTypeName, binaryInfo.ScanType) + } + if !IsBinaryGeospatialField(binaryField) { + t.Error("logical geospatial struct was not classified as binary") + } + + ordinaryStruct := arrow.Field{Name: "s", Type: geoStruct, Nullable: true} + ordinaryInfo := ColumnTypeInfoForField(ordinaryStruct) + if ordinaryInfo.DatabaseTypeName != "STRUCT" || IsBinaryGeospatialField(ordinaryStruct) { + t.Errorf("ordinary struct misclassified: %+v", ordinaryInfo) + } +} diff --git a/internal/backend/kernel/backend.go b/internal/backend/kernel/backend.go index a916cb16..3370f1e1 100644 --- a/internal/backend/kernel/backend.go +++ b/internal/backend/kernel/backend.go @@ -48,6 +48,11 @@ static inline KernelStatusCode go_kernel_set_driver_system_configuration( os_name, os_version, os_arch, client_app_name, locale_name, char_set_encoding, process_name); } + +static inline KernelStatusCode go_kernel_set_geospatial_as_string( + KernelSessionConfig* config, bool as_string) { + return kernel_session_config_set_geospatial_as_string(config, as_string); +} */ import "C" @@ -251,6 +256,9 @@ func (k *KernelBackend) OpenSession(ctx context.Context) error { if err := k.applyMaxConnections(cfg); err != nil { return err } + if err := k.applyGeospatialAsString(cfg); err != nil { + return err + } // Retry / backoff policy (WithRetries). See applyRetry. if err := k.applyRetry(cfg); err != nil { @@ -423,6 +431,22 @@ func (k *KernelBackend) applyMaxConnections(cfg *C.KernelSessionConfig) error { return nil } +// applyGeospatialAsString selects the kernel's client-side representation for +// GEOMETRY / GEOGRAPHY results. Nil deliberately skips the setter so the kernel +// owns its default; an explicit false must cross the ABI rather than collapsing +// into the zero value. +func (k *KernelBackend) applyGeospatialAsString(cfg *C.KernelSessionConfig) error { + if k.cfg.GeospatialAsString == nil { + return nil + } + if err := call(func() C.KernelStatusCode { + return C.go_kernel_set_geospatial_as_string(cfg, C.bool(*k.cfg.GeospatialAsString)) + }); err != nil { + return fmt.Errorf("kernel: set_geospatial_as_string: %w", toConnError(err)) + } + return nil +} + // applyRetry forwards the driver's HTTP retry / backoff policy to the session // config. A no-op when Config.Retry is nil, so the kernel's own default policy // (exponential backoff with jitter, 5 retries, 1s..60s, 900s budget) is preserved @@ -726,6 +750,18 @@ func trySetMaxConnections(cfg Config) error { return k.applyMaxConnections(c) } +// trySetGeospatialAsString exercises the geospatial representation setter +// without opening a network session. It is used only by tagged kernel tests. +func trySetGeospatialAsString(cfg Config) error { + var c *C.KernelSessionConfig + if err := call(func() C.KernelStatusCode { return C.kernel_session_config_new(&c) }); err != nil { + return fmt.Errorf("config_new: %w", err) + } + defer C.kernel_session_config_free(c) + k := &KernelBackend{cfg: cfg} + return k.applyGeospatialAsString(c) +} + // trySetTelemetry exercises the telemetry C setter without opening a network // session. It is used only by the tagged kernel tests. func trySetTelemetry(cfg Config) error { diff --git a/internal/backend/kernel/config.go b/internal/backend/kernel/config.go index 062030a5..0a267666 100644 --- a/internal/backend/kernel/config.go +++ b/internal/backend/kernel/config.go @@ -99,6 +99,11 @@ type Config struct { // the exact string (from WithKernelDecimalAsFloat). Kernel still sends Decimal128. DecimalAsFloat bool + // GeospatialAsString selects the local GEOMETRY / GEOGRAPHY result + // representation. Nil keeps the kernel default; true requests EWKT strings; + // false requests the canonical struct Arrow representation. + GeospatialAsString *bool + // TokenCacheEnabled controls the kernel's on-disk OAuth U2M token-cache persistence // (WithTokenCache / tokenCache DSN param). When false (the default), tokens are held // in memory only; when true, the refresh token is persisted encrypted to diff --git a/internal/backend/kernel/include/databricks_kernel.h b/internal/backend/kernel/include/databricks_kernel.h index 4cb46c60..85dd60cb 100644 --- a/internal/backend/kernel/include/databricks_kernel.h +++ b/internal/backend/kernel/include/databricks_kernel.h @@ -201,10 +201,11 @@ typedef struct kernel_session_config_t KernelSessionConfig; typedef struct kernel_session_t kernel_session_t; typedef struct kernel_statement_t kernel_statement_t; typedef struct kernel_executed_statement_t kernel_executed_statement_t; -/* Async-consumption executed handle. Reserved: the async submit path is - * deferred in v0 (see kernel_statement_submit), so no consumer receives a - * live instance yet. */ +/* Async-consumption executed handle returned by kernel_statement_submit. */ typedef struct kernel_executed_async_statement_t kernel_executed_async_statement_t; +/* Execution-scoped canceller detached from an async executed handle. */ +typedef struct kernel_executed_async_statement_canceller_t + kernel_executed_async_statement_canceller_t; typedef struct kernel_result_stream_t kernel_result_stream_t; /* Detached canceller for a sync-executing statement, returned by * kernel_statement_canceller_new. The returned handle is safe to use from @@ -343,6 +344,18 @@ KernelStatusCode kernel_session_config_set_session_conf(KernelSessionConfig* con const char* key, const char* value); +/* Select the physical Arrow representation for GEOMETRY / GEOGRAPHY results. + * `as_string = true` requests EWKT in Arrow UTF-8 values; `false` requests + * Arrow `struct` values. Omitting this setter uses + * the kernel string default. Binary mode requires the native Reyden Arrow + * path; legacy text-only results cannot be reconstructed as WKB. This + * client-side setting is never sent to SEA. + * A raw `geospatial_as_string` entry set through + * kernel_session_config_set_session_conf remains server-bound and is forwarded + * unchanged. */ +KernelStatusCode kernel_session_config_set_geospatial_as_string( + KernelSessionConfig* config, bool as_string); + /* Append one HTTP header sent on every request. Call once per header * (order preserved); `name` and `value` are both required. */ KernelStatusCode kernel_session_config_set_custom_header(KernelSessionConfig* config, @@ -735,13 +748,7 @@ KernelStatusCode kernel_statement_execute_with_timeout_ms( kernel_statement_t* stmt, uint64_t timeout_ms, kernel_executed_statement_t** out); -/* - * Submit-and-return (async). DEFERRED in v0: this always returns - * `KernelStatusCode_InvalidArgument` (with an explanatory last error) and - * writes nothing to `*out`. Use the synchronous `kernel_statement_execute` - * path. Declared so the symbol/contract is visible; wired when an ODBC / - * Go consumer needs caller-driven polling. - */ +/* Submit-and-return. On success, *out owns an async executed handle. */ KernelStatusCode kernel_statement_submit(kernel_statement_t* stmt, kernel_executed_async_statement_t** out); @@ -868,6 +875,50 @@ const char* kernel_executed_statement_query_id(const kernel_executed_statement_t KernelStatusCode kernel_executed_statement_close(kernel_executed_statement_t* exec); +/* ─── Executed statement (caller-driven async) ─────────────────── */ + +/* Perform at most one status request. Pending, Running, Succeeded, and Closed + * return Success. Failed returns SqlError; Cancelled returns Cancelled. On + * either non-success terminal outcome, *out is still set and the calling + * thread can retrieve the error information with kernel_get_last_error. */ +KernelStatusCode kernel_executed_async_statement_status( + kernel_executed_async_statement_t* exec, KernelStatementStatusKind* out); + +/* Materialise a successful terminal result as the existing owned result-stream + * handle. This is NON-BLOCKING: the caller must first poll + * kernel_executed_async_statement_status until it reports Succeeded (or Closed) + * before calling this. Invoking it before the statement reaches a terminal + * state returns KernelStatusCode_InvalidArgument rather than blocking (v0 + * collapses invalid-state into that code). This is single-consumption: + * a second successful take is rejected. Close the stream with + * kernel_result_stream_close. */ +KernelStatusCode kernel_executed_async_statement_get_result_stream( + kernel_executed_async_statement_t* exec, kernel_result_stream_t** out); + +/* Most recently observed DML count, or -1 when unavailable. */ +int64_t kernel_executed_async_statement_num_modified_rows( + const kernel_executed_async_statement_t* exec); + +/* Borrowed NUL-terminated server statement id, valid until async close. */ +const char* kernel_executed_async_statement_query_id( + const kernel_executed_async_statement_t* exec); + +/* Create an execution-scoped canceller. The returned handle owns no borrow of + * exec and may cancel concurrently with status or result work. Join any + * in-flight cancel before freeing this handle. */ +KernelStatusCode kernel_executed_async_statement_canceller_new( + const kernel_executed_async_statement_t* exec, + kernel_executed_async_statement_canceller_t** out); +KernelStatusCode kernel_executed_async_statement_canceller_cancel( + kernel_executed_async_statement_canceller_t* canceller); +KernelStatusCode kernel_executed_async_statement_canceller_free( + kernel_executed_async_statement_canceller_t* canceller); + +/* Close the server operation and consume the async handle. The handle is + * consumed even when close returns an error. */ +KernelStatusCode kernel_executed_async_statement_close( + kernel_executed_async_statement_t* exec); + /* ─── Result stream ─────────────────────────────────────────────────── * * The schema is exported once; each `next_batch` call exports one Arrow diff --git a/internal/backend/kernel/kernel_test.go b/internal/backend/kernel/kernel_test.go index 9039db2e..bb2ef4bc 100644 --- a/internal/backend/kernel/kernel_test.go +++ b/internal/backend/kernel/kernel_test.go @@ -255,6 +255,25 @@ func TestSetTokenCache(t *testing.T) { } } +func TestSetGeospatialAsString(t *testing.T) { + stringMode := true + binaryMode := false + for _, tc := range []struct { + name string + cfg Config + }{ + {name: "omitted is a no-op", cfg: Config{}}, + {name: "string", cfg: Config{GeospatialAsString: &stringMode}}, + {name: "binary", cfg: Config{GeospatialAsString: &binaryMode}}, + } { + t.Run(tc.name, func(t *testing.T) { + if err := trySetGeospatialAsString(tc.cfg); err != nil { + t.Errorf("applyGeospatialAsString(%s) = %v, want nil", tc.name, err) + } + }) + } +} + func TestSetRequestTimeout(t *testing.T) { for _, tc := range []struct { name string diff --git a/internal/backend/kernel/rows.go b/internal/backend/kernel/rows.go index 635fc10c..25299833 100644 --- a/internal/backend/kernel/rows.go +++ b/internal/backend/kernel/rows.go @@ -65,6 +65,7 @@ type kernelRows struct { cols []string colTypes []arrowscan.ColumnTypeInfo // per-column type metadata (PECOBLR-3692) + geoBinary []bool // top-level logical geo structs unwrap to WKB for database/sql schema *arrow.Schema // result-set schema, for GetArrowBatches().Schema() cur arrow.Record // current batch (nil until first Next) rowInCur int // next row index within cur @@ -118,9 +119,11 @@ func newKernelRows(ctx context.Context, op *kernelOp, stream *C.kernel_result_st // (PECOBLR-3692) with no per-call work. Kept in lockstep with the value scanner // (ScanCellCached) via the shared arrowscan.ColumnTypeInfoFor mapper. r.colTypes = make([]arrowscan.ColumnTypeInfo, len(fields)) + r.geoBinary = make([]bool, len(fields)) for i, f := range fields { r.cols[i] = f.Name - r.colTypes[i] = arrowscan.ColumnTypeInfoFor(f.Type) + r.colTypes[i] = arrowscan.ColumnTypeInfoForField(f) + r.geoBinary[i] = arrowscan.IsBinaryGeospatialField(f) } // Construction succeeded — now arm the close telemetry callback so a normal // Close() (after row iteration) records CLOSE_STATEMENT. @@ -232,7 +235,10 @@ func (r *kernelRows) next(dest []driver.Value) error { } rec := r.cur for c := 0; c < len(dest); c++ { - v, err := arrowscan.ScanCellCachedDecimalFloat(rec.Column(c), r.rowInCur, r.op.location, r.keyCache, r.op.decimalAsFloat) + v, err := arrowscan.ScanCellCachedWithOptions(rec.Column(c), r.rowInCur, r.op.location, r.keyCache, arrowscan.ScanOptions{ + DecimalAsFloat: r.op.decimalAsFloat, + GeospatialBinary: r.geoBinary[c], + }) if err != nil { return fmt.Errorf("kernel: scan col %d (%s): %w", c, r.cols[c], err) } diff --git a/internal/config/config.go b/internal/config/config.go index 75635927..0a171f7e 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -137,6 +137,12 @@ type KernelExperimentalConfig struct { // false TokenCacheEnabled calls the setter with enabled=false so the kernel does // NOT persist to disk by default. TokenCacheEnabled bool + + // GeospatialAsString selects the kernel's client-side GEOMETRY / GEOGRAPHY + // result representation. Nil keeps the kernel default (EWKT strings); true + // requests EWKT; false requests the canonical struct Arrow value. + // A pointer preserves omitted versus explicitly false. + GeospatialAsString *bool } // DeepCopy returns a deep copy of the experimental config, or nil for a nil @@ -158,6 +164,10 @@ func (k *KernelExperimentalConfig) DeepCopy() *KernelExperimentalConfig { DecimalAsFloat: k.DecimalAsFloat, TokenCacheEnabled: k.TokenCacheEnabled, } + if k.GeospatialAsString != nil { + asString := *k.GeospatialAsString + cp.GeospatialAsString = &asString + } if k.MaxConnections != nil { max := *k.MaxConnections cp.MaxConnections = &max @@ -302,6 +312,10 @@ type UserConfig struct { // unambiguous. False by default; when true, enables on-disk token-cache persistence // for U2M OAuth. DSN: tokenCache=true. TokenCacheEnabledDSN bool + // GeospatialAsStringDSN is a DSN-only carrier for + // geospatialAsString=true|false. Nil means omitted, preserving the kernel's + // default; the connector copies a present value into KernelExperimental. + GeospatialAsStringDSN *bool } // DeepCopy returns a true deep copy of UserConfig @@ -323,7 +337,7 @@ func (ucfg UserConfig) DeepCopy() UserConfig { } - return UserConfig{ + result := UserConfig{ Protocol: ucfg.Protocol, Host: ucfg.Host, Port: ucfg.Port, @@ -352,6 +366,11 @@ func (ucfg UserConfig) DeepCopy() UserConfig { WarehouseID: ucfg.WarehouseID, TokenCacheEnabledDSN: ucfg.TokenCacheEnabledDSN, } + if ucfg.GeospatialAsStringDSN != nil { + asString := *ucfg.GeospatialAsStringDSN + result.GeospatialAsStringDSN = &asString + } + return result } var defaultMaxRows = 100000 @@ -516,6 +535,12 @@ func ParseDSN(dsn string) (UserConfig, error) { } ucfg.TokenCacheEnabledDSN = tokenCache } + if geospatialAsString, ok, err := params.extractAsBool("geospatialAsString"); ok { + if err != nil { + return UserConfig{}, err + } + ucfg.GeospatialAsStringDSN = &geospatialAsString + } // Telemetry parameters if enableTelemetry, ok, err := params.extractAsBool("enableTelemetry"); ok { diff --git a/internal/config/config_test.go b/internal/config/config_test.go index b1e235b0..b0f9521b 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -693,6 +693,39 @@ func TestParseConfig(t *testing.T) { } } +func TestParseDSNGeospatialAsString(t *testing.T) { + base := "token:supersecret@example.cloud.databricks.com:443/sql/1.0/endpoints/abc" + for _, tc := range []struct { + name string + query string + want *bool + }{ + {name: "omitted", query: "", want: nil}, + {name: "string", query: "?geospatialAsString=true", want: func() *bool { v := true; return &v }()}, + {name: "binary", query: "?geospatialAsString=false", want: func() *bool { v := false; return &v }()}, + } { + t.Run(tc.name, func(t *testing.T) { + got, err := ParseDSN(base + tc.query) + if err != nil { + t.Fatal(err) + } + if tc.want == nil { + if got.GeospatialAsStringDSN != nil { + t.Errorf("GeospatialAsStringDSN = %v, want nil", *got.GeospatialAsStringDSN) + } + return + } + if got.GeospatialAsStringDSN == nil || *got.GeospatialAsStringDSN != *tc.want { + t.Errorf("GeospatialAsStringDSN = %v, want %v", got.GeospatialAsStringDSN, *tc.want) + } + }) + } + + if _, err := ParseDSN(base + "?geospatialAsString=notabool"); err == nil { + t.Error("malformed geospatialAsString should fail DSN parsing") + } +} + func TestEffectiveSessionParams(t *testing.T) { t.Run("no metric view leaves params unchanged", func(t *testing.T) { c := &Config{UserConfig: UserConfig{SessionParams: map[string]string{"QUERY_TAGS": "a:1"}}} diff --git a/kernel_config.go b/kernel_config.go index 2c461634..bcfde17b 100644 --- a/kernel_config.go +++ b/kernel_config.go @@ -209,6 +209,10 @@ func buildKernelConfig(cfg *config.Config, kauth kernel.Auth) kernel.Config { } // Client-side scan choice: lossy float64 decimals instead of exact strings. kc.DecimalAsFloat = ke.DecimalAsFloat + if ke.GeospatialAsString != nil { + asString := *ke.GeospatialAsString + kc.GeospatialAsString = &asString + } // U2M token-cache control (WithTokenCache / tokenCache DSN param): forward to kernel. kc.TokenCacheEnabled = ke.TokenCacheEnabled } diff --git a/kernel_config_test.go b/kernel_config_test.go index 5c1b1c3b..bfe54417 100644 --- a/kernel_config_test.go +++ b/kernel_config_test.go @@ -471,6 +471,7 @@ var kernelConfigFieldDisposition = map[string]string{ "TelemetryFlushInterval": "forwarded", "UseArrowNativeDecimalDSN": "inert", // DSN carrier; kernel renders decimals exactly regardless "TokenCacheEnabledDSN": "inert", // DSN carrier; forwarded to KernelExperimental.TokenCacheEnabled + "GeospatialAsStringDSN": "inert", // DSN carrier; forwarded to KernelExperimental.GeospatialAsString // Fields promoted from the embedded CloudFetchConfig. The kernel does // CloudFetch internally (below the C ABI), so none is forwarded — but each is @@ -559,6 +560,26 @@ func TestBuildKernelConfig(t *testing.T) { } }) + t.Run("geospatial representation preserves omitted true and false", func(t *testing.T) { + c := baseKernelConfig() + kc := buildKernelConfig(c, kernel.Auth{Mode: kernel.AuthPAT, Token: "dapi-x"}) + if kc.GeospatialAsString != nil { + t.Fatalf("omitted GeospatialAsString = %v, want nil", *kc.GeospatialAsString) + } + + WithKernelGeospatialAsString(true)(c) + kc = buildKernelConfig(c, kernel.Auth{Mode: kernel.AuthPAT, Token: "dapi-x"}) + if kc.GeospatialAsString == nil || !*kc.GeospatialAsString { + t.Fatalf("explicit string GeospatialAsString = %v, want true", kc.GeospatialAsString) + } + + WithKernelGeospatialAsString(false)(c) + kc = buildKernelConfig(c, kernel.Auth{Mode: kernel.AuthPAT, Token: "dapi-x"}) + if kc.GeospatialAsString == nil || *kc.GeospatialAsString { + t.Fatalf("explicit binary GeospatialAsString = %v, want false", kc.GeospatialAsString) + } + }) + t.Run("experimental TLS fields forwarded", func(t *testing.T) { c := baseKernelConfig() c.KernelExperimental = &config.KernelExperimentalConfig{ diff --git a/kernel_experimental_test.go b/kernel_experimental_test.go index 3f936a57..9ebda216 100644 --- a/kernel_experimental_test.go +++ b/kernel_experimental_test.go @@ -40,6 +40,7 @@ var kernelExperimentalFieldDisposition = map[string]string{ "MaxConnections": "forwarded", // set_max_connections "MaxChunksInMemory": "forwarded", // set_session_conf (cloudfetch_max_chunks_in_memory, client-only) "DecimalAsFloat": "forwarded", // kernel.Config.DecimalAsFloat → kernelOp → arrowscan (client-side scan choice) + "GeospatialAsString": "forwarded", // set_geospatial_as_string; false is preserved via pointer "TokenCacheEnabled": "forwarded", // set_u2m_token_cache_config (enabled, nil passphrase) } @@ -100,6 +101,9 @@ func TestWithKernelTLSOptionsSetExperimental(t *testing.T) { {"decimal as float", WithKernelDecimalAsFloat(true), func(k *config.KernelExperimentalConfig) bool { return k.DecimalAsFloat }}, + {"geospatial as binary", WithKernelGeospatialAsString(false), func(k *config.KernelExperimentalConfig) bool { + return k.GeospatialAsString != nil && !*k.GeospatialAsString + }}, {"token cache", WithTokenCache(true), func(k *config.KernelExperimentalConfig) bool { return k.TokenCacheEnabled }}, @@ -138,6 +142,7 @@ func TestWithKernelOptionsRejectedOnThriftPath(t *testing.T) { {"max connections", WithKernelMaxConnections(37)}, {"max chunks in memory", WithKernelMaxChunksInMemory(4)}, {"decimal as float", WithKernelDecimalAsFloat(true)}, + {"geospatial as binary", WithKernelGeospatialAsString(false)}, {"token cache", WithTokenCache(true)}, } for _, tc := range cases { @@ -218,6 +223,7 @@ func TestWithKernelClientCertificateCopiesPEM(t *testing.T) { // mutation reach another. func TestKernelExperimentalDeepCopy(t *testing.T) { maxConnections := 37 + geospatialAsString := false orig := &config.KernelExperimentalConfig{ TLSTrustedCertsPEM: []byte("ca-bundle"), TLSClientCertPEM: []byte("client-cert"), @@ -231,6 +237,7 @@ func TestKernelExperimentalDeepCopy(t *testing.T) { RetryOverallTimeout: 5 * time.Minute, MaxConnections: &maxConnections, MaxChunksInMemory: 4, + GeospatialAsString: &geospatialAsString, TokenCacheEnabled: true, } cp := orig.DeepCopy() @@ -252,6 +259,11 @@ func TestKernelExperimentalDeepCopy(t *testing.T) { if cp.MaxChunksInMemory != 4 { t.Errorf("DeepCopy lost MaxChunksInMemory: %v", cp.MaxChunksInMemory) } + if cp.GeospatialAsString == nil || *cp.GeospatialAsString { + t.Errorf("DeepCopy lost GeospatialAsString: %v", cp.GeospatialAsString) + } else if cp.GeospatialAsString == orig.GeospatialAsString { + t.Error("DeepCopy aliased GeospatialAsString pointer") + } if !cp.TokenCacheEnabled { t.Errorf("DeepCopy lost TokenCacheEnabled: %v", cp.TokenCacheEnabled) }