From b02d042388d9fb04a73b66caada1c2f5a8ab7efb Mon Sep 17 00:00:00 2001 From: "peco-engineer-bot[bot]" <3815206+peco-engineer-bot[bot]@users.noreply.github.com> Date: Wed, 21 Oct 2026 20:33:21 +0000 Subject: [PATCH 01/15] TIMESTAMP_NTZ is scanned as a zoned TIMESTAMP (#506) Signed-off-by: peco-engineer-bot[bot] <3815206+peco-engineer-bot[bot]@users.noreply.github.com> Signed-off-by: Vu Anh Phung --- doc.go | 5 +- driver_e2e_test.go | 82 +++++++++++++++++++ internal/arrowscan/arrowscan.go | 12 ++- internal/arrowscan/arrowscan_test.go | 7 +- internal/arrowscan/coltype.go | 8 ++ internal/arrowscan/coltype_ntz_test.go | 33 ++++++++ internal/arrowscan/timestamp_ntz_test.go | 67 +++++++++++++++ internal/rows/arrowbased/arrowRows.go | 26 +++++- .../rows/arrowbased/arrowscan_parity_test.go | 23 +++--- internal/rows/arrowbased/columnValues.go | 11 ++- .../rows/arrowbased/timestamp_ntz_test.go | 53 ++++++++++++ internal/rows/rows.go | 9 ++ internal/rows/rowscanner/rowScanner.go | 13 +++ zz_probe_test.go | 3 + 14 files changed, 334 insertions(+), 18 deletions(-) create mode 100644 internal/arrowscan/coltype_ntz_test.go create mode 100644 internal/arrowscan/timestamp_ntz_test.go create mode 100644 internal/rows/arrowbased/timestamp_ntz_test.go create mode 100644 zz_probe_test.go diff --git a/doc.go b/doc.go index e3523d69..34445ce9 100644 --- a/doc.go +++ b/doc.go @@ -290,8 +290,9 @@ pool and connection lifecycle. The Go wrapper telemetry interceptor is skipped o kernel path so it does not duplicate kernel-owned telemetry for the same connection and statements; `enableTelemetry`, `telemetry_batch_size`, and `telemetry_flush_interval` are forwarded to kernel-owned telemetry config. Result types render byte-for-byte identical to the -Thrift backend: scalars, DECIMAL (exact string), TIMESTAMP / TIMESTAMP_NTZ (shifted -into the session time zone), INTERVAL, nested Array/Map/Struct and VARIANT (as JSON), +Thrift backend: scalars, DECIMAL (exact string), TIMESTAMP (an instant shifted +into the session time zone) and TIMESTAMP_NTZ (a zone-less wall clock kept +as-is), INTERVAL, nested Array/Map/Struct and VARIANT (as JSON), and GEOMETRY / GEOGRAPHY (WKT). The server query id is surfaced on the success path, so a QueryIdCallback (see below) fires with the real id and EXECUTE_STATEMENT telemetry carries it. diff --git a/driver_e2e_test.go b/driver_e2e_test.go index 4bdaeb84..e69aba96 100644 --- a/driver_e2e_test.go +++ b/driver_e2e_test.go @@ -709,3 +709,85 @@ func TestE2ECloudFetchExactRowCount(t *testing.T) { require.Equal(t, int64(wantRows), rowCount, "CloudFetch must surface exactly the requested rows, with no Arrow padding") } + +// TestE2ETimestampNTZTypeName asserts that a zoned TIMESTAMP and a TIMESTAMP_NTZ +// in the same result set report distinct database type names. The Databricks +// TIMESTAMP_NTZ type is a wall-clock value with no zone; it must be reported as +// "TIMESTAMP_NTZ", distinct from the zoned "TIMESTAMP" instant type. +// See https://github.com/databricks/databricks-sql-go/issues/506 and +// https://docs.databricks.com/aws/en/sql/language-manual/data-types/timestamp-ntz-type. +func TestE2ETimestampNTZTypeName(t *testing.T) { + host, httpPath, token := pecoTestingCreds(t) + + connector, err := NewConnector( + WithServerHostname(host), + WithPort(443), + WithHTTPPath(httpPath), + WithAccessToken(token), + ) + require.NoError(t, err) + + db := sql.OpenDB(connector) + defer db.Close() //nolint:errcheck + + rows, err := db.QueryContext(context.Background(), + "SELECT TIMESTAMP '2024-06-15 09:15:00.123456' AS zoned, "+ + "TIMESTAMP_NTZ '2024-06-15 09:15:00.123456' AS ntz") + require.NoError(t, err) + defer rows.Close() //nolint:errcheck + + cts, err := rows.ColumnTypes() + require.NoError(t, err) + require.Len(t, cts, 2) + + assert.Equal(t, "TIMESTAMP", cts[0].DatabaseTypeName(), "zoned column should report TIMESTAMP") + assert.Equal(t, "TIMESTAMP_NTZ", cts[1].DatabaseTypeName(), "zone-less column should report TIMESTAMP_NTZ") +} + +// TestE2ETimestampNTZPreservesWallClock asserts that at a non-UTC session +// timezone a TIMESTAMP_NTZ value keeps the exact wall clock it was given, with +// no location shift, while a zoned TIMESTAMP remains the correct instant. +// TIMESTAMP_NTZ carries no zone, so applying a session-tz conversion to it is a +// bug. See https://github.com/databricks/databricks-sql-go/issues/506. +func TestE2ETimestampNTZPreservesWallClock(t *testing.T) { + host, httpPath, token := pecoTestingCreds(t) + + connector, err := NewConnector( + WithServerHostname(host), + WithPort(443), + WithHTTPPath(httpPath), + WithAccessToken(token), + WithSessionParams(map[string]string{"timezone": "America/New_York"}), + ) + require.NoError(t, err) + + db := sql.OpenDB(connector) + defer db.Close() //nolint:errcheck + + rows, err := db.QueryContext(context.Background(), + "SELECT TIMESTAMP '2024-06-15 09:15:00.123456' AS zoned, "+ + "TIMESTAMP_NTZ '2024-06-15 09:15:00.123456' AS ntz") + require.NoError(t, err) + defer rows.Close() //nolint:errcheck + + require.True(t, rows.Next(), "expected one row") + var zoned, ntz time.Time + require.NoError(t, rows.Scan(&zoned, &ntz)) + + // The NTZ value must preserve the wall clock it was given, with no shift + // from the session timezone. We compare the wall-clock fields directly so + // the assertion is independent of whatever *time.Location the driver tags + // the zone-less value with. + assert.Equal(t, 2024, ntz.Year()) + assert.Equal(t, time.June, ntz.Month()) + assert.Equal(t, 15, ntz.Day()) + assert.Equal(t, 9, ntz.Hour(), "TIMESTAMP_NTZ wall-clock hour must not be shifted by the session timezone") + assert.Equal(t, 15, ntz.Minute()) + assert.Equal(t, 0, ntz.Second()) + assert.Equal(t, 123456000, ntz.Nanosecond()) + + // The zoned TIMESTAMP is an instant: the literal is interpreted in the + // session timezone (America/New_York), so its UTC instant is 13:15:00.123456Z. + assert.Equal(t, time.Date(2024, time.June, 15, 13, 15, 0, 123456000, time.UTC), zoned.UTC(), + "zoned TIMESTAMP must remain the correct instant") +} diff --git a/internal/arrowscan/arrowscan.go b/internal/arrowscan/arrowscan.go index c6705955..72c43454 100644 --- a/internal/arrowscan/arrowscan.go +++ b/internal/arrowscan/arrowscan.go @@ -213,7 +213,17 @@ func scanCell(col arrow.Array, row int, loc *time.Location, keys *StructKeyCache // UnixMilli/Unix split into (sec, nsec) internally, so the full 0001–9999 range // round-trips. The Thrift default path never hits this — it receives TIMESTAMP // as a preformatted server string — so this divergence is kernel-only. - return inLocation(timestampToTime(int64(c.Value(row)), dt.Unit), loc), nil + // + // A zone-less TIMESTAMP_NTZ (empty TimeZone) is a wall clock with no zone, + // so it must NOT be shifted into the session location the way a zoned + // TIMESTAMP instant is: applying .In(loc) would corrupt the wall-clock + // fields the server sent (databricks/databricks-sql-go#506). Only a zoned + // TIMESTAMP is rendered in loc. + ts := timestampToTime(int64(c.Value(row)), dt.Unit) + if dt.TimeZone == "" { + return ts, nil + } + return inLocation(ts, loc), nil case *array.Decimal128: dt := col.DataType().(*arrow.Decimal128Type) if options.DecimalAsFloat { diff --git a/internal/arrowscan/arrowscan_test.go b/internal/arrowscan/arrowscan_test.go index e68745eb..52fecbd3 100644 --- a/internal/arrowscan/arrowscan_test.go +++ b/internal/arrowscan/arrowscan_test.go @@ -234,9 +234,12 @@ func TestScanCellTimestampLocation(t *testing.T) { t.Skipf("tz database unavailable: %v", err) } - // 2026-07-09T12:00:00Z as microseconds since epoch. + // 2026-07-09T12:00:00Z as microseconds since epoch. A ZONED timestamp + // (non-empty TimeZone) is an instant that renders .In(location); a zone-less + // TIMESTAMP_NTZ deliberately does NOT shift (issue #506), so the location + // path must be exercised with a zoned type. utcTS := time.Date(2026, time.July, 9, 12, 0, 0, 0, time.UTC) - b := array.NewTimestampBuilder(pool, &arrow.TimestampType{Unit: arrow.Microsecond}) + b := array.NewTimestampBuilder(pool, &arrow.TimestampType{Unit: arrow.Microsecond, TimeZone: "UTC"}) defer b.Release() b.Append(arrow.Timestamp(utcTS.UnixMicro())) arr := b.NewArray() diff --git a/internal/arrowscan/coltype.go b/internal/arrowscan/coltype.go index d1e8a750..a77c876a 100644 --- a/internal/arrowscan/coltype.go +++ b/internal/arrowscan/coltype.go @@ -78,6 +78,14 @@ func ColumnTypeInfoFor(dt arrow.DataType) ColumnTypeInfo { case arrow.DATE32, arrow.DATE64: return ColumnTypeInfo{DatabaseTypeName: "DATE", ScanType: scanTypeDateTime} case arrow.TIMESTAMP: + // TIMESTAMP and TIMESTAMP_NTZ both arrive as Arrow timestamp[us] and are + // distinguishable only by the schema's TimeZone field: a zoned TIMESTAMP + // carries a zone, a zone-less TIMESTAMP_NTZ carries an empty TimeZone. The + // database/sql type name must reflect that distinction so a wall clock is + // not mistaken for an instant (databricks/databricks-sql-go#506). + if ts, ok := dt.(*arrow.TimestampType); ok && ts.TimeZone == "" { + return ColumnTypeInfo{DatabaseTypeName: "TIMESTAMP_NTZ", ScanType: scanTypeDateTime} + } return ColumnTypeInfo{DatabaseTypeName: "TIMESTAMP", ScanType: scanTypeDateTime} case arrow.DECIMAL128: // Match Thrift's sql.RawBytes scan type even though the kernel renders the diff --git a/internal/arrowscan/coltype_ntz_test.go b/internal/arrowscan/coltype_ntz_test.go new file mode 100644 index 00000000..3bd0f02a --- /dev/null +++ b/internal/arrowscan/coltype_ntz_test.go @@ -0,0 +1,33 @@ +package arrowscan + +import ( + "testing" + + "github.com/apache/arrow/go/v12/arrow" +) + +// TestColumnTypeInfoForTimestampNTZ pins the TIMESTAMP vs TIMESTAMP_NTZ type-name +// distinction. Over Arrow both arrive as timestamp[us], differing only in the +// schema's TimestampType.TimeZone field: a zoned TIMESTAMP carries a zone, a +// TIMESTAMP_NTZ carries an empty TimeZone. database/sql's +// ColumnTypeDatabaseTypeName must report them distinctly ("TIMESTAMP" vs +// "TIMESTAMP_NTZ") so a zone-less wall clock is not mistaken for an instant. +// +// External authority: the Databricks TIMESTAMP_NTZ type is a distinct SQL type +// (https://docs.databricks.com/aws/en/sql/language-manual/data-types/timestamp-ntz-type) +// and issue #506 requires it be surfaced as such. +func TestColumnTypeInfoForTimestampNTZ(t *testing.T) { + t.Run("zoned timestamp reports TIMESTAMP", func(t *testing.T) { + dt := &arrow.TimestampType{Unit: arrow.Microsecond, TimeZone: "UTC"} + if got := ColumnTypeInfoFor(dt).DatabaseTypeName; got != "TIMESTAMP" { + t.Errorf("DatabaseTypeName = %q, want %q", got, "TIMESTAMP") + } + }) + + t.Run("zone-less timestamp reports TIMESTAMP_NTZ", func(t *testing.T) { + dt := &arrow.TimestampType{Unit: arrow.Microsecond, TimeZone: ""} + if got := ColumnTypeInfoFor(dt).DatabaseTypeName; got != "TIMESTAMP_NTZ" { + t.Errorf("DatabaseTypeName = %q, want %q", got, "TIMESTAMP_NTZ") + } + }) +} diff --git a/internal/arrowscan/timestamp_ntz_test.go b/internal/arrowscan/timestamp_ntz_test.go new file mode 100644 index 00000000..3f807687 --- /dev/null +++ b/internal/arrowscan/timestamp_ntz_test.go @@ -0,0 +1,67 @@ +package arrowscan + +import ( + "testing" + "time" + + "github.com/apache/arrow/go/v12/arrow" + "github.com/apache/arrow/go/v12/arrow/array" + "github.com/apache/arrow/go/v12/arrow/memory" +) + +// TestScanCellTimestampNTZPreservesWallClock pins the value contract for +// TIMESTAMP_NTZ: a zone-less timestamp (empty TimestampType.TimeZone) is a wall +// clock with no zone, so it must be returned with the exact wall-clock fields it +// was given, WITHOUT the session-timezone .In(loc) conversion that is correct for +// a zoned TIMESTAMP instant. Applying the location shift to an NTZ value (the +// pre-#506 behavior) corrupts the wall clock. +// +// External authority: Databricks TIMESTAMP_NTZ +// (https://docs.databricks.com/aws/en/sql/language-manual/data-types/timestamp-ntz-type) +// "does not ... apply a time zone"; issue #506 requires the wall clock be kept. +func TestScanCellTimestampNTZPreservesWallClock(t *testing.T) { + pool := memory.NewGoAllocator() + loc, err := time.LoadLocation("America/New_York") + if err != nil { + t.Skipf("tz database unavailable: %v", err) + } + + // The wall clock the TIMESTAMP_NTZ literal carries. On the wire the server + // encodes it as microseconds-since-epoch interpreted as if UTC. + wall := time.Date(2024, time.June, 15, 9, 15, 0, 123456000, time.UTC) + + build := func(tz string) arrow.Array { + b := array.NewTimestampBuilder(pool, &arrow.TimestampType{Unit: arrow.Microsecond, TimeZone: tz}) + b.Append(arrow.Timestamp(wall.UnixMicro())) + return b.NewArray() + } + + t.Run("ntz empty zone keeps wall clock", func(t *testing.T) { + arr := build("") + defer arr.Release() + v, err := ScanCell(arr, 0, loc) + if err != nil { + t.Fatal(err) + } + got := v.(time.Time) + if got.Year() != 2024 || got.Month() != time.June || got.Day() != 15 || + got.Hour() != 9 || got.Minute() != 15 || got.Second() != 0 || got.Nanosecond() != 123456000 { + t.Errorf("NTZ wall clock not preserved: got %v, want wall 2024-06-15 09:15:00.123456", got) + } + }) + + t.Run("zoned keeps instant via location", func(t *testing.T) { + arr := build("UTC") + defer arr.Release() + v, err := ScanCell(arr, 0, loc) + if err != nil { + t.Fatal(err) + } + got := v.(time.Time) + // A zoned TIMESTAMP is an instant: rendered in loc, it is the same + // instant as the UTC wall value. + if !got.Equal(wall) { + t.Errorf("zoned instant changed: got %v, want %v", got.UTC(), wall) + } + }) +} diff --git a/internal/rows/arrowbased/arrowRows.go b/internal/rows/arrowbased/arrowRows.go index 58970d85..8d6f3d7f 100644 --- a/internal/rows/arrowbased/arrowRows.go +++ b/internal/rows/arrowbased/arrowRows.go @@ -153,6 +153,27 @@ func NewArrowRowScanner(resultSetMetadata *cli_service.TGetResultSetMetadataResp return rs, nil } +// DatabaseTypeNameOverride implements rowscanner.DatabaseTypeNameOverrider. Over +// Thrift, a zoned TIMESTAMP and a zone-less TIMESTAMP_NTZ share the same +// primitive type id (TIMESTAMP_TYPE), so GetDBTypeName reports "TIMESTAMP" for +// both. The distinction survives only in the Arrow schema's +// TimestampType.TimeZone: a zone-less column is TIMESTAMP_NTZ +// (databricks/databricks-sql-go#506). Non-timestamp columns and zoned +// timestamps fall through to the Thrift-derived name. +func (ars *arrowRowScanner) DatabaseTypeNameOverride(index int) (string, bool) { + if ars == nil || ars.arrowSchema == nil { + return "", false + } + fields := ars.arrowSchema.Fields() + if index < 0 || index >= len(fields) { + return "", false + } + if ts, ok := fields[index].Type.(*arrow.TimestampType); ok && ts.TimeZone == "" { + return "TIMESTAMP_NTZ", true + } + return "", false +} + // Close is called when the Rows instance is closed. func (ars *arrowRowScanner) Close() { if ars.rowValues != nil { @@ -649,7 +670,10 @@ func (vcm *arrowValueContainerMaker) makeColumnValueContainer(t arrow.DataType, return &dateValueContainer{location: location}, nil case *arrow.TimestampType: - return ×tampValueContainer{location: location, toTimestampFn: toTimestampFn}, nil + // A zone-less TIMESTAMP_NTZ (empty TimeZone) carries a wall clock that + // must not be shifted into the session location (#506); a zoned + // TIMESTAMP is an instant rendered .In(location) as before. + return ×tampValueContainer{location: location, toTimestampFn: toTimestampFn, wallClock: t.TimeZone == ""}, nil case *arrow.BinaryType: return &columnValuesTyped[*array.Binary, []byte]{}, nil diff --git a/internal/rows/arrowbased/arrowscan_parity_test.go b/internal/rows/arrowbased/arrowscan_parity_test.go index a42856d2..aeaaab30 100644 --- a/internal/rows/arrowbased/arrowscan_parity_test.go +++ b/internal/rows/arrowbased/arrowscan_parity_test.go @@ -402,18 +402,19 @@ func TestArrowbasedKernelTopLevelScalarParity(t *testing.T) { } } -// TestArrowbasedKernelTimestampTZParity pins the TIMESTAMP vs TIMESTAMP_NTZ -// rendering contract: over Arrow, a TIMESTAMP arrives with TimeZone "UTC" and a -// TIMESTAMP_NTZ with an empty TimeZone (kernel json.rs:171 — "TIMESTAMP carries a -// tz on the Arrow side; TIMESTAMP_NTZ does not"). Both backends deliberately IGNORE -// that tz field: each renders via ToTime + .In(loc), so the LTZ-vs-NTZ difference is -// carried entirely by the instant value the server sends, NOT by the client -// inspecting the tz. Verified live on both backends (NY + Kolkata, incl. a DST-gap -// literal): kernel == Thrift byte-for-byte for both types. +// TestArrowbasedKernelTimestampTZParity pins the CROSS-BACKEND parity of the +// TIMESTAMP vs TIMESTAMP_NTZ rendering contract: over Arrow, a TIMESTAMP arrives +// with TimeZone "UTC" and a TIMESTAMP_NTZ with an empty TimeZone (kernel +// json.rs:171 — "TIMESTAMP carries a tz on the Arrow side; TIMESTAMP_NTZ does +// not"). Per the TIMESTAMP_NTZ spec and issue #506, a zone-less column is a wall +// clock and must NOT be shifted into the session location, while a zoned +// TIMESTAMP is an instant rendered via ToTime + .In(loc). Both backends now read +// the tz field and agree on that split, so the kernel and Thrift paths stay +// byte-for-byte identical for both types. // -// This test locks that in so a future "semantically-correct" change that skips -// .In(loc) for TimeZone=="" (which would look right in isolation) fails CI — it -// would make the kernel diverge from the Thrift path, which shifts NTZ too. +// This test locks the two backends together: a change that alters the NTZ / +// zoned handling on only one of them (making the kernel diverge from the Thrift +// path, or vice versa) fails CI. func TestArrowbasedKernelTimestampTZParity(t *testing.T) { pool := memory.NewGoAllocator() loc, err := time.LoadLocation("America/New_York") diff --git a/internal/rows/arrowbased/columnValues.go b/internal/rows/arrowbased/columnValues.go index 73aeb01f..f525478c 100644 --- a/internal/rows/arrowbased/columnValues.go +++ b/internal/rows/arrowbased/columnValues.go @@ -477,13 +477,22 @@ type timestampValueContainer struct { timeArray *array.Timestamp location *time.Location toTimestampFn func(arrow.Timestamp) time.Time + // wallClock is true for a zone-less TIMESTAMP_NTZ column (empty Arrow + // TimeZone): its value is a wall clock with no zone and must be returned + // with the exact fields the server sent, WITHOUT the session .In(location) + // shift that is correct only for a zoned TIMESTAMP instant + // (databricks/databricks-sql-go#506). + wallClock bool } var _ columnValues = (*timestampValueContainer)(nil) func (tvc *timestampValueContainer) Value(i int) (any, error) { ats := tvc.timeArray.Value(i) - val := tvc.toTimestampFn(ats).In(tvc.location) + val := tvc.toTimestampFn(ats) + if !tvc.wallClock { + val = val.In(tvc.location) + } return val, nil } diff --git a/internal/rows/arrowbased/timestamp_ntz_test.go b/internal/rows/arrowbased/timestamp_ntz_test.go new file mode 100644 index 00000000..164123f7 --- /dev/null +++ b/internal/rows/arrowbased/timestamp_ntz_test.go @@ -0,0 +1,53 @@ +package arrowbased + +import ( + "testing" + "time" + + "github.com/apache/arrow/go/v12/arrow" + "github.com/apache/arrow/go/v12/arrow/array" + "github.com/apache/arrow/go/v12/arrow/memory" +) + +// TestThriftTimestampNTZPreservesWallClock pins the Thrift (arrowbased) backend's +// value contract for TIMESTAMP_NTZ: a zone-less timestamp (empty +// TimestampType.TimeZone) is a wall clock and must be returned with the exact +// wall-clock fields it was given, WITHOUT the session-tz .In(location) shift that +// is correct only for a zoned TIMESTAMP instant. +// +// External authority: Databricks TIMESTAMP_NTZ +// (https://docs.databricks.com/aws/en/sql/language-manual/data-types/timestamp-ntz-type) +// and issue #506. +func TestThriftTimestampNTZPreservesWallClock(t *testing.T) { + pool := memory.NewGoAllocator() + loc, err := time.LoadLocation("America/New_York") + if err != nil { + t.Skipf("tz database unavailable: %v", err) + } + + wall := time.Date(2024, time.June, 15, 9, 15, 0, 123456000, time.UTC) + build := func(tz string) arrow.Array { + b := array.NewTimestampBuilder(pool, &arrow.TimestampType{Unit: arrow.Microsecond, TimeZone: tz}) + b.Append(arrow.Timestamp(wall.UnixMicro())) + return b.NewArray() + } + + t.Run("ntz empty zone keeps wall clock", func(t *testing.T) { + arr := build("") + defer arr.Release() + got := renderViaArrowbased(t, arr, 0, loc).(time.Time) + if got.Year() != 2024 || got.Month() != time.June || got.Day() != 15 || + got.Hour() != 9 || got.Minute() != 15 || got.Second() != 0 || got.Nanosecond() != 123456000 { + t.Errorf("NTZ wall clock not preserved: got %v, want wall 2024-06-15 09:15:00.123456", got) + } + }) + + t.Run("zoned keeps instant via location", func(t *testing.T) { + arr := build("UTC") + defer arr.Release() + got := renderViaArrowbased(t, arr, 0, loc).(time.Time) + if !got.Equal(wall) { + t.Errorf("zoned instant changed: got %v, want %v", got.UTC(), wall) + } + }) +} diff --git a/internal/rows/rows.go b/internal/rows/rows.go index 1ba3bb7c..8ddb33f9 100644 --- a/internal/rows/rows.go +++ b/internal/rows/rows.go @@ -364,6 +364,15 @@ func (r *rows) ColumnTypeDatabaseTypeName(index int) string { return "" } + // The Thrift primitive type id cannot tell a zoned TIMESTAMP from a + // zone-less TIMESTAMP_NTZ; that distinction lives in the Arrow schema, which + // the arrow row scanner can expose via DatabaseTypeNameOverrider (#506). + if o, ok := r.RowScanner.(rowscanner.DatabaseTypeNameOverrider); ok && o != nil { + if name, ok := o.DatabaseTypeNameOverride(index); ok { + return name + } + } + dbtype := rowscanner.GetDBTypeName(column) return dbtype diff --git a/internal/rows/rowscanner/rowScanner.go b/internal/rows/rowscanner/rowScanner.go index ebb25bf6..dd332be4 100644 --- a/internal/rows/rowscanner/rowScanner.go +++ b/internal/rows/rowscanner/rowScanner.go @@ -35,6 +35,19 @@ type RowScanner interface { GetArrowIPCStreams(ctx context.Context, cfg config.Config, rpi ResultPageIterator) (dbsqlrows.ArrowIPCStreamIterator, error) } +// DatabaseTypeNameOverrider is an optional interface a RowScanner may implement +// to refine the database type name reported for a column beyond what the Thrift +// TColumnDesc carries. It exists for TIMESTAMP_NTZ: the Thrift primitive type id +// is TIMESTAMP_TYPE for both zoned TIMESTAMP and zone-less TIMESTAMP_NTZ, so the +// distinction survives only in the Arrow schema's TimestampType.TimeZone field, +// which only the arrow row scanner can see (databricks/databricks-sql-go#506). +type DatabaseTypeNameOverrider interface { + // DatabaseTypeNameOverride returns the refined type name for the column at + // index and ok=true when it overrides the Thrift-derived name; ok=false + // means the caller should fall back to GetDBTypeName. + DatabaseTypeNameOverride(index int) (name string, ok bool) +} + // Expected formats for TIMESTAMP and DATE types when represented by a string value var DateTimeFormats map[string]string = map[string]string{ "TIMESTAMP": "2006-01-02 15:04:05.999999999", diff --git a/zz_probe_test.go b/zz_probe_test.go new file mode 100644 index 00000000..c66404fc --- /dev/null +++ b/zz_probe_test.go @@ -0,0 +1,3 @@ +//go:build ignore_probe_scratch + +package dbsql From f4039523fbfd4dcb874a2d4c102f26b87c47f532 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 22 Oct 2026 21:21:14 +0000 Subject: [PATCH 02/15] chore: remove unused probe test file Signed-off-by: Vu Anh Phung --- zz_probe_test.go | 3 --- 1 file changed, 3 deletions(-) delete mode 100644 zz_probe_test.go diff --git a/zz_probe_test.go b/zz_probe_test.go deleted file mode 100644 index c66404fc..00000000 --- a/zz_probe_test.go +++ /dev/null @@ -1,3 +0,0 @@ -//go:build ignore_probe_scratch - -package dbsql From 9079b5704d15cfe6489095a2f52df873b2a14a5a Mon Sep 17 00:00:00 2001 From: "peco-engineer-bot[bot]" Date: Thu, 22 Oct 2026 21:34:08 +0000 Subject: [PATCH 03/15] ai: apply changes for #510 (4 review threads) Addresses: - #4353704768 at internal/arrowscan/arrowscan.go:223 - #4353704778 at doc.go:293 - #4353704786 at internal/rows/arrowbased/arrowRows.go:163 - #4353704797 at internal/rows/rows.go:370 Signed-off-by: peco-engineer-bot[bot] --- CHANGELOG.md | 3 ++ CONNECTION_PARAMETERS.md | 3 +- README.md | 5 +- .../rows/arrowbased/timestamp_ntz_test.go | 39 ++++++++++++++ internal/rows/rows_test.go | 51 +++++++++++++++++++ internal/rows/rowscanner/rowScanner.go | 6 +++ kernel_e2e_test.go | 30 ++++++----- 7 files changed, 120 insertions(+), 17 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index dbae70fa..b7a7cb0e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,8 @@ # Release History +## Unreleased +- **Behavior change:** TIMESTAMP_NTZ values are no longer shifted into the session time zone (databricks/databricks-sql-go#506). A zone-less TIMESTAMP_NTZ is now returned as the exact wall clock stored, with `Location()` == `time.UTC`, on both backends (including the default `UseArrowNativeTimestamp=true` Arrow path). Previously both the instant and the location changed with the session time zone; callers that compensated for the old shift (e.g. by calling `.UTC()` or re-interpreting the value in the session location) must drop that workaround. Zoned TIMESTAMP values are unchanged and still render in the session time zone. + ## v1.17.0 (2026-10-20) - Upgrade the SQL kernel bindings to v1.2.0. - Add opt-in GEOMETRY / GEOGRAPHY support on the kernel backend with `WithEnableGeoSpatialSupport` or `enableGeoSpatialSupport=true`. Enabled mode returns EWKB bytes with per-row SRIDs and geo column metadata; the default remains WKT / EWKT strings. Arrow batches retain the `{srid,wkb}` struct in enabled mode. diff --git a/CONNECTION_PARAMETERS.md b/CONNECTION_PARAMETERS.md index 422eaf5e..80df2502 100644 --- a/CONNECTION_PARAMETERS.md +++ b/CONNECTION_PARAMETERS.md @@ -159,7 +159,8 @@ ignored with a warning and still stripped from the SEA wire. Apart from the geospatial representation selected above, results render **byte-for-byte identically** on both backends (scalars, -DECIMAL as exact string, TIMESTAMP / TIMESTAMP_NTZ shifted into the session time zone, +DECIMAL as exact string, TIMESTAMP as an instant shifted into the session time zone, +TIMESTAMP_NTZ as a zone-less wall clock kept as-is (`Location()` == `time.UTC`), INTERVAL, nested ARRAY / MAP / STRUCT and VARIANT as JSON, and BINARY as `sql.RawBytes`). The default `enableGeoSpatialSupport=false` on the kernel path matches the Thrift path's WKT / EWKT geospatial values. diff --git a/README.md b/README.md index d59054af..110eef25 100644 --- a/README.md +++ b/README.md @@ -360,8 +360,9 @@ The Thrift and kernel backends both honor the standard `HTTP_PROXY` / `HTTPS_PRO ## Data types Apart from the configurable geospatial representation, results render **byte-for-byte -identically** on both backends. Scalars, DECIMAL (exact string), TIMESTAMP / -TIMESTAMP_NTZ (shifted into the session time zone), INTERVAL, and nested ARRAY / MAP / +identically** on both backends. Scalars, DECIMAL (exact string), TIMESTAMP (an +instant shifted into the session time zone), TIMESTAMP_NTZ (a zone-less wall clock kept +as-is, returned with `Location()` == `time.UTC`), INTERVAL, and nested ARRAY / MAP / STRUCT and VARIANT (as JSON) are supported. BINARY is returned as `sql.RawBytes`. Both backends return GEOMETRY / GEOGRAPHY as WKT / EWKT by default. Set `enableGeoSpatialSupport=true` on the kernel path to opt into EWKB. diff --git a/internal/rows/arrowbased/timestamp_ntz_test.go b/internal/rows/arrowbased/timestamp_ntz_test.go index 164123f7..95683be3 100644 --- a/internal/rows/arrowbased/timestamp_ntz_test.go +++ b/internal/rows/arrowbased/timestamp_ntz_test.go @@ -51,3 +51,42 @@ func TestThriftTimestampNTZPreservesWallClock(t *testing.T) { } }) } + +// TestArrowRowScannerDatabaseTypeNameOverride pins the Thrift-path type-name +// refinement (#506): only a zone-less Arrow timestamp column is reported as +// TIMESTAMP_NTZ; zoned timestamps, non-timestamp columns, out-of-range indexes +// and a missing schema all fall through (ok=false) to the Thrift-derived name. +func TestArrowRowScannerDatabaseTypeNameOverride(t *testing.T) { + schema := arrow.NewSchema([]arrow.Field{ + {Name: "ts", Type: &arrow.TimestampType{Unit: arrow.Microsecond, TimeZone: "UTC"}}, + {Name: "ts_ntz", Type: &arrow.TimestampType{Unit: arrow.Microsecond}}, + {Name: "i", Type: arrow.PrimitiveTypes.Int32}, + }, nil) + ars := &arrowRowScanner{arrowSchema: schema} + + cases := []struct { + index int + wantName string + wantOK bool + }{ + {0, "", false}, + {1, "TIMESTAMP_NTZ", true}, + {2, "", false}, + {-1, "", false}, + {3, "", false}, + } + for _, c := range cases { + name, ok := ars.DatabaseTypeNameOverride(c.index) + if name != c.wantName || ok != c.wantOK { + t.Errorf("index %d: got (%q, %v), want (%q, %v)", c.index, name, ok, c.wantName, c.wantOK) + } + } + + if name, ok := (&arrowRowScanner{}).DatabaseTypeNameOverride(0); name != "" || ok { + t.Errorf("nil schema: got (%q, %v), want (\"\", false)", name, ok) + } + var nilScanner *arrowRowScanner + if name, ok := nilScanner.DatabaseTypeNameOverride(0); name != "" || ok { + t.Errorf("nil scanner: got (%q, %v), want (\"\", false)", name, ok) + } +} diff --git a/internal/rows/rows_test.go b/internal/rows/rows_test.go index e04f9ef5..100b694c 100644 --- a/internal/rows/rows_test.go +++ b/internal/rows/rows_test.go @@ -708,6 +708,57 @@ func TestColumnTypeDatabaseTypeName(t *testing.T) { assert.Equal(t, expectedScanTypes, scanTypes) } +// typeNameOverrideScanner is a RowScanner stub that only implements +// rowscanner.DatabaseTypeNameOverrider; other methods are never called. +type typeNameOverrideScanner struct { + rowscanner.RowScanner + overrides map[int]string +} + +func (s *typeNameOverrideScanner) DatabaseTypeNameOverride(index int) (string, bool) { + name, ok := s.overrides[index] + return name, ok +} + +// TestColumnTypeDatabaseTypeNameOverride checks that ColumnTypeDatabaseTypeName +// consults a DatabaseTypeNameOverrider row scanner (#506): Thrift reports +// TIMESTAMP_TYPE for both zoned and zone-less timestamps, so the override is +// what surfaces TIMESTAMP_NTZ, and columns without an override keep the +// Thrift-derived name. +func TestColumnTypeDatabaseTypeNameOverride(t *testing.T) { + var getMetadataCount, fetchResultsCount int + + col := func(name string, typ cli_service.TTypeId) *cli_service.TColumnDesc { + return &cli_service.TColumnDesc{ + ColumnName: name, + TypeDesc: &cli_service.TTypeDesc{ + Types: []*cli_service.TTypeEntry{ + {PrimitiveEntry: &cli_service.TPrimitiveTypeEntry{Type: typ}}, + }, + }, + } + } + + rowSet := &rows{ + client: getRowsTestSimpleClient(&getMetadataCount, &fetchResultsCount), + schema: &cli_service.TTableSchema{Columns: []*cli_service.TColumnDesc{ + col("ts", cli_service.TTypeId_TIMESTAMP_TYPE), + col("ts_ntz", cli_service.TTypeId_TIMESTAMP_TYPE), + col("i", cli_service.TTypeId_INT_TYPE), + }}, + RowScanner: &typeNameOverrideScanner{overrides: map[int]string{1: "TIMESTAMP_NTZ"}}, + } + + assert.Equal(t, "TIMESTAMP", rowSet.ColumnTypeDatabaseTypeName(0)) + assert.Equal(t, "TIMESTAMP_NTZ", rowSet.ColumnTypeDatabaseTypeName(1)) + assert.Equal(t, "INT", rowSet.ColumnTypeDatabaseTypeName(2)) + assert.Equal(t, "", rowSet.ColumnTypeDatabaseTypeName(3)) + + // Without an overrider (or with no row scanner yet) the Thrift name is used. + rowSet.RowScanner = nil + assert.Equal(t, "TIMESTAMP", rowSet.ColumnTypeDatabaseTypeName(1)) +} + func TestRowsCloseOptimization(t *testing.T) { var closeCount int client := &client.TestClient{ diff --git a/internal/rows/rowscanner/rowScanner.go b/internal/rows/rowscanner/rowScanner.go index dd332be4..d1a75091 100644 --- a/internal/rows/rowscanner/rowScanner.go +++ b/internal/rows/rowscanner/rowScanner.go @@ -41,6 +41,12 @@ type RowScanner interface { // is TIMESTAMP_TYPE for both zoned TIMESTAMP and zone-less TIMESTAMP_NTZ, so the // distinction survives only in the Arrow schema's TimestampType.TimeZone field, // which only the arrow row scanner can see (databricks/databricks-sql-go#506). +// +// Limitation: the override applies only when the Arrow schema carries a native +// *arrow.TimestampType. With UseArrowNativeTimestamp=false (timestamps sent as +// strings) or with the column-based (non-Arrow) row scanner, there is no zone +// information to inspect, so a TIMESTAMP_NTZ column still reports "TIMESTAMP" +// and its string value is parsed in the session location. type DatabaseTypeNameOverrider interface { // DatabaseTypeNameOverride returns the refined type name for the column at // index and ok=true when it overrides the Thrift-derived name; ok=false diff --git a/kernel_e2e_test.go b/kernel_e2e_test.go index 18eea660..b1fb77b8 100644 --- a/kernel_e2e_test.go +++ b/kernel_e2e_test.go @@ -123,14 +123,13 @@ func TestKernelE2ETimeZone(t *testing.T) { } } -// TestKernelE2ETimestampNTZ proves TIMESTAMP and TIMESTAMP_NTZ round-trip -// byte-identically to the Thrift path at a non-UTC session tz. The kernel delivers -// TIMESTAMP with an Arrow tz and TIMESTAMP_NTZ without one, but — like the Thrift -// path — the driver ignores that field and renders via .In(loc). So the naive -// wall-clock of an NTZ literal is treated as a UTC instant and shifted into the -// session tz (12:00 NTZ in America/New_York → the -04:00 local of 12:00Z = 08:00), -// exactly as Thrift does. Verified live on both backends; this pins the round-trip -// so a one-sided "don't shift NTZ" change is caught. +// TestKernelE2ETimestampNTZ pins how TIMESTAMP and TIMESTAMP_NTZ scan at a +// non-UTC session tz. The kernel delivers TIMESTAMP with an Arrow tz and +// TIMESTAMP_NTZ without one, and the driver honors that field: a zoned TIMESTAMP +// is rendered via .In(loc), while a zone-less TIMESTAMP_NTZ is a wall clock and is +// NOT shifted into the session tz (databricks/databricks-sql-go#506). So 12:00 NTZ +// in America/New_York comes back as the 12:00 wall clock, tagged UTC — not the +// -04:00 local of 12:00Z (08:00) that the old .In(loc) behavior produced. func TestKernelE2ETimestampNTZ(t *testing.T) { // Reyden shifts a zoned TIMESTAMP literal like an NTZ (12:00 → 08:00 -0400) where a // normal warehouse keeps the 12:00 wall-clock — an engine timestamp-tz divergence, @@ -161,13 +160,16 @@ func TestKernelE2ETimestampNTZ(t *testing.T) { "SELECT CAST('2026-07-09 12:00:00' AS TIMESTAMP_NTZ)").Scan(&ntz); err != nil { t.Fatalf("TIMESTAMP_NTZ query: %v", err) } - // The NTZ wall-clock is treated as a UTC instant then shifted into the session - // tz: 12:00Z == 08:00 in America/New_York (-04:00 in July). - if want := time.Date(2026, 7, 9, 8, 0, 0, 0, loc); !ntz.Equal(want) { - t.Errorf("TIMESTAMP_NTZ = %s, want %s (UTC instant %s)", ntz, want, ntz.UTC()) + // The NTZ wall clock is preserved as sent (12:00) and NOT shifted into the + // session tz; it carries no zone, so it is tagged UTC. + if want := time.Date(2026, 7, 9, 12, 0, 0, 0, time.UTC); !ntz.Equal(want) { + t.Errorf("TIMESTAMP_NTZ = %s, want %s", ntz, want) } - if ntz.Location().String() != tz { - t.Errorf("TIMESTAMP_NTZ location = %q, want %q", ntz.Location(), tz) + if h, m := ntz.Hour(), ntz.Minute(); h != 12 || m != 0 { + t.Errorf("TIMESTAMP_NTZ wall clock = %02d:%02d, want 12:00", h, m) + } + if ntz.Location() != time.UTC { + t.Errorf("TIMESTAMP_NTZ location = %q, want %q", ntz.Location(), time.UTC) } } From f1269600109d529a3d41823603d0045001b9ec2b Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 22 Oct 2026 21:38:05 +0000 Subject: [PATCH 04/15] test: tighten TIMESTAMP_NTZ coverage and docs Signed-off-by: Vu Anh Phung --- CHANGELOG.md | 3 +- CONNECTION_PARAMETERS.md | 4 +- README.md | 6 +-- doc.go | 6 +-- internal/arrowscan/arrowscan.go | 3 +- internal/arrowscan/timestamp_ntz_test.go | 6 +++ .../rows/arrowbased/timestamp_ntz_test.go | 6 +++ internal/rows/rows_test.go | 48 +++++++++++++++++++ internal/rows/rowscanner/rowScanner.go | 16 ++----- kernel_e2e_test.go | 20 ++------ 10 files changed, 80 insertions(+), 38 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b7a7cb0e..071795ef 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,7 +1,8 @@ # Release History ## Unreleased -- **Behavior change:** TIMESTAMP_NTZ values are no longer shifted into the session time zone (databricks/databricks-sql-go#506). A zone-less TIMESTAMP_NTZ is now returned as the exact wall clock stored, with `Location()` == `time.UTC`, on both backends (including the default `UseArrowNativeTimestamp=true` Arrow path). Previously both the instant and the location changed with the session time zone; callers that compensated for the old shift (e.g. by calling `.UTC()` or re-interpreting the value in the session location) must drop that workaround. Zoned TIMESTAMP values are unchanged and still render in the session time zone. + +- **Behavior change:** Native Arrow TIMESTAMP_NTZ results now preserve their wall-clock fields and report `TIMESTAMP_NTZ` on both backends. Returned `time.Time` values use UTC instead of the session location; their represented instant and `.UTC()` value are unchanged (databricks/databricks-sql-go#506). ## v1.17.0 (2026-10-20) - Upgrade the SQL kernel bindings to v1.2.0. diff --git a/CONNECTION_PARAMETERS.md b/CONNECTION_PARAMETERS.md index 80df2502..4be1cfe9 100644 --- a/CONNECTION_PARAMETERS.md +++ b/CONNECTION_PARAMETERS.md @@ -159,8 +159,8 @@ ignored with a warning and still stripped from the SEA wire. Apart from the geospatial representation selected above, results render **byte-for-byte identically** on both backends (scalars, -DECIMAL as exact string, TIMESTAMP as an instant shifted into the session time zone, -TIMESTAMP_NTZ as a zone-less wall clock kept as-is (`Location()` == `time.UTC`), +DECIMAL as exact string, TIMESTAMP rendered in the session time zone, +TIMESTAMP_NTZ with wall-clock fields preserved in a `time.Time` with a UTC location, INTERVAL, nested ARRAY / MAP / STRUCT and VARIANT as JSON, and BINARY as `sql.RawBytes`). The default `enableGeoSpatialSupport=false` on the kernel path matches the Thrift path's WKT / EWKT geospatial values. diff --git a/README.md b/README.md index 110eef25..11a594ea 100644 --- a/README.md +++ b/README.md @@ -360,9 +360,9 @@ The Thrift and kernel backends both honor the standard `HTTP_PROXY` / `HTTPS_PRO ## Data types Apart from the configurable geospatial representation, results render **byte-for-byte -identically** on both backends. Scalars, DECIMAL (exact string), TIMESTAMP (an -instant shifted into the session time zone), TIMESTAMP_NTZ (a zone-less wall clock kept -as-is, returned with `Location()` == `time.UTC`), INTERVAL, and nested ARRAY / MAP / +identically** on both backends. Scalars, DECIMAL (exact string), TIMESTAMP +(rendered in the session time zone), TIMESTAMP_NTZ (wall-clock fields preserved +in a `time.Time` with a UTC location), INTERVAL, and nested ARRAY / MAP / STRUCT and VARIANT (as JSON) are supported. BINARY is returned as `sql.RawBytes`. Both backends return GEOMETRY / GEOGRAPHY as WKT / EWKT by default. Set `enableGeoSpatialSupport=true` on the kernel path to opt into EWKB. diff --git a/doc.go b/doc.go index 34445ce9..2ff02ca2 100644 --- a/doc.go +++ b/doc.go @@ -290,9 +290,9 @@ pool and connection lifecycle. The Go wrapper telemetry interceptor is skipped o kernel path so it does not duplicate kernel-owned telemetry for the same connection and statements; `enableTelemetry`, `telemetry_batch_size`, and `telemetry_flush_interval` are forwarded to kernel-owned telemetry config. Result types render byte-for-byte identical to the -Thrift backend: scalars, DECIMAL (exact string), TIMESTAMP (an instant shifted -into the session time zone) and TIMESTAMP_NTZ (a zone-less wall clock kept -as-is), INTERVAL, nested Array/Map/Struct and VARIANT (as JSON), +Thrift backend: scalars, DECIMAL (exact string), TIMESTAMP (rendered in the +session time zone) and TIMESTAMP_NTZ (wall-clock fields preserved in a time.Time +with a UTC location), INTERVAL, nested Array/Map/Struct and VARIANT (as JSON), and GEOMETRY / GEOGRAPHY (WKT). The server query id is surfaced on the success path, so a QueryIdCallback (see below) fires with the real id and EXECUTE_STATEMENT telemetry carries it. diff --git a/internal/arrowscan/arrowscan.go b/internal/arrowscan/arrowscan.go index 72c43454..b996419a 100644 --- a/internal/arrowscan/arrowscan.go +++ b/internal/arrowscan/arrowscan.go @@ -211,8 +211,7 @@ func scanCell(col arrow.Array, row int, loc *time.Location, keys *StructKeyCache // which overflows for TIMESTAMPs outside ~1678–2262 and silently wraps a valid // instant to a wrong one (e.g. TIMESTAMP '0001-01-01' → 1754). time.UnixMicro/ // UnixMilli/Unix split into (sec, nsec) internally, so the full 0001–9999 range - // round-trips. The Thrift default path never hits this — it receives TIMESTAMP - // as a preformatted server string — so this divergence is kernel-only. + // round-trips. The Thrift renderer still uses Arrow's ToTime conversion. // // A zone-less TIMESTAMP_NTZ (empty TimeZone) is a wall clock with no zone, // so it must NOT be shifted into the session location the way a zoned diff --git a/internal/arrowscan/timestamp_ntz_test.go b/internal/arrowscan/timestamp_ntz_test.go index 3f807687..c232f3cc 100644 --- a/internal/arrowscan/timestamp_ntz_test.go +++ b/internal/arrowscan/timestamp_ntz_test.go @@ -48,6 +48,12 @@ func TestScanCellTimestampNTZPreservesWallClock(t *testing.T) { got.Hour() != 9 || got.Minute() != 15 || got.Second() != 0 || got.Nanosecond() != 123456000 { t.Errorf("NTZ wall clock not preserved: got %v, want wall 2024-06-15 09:15:00.123456", got) } + if got.Location() != time.UTC { + t.Errorf("NTZ location = %q, want UTC", got.Location()) + } + if !got.Equal(wall) { + t.Errorf("NTZ instant changed: got %v, want %v", got.UTC(), wall) + } }) t.Run("zoned keeps instant via location", func(t *testing.T) { diff --git a/internal/rows/arrowbased/timestamp_ntz_test.go b/internal/rows/arrowbased/timestamp_ntz_test.go index 95683be3..9878c757 100644 --- a/internal/rows/arrowbased/timestamp_ntz_test.go +++ b/internal/rows/arrowbased/timestamp_ntz_test.go @@ -40,6 +40,12 @@ func TestThriftTimestampNTZPreservesWallClock(t *testing.T) { got.Hour() != 9 || got.Minute() != 15 || got.Second() != 0 || got.Nanosecond() != 123456000 { t.Errorf("NTZ wall clock not preserved: got %v, want wall 2024-06-15 09:15:00.123456", got) } + if got.Location() != time.UTC { + t.Errorf("NTZ location = %q, want UTC", got.Location()) + } + if !got.Equal(wall) { + t.Errorf("NTZ instant changed: got %v, want %v", got.UTC(), wall) + } }) t.Run("zoned keeps instant via location", func(t *testing.T) { diff --git a/internal/rows/rows_test.go b/internal/rows/rows_test.go index 100b694c..e5ab93f9 100644 --- a/internal/rows/rows_test.go +++ b/internal/rows/rows_test.go @@ -1,6 +1,7 @@ package rows import ( + "bytes" "context" "database/sql/driver" "encoding/json" @@ -13,6 +14,8 @@ import ( "testing" "time" + "github.com/apache/arrow/go/v12/arrow" + "github.com/apache/arrow/go/v12/arrow/ipc" "github.com/databricks/databricks-sql-go/driverctx" "github.com/databricks/databricks-sql-go/internal/cli_service" "github.com/databricks/databricks-sql-go/internal/client" @@ -21,6 +24,7 @@ import ( dbsqlrows "github.com/databricks/databricks-sql-go/rows" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) func TestRowsNextRowInPage(t *testing.T) { @@ -759,6 +763,50 @@ func TestColumnTypeDatabaseTypeNameOverride(t *testing.T) { assert.Equal(t, "TIMESTAMP", rowSet.ColumnTypeDatabaseTypeName(1)) } +func TestColumnTypeDatabaseTypeNameArrowTimestamps(t *testing.T) { + arrowSchema := arrow.NewSchema([]arrow.Field{ + {Name: "zoned", Type: &arrow.TimestampType{Unit: arrow.Microsecond, TimeZone: "UTC"}}, + {Name: "ntz", Type: &arrow.TimestampType{Unit: arrow.Microsecond}}, + {Name: "number", Type: arrow.PrimitiveTypes.Int64}, + }, nil) + var schemaBytes bytes.Buffer + writer := ipc.NewWriter(&schemaBytes, ipc.WithSchema(arrowSchema)) + require.NoError(t, writer.Close()) + + schema := &cli_service.TTableSchema{} + for i, typeID := range []cli_service.TTypeId{ + cli_service.TTypeId_TIMESTAMP_TYPE, + cli_service.TTypeId_TIMESTAMP_TYPE, + cli_service.TTypeId_BIGINT_TYPE, + } { + schema.Columns = append(schema.Columns, &cli_service.TColumnDesc{ + ColumnName: arrowSchema.Field(i).Name, + TypeDesc: &cli_service.TTypeDesc{Types: []*cli_service.TTypeEntry{{ + PrimitiveEntry: &cli_service.TPrimitiveTypeEntry{Type: typeID}, + }}}, + }) + } + rowSet := &rows{ + client: &client.TestClient{}, + config: config.WithDefaults(), + ctx: context.Background(), + schema: schema, + resultSetMetadata: &cli_service.TGetResultSetMetadataResp{ + Schema: schema, ArrowSchema: schemaBytes.Bytes(), + }, + } + require.NoError(t, rowSet.makeRowScanner(&cli_service.TFetchResultsResp{ + Results: &cli_service.TRowSet{ArrowBatches: []*cli_service.TSparkArrowBatch{}}, + })) + defer rowSet.RowScanner.Close() + + for i, want := range []string{"TIMESTAMP", "TIMESTAMP_NTZ", "BIGINT"} { + assert.Equal(t, want, rowSet.ColumnTypeDatabaseTypeName(i), "column %d", i) + } + assert.Empty(t, rowSet.ColumnTypeDatabaseTypeName(-1)) + assert.Empty(t, rowSet.ColumnTypeDatabaseTypeName(3)) +} + func TestRowsCloseOptimization(t *testing.T) { var closeCount int client := &client.TestClient{ diff --git a/internal/rows/rowscanner/rowScanner.go b/internal/rows/rowscanner/rowScanner.go index d1a75091..9c5d7aeb 100644 --- a/internal/rows/rowscanner/rowScanner.go +++ b/internal/rows/rowscanner/rowScanner.go @@ -35,18 +35,10 @@ type RowScanner interface { GetArrowIPCStreams(ctx context.Context, cfg config.Config, rpi ResultPageIterator) (dbsqlrows.ArrowIPCStreamIterator, error) } -// DatabaseTypeNameOverrider is an optional interface a RowScanner may implement -// to refine the database type name reported for a column beyond what the Thrift -// TColumnDesc carries. It exists for TIMESTAMP_NTZ: the Thrift primitive type id -// is TIMESTAMP_TYPE for both zoned TIMESTAMP and zone-less TIMESTAMP_NTZ, so the -// distinction survives only in the Arrow schema's TimestampType.TimeZone field, -// which only the arrow row scanner can see (databricks/databricks-sql-go#506). -// -// Limitation: the override applies only when the Arrow schema carries a native -// *arrow.TimestampType. With UseArrowNativeTimestamp=false (timestamps sent as -// strings) or with the column-based (non-Arrow) row scanner, there is no zone -// information to inspect, so a TIMESTAMP_NTZ column still reports "TIMESTAMP" -// and its string value is parsed in the session location. +// DatabaseTypeNameOverrider refines Thrift column type names using the Arrow +// schema, which distinguishes TIMESTAMP_NTZ from TIMESTAMP. String timestamps +// and non-Arrow results lack this distinction and retain the legacy TIMESTAMP +// name and session-location parsing. type DatabaseTypeNameOverrider interface { // DatabaseTypeNameOverride returns the refined type name for the column at // index and ok=true when it overrides the Thrift-derived name; ok=false diff --git a/kernel_e2e_test.go b/kernel_e2e_test.go index b1fb77b8..7eef7c93 100644 --- a/kernel_e2e_test.go +++ b/kernel_e2e_test.go @@ -123,13 +123,8 @@ func TestKernelE2ETimeZone(t *testing.T) { } } -// TestKernelE2ETimestampNTZ pins how TIMESTAMP and TIMESTAMP_NTZ scan at a -// non-UTC session tz. The kernel delivers TIMESTAMP with an Arrow tz and -// TIMESTAMP_NTZ without one, and the driver honors that field: a zoned TIMESTAMP -// is rendered via .In(loc), while a zone-less TIMESTAMP_NTZ is a wall clock and is -// NOT shifted into the session tz (databricks/databricks-sql-go#506). So 12:00 NTZ -// in America/New_York comes back as the 12:00 wall clock, tagged UTC — not the -// -04:00 local of 12:00Z (08:00) that the old .In(loc) behavior produced. +// TestKernelE2ETimestampNTZ checks that TIMESTAMP uses the session location, +// while TIMESTAMP_NTZ preserves its wall clock with a UTC location. func TestKernelE2ETimestampNTZ(t *testing.T) { // Reyden shifts a zoned TIMESTAMP literal like an NTZ (12:00 → 08:00 -0400) where a // normal warehouse keeps the 12:00 wall-clock — an engine timestamp-tz divergence, @@ -160,16 +155,11 @@ func TestKernelE2ETimestampNTZ(t *testing.T) { "SELECT CAST('2026-07-09 12:00:00' AS TIMESTAMP_NTZ)").Scan(&ntz); err != nil { t.Fatalf("TIMESTAMP_NTZ query: %v", err) } - // The NTZ wall clock is preserved as sent (12:00) and NOT shifted into the - // session tz; it carries no zone, so it is tagged UTC. - if want := time.Date(2026, 7, 9, 12, 0, 0, 0, time.UTC); !ntz.Equal(want) { - t.Errorf("TIMESTAMP_NTZ = %s, want %s", ntz, want) - } - if h, m := ntz.Hour(), ntz.Minute(); h != 12 || m != 0 { - t.Errorf("TIMESTAMP_NTZ wall clock = %02d:%02d, want 12:00", h, m) + if want := "2026-07-09 12:00:00"; ntz.Format(time.DateTime) != want { + t.Errorf("TIMESTAMP_NTZ wall clock = %s, want %s", ntz.Format(time.DateTime), want) } if ntz.Location() != time.UTC { - t.Errorf("TIMESTAMP_NTZ location = %q, want %q", ntz.Location(), time.UTC) + t.Errorf("TIMESTAMP_NTZ location = %q, want UTC", ntz.Location()) } } From 91ac281a1e1644446d050bca00d12413ccd71a4b Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 22 Oct 2026 21:54:01 +0000 Subject: [PATCH 05/15] fix: preserve NTZ metadata across result iteration and kernel schemas Signed-off-by: Vu Anh Phung --- CHANGELOG.md | 2 +- CONNECTION_PARAMETERS.md | 5 + README.md | 5 + doc.go | 5 + driver_e2e_test.go | 1 + internal/arrowscan/coltype.go | 12 +- internal/arrowscan/coltype_ntz_test.go | 23 +++ .../timestamp_ntz_kernel_schema.arrow | Bin 0 -> 736 bytes .../testdata/timestamp_ntz_kernel_schema.md | 5 + internal/backend/kernel/batchiterator_test.go | 44 +++++ internal/backend/kernel/rows.go | 7 +- internal/rows/arrowbased/arrowRows.go | 21 --- .../rows/arrowbased/timestamp_ntz_test.go | 39 ----- internal/rows/rows.go | 41 ++++- internal/rows/rows_test.go | 157 +++++++++--------- internal/rows/rowscanner/rowScanner.go | 11 -- .../testdata/timestamp_ntz_arrow_schema.arrow | Bin 0 -> 664 bytes .../testdata/timestamp_ntz_arrow_schema.md | 5 + kernel_e2e_test.go | 33 +++- 19 files changed, 246 insertions(+), 170 deletions(-) create mode 100644 internal/arrowscan/testdata/timestamp_ntz_kernel_schema.arrow create mode 100644 internal/arrowscan/testdata/timestamp_ntz_kernel_schema.md create mode 100644 internal/rows/testdata/timestamp_ntz_arrow_schema.arrow create mode 100644 internal/rows/testdata/timestamp_ntz_arrow_schema.md diff --git a/CHANGELOG.md b/CHANGELOG.md index 071795ef..444cf6dc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,7 +2,7 @@ ## Unreleased -- **Behavior change:** Native Arrow TIMESTAMP_NTZ results now preserve their wall-clock fields and report `TIMESTAMP_NTZ` on both backends. Returned `time.Time` values use UTC instead of the session location; their represented instant and `.UTC()` value are unchanged (databricks/databricks-sql-go#506). +- **Behavior change:** Native Arrow TIMESTAMP_NTZ results now preserve their wall-clock fields and report `TIMESTAMP_NTZ` on both backends. Returned `time.Time` values use UTC instead of the session location; their represented instant and `.UTC()` value are unchanged (databricks/databricks-sql-go#506). On Thrift, this requires a server-supplied Arrow schema with native timestamps; without it, NTZ still reports `TIMESTAMP` and uses session-location handling. String timestamps (`UseArrowNativeTimestamp=false`, internal) and non-Arrow values also retain legacy parsing. ## v1.17.0 (2026-10-20) - Upgrade the SQL kernel bindings to v1.2.0. diff --git a/CONNECTION_PARAMETERS.md b/CONNECTION_PARAMETERS.md index 4be1cfe9..35a4ad0e 100644 --- a/CONNECTION_PARAMETERS.md +++ b/CONNECTION_PARAMETERS.md @@ -165,6 +165,11 @@ INTERVAL, nested ARRAY / MAP / STRUCT and VARIANT as JSON, and BINARY as `sql.RawBytes`). The default `enableGeoSpatialSupport=false` on the kernel path matches the Thrift path's WKT / EWKT geospatial values. +On Thrift, TIMESTAMP_NTZ handling requires native timestamp fields in the +server-supplied Arrow schema. Without that schema, NTZ still reports `TIMESTAMP` +and uses session-location handling. String timestamps (`UseArrowNativeTimestamp=false`, +an internal setting) and non-Arrow values also retain legacy session-location parsing. + ## Cloud Fetch | DSN parameter | Connector option | Thrift | Kernel | Default | Notes | diff --git a/README.md b/README.md index 11a594ea..504ec79a 100644 --- a/README.md +++ b/README.md @@ -367,6 +367,11 @@ STRUCT and VARIANT (as JSON) are supported. BINARY is returned as `sql.RawBytes` Both backends return GEOMETRY / GEOGRAPHY as WKT / EWKT by default. Set `enableGeoSpatialSupport=true` on the kernel path to opt into EWKB. +On Thrift, TIMESTAMP_NTZ handling requires native timestamp fields in the +server-supplied Arrow schema. Without that schema, NTZ still reports `TIMESTAMP` +and uses session-location handling. String timestamps (`UseArrowNativeTimestamp=false`, +an internal setting) and non-Arrow values also retain legacy session-location parsing. + Metadata is reached through SQL (`SHOW`, `DESCRIBE`, `information_schema`) on both backends — the driver exposes no `GetCatalogs`/`GetSchemas`/`GetTables`/`GetColumns` API (a `database/sql` limitation, not backend-specific). diff --git a/doc.go b/doc.go index 2ff02ca2..e78f5a61 100644 --- a/doc.go +++ b/doc.go @@ -297,6 +297,11 @@ and GEOMETRY / GEOGRAPHY (WKT). The server query id is surfaced on the success p QueryIdCallback (see below) fires with the real id and EXECUTE_STATEMENT telemetry carries it. +On Thrift, TIMESTAMP_NTZ handling requires native timestamps in the server-supplied +Arrow schema. Without that schema, NTZ still reports TIMESTAMP and uses +session-location handling. String timestamps (internal UseArrowNativeTimestamp=false) +and non-Arrow values also retain legacy session-location parsing. + On the read path, context cancellation is honored at result-batch boundaries, not mid-fetch: an in-flight CloudFetch batch runs to completion before the cancel takes effect. diff --git a/driver_e2e_test.go b/driver_e2e_test.go index e69aba96..f07837b0 100644 --- a/driver_e2e_test.go +++ b/driver_e2e_test.go @@ -785,6 +785,7 @@ func TestE2ETimestampNTZPreservesWallClock(t *testing.T) { assert.Equal(t, 15, ntz.Minute()) assert.Equal(t, 0, ntz.Second()) assert.Equal(t, 123456000, ntz.Nanosecond()) + assert.Same(t, time.UTC, ntz.Location()) // The zoned TIMESTAMP is an instant: the literal is interpreted in the // session timezone (America/New_York), so its UTC instant is 13:15:00.123456Z. diff --git a/internal/arrowscan/coltype.go b/internal/arrowscan/coltype.go index a77c876a..acb3121a 100644 --- a/internal/arrowscan/coltype.go +++ b/internal/arrowscan/coltype.go @@ -148,11 +148,15 @@ func IsBinaryGeospatialField(field arrow.Field) bool { return geospatial && field.Type.ID() == arrow.STRUCT } -// ColumnTypeInfoForField reports logical GEOMETRY / GEOGRAPHY metadata for the -// native struct representation. String mode follows the existing driver -// contract and reports STRING, matching its UTF-8 physical representation. -// Other fields retain the Arrow-only mapping. +// ColumnTypeInfoForField preserves logical timestamp and geospatial types that +// the kernel's physical Arrow representation alone cannot distinguish. func ColumnTypeInfoForField(field arrow.Field) ColumnTypeInfo { + if field.Type.ID() == arrow.TIMESTAMP { + name, _ := field.Metadata.GetValue(databricksTypeNameMetadataKey) + if strings.EqualFold(strings.TrimSpace(name), "TIMESTAMP_NTZ") { + return ColumnTypeInfo{DatabaseTypeName: "TIMESTAMP_NTZ", ScanType: scanTypeDateTime} + } + } name, geospatial := LogicalGeospatialType(field) if !geospatial { return ColumnTypeInfoFor(field.Type) diff --git a/internal/arrowscan/coltype_ntz_test.go b/internal/arrowscan/coltype_ntz_test.go index 3bd0f02a..0b1c6621 100644 --- a/internal/arrowscan/coltype_ntz_test.go +++ b/internal/arrowscan/coltype_ntz_test.go @@ -31,3 +31,26 @@ func TestColumnTypeInfoForTimestampNTZ(t *testing.T) { } }) } + +func TestColumnTypeInfoForFieldTimestampNTZ(t *testing.T) { + for _, tc := range []struct { + name string + typeName string + dataType arrow.DataType + want string + }{ + {"kernel NTZ with UTC zone", "TIMESTAMP_NTZ", arrow.FixedWidthTypes.Timestamp_us, "TIMESTAMP_NTZ"}, + {"zoned timestamp", "TIMESTAMP", arrow.FixedWidthTypes.Timestamp_us, "TIMESTAMP"}, + {"no logical metadata", "", arrow.FixedWidthTypes.Timestamp_us, "TIMESTAMP"}, + {"string representation", "TIMESTAMP_NTZ", arrow.BinaryTypes.String, "STRING"}, + } { + t.Run(tc.name, func(t *testing.T) { + field := arrow.Field{Type: tc.dataType, Metadata: arrow.MetadataFrom(map[string]string{ + databricksTypeNameMetadataKey: tc.typeName, + })} + if got := ColumnTypeInfoForField(field).DatabaseTypeName; got != tc.want { + t.Errorf("DatabaseTypeName = %q, want %q", got, tc.want) + } + }) + } +} diff --git a/internal/arrowscan/testdata/timestamp_ntz_kernel_schema.arrow b/internal/arrowscan/testdata/timestamp_ntz_kernel_schema.arrow new file mode 100644 index 0000000000000000000000000000000000000000..2115c9716650e06adf4a6fbfd849f5358423b4e9 GIT binary patch literal 736 zcmbVJ!AiqG6nxc+X{;!8SniJf_Kd?F-~r5akjom|z#34~{W!|BiqkOBh>Pv)0^ zFHbBt>&JAa`FE;Q O$HOG*A#$Nr!}tym%a>mO literal 0 HcmV?d00001 diff --git a/internal/arrowscan/testdata/timestamp_ntz_kernel_schema.md b/internal/arrowscan/testdata/timestamp_ntz_kernel_schema.md new file mode 100644 index 00000000..78bf445e --- /dev/null +++ b/internal/arrowscan/testdata/timestamp_ntz_kernel_schema.md @@ -0,0 +1,5 @@ +Captured from a live test warehouse on 2026-10-22. The `.arrow` file contains the Arrow IPC schema for: + +```sql +SELECT TIMESTAMP '2024-06-15 09:15:00.123456' AS zoned, TIMESTAMP_NTZ '2024-06-15 09:15:00.123456' AS ntz +``` diff --git a/internal/backend/kernel/batchiterator_test.go b/internal/backend/kernel/batchiterator_test.go index 4eb22c8f..bb471330 100644 --- a/internal/backend/kernel/batchiterator_test.go +++ b/internal/backend/kernel/batchiterator_test.go @@ -3,15 +3,24 @@ package kernel import ( + "bytes" "context" + "database/sql/driver" "errors" "io" + "os" "testing" + "time" "github.com/apache/arrow/go/v12/arrow" + "github.com/apache/arrow/go/v12/arrow/array" + "github.com/apache/arrow/go/v12/arrow/ipc" + "github.com/apache/arrow/go/v12/arrow/memory" dbsqlerr "github.com/databricks/databricks-sql-go/errors" + "github.com/databricks/databricks-sql-go/internal/arrowscan" dbsqlrows "github.com/databricks/databricks-sql-go/internal/rows" dbrows "github.com/databricks/databricks-sql-go/rows" + "github.com/stretchr/testify/require" ) // TestGetArrowIPCStreamsRejected pins the discoverable sentinel so callers can @@ -111,3 +120,38 @@ func TestBatchIteratorFetchErrorTerminates(t *testing.T) { // compile-time guard: kernelBatchIterator must satisfy the public iterator surface. var _ dbrows.ArrowBatchIterator = (*kernelBatchIterator)(nil) + +func TestRowsTimestampNTZWithServerSchema(t *testing.T) { + data, err := os.ReadFile("../../arrowscan/testdata/timestamp_ntz_kernel_schema.arrow") + require.NoError(t, err) + reader, err := ipc.NewReader(bytes.NewReader(data)) + require.NoError(t, err) + defer reader.Release() + schema := reader.Schema() + require.Len(t, schema.Fields(), 2) + + location, err := time.LoadLocation("America/New_York") + require.NoError(t, err) + wall := time.Date(2024, time.June, 15, 9, 15, 0, 123456000, time.UTC) + builder := array.NewRecordBuilder(memory.NewGoAllocator(), schema) + defer builder.Release() + builder.Field(0).(*array.TimestampBuilder).Append(arrow.Timestamp(wall.Add(4 * time.Hour).UnixMicro())) + builder.Field(1).(*array.TimestampBuilder).Append(arrow.Timestamp(wall.UnixMicro())) + record := builder.NewRecord() + defer record.Release() + rowSet := &kernelRows{ + op: &kernelOp{location: location}, cur: record, + cols: []string{"zoned", "ntz"}, geoBinary: make([]bool, 2), + colTypes: []arrowscan.ColumnTypeInfo{ + arrowscan.ColumnTypeInfoForField(schema.Field(0)), + arrowscan.ColumnTypeInfoForField(schema.Field(1)), + }, + } + require.Equal(t, "TIMESTAMP", rowSet.ColumnTypeDatabaseTypeName(0)) + require.Equal(t, "TIMESTAMP_NTZ", rowSet.ColumnTypeDatabaseTypeName(1)) + values := make([]driver.Value, 2) + require.NoError(t, rowSet.Next(values)) + require.Equal(t, wall.Add(4*time.Hour).In(location), values[0]) + require.Equal(t, wall, values[1]) + require.Same(t, time.UTC, values[1].(time.Time).Location()) +} diff --git a/internal/backend/kernel/rows.go b/internal/backend/kernel/rows.go index d55f77ea..c7ca72b5 100644 --- a/internal/backend/kernel/rows.go +++ b/internal/backend/kernel/rows.go @@ -235,7 +235,12 @@ func (r *kernelRows) next(dest []driver.Value) error { } rec := r.cur for c := 0; c < len(dest); c++ { - v, err := arrowscan.ScanCellCachedWithOptions(rec.Column(c), r.rowInCur, r.op.location, r.keyCache, arrowscan.ScanOptions{ + location := r.op.location + // The kernel may tag NTZ with UTC in Arrow while preserving its logical type. + if r.colTypes[c].DatabaseTypeName == "TIMESTAMP_NTZ" { + location = time.UTC + } + v, err := arrowscan.ScanCellCachedWithOptions(rec.Column(c), r.rowInCur, location, r.keyCache, arrowscan.ScanOptions{ DecimalAsFloat: r.op.decimalAsFloat, GeospatialBinary: r.geoBinary[c], }) diff --git a/internal/rows/arrowbased/arrowRows.go b/internal/rows/arrowbased/arrowRows.go index 8d6f3d7f..2c378b83 100644 --- a/internal/rows/arrowbased/arrowRows.go +++ b/internal/rows/arrowbased/arrowRows.go @@ -153,27 +153,6 @@ func NewArrowRowScanner(resultSetMetadata *cli_service.TGetResultSetMetadataResp return rs, nil } -// DatabaseTypeNameOverride implements rowscanner.DatabaseTypeNameOverrider. Over -// Thrift, a zoned TIMESTAMP and a zone-less TIMESTAMP_NTZ share the same -// primitive type id (TIMESTAMP_TYPE), so GetDBTypeName reports "TIMESTAMP" for -// both. The distinction survives only in the Arrow schema's -// TimestampType.TimeZone: a zone-less column is TIMESTAMP_NTZ -// (databricks/databricks-sql-go#506). Non-timestamp columns and zoned -// timestamps fall through to the Thrift-derived name. -func (ars *arrowRowScanner) DatabaseTypeNameOverride(index int) (string, bool) { - if ars == nil || ars.arrowSchema == nil { - return "", false - } - fields := ars.arrowSchema.Fields() - if index < 0 || index >= len(fields) { - return "", false - } - if ts, ok := fields[index].Type.(*arrow.TimestampType); ok && ts.TimeZone == "" { - return "TIMESTAMP_NTZ", true - } - return "", false -} - // Close is called when the Rows instance is closed. func (ars *arrowRowScanner) Close() { if ars.rowValues != nil { diff --git a/internal/rows/arrowbased/timestamp_ntz_test.go b/internal/rows/arrowbased/timestamp_ntz_test.go index 9878c757..8761e66c 100644 --- a/internal/rows/arrowbased/timestamp_ntz_test.go +++ b/internal/rows/arrowbased/timestamp_ntz_test.go @@ -57,42 +57,3 @@ func TestThriftTimestampNTZPreservesWallClock(t *testing.T) { } }) } - -// TestArrowRowScannerDatabaseTypeNameOverride pins the Thrift-path type-name -// refinement (#506): only a zone-less Arrow timestamp column is reported as -// TIMESTAMP_NTZ; zoned timestamps, non-timestamp columns, out-of-range indexes -// and a missing schema all fall through (ok=false) to the Thrift-derived name. -func TestArrowRowScannerDatabaseTypeNameOverride(t *testing.T) { - schema := arrow.NewSchema([]arrow.Field{ - {Name: "ts", Type: &arrow.TimestampType{Unit: arrow.Microsecond, TimeZone: "UTC"}}, - {Name: "ts_ntz", Type: &arrow.TimestampType{Unit: arrow.Microsecond}}, - {Name: "i", Type: arrow.PrimitiveTypes.Int32}, - }, nil) - ars := &arrowRowScanner{arrowSchema: schema} - - cases := []struct { - index int - wantName string - wantOK bool - }{ - {0, "", false}, - {1, "TIMESTAMP_NTZ", true}, - {2, "", false}, - {-1, "", false}, - {3, "", false}, - } - for _, c := range cases { - name, ok := ars.DatabaseTypeNameOverride(c.index) - if name != c.wantName || ok != c.wantOK { - t.Errorf("index %d: got (%q, %v), want (%q, %v)", c.index, name, ok, c.wantName, c.wantOK) - } - } - - if name, ok := (&arrowRowScanner{}).DatabaseTypeNameOverride(0); name != "" || ok { - t.Errorf("nil schema: got (%q, %v), want (\"\", false)", name, ok) - } - var nilScanner *arrowRowScanner - if name, ok := nilScanner.DatabaseTypeNameOverride(0); name != "" || ok { - t.Errorf("nil scanner: got (%q, %v), want (\"\", false)", name, ok) - } -} diff --git a/internal/rows/rows.go b/internal/rows/rows.go index 8ddb33f9..34438b0c 100644 --- a/internal/rows/rows.go +++ b/internal/rows/rows.go @@ -1,6 +1,7 @@ package rows import ( + "bytes" "context" "database/sql" "database/sql/driver" @@ -9,6 +10,8 @@ import ( "reflect" "time" + "github.com/apache/arrow/go/v12/arrow" + "github.com/apache/arrow/go/v12/arrow/ipc" "github.com/databricks/databricks-sql-go/driverctx" dbsqlerr "github.com/databricks/databricks-sql-go/errors" "github.com/databricks/databricks-sql-go/internal/cli_service" @@ -43,8 +46,9 @@ type rows struct { location *time.Location // Metadata for result set - resultSetMetadata *cli_service.TGetResultSetMetadataResp - schema *cli_service.TTableSchema + resultSetMetadata *cli_service.TGetResultSetMetadataResp + schema *cli_service.TTableSchema + databaseTypeNameOverrides map[int]string config *config.Config @@ -364,13 +368,9 @@ func (r *rows) ColumnTypeDatabaseTypeName(index int) string { return "" } - // The Thrift primitive type id cannot tell a zoned TIMESTAMP from a - // zone-less TIMESTAMP_NTZ; that distinction lives in the Arrow schema, which - // the arrow row scanner can expose via DatabaseTypeNameOverrider (#506). - if o, ok := r.RowScanner.(rowscanner.DatabaseTypeNameOverrider); ok && o != nil { - if name, ok := o.DatabaseTypeNameOverride(index); ok { - return name - } + r.cacheDatabaseTypeNameOverrides() + if name, ok := r.databaseTypeNameOverrides[index]; ok { + return name } dbtype := rowscanner.GetDBTypeName(column) @@ -378,6 +378,29 @@ func (r *rows) ColumnTypeDatabaseTypeName(index int) string { return dbtype } +// Thrift shares TIMESTAMP_TYPE between TIMESTAMP and TIMESTAMP_NTZ. Cache the +// distinction from server metadata so it survives result-page scanner changes. +func (r *rows) cacheDatabaseTypeNameOverrides() { + if r.databaseTypeNameOverrides != nil { + return + } + r.databaseTypeNameOverrides = make(map[int]string) + if r.resultSetMetadata == nil || len(r.resultSetMetadata.ArrowSchema) == 0 { + return + } + + reader, err := ipc.NewReader(bytes.NewReader(r.resultSetMetadata.ArrowSchema)) + if err != nil { + return + } + defer reader.Release() + for i, field := range reader.Schema().Fields() { + if ts, ok := field.Type.(*arrow.TimestampType); ok && ts.TimeZone == "" { + r.databaseTypeNameOverrides[i] = "TIMESTAMP_NTZ" + } + } +} + // ColumnTypeNullable returns a flag indicating whether the column is nullable // and an ok value of true if the status of the column is known. Otherwise // a value of false is returned for ok. diff --git a/internal/rows/rows_test.go b/internal/rows/rows_test.go index e5ab93f9..5dbb7af5 100644 --- a/internal/rows/rows_test.go +++ b/internal/rows/rows_test.go @@ -712,66 +712,16 @@ func TestColumnTypeDatabaseTypeName(t *testing.T) { assert.Equal(t, expectedScanTypes, scanTypes) } -// typeNameOverrideScanner is a RowScanner stub that only implements -// rowscanner.DatabaseTypeNameOverrider; other methods are never called. -type typeNameOverrideScanner struct { - rowscanner.RowScanner - overrides map[int]string -} - -func (s *typeNameOverrideScanner) DatabaseTypeNameOverride(index int) (string, bool) { - name, ok := s.overrides[index] - return name, ok -} - -// TestColumnTypeDatabaseTypeNameOverride checks that ColumnTypeDatabaseTypeName -// consults a DatabaseTypeNameOverrider row scanner (#506): Thrift reports -// TIMESTAMP_TYPE for both zoned and zone-less timestamps, so the override is -// what surfaces TIMESTAMP_NTZ, and columns without an override keep the -// Thrift-derived name. -func TestColumnTypeDatabaseTypeNameOverride(t *testing.T) { - var getMetadataCount, fetchResultsCount int - - col := func(name string, typ cli_service.TTypeId) *cli_service.TColumnDesc { - return &cli_service.TColumnDesc{ - ColumnName: name, - TypeDesc: &cli_service.TTypeDesc{ - Types: []*cli_service.TTypeEntry{ - {PrimitiveEntry: &cli_service.TPrimitiveTypeEntry{Type: typ}}, - }, - }, - } - } - - rowSet := &rows{ - client: getRowsTestSimpleClient(&getMetadataCount, &fetchResultsCount), - schema: &cli_service.TTableSchema{Columns: []*cli_service.TColumnDesc{ - col("ts", cli_service.TTypeId_TIMESTAMP_TYPE), - col("ts_ntz", cli_service.TTypeId_TIMESTAMP_TYPE), - col("i", cli_service.TTypeId_INT_TYPE), - }}, - RowScanner: &typeNameOverrideScanner{overrides: map[int]string{1: "TIMESTAMP_NTZ"}}, - } - - assert.Equal(t, "TIMESTAMP", rowSet.ColumnTypeDatabaseTypeName(0)) - assert.Equal(t, "TIMESTAMP_NTZ", rowSet.ColumnTypeDatabaseTypeName(1)) - assert.Equal(t, "INT", rowSet.ColumnTypeDatabaseTypeName(2)) - assert.Equal(t, "", rowSet.ColumnTypeDatabaseTypeName(3)) - - // Without an overrider (or with no row scanner yet) the Thrift name is used. - rowSet.RowScanner = nil - assert.Equal(t, "TIMESTAMP", rowSet.ColumnTypeDatabaseTypeName(1)) -} - func TestColumnTypeDatabaseTypeNameArrowTimestamps(t *testing.T) { - arrowSchema := arrow.NewSchema([]arrow.Field{ - {Name: "zoned", Type: &arrow.TimestampType{Unit: arrow.Microsecond, TimeZone: "UTC"}}, - {Name: "ntz", Type: &arrow.TimestampType{Unit: arrow.Microsecond}}, - {Name: "number", Type: arrow.PrimitiveTypes.Int64}, - }, nil) - var schemaBytes bytes.Buffer - writer := ipc.NewWriter(&schemaBytes, ipc.WithSchema(arrowSchema)) - require.NoError(t, writer.Close()) + metadata := &cli_service.TGetResultSetMetadataResp{} + schemaBytes, err := os.ReadFile("testdata/timestamp_ntz_arrow_schema.arrow") + require.NoError(t, err) + metadata.ArrowSchema = schemaBytes + require.NotEmpty(t, metadata.ArrowSchema) + reader, err := ipc.NewReader(bytes.NewReader(metadata.ArrowSchema)) + require.NoError(t, err) + defer reader.Release() + arrowSchema := reader.Schema() schema := &cli_service.TTableSchema{} for i, typeID := range []cli_service.TTypeId{ @@ -786,25 +736,82 @@ func TestColumnTypeDatabaseTypeNameArrowTimestamps(t *testing.T) { }}}, }) } - rowSet := &rows{ - client: &client.TestClient{}, - config: config.WithDefaults(), - ctx: context.Background(), - schema: schema, - resultSetMetadata: &cli_service.TGetResultSetMetadataResp{ - Schema: schema, ArrowSchema: schemaBytes.Bytes(), - }, + metadata.Schema = schema + for _, fetchMetadata := range []bool{false, true} { + t.Run(fmt.Sprintf("fetchMetadata=%t", fetchMetadata), func(t *testing.T) { + getMetadataCount := 0 + testClient := &client.TestClient{ + FnGetResultSetMetadata: func(context.Context, *cli_service.TGetResultSetMetadataReq) (*cli_service.TGetResultSetMetadataResp, error) { + getMetadataCount++ + return metadata, nil + }, + } + rowSet := &rows{ + client: testClient, + config: config.WithDefaults(), + ctx: context.Background(), + } + if !fetchMetadata { + rowSet.schema = schema + rowSet.resultSetMetadata = metadata + } + rowSet.ResultPageIterator = rowscanner.NewResultPageIterator( + rowSet.ctx, rowscanner.NewDelimiter(0, 0), 1, nil, true, testClient, rowSet.logger()) + + checkTypes := func() { + t.Helper() + for i, want := range []string{"TIMESTAMP", "TIMESTAMP_NTZ", "BIGINT"} { + assert.Equal(t, want, rowSet.ColumnTypeDatabaseTypeName(i), "column %d", i) + } + assert.Empty(t, rowSet.ColumnTypeDatabaseTypeName(-1)) + assert.Empty(t, rowSet.ColumnTypeDatabaseTypeName(3)) + } + require.Nil(t, rowSet.RowScanner) + checkTypes() // Before the first Next, including metadata fetched on demand. + + require.NoError(t, rowSet.makeRowScanner(&cli_service.TFetchResultsResp{ + Results: &cli_service.TRowSet{ArrowBatches: []*cli_service.TSparkArrowBatch{}}, + })) + checkTypes() + require.NoError(t, rowSet.makeRowScanner(&cli_service.TFetchResultsResp{ + Results: &cli_service.TRowSet{Columns: []*cli_service.TColumn{}}, + })) + checkTypes() + + require.ErrorIs(t, rowSet.Next(make([]driver.Value, 3)), io.EOF) + require.Nil(t, rowSet.RowScanner) + checkTypes() // EOF releases the scanner, but not the result metadata. + assert.Equal(t, map[bool]int{false: 0, true: 1}[fetchMetadata], getMetadataCount) + }) } - require.NoError(t, rowSet.makeRowScanner(&cli_service.TFetchResultsResp{ - Results: &cli_service.TRowSet{ArrowBatches: []*cli_service.TSparkArrowBatch{}}, - })) - defer rowSet.RowScanner.Close() +} + +func TestColumnTypeDatabaseTypeNameWithoutNativeArrowSchema(t *testing.T) { + var stringSchema bytes.Buffer + writer := ipc.NewWriter(&stringSchema, ipc.WithSchema(arrow.NewSchema([]arrow.Field{ + {Name: "ts", Type: arrow.BinaryTypes.String}, + }, nil))) + require.NoError(t, writer.Close()) - for i, want := range []string{"TIMESTAMP", "TIMESTAMP_NTZ", "BIGINT"} { - assert.Equal(t, want, rowSet.ColumnTypeDatabaseTypeName(i), "column %d", i) + for name, schemaBytes := range map[string][]byte{ + "missing": nil, + "invalid": []byte("invalid arrow schema"), + "string timestamp": stringSchema.Bytes(), + } { + t.Run(name, func(t *testing.T) { + schema := &cli_service.TTableSchema{Columns: []*cli_service.TColumnDesc{{ + ColumnName: "ts", + TypeDesc: &cli_service.TTypeDesc{Types: []*cli_service.TTypeEntry{{ + PrimitiveEntry: &cli_service.TPrimitiveTypeEntry{Type: cli_service.TTypeId_TIMESTAMP_TYPE}, + }}}, + }}} + rowSet := &rows{ + client: &client.TestClient{}, schema: schema, + resultSetMetadata: &cli_service.TGetResultSetMetadataResp{Schema: schema, ArrowSchema: schemaBytes}, + } + assert.Equal(t, "TIMESTAMP", rowSet.ColumnTypeDatabaseTypeName(0)) + }) } - assert.Empty(t, rowSet.ColumnTypeDatabaseTypeName(-1)) - assert.Empty(t, rowSet.ColumnTypeDatabaseTypeName(3)) } func TestRowsCloseOptimization(t *testing.T) { diff --git a/internal/rows/rowscanner/rowScanner.go b/internal/rows/rowscanner/rowScanner.go index 9c5d7aeb..ebb25bf6 100644 --- a/internal/rows/rowscanner/rowScanner.go +++ b/internal/rows/rowscanner/rowScanner.go @@ -35,17 +35,6 @@ type RowScanner interface { GetArrowIPCStreams(ctx context.Context, cfg config.Config, rpi ResultPageIterator) (dbsqlrows.ArrowIPCStreamIterator, error) } -// DatabaseTypeNameOverrider refines Thrift column type names using the Arrow -// schema, which distinguishes TIMESTAMP_NTZ from TIMESTAMP. String timestamps -// and non-Arrow results lack this distinction and retain the legacy TIMESTAMP -// name and session-location parsing. -type DatabaseTypeNameOverrider interface { - // DatabaseTypeNameOverride returns the refined type name for the column at - // index and ok=true when it overrides the Thrift-derived name; ok=false - // means the caller should fall back to GetDBTypeName. - DatabaseTypeNameOverride(index int) (name string, ok bool) -} - // Expected formats for TIMESTAMP and DATE types when represented by a string value var DateTimeFormats map[string]string = map[string]string{ "TIMESTAMP": "2006-01-02 15:04:05.999999999", diff --git a/internal/rows/testdata/timestamp_ntz_arrow_schema.arrow b/internal/rows/testdata/timestamp_ntz_arrow_schema.arrow new file mode 100644 index 0000000000000000000000000000000000000000..3cd1798e1f057eb23f4f83932c753a613086ce8a GIT binary patch literal 664 zcmaixJx;?w5QWDfB#R+{0s?AZfEGm*{3L)8v5>qpG?4)dA>tnrQ;>>+f`TFiqT&b~ zfFp1y#PVKZ1u00$NZ-!vdS~~&wK2x@3nI(Huw_w7QW1rUQ5>#Ywk%+#Fz33;L6RBp zY>9j#7s1<)F{{W=-~}KGeGXt7aL)<$@C@Er0!}-LuN{$f)P5RtZfo@*3-m)8){eVL z%t9W=diLBy*V7_v|9kwqwiiSpJ|Gm7D=7~&V`}YS;GOY(bP;w$4&a}9U%@*-*v!;h z2h{tW-@So)2{_qJ6n3*9N?UQ(b0+HE!uj;d-U`rev*GLA=4s2*=a^5^PqhkAGga}u zaVf?s&ZP2{(Z8u;HtyvEYxHsi=;bTd!+rDM&&y;V!*?tW-jROTKL Date: Thu, 22 Oct 2026 22:04:26 +0000 Subject: [PATCH 06/15] ai: apply changes for #510 (2 review threads) Addresses: - #4353955766 at internal/backend/kernel/rows.go:240 - #4353955769 at internal/rows/rows.go:388 Signed-off-by: peco-engineer-bot[bot] --- internal/arrowscan/arrowscan.go | 25 ++++-- internal/arrowscan/coltype.go | 29 +++++-- .../arrowscan/timestamp_ntz_nested_test.go | 80 +++++++++++++++++++ internal/rows/rows.go | 13 +++ internal/rows/rows_test.go | 44 ++++++++++ 5 files changed, 178 insertions(+), 13 deletions(-) create mode 100644 internal/arrowscan/timestamp_ntz_nested_test.go diff --git a/internal/arrowscan/arrowscan.go b/internal/arrowscan/arrowscan.go index b996419a..94cf0f9e 100644 --- a/internal/arrowscan/arrowscan.go +++ b/internal/arrowscan/arrowscan.go @@ -426,16 +426,19 @@ func writeJSON(b *strings.Builder, col arrow.Array, row int, loc *time.Location, // re-slices preserving data.offset) must index offsets[row+data.offset]. // ValueOffsets does that; this mirrors the Thrift path's ValueOffsets use. s, e := c.ValueOffsets(row) - return writeListJSON(b, c.ListValues(), int(s), int(e), loc, keys) + elemLoc := childLocation(c.DataType().(*arrow.ListType).ElemField(), loc) + return writeListJSON(b, c.ListValues(), int(s), int(e), elemLoc, keys) case *array.LargeList: // arrow-go's LargeList.ValueOffsets (unlike List's) does NOT add data.offset, // so add it by hand to stay offset-correct. off := c.Data().Offset() - return writeListJSON(b, c.ListValues(), int(c.Offsets()[row+off]), int(c.Offsets()[row+off+1]), loc, keys) + elemLoc := childLocation(c.DataType().(*arrow.LargeListType).ElemField(), loc) + return writeListJSON(b, c.ListValues(), int(c.Offsets()[row+off]), int(c.Offsets()[row+off+1]), elemLoc, keys) case *array.FixedSizeList: - n := int(c.DataType().(*arrow.FixedSizeListType).Len()) + fdt := c.DataType().(*arrow.FixedSizeListType) + n := int(fdt.Len()) base := (row + c.Data().Offset()) * n - return writeListJSON(b, c.ListValues(), base, base+n, loc, keys) + return writeListJSON(b, c.ListValues(), base, base+n, childLocation(fdt.ElemField(), loc), keys) case *array.Map: return writeMapJSON(b, c, row, loc, keys) case *array.Struct: @@ -490,19 +493,22 @@ func writeMapJSON(b *strings.Builder, m *array.Map, row int, loc *time.Location, start, end := int(s), int(e) mapKeys := m.Keys() items := m.Items() + mdt := m.DataType().(*arrow.MapType) + keyLoc := childLocation(mdt.KeyField(), loc) + itemLoc := childLocation(mdt.ItemField(), loc) b.WriteByte('{') for i := start; i < end; i++ { if i > start { b.WriteByte(',') } // JSON object keys must be strings; stringify the key value. - kv, err := scalarForJSON(mapKeys, i, loc) + kv, err := scalarForJSON(mapKeys, i, keyLoc) if err != nil { return err } writeJSONKey(b, kv) b.WriteByte(':') - if err := writeJSON(b, items, i, loc, keys); err != nil { + if err := writeJSON(b, items, i, itemLoc, keys); err != nil { return err } } @@ -514,14 +520,17 @@ func writeStructJSON(b *strings.Builder, s *array.Struct, row int, loc *time.Loc // Field-name keys are constant across rows; keys.keyPrefixes memoizes them per // result set so this per-row/per-cell path doesn't re-marshal them (a nil cache // recomputes inline — correct, just not memoized). - prefixes := keys.keyPrefixes(s.DataType().(*arrow.StructType)) + st := s.DataType().(*arrow.StructType) + prefixes := keys.keyPrefixes(st) b.WriteByte('{') for f := 0; f < s.NumField(); f++ { if f > 0 { b.WriteByte(',') } b.WriteString(prefixes[f]) // pre-escaped `"name":` - if err := writeJSON(b, s.Field(f), row, loc, keys); err != nil { + // A kernel-labeled TIMESTAMP_NTZ child may carry a physical TimeZone; keep + // its wall clock unshifted, like the top-level NTZ column. + if err := writeJSON(b, s.Field(f), row, childLocation(st.Field(f), loc), keys); err != nil { return err } } diff --git a/internal/arrowscan/coltype.go b/internal/arrowscan/coltype.go index acb3121a..fa50cc9a 100644 --- a/internal/arrowscan/coltype.go +++ b/internal/arrowscan/coltype.go @@ -151,11 +151,8 @@ func IsBinaryGeospatialField(field arrow.Field) bool { // ColumnTypeInfoForField preserves logical timestamp and geospatial types that // the kernel's physical Arrow representation alone cannot distinguish. func ColumnTypeInfoForField(field arrow.Field) ColumnTypeInfo { - if field.Type.ID() == arrow.TIMESTAMP { - name, _ := field.Metadata.GetValue(databricksTypeNameMetadataKey) - if strings.EqualFold(strings.TrimSpace(name), "TIMESTAMP_NTZ") { - return ColumnTypeInfo{DatabaseTypeName: "TIMESTAMP_NTZ", ScanType: scanTypeDateTime} - } + if isLogicalNTZField(field) { + return ColumnTypeInfo{DatabaseTypeName: "TIMESTAMP_NTZ", ScanType: scanTypeDateTime} } name, geospatial := LogicalGeospatialType(field) if !geospatial { @@ -172,6 +169,28 @@ func ColumnTypeInfoForField(field arrow.Field) ColumnTypeInfo { return info } +// isLogicalNTZField reports whether field is an Arrow timestamp the kernel has +// labeled TIMESTAMP_NTZ via field metadata. The kernel may give such a field a +// physical TimeZone (e.g. "UTC"), so the metadata is the authoritative signal. +func isLogicalNTZField(field arrow.Field) bool { + if field.Type == nil || field.Type.ID() != arrow.TIMESTAMP { + return false + } + name, _ := field.Metadata.GetValue(databricksTypeNameMetadataKey) + return strings.EqualFold(strings.TrimSpace(name), "TIMESTAMP_NTZ") +} + +// childLocation returns the location to render a nested child field's values +// in: a metadata-labeled TIMESTAMP_NTZ child is a wall clock and must not be +// shifted into the session location, so it renders in UTC (the stored wall +// clock); every other child inherits loc. +func childLocation(field arrow.Field, loc *time.Location) *time.Location { + if isLogicalNTZField(field) { + return time.UTC + } + return loc +} + // 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/timestamp_ntz_nested_test.go b/internal/arrowscan/timestamp_ntz_nested_test.go new file mode 100644 index 00000000..2576a184 --- /dev/null +++ b/internal/arrowscan/timestamp_ntz_nested_test.go @@ -0,0 +1,80 @@ +package arrowscan + +import ( + "testing" + "time" + + "github.com/apache/arrow/go/v12/arrow" + "github.com/apache/arrow/go/v12/arrow/array" + "github.com/apache/arrow/go/v12/arrow/memory" +) + +// TestScanCellNestedKernelNTZKeepsWallClock pins that a nested child the kernel +// labels TIMESTAMP_NTZ via field metadata — while giving it a physical +// TimeZone "UTC" — renders its wall clock unshifted, exactly like a top-level +// NTZ column, whereas a zoned TIMESTAMP sibling is still shifted into loc. +func TestScanCellNestedKernelNTZKeepsWallClock(t *testing.T) { + pool := memory.NewGoAllocator() + loc, err := time.LoadLocation("America/New_York") + if err != nil { + t.Skipf("tz database unavailable: %v", err) + } + wall := time.Date(2024, time.June, 15, 9, 15, 0, 123456000, time.UTC) + tsUTC := &arrow.TimestampType{Unit: arrow.Microsecond, TimeZone: "UTC"} + ntzMeta := arrow.NewMetadata([]string{databricksTypeNameMetadataKey}, []string{"TIMESTAMP_NTZ"}) + ntzField := arrow.Field{Name: "n", Type: tsUTC, Nullable: true, Metadata: ntzMeta} + zonedField := arrow.Field{Name: "z", Type: tsUTC, Nullable: true} + + const ntzStr = `"2024-06-15 09:15:00.123456 +0000 UTC"` + const zonedStr = `"2024-06-15 05:15:00.123456 -0400 EDT"` + + t.Run("struct", func(t *testing.T) { + b := array.NewStructBuilder(pool, arrow.StructOf(ntzField, zonedField)) + defer b.Release() + b.Append(true) + b.FieldBuilder(0).(*array.TimestampBuilder).Append(arrow.Timestamp(wall.UnixMicro())) + b.FieldBuilder(1).(*array.TimestampBuilder).Append(arrow.Timestamp(wall.UnixMicro())) + arr := b.NewArray() + defer arr.Release() + v, err := ScanCellCached(arr, 0, loc, NewStructKeyCache()) + if err != nil { + t.Fatal(err) + } + want := `{"n":` + ntzStr + `,"z":` + zonedStr + `}` + if v != want { + t.Errorf("got %v, want %v", v, want) + } + }) + + t.Run("list", func(t *testing.T) { + b := array.NewBuilder(pool, arrow.ListOfField(arrow.Field{Name: "item", Type: tsUTC, Nullable: true, Metadata: ntzMeta})).(*array.ListBuilder) + defer b.Release() + b.Append(true) + b.ValueBuilder().(*array.TimestampBuilder).Append(arrow.Timestamp(wall.UnixMicro())) + arr := b.NewArray() + defer arr.Release() + v, err := ScanCellCached(arr, 0, loc, NewStructKeyCache()) + if err != nil { + t.Fatal(err) + } + if want := `[` + ntzStr + `]`; v != want { + t.Errorf("got %v, want %v", v, want) + } + }) + + t.Run("list without NTZ metadata still shifts", func(t *testing.T) { + b := array.NewListBuilder(pool, tsUTC) + defer b.Release() + b.Append(true) + b.ValueBuilder().(*array.TimestampBuilder).Append(arrow.Timestamp(wall.UnixMicro())) + arr := b.NewArray() + defer arr.Release() + v, err := ScanCellCached(arr, 0, loc, NewStructKeyCache()) + if err != nil { + t.Fatal(err) + } + if want := `[` + zonedStr + `]`; v != want { + t.Errorf("got %v, want %v", v, want) + } + }) +} diff --git a/internal/rows/rows.go b/internal/rows/rows.go index 34438b0c..2a8e5bcd 100644 --- a/internal/rows/rows.go +++ b/internal/rows/rows.go @@ -380,6 +380,9 @@ func (r *rows) ColumnTypeDatabaseTypeName(index int) string { // Thrift shares TIMESTAMP_TYPE between TIMESTAMP and TIMESTAMP_NTZ. Cache the // distinction from server metadata so it survives result-page scanner changes. +// The override is only applied when the result set is Arrow-decoded: columnar +// (TColumn) results keep legacy session-location timestamp parsing, so they +// keep the legacy TIMESTAMP type name too and the name matches the values. func (r *rows) cacheDatabaseTypeNameOverrides() { if r.databaseTypeNameOverrides != nil { return @@ -388,6 +391,16 @@ func (r *rows) cacheDatabaseTypeNameOverrides() { if r.resultSetMetadata == nil || len(r.resultSetMetadata.ArrowSchema) == 0 { return } + if r.config != nil && !r.config.UseArrowBatches { + return + } + if r.resultSetMetadata.IsSetResultFormat() { + switch r.resultSetMetadata.GetResultFormat() { + case cli_service.TSparkRowSetType_ARROW_BASED_SET, cli_service.TSparkRowSetType_URL_BASED_SET: + default: + return + } + } reader, err := ipc.NewReader(bytes.NewReader(r.resultSetMetadata.ArrowSchema)) if err != nil { diff --git a/internal/rows/rows_test.go b/internal/rows/rows_test.go index 5dbb7af5..95688399 100644 --- a/internal/rows/rows_test.go +++ b/internal/rows/rows_test.go @@ -814,6 +814,50 @@ func TestColumnTypeDatabaseTypeNameWithoutNativeArrowSchema(t *testing.T) { } } +func TestColumnTypeDatabaseTypeNameNonArrowResults(t *testing.T) { + schemaBytes, err := os.ReadFile("testdata/timestamp_ntz_arrow_schema.arrow") + require.NoError(t, err) + schema := &cli_service.TTableSchema{} + for _, typeID := range []cli_service.TTypeId{ + cli_service.TTypeId_TIMESTAMP_TYPE, + cli_service.TTypeId_TIMESTAMP_TYPE, + cli_service.TTypeId_BIGINT_TYPE, + } { + schema.Columns = append(schema.Columns, &cli_service.TColumnDesc{ + TypeDesc: &cli_service.TTypeDesc{Types: []*cli_service.TTypeEntry{{ + PrimitiveEntry: &cli_service.TPrimitiveTypeEntry{Type: typeID}, + }}}, + }) + } + columnFormat := cli_service.TSparkRowSetType_COLUMN_BASED_SET + noArrowConfig := config.WithDefaults() + noArrowConfig.UseArrowBatches = false + + for name, tc := range map[string]struct { + cfg *config.Config + metadata *cli_service.TGetResultSetMetadataResp + }{ + "column based result format": { + cfg: config.WithDefaults(), + metadata: &cli_service.TGetResultSetMetadataResp{Schema: schema, ArrowSchema: schemaBytes, ResultFormat: &columnFormat}, + }, + "arrow batches disabled": { + cfg: noArrowConfig, + metadata: &cli_service.TGetResultSetMetadataResp{Schema: schema, ArrowSchema: schemaBytes}, + }, + } { + t.Run(name, func(t *testing.T) { + rowSet := &rows{ + client: &client.TestClient{}, schema: schema, config: tc.cfg, + resultSetMetadata: tc.metadata, + } + for i, want := range []string{"TIMESTAMP", "TIMESTAMP", "BIGINT"} { + assert.Equal(t, want, rowSet.ColumnTypeDatabaseTypeName(i), "column %d", i) + } + }) + } +} + func TestRowsCloseOptimization(t *testing.T) { var closeCount int client := &client.TestClient{ From b8113f0c5a4eb6c85349127444e689686e0c7136 Mon Sep 17 00:00:00 2001 From: "peco-engineer-bot[bot]" Date: Thu, 22 Oct 2026 22:11:55 +0000 Subject: [PATCH 07/15] ai: apply changes for #510 (2 review threads) Addresses: - #4354033451 at internal/rows/rows.go:405 - #4354033458 at internal/backend/kernel/rows.go:240 Signed-off-by: peco-engineer-bot[bot] --- internal/backend/kernel/batchiterator_test.go | 2 +- internal/backend/kernel/rows.go | 9 +++++++-- internal/rows/rows.go | 4 ++++ 3 files changed, 12 insertions(+), 3 deletions(-) diff --git a/internal/backend/kernel/batchiterator_test.go b/internal/backend/kernel/batchiterator_test.go index bb471330..28981560 100644 --- a/internal/backend/kernel/batchiterator_test.go +++ b/internal/backend/kernel/batchiterator_test.go @@ -141,7 +141,7 @@ func TestRowsTimestampNTZWithServerSchema(t *testing.T) { defer record.Release() rowSet := &kernelRows{ op: &kernelOp{location: location}, cur: record, - cols: []string{"zoned", "ntz"}, geoBinary: make([]bool, 2), + cols: []string{"zoned", "ntz"}, geoBinary: make([]bool, 2), ntz: []bool{false, true}, colTypes: []arrowscan.ColumnTypeInfo{ arrowscan.ColumnTypeInfoForField(schema.Field(0)), arrowscan.ColumnTypeInfoForField(schema.Field(1)), diff --git a/internal/backend/kernel/rows.go b/internal/backend/kernel/rows.go index c7ca72b5..87dcd639 100644 --- a/internal/backend/kernel/rows.go +++ b/internal/backend/kernel/rows.go @@ -66,6 +66,7 @@ type kernelRows struct { cols []string colTypes []arrowscan.ColumnTypeInfo // per-column type metadata (PECOBLR-3692) geoBinary []bool // top-level logical geo structs become EWKB for database/sql + ntz []bool // top-level TIMESTAMP_NTZ columns scan in UTC, not the session location 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 @@ -120,10 +121,15 @@ func newKernelRows(ctx context.Context, op *kernelOp, stream *C.kernel_result_st // (ScanCellCached) via the shared arrowscan.ColumnTypeInfoFor mapper. r.colTypes = make([]arrowscan.ColumnTypeInfo, len(fields)) r.geoBinary = make([]bool, len(fields)) + r.ntz = make([]bool, len(fields)) for i, f := range fields { r.cols[i] = f.Name r.colTypes[i] = arrowscan.ColumnTypeInfoForField(f) r.geoBinary[i] = arrowscan.IsBinaryGeospatialField(f) + // The kernel may tag NTZ with UTC in Arrow while preserving its logical + // type; decide once here (from the same mapper as ColumnTypeDatabaseTypeName) + // rather than per cell in next(). + r.ntz[i] = r.colTypes[i].DatabaseTypeName == "TIMESTAMP_NTZ" } // Construction succeeded — now arm the close telemetry callback so a normal // Close() (after row iteration) records CLOSE_STATEMENT. @@ -236,8 +242,7 @@ func (r *kernelRows) next(dest []driver.Value) error { rec := r.cur for c := 0; c < len(dest); c++ { location := r.op.location - // The kernel may tag NTZ with UTC in Arrow while preserving its logical type. - if r.colTypes[c].DatabaseTypeName == "TIMESTAMP_NTZ" { + if r.ntz[c] { location = time.UTC } v, err := arrowscan.ScanCellCachedWithOptions(rec.Column(c), r.rowInCur, location, r.keyCache, arrowscan.ScanOptions{ diff --git a/internal/rows/rows.go b/internal/rows/rows.go index 2a8e5bcd..eb3c280a 100644 --- a/internal/rows/rows.go +++ b/internal/rows/rows.go @@ -404,6 +404,10 @@ func (r *rows) cacheDatabaseTypeNameOverrides() { reader, err := ipc.NewReader(bytes.NewReader(r.resultSetMetadata.ArrowSchema)) if err != nil { + // Leave the cache empty (and cached): result set metadata never changes, + // so retrying would fail the same way. Columns fall back to the legacy + // TIMESTAMP name; the Arrow row scanner surfaces the parse error itself. + r.logger().Debug().Err(err).Msg("databricks: failed to parse arrow schema for column type name overrides; falling back to thrift type names") return } defer reader.Release() From a4ec80285af5eff8d48248c2d611fb15dbd7744e Mon Sep 17 00:00:00 2001 From: "peco-engineer-bot[bot]" Date: Thu, 22 Oct 2026 22:20:41 +0000 Subject: [PATCH 08/15] ai: apply changes for #510 (1 review thread) Addresses: - #4354086520 at internal/rows/arrowbased/arrowRows.go:655 Signed-off-by: peco-engineer-bot[bot] --- internal/rows/arrowbased/arrowRows_test.go | 54 ++++++++++++++++++++++ 1 file changed, 54 insertions(+) diff --git a/internal/rows/arrowbased/arrowRows_test.go b/internal/rows/arrowbased/arrowRows_test.go index 713ab314..05b5616b 100644 --- a/internal/rows/arrowbased/arrowRows_test.go +++ b/internal/rows/arrowbased/arrowRows_test.go @@ -2398,3 +2398,57 @@ func TestDecimalInComplexTypes(t *testing.T) { assert.Equal(t, `{"col2":null}`, v) }) } + +// TestDerivedSchemaTimestampKeepsSessionLocation pins that when the server sends +// no ArrowSchema, the schema derived from the Thrift TTableSchema types +// TIMESTAMP_TYPE as a *zoned* Arrow timestamp. The zone survives the +// serialize/deserialize round-trip that also feeds GetArrowBatches. Values from +// that column are therefore still rendered in the session location rather than +// treated as a zone-less TIMESTAMP_NTZ wall clock (#506): without the server +// schema, NTZ and TIMESTAMP cannot be told apart, so the zoned behaviour stays. +func TestDerivedSchemaTimestampKeepsSessionLocation(t *testing.T) { + const tsCol = 8 // timestamp column in getAllTypesSchema + ctx := context.Background() + arrowConfig := config.ArrowConfig{UseArrowNativeTimestamp: true} + md := &cli_service.TGetResultSetMetadataResp{Schema: getAllTypesSchema()} + require.Nil(t, md.ArrowSchema) + + schemaBytes, derived, dErr := tGetResultSetMetadataRespToArrowSchema(md, arrowConfig, ctx, dbsqllog.Logger) + require.Nil(t, dErr) + + // Round-trip the serialized bytes the way a native-schema response is read. + _, roundTripped, rErr := tGetResultSetMetadataRespToArrowSchema( + &cli_service.TGetResultSetMetadataResp{Schema: md.Schema, ArrowSchema: schemaBytes}, arrowConfig, ctx, dbsqllog.Logger) + require.Nil(t, rErr) + + loc, err := time.LoadLocation("America/New_York") + require.NoError(t, err) + ttsf, err := arrow.FixedWidthTypes.Timestamp_us.(*arrow.TimestampType).GetToTimeFunc() + require.NoError(t, err) + instant := time.Date(2021, 7, 1, 12, 30, 0, 0, time.UTC) + + for name, s := range map[string]*arrow.Schema{"derived": derived, "round-tripped": roundTripped} { + t.Run(name, func(t *testing.T) { + tsType, ok := s.Field(tsCol).Type.(*arrow.TimestampType) + require.True(t, ok, "expected timestamp type, got %s", s.Field(tsCol).Type) + assert.NotEmpty(t, tsType.TimeZone, "derived TIMESTAMP must be zoned, not NTZ") + + bldr := array.NewTimestampBuilder(memory.DefaultAllocator, tsType) + defer bldr.Release() + bldr.Append(arrow.Timestamp(instant.UnixMicro())) + arr := bldr.NewTimestampArray() + defer arr.Release() + + c, err := (&arrowValueContainerMaker{}).makeColumnValueContainer(tsType, loc, ttsf, nil) + require.NoError(t, err) + require.NoError(t, c.SetValueArray(arr.Data())) + + v, err := c.Value(0) + require.NoError(t, err) + got, ok := v.(time.Time) + require.True(t, ok) + assert.True(t, instant.Equal(got)) + assert.Equal(t, loc, got.Location(), "zoned TIMESTAMP must render in the session location") + }) + } +} From 332032f6b59d5e6984acd2e8da1eb1cc817c5195 Mon Sep 17 00:00:00 2001 From: "peco-engineer-bot[bot]" Date: Thu, 22 Oct 2026 22:27:12 +0000 Subject: [PATCH 09/15] ai: apply changes for #510 (1 review thread) Addresses: - #4354148614 at internal/rows/rows.go:397 Signed-off-by: peco-engineer-bot[bot] --- internal/rows/rows.go | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/internal/rows/rows.go b/internal/rows/rows.go index eb3c280a..f442b488 100644 --- a/internal/rows/rows.go +++ b/internal/rows/rows.go @@ -380,9 +380,16 @@ func (r *rows) ColumnTypeDatabaseTypeName(index int) string { // Thrift shares TIMESTAMP_TYPE between TIMESTAMP and TIMESTAMP_NTZ. Cache the // distinction from server metadata so it survives result-page scanner changes. -// The override is only applied when the result set is Arrow-decoded: columnar -// (TColumn) results keep legacy session-location timestamp parsing, so they -// keep the legacy TIMESTAMP type name too and the name matches the values. +// The override is only applied when the result set metadata indicates Arrow +// decoding: columnar (TColumn) results keep legacy session-location timestamp +// parsing, so they keep the legacy TIMESTAMP type name too. +// +// The decision is made from result set metadata alone (it is cached before the +// first page arrives), not from the format of each fetched page. A result set +// with an ArrowSchema, Arrow batches enabled, and ResultFormat unset is treated +// as Arrow. If a server sent such metadata but then returned TColumn pages, the +// column would still report TIMESTAMP_NTZ while its values went through the +// legacy column scanner. func (r *rows) cacheDatabaseTypeNameOverrides() { if r.databaseTypeNameOverrides != nil { return From 9f08fef17c69c16240ab64b47baed6df7601a102 Mon Sep 17 00:00:00 2001 From: "peco-engineer-bot[bot]" Date: Thu, 22 Oct 2026 23:48:53 +0000 Subject: [PATCH 10/15] engineer-bot: address @mention on PR #510 Signed-off-by: peco-engineer-bot[bot] --- internal/arrowscan/arrowscan.go | 11 ++-- internal/arrowscan/coltype.go | 59 +++++++++++++++------- internal/arrowscan/coltype_ntz_test.go | 41 +++++++++++++++ internal/rows/arrowbased/arrowRows.go | 31 ++++++++---- internal/rows/arrowbased/arrowRows_test.go | 50 ++++++++++++++++++ internal/rows/rows.go | 6 ++- 6 files changed, 164 insertions(+), 34 deletions(-) diff --git a/internal/arrowscan/arrowscan.go b/internal/arrowscan/arrowscan.go index 94cf0f9e..b1ac260d 100644 --- a/internal/arrowscan/arrowscan.go +++ b/internal/arrowscan/arrowscan.go @@ -217,9 +217,12 @@ func scanCell(col arrow.Array, row int, loc *time.Location, keys *StructKeyCache // so it must NOT be shifted into the session location the way a zoned // TIMESTAMP instant is: applying .In(loc) would corrupt the wall-clock // fields the server sent (databricks/databricks-sql-go#506). Only a zoned - // TIMESTAMP is rendered in loc. + // TIMESTAMP is rendered in loc. ScanCell sees only the data type, so this + // uses IsTimestampNTZField's TimeZone fallback; metadata-labeled NTZ + // columns and children reach here with loc already set to UTC (the + // kernel's top-level ntz flag and childLocation). ts := timestampToTime(int64(c.Value(row)), dt.Unit) - if dt.TimeZone == "" { + if IsTimestampNTZField(arrow.Field{Type: dt}) { return ts, nil } return inLocation(ts, loc), nil @@ -528,8 +531,8 @@ func writeStructJSON(b *strings.Builder, s *array.Struct, row int, loc *time.Loc b.WriteByte(',') } b.WriteString(prefixes[f]) // pre-escaped `"name":` - // A kernel-labeled TIMESTAMP_NTZ child may carry a physical TimeZone; keep - // its wall clock unshifted, like the top-level NTZ column. + // A TIMESTAMP_NTZ child (IsTimestampNTZField) may carry a physical + // TimeZone; keep its wall clock unshifted, like the top-level NTZ column. if err := writeJSON(b, s.Field(f), row, childLocation(st.Field(f), loc), keys); err != nil { return err } diff --git a/internal/arrowscan/coltype.go b/internal/arrowscan/coltype.go index fa50cc9a..b746a40f 100644 --- a/internal/arrowscan/coltype.go +++ b/internal/arrowscan/coltype.go @@ -82,11 +82,9 @@ func ColumnTypeInfoFor(dt arrow.DataType) ColumnTypeInfo { // distinguishable only by the schema's TimeZone field: a zoned TIMESTAMP // carries a zone, a zone-less TIMESTAMP_NTZ carries an empty TimeZone. The // database/sql type name must reflect that distinction so a wall clock is - // not mistaken for an instant (databricks/databricks-sql-go#506). - if ts, ok := dt.(*arrow.TimestampType); ok && ts.TimeZone == "" { - return ColumnTypeInfo{DatabaseTypeName: "TIMESTAMP_NTZ", ScanType: scanTypeDateTime} - } - return ColumnTypeInfo{DatabaseTypeName: "TIMESTAMP", ScanType: scanTypeDateTime} + // not mistaken for an instant (databricks/databricks-sql-go#506). With no + // field metadata, IsTimestampNTZField falls back to that TimeZone check. + return timestampColumnTypeInfo(arrow.Field{Type: dt}) case arrow.DECIMAL128: // Match Thrift's sql.RawBytes scan type even though the kernel renders the // value as an exact string — both convert cleanly to a caller's *string/*[]byte. @@ -151,8 +149,8 @@ func IsBinaryGeospatialField(field arrow.Field) bool { // ColumnTypeInfoForField preserves logical timestamp and geospatial types that // the kernel's physical Arrow representation alone cannot distinguish. func ColumnTypeInfoForField(field arrow.Field) ColumnTypeInfo { - if isLogicalNTZField(field) { - return ColumnTypeInfo{DatabaseTypeName: "TIMESTAMP_NTZ", ScanType: scanTypeDateTime} + if field.Type != nil && field.Type.ID() == arrow.TIMESTAMP { + return timestampColumnTypeInfo(field) } name, geospatial := LogicalGeospatialType(field) if !geospatial { @@ -169,28 +167,53 @@ func ColumnTypeInfoForField(field arrow.Field) ColumnTypeInfo { return info } -// isLogicalNTZField reports whether field is an Arrow timestamp the kernel has -// labeled TIMESTAMP_NTZ via field metadata. The kernel may give such a field a -// physical TimeZone (e.g. "UTC"), so the metadata is the authoritative signal. -func isLogicalNTZField(field arrow.Field) bool { - if field.Type == nil || field.Type.ID() != arrow.TIMESTAMP { +// sparkSQLNameMetadataKey is the Arrow field-metadata key the Thrift server sets +// on its serialized result schema with the Spark SQL type name (e.g. +// "TIMESTAMP_NTZ"). The kernel copies it too. +const sparkSQLNameMetadataKey = "Spark:DataType:SqlName" + +// IsTimestampNTZField reports whether field is an Arrow timestamp carrying a +// TIMESTAMP_NTZ wall clock. Every backend uses this one check, for both the +// reported type name and value handling, so the two always agree. +// +// Explicit logical-type metadata wins: databricks.type_name (kernel) first, then +// Spark:DataType:SqlName (Thrift server schema; the kernel copies it). Metadata +// is authoritative because the kernel may give an NTZ field a physical TimeZone +// such as "UTC". Without either key, a zone-less timestamp (empty TimeZone) is +// treated as NTZ and a zoned one as TIMESTAMP. +func IsTimestampNTZField(field arrow.Field) bool { + ts, ok := field.Type.(*arrow.TimestampType) + if !ok { return false } - name, _ := field.Metadata.GetValue(databricksTypeNameMetadataKey) - return strings.EqualFold(strings.TrimSpace(name), "TIMESTAMP_NTZ") + for _, key := range [...]string{databricksTypeNameMetadataKey, sparkSQLNameMetadataKey} { + if name, ok := field.Metadata.GetValue(key); ok && strings.TrimSpace(name) != "" { + return strings.EqualFold(strings.TrimSpace(name), "TIMESTAMP_NTZ") + } + } + return ts.TimeZone == "" } // childLocation returns the location to render a nested child field's values -// in: a metadata-labeled TIMESTAMP_NTZ child is a wall clock and must not be -// shifted into the session location, so it renders in UTC (the stored wall -// clock); every other child inherits loc. +// in: a TIMESTAMP_NTZ child (see IsTimestampNTZField) is a wall clock and must +// not be shifted into the session location, so it renders in UTC (the stored +// wall clock); every other child inherits loc. func childLocation(field arrow.Field, loc *time.Location) *time.Location { - if isLogicalNTZField(field) { + if IsTimestampNTZField(field) { return time.UTC } return loc } +// timestampColumnTypeInfo reports TIMESTAMP_NTZ or TIMESTAMP for an Arrow +// timestamp field, using the same IsTimestampNTZField check as value scanning. +func timestampColumnTypeInfo(field arrow.Field) ColumnTypeInfo { + if IsTimestampNTZField(field) { + return ColumnTypeInfo{DatabaseTypeName: "TIMESTAMP_NTZ", ScanType: scanTypeDateTime} + } + return ColumnTypeInfo{DatabaseTypeName: "TIMESTAMP", ScanType: scanTypeDateTime} +} + // 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_ntz_test.go b/internal/arrowscan/coltype_ntz_test.go index 0b1c6621..f3336829 100644 --- a/internal/arrowscan/coltype_ntz_test.go +++ b/internal/arrowscan/coltype_ntz_test.go @@ -32,6 +32,47 @@ func TestColumnTypeInfoForTimestampNTZ(t *testing.T) { }) } +// TestIsTimestampNTZField pins the single NTZ check shared by type reporting and +// value handling on both backends: explicit metadata (databricks.type_name, then +// Spark:DataType:SqlName) wins over the physical TimeZone, which is only the +// fallback when no logical type name is present. +func TestIsTimestampNTZField(t *testing.T) { + zoned := arrow.FixedWidthTypes.Timestamp_us // tz=UTC + zoneless := &arrow.TimestampType{Unit: arrow.Microsecond} + for _, tc := range []struct { + name string + dt arrow.DataType + meta map[string]string + want bool + }{ + {"zone-less, no metadata", zoneless, nil, true}, + {"zoned, no metadata", zoned, nil, false}, + {"kernel type_name NTZ on UTC zone", zoned, map[string]string{databricksTypeNameMetadataKey: "TIMESTAMP_NTZ"}, true}, + {"thrift SqlName NTZ on zone-less", zoneless, map[string]string{sparkSQLNameMetadataKey: "TIMESTAMP_NTZ"}, true}, + {"SqlName NTZ on UTC zone", zoned, map[string]string{sparkSQLNameMetadataKey: "timestamp_ntz"}, true}, + {"SqlName TIMESTAMP overrides zone-less fallback", zoneless, map[string]string{sparkSQLNameMetadataKey: "TIMESTAMP"}, false}, + {"type_name wins over SqlName", zoned, map[string]string{databricksTypeNameMetadataKey: "TIMESTAMP", sparkSQLNameMetadataKey: "TIMESTAMP_NTZ"}, false}, + {"blank metadata falls back to zone", zoneless, map[string]string{databricksTypeNameMetadataKey: " "}, true}, + {"non-timestamp labeled NTZ", arrow.BinaryTypes.String, map[string]string{sparkSQLNameMetadataKey: "TIMESTAMP_NTZ"}, false}, + } { + t.Run(tc.name, func(t *testing.T) { + field := arrow.Field{Type: tc.dt, Metadata: arrow.MetadataFrom(tc.meta)} + if got := IsTimestampNTZField(field); got != tc.want { + t.Errorf("IsTimestampNTZField = %v, want %v", got, tc.want) + } + if tc.dt.ID() == arrow.TIMESTAMP { + want := "TIMESTAMP" + if tc.want { + want = "TIMESTAMP_NTZ" + } + if got := ColumnTypeInfoForField(field).DatabaseTypeName; got != want { + t.Errorf("DatabaseTypeName = %q, want %q", got, want) + } + } + }) + } +} + func TestColumnTypeInfoForFieldTimestampNTZ(t *testing.T) { for _, tc := range []struct { name string diff --git a/internal/rows/arrowbased/arrowRows.go b/internal/rows/arrowbased/arrowRows.go index 2c378b83..144b6909 100644 --- a/internal/rows/arrowbased/arrowRows.go +++ b/internal/rows/arrowbased/arrowRows.go @@ -12,6 +12,7 @@ import ( "github.com/apache/arrow/go/v12/arrow/array" "github.com/apache/arrow/go/v12/arrow/ipc" dbsqlerr "github.com/databricks/databricks-sql-go/errors" + "github.com/databricks/databricks-sql-go/internal/arrowscan" "github.com/databricks/databricks-sql-go/internal/cli_service" "github.com/databricks/databricks-sql-go/internal/config" dbsqlerrint "github.com/databricks/databricks-sql-go/internal/errors" @@ -592,7 +593,7 @@ func (vcm *arrowValueContainerMaker) makeColumnValuesContainers(ars *arrowRowSca if ars.rowValues == nil { columnValueHolders := make([]columnValues, len(ars.colInfo)) for i, field := range ars.arrowSchema.Fields() { - holder, err := vcm.makeColumnValueContainer(field.Type, ars.location, ars.toTimestampFn, &ars.colInfo[i]) + holder, err := vcm.makeFieldValueContainer(field, ars.location, ars.toTimestampFn, &ars.colInfo[i]) if err != nil { ars.Error().Msg(err.Error()) return err @@ -607,12 +608,21 @@ func (vcm *arrowValueContainerMaker) makeColumnValuesContainers(ars *arrowRowSca return nil } +// makeColumnValueContainer builds a container for a bare Arrow type with no field +// metadata. See makeFieldValueContainer. func (vcm *arrowValueContainerMaker) makeColumnValueContainer(t arrow.DataType, location *time.Location, toTimestampFn timeStampFn, colInfo *colInfo) (columnValues, error) { + return vcm.makeFieldValueContainer(arrow.Field{Type: t}, location, toTimestampFn, colInfo) +} + +// makeFieldValueContainer builds a container for an Arrow field. It takes the +// whole field, not just its type, so field metadata (e.g. the logical +// TIMESTAMP_NTZ type name) reaches top-level columns and nested children. +func (vcm *arrowValueContainerMaker) makeFieldValueContainer(field arrow.Field, location *time.Location, toTimestampFn timeStampFn, colInfo *colInfo) (columnValues, error) { if location == nil { location = time.UTC } - switch t := t.(type) { + switch t := field.Type.(type) { case *arrow.BooleanType: return &columnValuesTyped[*array.Boolean, bool]{}, nil @@ -649,10 +659,11 @@ func (vcm *arrowValueContainerMaker) makeColumnValueContainer(t arrow.DataType, return &dateValueContainer{location: location}, nil case *arrow.TimestampType: - // A zone-less TIMESTAMP_NTZ (empty TimeZone) carries a wall clock that - // must not be shifted into the session location (#506); a zoned - // TIMESTAMP is an instant rendered .In(location) as before. - return ×tampValueContainer{location: location, toTimestampFn: toTimestampFn, wallClock: t.TimeZone == ""}, nil + // A TIMESTAMP_NTZ carries a wall clock that must not be shifted into the + // session location (#506); a zoned TIMESTAMP is an instant rendered + // .In(location) as before. IsTimestampNTZField is the same check that + // produces the TIMESTAMP_NTZ type name: metadata first, then TimeZone. + return ×tampValueContainer{location: location, toTimestampFn: toTimestampFn, wallClock: arrowscan.IsTimestampNTZField(field)}, nil case *arrow.BinaryType: return &columnValuesTyped[*array.Binary, []byte]{}, nil @@ -660,7 +671,7 @@ func (vcm *arrowValueContainerMaker) makeColumnValueContainer(t arrow.DataType, case *arrow.ListType: lvc := &listValueContainer{listArrayType: t} var err error - lvc.values, err = vcm.makeColumnValueContainer(t.Elem(), location, toTimestampFn, nil) + lvc.values, err = vcm.makeFieldValueContainer(t.ElemField(), location, toTimestampFn, nil) if err != nil { return nil, err } @@ -673,11 +684,11 @@ func (vcm *arrowValueContainerMaker) makeColumnValueContainer(t arrow.DataType, case *arrow.MapType: mvc := &mapValueContainer{mapArrayType: t} var err error - mvc.values, err = vcm.makeColumnValueContainer(t.ItemType(), location, toTimestampFn, nil) + mvc.values, err = vcm.makeFieldValueContainer(t.ItemField(), location, toTimestampFn, nil) if err != nil { return nil, err } - mvc.keys, err = vcm.makeColumnValueContainer(t.KeyType(), location, toTimestampFn, nil) + mvc.keys, err = vcm.makeFieldValueContainer(t.KeyField(), location, toTimestampFn, nil) if err != nil { return nil, err } @@ -700,7 +711,7 @@ func (vcm *arrowValueContainerMaker) makeColumnValueContainer(t arrow.DataType, // string never errors. keyBytes, _ := json.Marshal(f.Name) svc.fieldKeys[i] = string(keyBytes) + ":" - c, err := vcm.makeColumnValueContainer(f.Type, location, toTimestampFn, nil) + c, err := vcm.makeFieldValueContainer(f, location, toTimestampFn, nil) if err != nil { return nil, err } diff --git a/internal/rows/arrowbased/arrowRows_test.go b/internal/rows/arrowbased/arrowRows_test.go index 05b5616b..77caff7c 100644 --- a/internal/rows/arrowbased/arrowRows_test.go +++ b/internal/rows/arrowbased/arrowRows_test.go @@ -2399,6 +2399,56 @@ func TestDecimalInComplexTypes(t *testing.T) { }) } +// TestFieldValueContainerUsesNTZMetadata pins that the Thrift container maker +// decides wall-clock handling from arrowscan.IsTimestampNTZField on the whole +// field, the same check that produces the TIMESTAMP_NTZ type name. Explicit +// Spark:DataType:SqlName metadata beats the physical TimeZone, at top level and +// on nested children. +func TestFieldValueContainerUsesNTZMetadata(t *testing.T) { + ntzMeta := arrow.NewMetadata([]string{"Spark:DataType:SqlName"}, []string{"TIMESTAMP_NTZ"}) + tsMeta := arrow.NewMetadata([]string{"Spark:DataType:SqlName"}, []string{"TIMESTAMP"}) + zoned := arrow.FixedWidthTypes.Timestamp_us + zoneless := &arrow.TimestampType{Unit: arrow.Microsecond} + vcm := &arrowValueContainerMaker{} + + wallClock := func(t *testing.T, c columnValues) bool { + tvc, ok := c.(*timestampValueContainer) + require.True(t, ok, "expected timestampValueContainer, got %T", c) + return tvc.wallClock + } + + t.Run("top level", func(t *testing.T) { + for _, tc := range []struct { + name string + field arrow.Field + want bool + }{ + {"SqlName NTZ on UTC zone", arrow.Field{Type: zoned, Metadata: ntzMeta}, true}, + {"SqlName TIMESTAMP on zone-less", arrow.Field{Type: zoneless, Metadata: tsMeta}, false}, + {"no metadata, zone-less", arrow.Field{Type: zoneless}, true}, + {"no metadata, zoned", arrow.Field{Type: zoned}, false}, + } { + t.Run(tc.name, func(t *testing.T) { + c, err := vcm.makeFieldValueContainer(tc.field, time.UTC, nil, nil) + require.NoError(t, err) + assert.Equal(t, tc.want, wallClock(t, c)) + }) + } + }) + + t.Run("nested children", func(t *testing.T) { + st := arrow.StructOf(arrow.Field{Name: "n", Type: zoned, Nullable: true, Metadata: ntzMeta}) + c, err := vcm.makeFieldValueContainer(arrow.Field{Type: st}, time.UTC, nil, nil) + require.NoError(t, err) + assert.True(t, wallClock(t, c.(*structValueContainer).fieldValues[0]), "struct child") + + lt := arrow.ListOfField(arrow.Field{Name: "element", Type: zoned, Nullable: true, Metadata: ntzMeta}) + c, err = vcm.makeFieldValueContainer(arrow.Field{Type: lt}, time.UTC, nil, nil) + require.NoError(t, err) + assert.True(t, wallClock(t, c.(*listValueContainer).values), "list element") + }) +} + // TestDerivedSchemaTimestampKeepsSessionLocation pins that when the server sends // no ArrowSchema, the schema derived from the Thrift TTableSchema types // TIMESTAMP_TYPE as a *zoned* Arrow timestamp. The zone survives the diff --git a/internal/rows/rows.go b/internal/rows/rows.go index f442b488..c0b090ac 100644 --- a/internal/rows/rows.go +++ b/internal/rows/rows.go @@ -10,10 +10,10 @@ import ( "reflect" "time" - "github.com/apache/arrow/go/v12/arrow" "github.com/apache/arrow/go/v12/arrow/ipc" "github.com/databricks/databricks-sql-go/driverctx" dbsqlerr "github.com/databricks/databricks-sql-go/errors" + "github.com/databricks/databricks-sql-go/internal/arrowscan" "github.com/databricks/databricks-sql-go/internal/cli_service" dbsqlclient "github.com/databricks/databricks-sql-go/internal/client" context2 "github.com/databricks/databricks-sql-go/internal/compat/context" @@ -419,7 +419,9 @@ func (r *rows) cacheDatabaseTypeNameOverrides() { } defer reader.Release() for i, field := range reader.Schema().Fields() { - if ts, ok := field.Type.(*arrow.TimestampType); ok && ts.TimeZone == "" { + // Same check the Arrow value scanner uses (metadata first, then + // TimeZone), so the type name and the value handling stay in sync. + if arrowscan.IsTimestampNTZField(field) { r.databaseTypeNameOverrides[i] = "TIMESTAMP_NTZ" } } From 52d12db2edb8d0d886fee29956271135db1fbdbd Mon Sep 17 00:00:00 2001 From: "peco-engineer-bot[bot]" Date: Thu, 22 Oct 2026 23:58:25 +0000 Subject: [PATCH 11/15] ai: apply changes for #510 (1 review thread) Addresses: - #4354729938 at internal/arrowscan/arrowscan.go:225 Signed-off-by: peco-engineer-bot[bot] --- internal/arrowscan/arrowscan.go | 47 ++++++++++++------- .../arrowscan/timestamp_ntz_nested_test.go | 21 +++++++++ internal/arrowscan/timestamp_ntz_test.go | 16 +++++++ 3 files changed, 67 insertions(+), 17 deletions(-) diff --git a/internal/arrowscan/arrowscan.go b/internal/arrowscan/arrowscan.go index b1ac260d..8ff76aba 100644 --- a/internal/arrowscan/arrowscan.go +++ b/internal/arrowscan/arrowscan.go @@ -108,14 +108,26 @@ 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, ScanOptions{}) + return scanCell(col, row, fieldlessLocation(col, loc), keys, ScanOptions{}) +} + +// fieldlessLocation applies IsTimestampNTZField's zone-less fallback for the +// entry points that receive only an array, not its schema Field: a top-level +// timestamp with an empty TimeZone is treated as TIMESTAMP_NTZ and scanned in +// UTC. Callers that know the Field (the kernel rows via +// ScanCellCachedWithOptions) decide NTZ from its metadata and pass loc directly. +func fieldlessLocation(col arrow.Array, loc *time.Location) *time.Location { + if IsTimestampNTZField(arrow.Field{Type: col.DataType()}) { + return time.UTC + } + return loc } // 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, ScanOptions{DecimalAsFloat: decimalAsFloat}) + return scanCell(col, row, fieldlessLocation(col, loc), keys, ScanOptions{DecimalAsFloat: decimalAsFloat}) } // ScanOptions controls top-level result adaptations that cannot be inferred from @@ -129,7 +141,9 @@ type ScanOptions struct { // 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. +// should continue to use ScanCellCached. Unlike the field-less entry points, loc +// is authoritative here: the caller decides TIMESTAMP_NTZ from the schema Field +// (IsTimestampNTZField) and passes time.UTC for an NTZ column. func ScanCellCachedWithOptions(col arrow.Array, row int, loc *time.Location, keys *StructKeyCache, options ScanOptions) (driver.Value, error) { return scanCell(col, row, loc, keys, options) } @@ -213,19 +227,16 @@ func scanCell(col arrow.Array, row int, loc *time.Location, keys *StructKeyCache // UnixMilli/Unix split into (sec, nsec) internally, so the full 0001–9999 range // round-trips. The Thrift renderer still uses Arrow's ToTime conversion. // - // A zone-less TIMESTAMP_NTZ (empty TimeZone) is a wall clock with no zone, - // so it must NOT be shifted into the session location the way a zoned - // TIMESTAMP instant is: applying .In(loc) would corrupt the wall-clock - // fields the server sent (databricks/databricks-sql-go#506). Only a zoned - // TIMESTAMP is rendered in loc. ScanCell sees only the data type, so this - // uses IsTimestampNTZField's TimeZone fallback; metadata-labeled NTZ - // columns and children reach here with loc already set to UTC (the - // kernel's top-level ntz flag and childLocation). - ts := timestampToTime(int64(c.Value(row)), dt.Unit) - if IsTimestampNTZField(arrow.Field{Type: dt}) { - return ts, nil - } - return inLocation(ts, loc), nil + // A TIMESTAMP_NTZ is a wall clock with no zone, so it must NOT be shifted + // into the session location the way a zoned TIMESTAMP instant is + // (databricks/databricks-sql-go#506). That decision needs the enclosing + // Field's metadata (IsTimestampNTZField), which this arm cannot see, so the + // caller's loc is authoritative: NTZ fields arrive with loc = UTC (the + // kernel's top-level ntz flag, childLocation for nested children, and + // fieldlessLocation for the field-less ScanCell entry points). A UTC value + // rendered in UTC is a no-op, so the NTZ wall clock is preserved, while a + // zone-less field the metadata labels TIMESTAMP is still rendered in loc. + return inLocation(timestampToTime(int64(c.Value(row)), dt.Unit), loc), nil case *array.Decimal128: dt := col.DataType().(*arrow.Decimal128Type) if options.DecimalAsFloat { @@ -597,5 +608,7 @@ func scalarForJSON(col arrow.Array, row int, loc *time.Location) (any, error) { if col.IsNull(row) { return nil, nil } - return ScanCell(col, row, loc) + // scanCell, not ScanCell: keyLoc already came from childLocation on the key + // Field, so the field-less zone-less fallback must not override it. + return scanCell(col, row, loc, nil, ScanOptions{}) } diff --git a/internal/arrowscan/timestamp_ntz_nested_test.go b/internal/arrowscan/timestamp_ntz_nested_test.go index 2576a184..c8ac6f55 100644 --- a/internal/arrowscan/timestamp_ntz_nested_test.go +++ b/internal/arrowscan/timestamp_ntz_nested_test.go @@ -62,6 +62,27 @@ func TestScanCellNestedKernelNTZKeepsWallClock(t *testing.T) { } }) + // A zone-less child the metadata labels TIMESTAMP is a TIMESTAMP (reported as + // such by ColumnTypeInfoForField), so it must be rendered in loc, not left + // in UTC by a data-type-only zone-less fallback at the leaf. + t.Run("zone-less child tagged TIMESTAMP shifts", func(t *testing.T) { + tsMeta := arrow.NewMetadata([]string{databricksTypeNameMetadataKey}, []string{"TIMESTAMP"}) + zoneless := &arrow.TimestampType{Unit: arrow.Microsecond} + b := array.NewStructBuilder(pool, arrow.StructOf(arrow.Field{Name: "z", Type: zoneless, Nullable: true, Metadata: tsMeta})) + defer b.Release() + b.Append(true) + b.FieldBuilder(0).(*array.TimestampBuilder).Append(arrow.Timestamp(wall.UnixMicro())) + arr := b.NewArray() + defer arr.Release() + v, err := ScanCellCached(arr, 0, loc, NewStructKeyCache()) + if err != nil { + t.Fatal(err) + } + if want := `{"z":` + zonedStr + `}`; v != want { + t.Errorf("got %v, want %v", v, want) + } + }) + t.Run("list without NTZ metadata still shifts", func(t *testing.T) { b := array.NewListBuilder(pool, tsUTC) defer b.Release() diff --git a/internal/arrowscan/timestamp_ntz_test.go b/internal/arrowscan/timestamp_ntz_test.go index c232f3cc..0cc8bc07 100644 --- a/internal/arrowscan/timestamp_ntz_test.go +++ b/internal/arrowscan/timestamp_ntz_test.go @@ -56,6 +56,22 @@ func TestScanCellTimestampNTZPreservesWallClock(t *testing.T) { } }) + // The kernel rows path (ScanCellCachedWithOptions) decides NTZ from the + // schema Field and passes loc accordingly, so a zone-less column whose + // metadata says TIMESTAMP (r.ntz false → session loc) must render in loc. + t.Run("kernel path honors caller loc for zone-less TIMESTAMP", func(t *testing.T) { + arr := build("") + defer arr.Release() + v, err := ScanCellCachedWithOptions(arr, 0, loc, nil, ScanOptions{}) + if err != nil { + t.Fatal(err) + } + got := v.(time.Time) + if got.Location() != loc || !got.Equal(wall) { + t.Errorf("got %v (%s), want instant %v in %s", got, got.Location(), wall, loc) + } + }) + t.Run("zoned keeps instant via location", func(t *testing.T) { arr := build("UTC") defer arr.Release() From 935738908782999f767337dcc7a0d4f5d206488d Mon Sep 17 00:00:00 2001 From: "peco-engineer-bot[bot]" Date: Fri, 23 Oct 2026 00:06:43 +0000 Subject: [PATCH 12/15] ai: apply changes for #510 (2 review threads) Addresses: - #4354787992 at internal/backend/kernel/batchiterator_test.go:144 - #4354787998 at internal/rows/arrowbased/timestamp_ntz_test.go:52 Signed-off-by: peco-engineer-bot[bot] --- internal/arrowscan/timestamp_ntz_test.go | 3 +++ internal/backend/kernel/batchiterator_test.go | 3 ++- internal/backend/kernel/rows.go | 20 ++++++++++++++----- .../rows/arrowbased/timestamp_ntz_test.go | 3 +++ 4 files changed, 23 insertions(+), 6 deletions(-) diff --git a/internal/arrowscan/timestamp_ntz_test.go b/internal/arrowscan/timestamp_ntz_test.go index 0cc8bc07..cb71b1d5 100644 --- a/internal/arrowscan/timestamp_ntz_test.go +++ b/internal/arrowscan/timestamp_ntz_test.go @@ -85,5 +85,8 @@ func TestScanCellTimestampNTZPreservesWallClock(t *testing.T) { if !got.Equal(wall) { t.Errorf("zoned instant changed: got %v, want %v", got.UTC(), wall) } + if got.Location() != loc { + t.Errorf("zoned location = %q, want %q", got.Location(), loc) + } }) } diff --git a/internal/backend/kernel/batchiterator_test.go b/internal/backend/kernel/batchiterator_test.go index 28981560..6429e8e1 100644 --- a/internal/backend/kernel/batchiterator_test.go +++ b/internal/backend/kernel/batchiterator_test.go @@ -141,12 +141,13 @@ func TestRowsTimestampNTZWithServerSchema(t *testing.T) { defer record.Release() rowSet := &kernelRows{ op: &kernelOp{location: location}, cur: record, - cols: []string{"zoned", "ntz"}, geoBinary: make([]bool, 2), ntz: []bool{false, true}, + cols: []string{"zoned", "ntz"}, geoBinary: make([]bool, 2), ntz: timestampNTZColumns(schema.Fields()), colTypes: []arrowscan.ColumnTypeInfo{ arrowscan.ColumnTypeInfoForField(schema.Field(0)), arrowscan.ColumnTypeInfoForField(schema.Field(1)), }, } + require.Equal(t, []bool{false, true}, rowSet.ntz) require.Equal(t, "TIMESTAMP", rowSet.ColumnTypeDatabaseTypeName(0)) require.Equal(t, "TIMESTAMP_NTZ", rowSet.ColumnTypeDatabaseTypeName(1)) values := make([]driver.Value, 2) diff --git a/internal/backend/kernel/rows.go b/internal/backend/kernel/rows.go index 87dcd639..697dd71f 100644 --- a/internal/backend/kernel/rows.go +++ b/internal/backend/kernel/rows.go @@ -121,15 +121,11 @@ func newKernelRows(ctx context.Context, op *kernelOp, stream *C.kernel_result_st // (ScanCellCached) via the shared arrowscan.ColumnTypeInfoFor mapper. r.colTypes = make([]arrowscan.ColumnTypeInfo, len(fields)) r.geoBinary = make([]bool, len(fields)) - r.ntz = make([]bool, len(fields)) + r.ntz = timestampNTZColumns(fields) for i, f := range fields { r.cols[i] = f.Name r.colTypes[i] = arrowscan.ColumnTypeInfoForField(f) r.geoBinary[i] = arrowscan.IsBinaryGeospatialField(f) - // The kernel may tag NTZ with UTC in Arrow while preserving its logical - // type; decide once here (from the same mapper as ColumnTypeDatabaseTypeName) - // rather than per cell in next(). - r.ntz[i] = r.colTypes[i].DatabaseTypeName == "TIMESTAMP_NTZ" } // Construction succeeded — now arm the close telemetry callback so a normal // Close() (after row iteration) records CLOSE_STATEMENT. @@ -138,6 +134,20 @@ func newKernelRows(ctx context.Context, op *kernelOp, stream *C.kernel_result_st return r, nil } +// timestampNTZColumns reports, per top-level field, whether it is a +// TIMESTAMP_NTZ wall clock that must scan in UTC rather than the session +// location. The kernel may tag NTZ with a UTC zone in Arrow while preserving its +// logical type in metadata, so this uses arrowscan.IsTimestampNTZField — the same +// check ColumnTypeDatabaseTypeName is derived from — decided once per result set +// rather than per cell in next(). +func timestampNTZColumns(fields []arrow.Field) []bool { + ntz := make([]bool, len(fields)) + for i, f := range fields { + ntz[i] = arrowscan.IsTimestampNTZField(f) + } + return ntz +} + // Columns returns the result-set column names. func (r *kernelRows) Columns() []string { return r.cols } diff --git a/internal/rows/arrowbased/timestamp_ntz_test.go b/internal/rows/arrowbased/timestamp_ntz_test.go index 8761e66c..938b41cd 100644 --- a/internal/rows/arrowbased/timestamp_ntz_test.go +++ b/internal/rows/arrowbased/timestamp_ntz_test.go @@ -55,5 +55,8 @@ func TestThriftTimestampNTZPreservesWallClock(t *testing.T) { if !got.Equal(wall) { t.Errorf("zoned instant changed: got %v, want %v", got.UTC(), wall) } + if got.Location() != loc { + t.Errorf("zoned location = %q, want %q", got.Location(), loc) + } }) } From 985362d022ac672a16eb8f7b43dadb2573881c9a Mon Sep 17 00:00:00 2001 From: "peco-engineer-bot[bot]" Date: Fri, 23 Oct 2026 00:13:43 +0000 Subject: [PATCH 13/15] ai: apply changes for #510 (1 review thread) Addresses: - #4354828451 at internal/rows/rows.go:394 Signed-off-by: peco-engineer-bot[bot] --- internal/rows/rows.go | 30 +++++++++++++++++++++--------- 1 file changed, 21 insertions(+), 9 deletions(-) diff --git a/internal/rows/rows.go b/internal/rows/rows.go index c0b090ac..f7a7a0f4 100644 --- a/internal/rows/rows.go +++ b/internal/rows/rows.go @@ -8,6 +8,7 @@ import ( "io" "math" "reflect" + "sync" "time" "github.com/apache/arrow/go/v12/arrow/ipc" @@ -49,6 +50,9 @@ type rows struct { resultSetMetadata *cli_service.TGetResultSetMetadataResp schema *cli_service.TTableSchema databaseTypeNameOverrides map[int]string + // databaseTypeNameOverridesOnce guards the lazy, one-time build of + // databaseTypeNameOverrides (ColumnTypes may be called concurrently). + databaseTypeNameOverridesOnce sync.Once config *config.Config @@ -390,22 +394,29 @@ func (r *rows) ColumnTypeDatabaseTypeName(index int) string { // as Arrow. If a server sent such metadata but then returned TColumn pages, the // column would still report TIMESTAMP_NTZ while its values went through the // legacy column scanner. +// +// database/sql's Rows.ColumnTypes() only holds a read lock, so this may be +// reached concurrently. The map is built in a local variable and published +// exactly once via sync.Once, so readers never observe a partially filled map. func (r *rows) cacheDatabaseTypeNameOverrides() { - if r.databaseTypeNameOverrides != nil { - return - } - r.databaseTypeNameOverrides = make(map[int]string) + r.databaseTypeNameOverridesOnce.Do(func() { + r.databaseTypeNameOverrides = r.buildDatabaseTypeNameOverrides() + }) +} + +func (r *rows) buildDatabaseTypeNameOverrides() map[int]string { + overrides := make(map[int]string) if r.resultSetMetadata == nil || len(r.resultSetMetadata.ArrowSchema) == 0 { - return + return overrides } if r.config != nil && !r.config.UseArrowBatches { - return + return overrides } if r.resultSetMetadata.IsSetResultFormat() { switch r.resultSetMetadata.GetResultFormat() { case cli_service.TSparkRowSetType_ARROW_BASED_SET, cli_service.TSparkRowSetType_URL_BASED_SET: default: - return + return overrides } } @@ -415,16 +426,17 @@ func (r *rows) cacheDatabaseTypeNameOverrides() { // so retrying would fail the same way. Columns fall back to the legacy // TIMESTAMP name; the Arrow row scanner surfaces the parse error itself. r.logger().Debug().Err(err).Msg("databricks: failed to parse arrow schema for column type name overrides; falling back to thrift type names") - return + return overrides } defer reader.Release() for i, field := range reader.Schema().Fields() { // Same check the Arrow value scanner uses (metadata first, then // TimeZone), so the type name and the value handling stay in sync. if arrowscan.IsTimestampNTZField(field) { - r.databaseTypeNameOverrides[i] = "TIMESTAMP_NTZ" + overrides[i] = "TIMESTAMP_NTZ" } } + return overrides } // ColumnTypeNullable returns a flag indicating whether the column is nullable From 95d217ecdbe54a8f7fee4261ca3103882f7bca8e Mon Sep 17 00:00:00 2001 From: "peco-engineer-bot[bot]" Date: Fri, 23 Oct 2026 00:19:58 +0000 Subject: [PATCH 14/15] ai: apply changes for #510 (1 review thread) Addresses: - #4354869375 at internal/rows/arrowbased/columnValues.go:480 Signed-off-by: peco-engineer-bot[bot] --- internal/arrowscan/coltype.go | 13 +++++++------ internal/rows/arrowbased/columnValues.go | 6 ++++-- 2 files changed, 11 insertions(+), 8 deletions(-) diff --git a/internal/arrowscan/coltype.go b/internal/arrowscan/coltype.go index b746a40f..c46c1cfb 100644 --- a/internal/arrowscan/coltype.go +++ b/internal/arrowscan/coltype.go @@ -78,12 +78,13 @@ func ColumnTypeInfoFor(dt arrow.DataType) ColumnTypeInfo { case arrow.DATE32, arrow.DATE64: return ColumnTypeInfo{DatabaseTypeName: "DATE", ScanType: scanTypeDateTime} case arrow.TIMESTAMP: - // TIMESTAMP and TIMESTAMP_NTZ both arrive as Arrow timestamp[us] and are - // distinguishable only by the schema's TimeZone field: a zoned TIMESTAMP - // carries a zone, a zone-less TIMESTAMP_NTZ carries an empty TimeZone. The - // database/sql type name must reflect that distinction so a wall clock is - // not mistaken for an instant (databricks/databricks-sql-go#506). With no - // field metadata, IsTimestampNTZField falls back to that TimeZone check. + // TIMESTAMP and TIMESTAMP_NTZ both arrive as Arrow timestamp[us]. NTZ is + // decided by IsTimestampNTZField: the field's logical-type metadata first + // (which may mark a UTC-zoned field as NTZ, as the kernel does), then an + // empty TimeZone as the fallback. The database/sql type name must reflect + // that distinction so a wall clock is not mistaken for an instant + // (databricks/databricks-sql-go#506). Only the DataType is available here, + // so with no field metadata this resolves via the empty-TimeZone fallback. return timestampColumnTypeInfo(arrow.Field{Type: dt}) case arrow.DECIMAL128: // Match Thrift's sql.RawBytes scan type even though the kernel renders the diff --git a/internal/rows/arrowbased/columnValues.go b/internal/rows/arrowbased/columnValues.go index f525478c..f6d877d0 100644 --- a/internal/rows/arrowbased/columnValues.go +++ b/internal/rows/arrowbased/columnValues.go @@ -477,8 +477,10 @@ type timestampValueContainer struct { timeArray *array.Timestamp location *time.Location toTimestampFn func(arrow.Timestamp) time.Time - // wallClock is true for a zone-less TIMESTAMP_NTZ column (empty Arrow - // TimeZone): its value is a wall clock with no zone and must be returned + // wallClock is true for a TIMESTAMP_NTZ column, as decided by + // arrowscan.IsTimestampNTZField (logical-type metadata first, then an empty + // Arrow TimeZone; a UTC-zoned field may still be NTZ by metadata): its + // value is a wall clock with no zone and must be returned // with the exact fields the server sent, WITHOUT the session .In(location) // shift that is correct only for a zoned TIMESTAMP instant // (databricks/databricks-sql-go#506). From 3612b9d034a79eeb42df095130b00f76f196e314 Mon Sep 17 00:00:00 2001 From: "peco-engineer-bot[bot]" Date: Fri, 23 Oct 2026 00:25:08 +0000 Subject: [PATCH 15/15] ai: apply changes for #510 (1 review thread) Addresses: - #4354900157 at driver_e2e_test.go:787 Signed-off-by: peco-engineer-bot[bot] --- driver_e2e_test.go | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/driver_e2e_test.go b/driver_e2e_test.go index f07837b0..bf8d4c8c 100644 --- a/driver_e2e_test.go +++ b/driver_e2e_test.go @@ -775,9 +775,10 @@ func TestE2ETimestampNTZPreservesWallClock(t *testing.T) { require.NoError(t, rows.Scan(&zoned, &ntz)) // The NTZ value must preserve the wall clock it was given, with no shift - // from the session timezone. We compare the wall-clock fields directly so - // the assertion is independent of whatever *time.Location the driver tags - // the zone-less value with. + // from the session timezone. We check the wall-clock fields directly, and + // separately pin the location: returning zone-less TIMESTAMP_NTZ values + // tagged with time.UTC is part of the documented contract (see CHANGELOG + // and README), not an incidental detail, so do not loosen that check. assert.Equal(t, 2024, ntz.Year()) assert.Equal(t, time.June, ntz.Month()) assert.Equal(t, 15, ntz.Day())