diff --git a/CHANGELOG.md b/CHANGELOG.md index dbae70fa..444cf6dc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,9 @@ # Release History +## 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). 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. - 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..35a4ad0e 100644 --- a/CONNECTION_PARAMETERS.md +++ b/CONNECTION_PARAMETERS.md @@ -159,11 +159,17 @@ 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 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. +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 d59054af..504ec79a 100644 --- a/README.md +++ b/README.md @@ -360,12 +360,18 @@ 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 +(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. +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 e3523d69..e78f5a61 100644 --- a/doc.go +++ b/doc.go @@ -290,12 +290,18 @@ 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 (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. +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 4bdaeb84..bf8d4c8c 100644 --- a/driver_e2e_test.go +++ b/driver_e2e_test.go @@ -709,3 +709,87 @@ 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 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()) + 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()) + 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. + 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..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) } @@ -211,8 +225,17 @@ 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 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) @@ -417,16 +440,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: @@ -481,19 +507,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 } } @@ -505,14 +534,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 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 } } @@ -576,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/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..c46c1cfb 100644 --- a/internal/arrowscan/coltype.go +++ b/internal/arrowscan/coltype.go @@ -78,7 +78,14 @@ func ColumnTypeInfoFor(dt arrow.DataType) ColumnTypeInfo { case arrow.DATE32, arrow.DATE64: return ColumnTypeInfo{DatabaseTypeName: "DATE", ScanType: scanTypeDateTime} case arrow.TIMESTAMP: - return ColumnTypeInfo{DatabaseTypeName: "TIMESTAMP", ScanType: scanTypeDateTime} + // 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 // value as an exact string — both convert cleanly to a caller's *string/*[]byte. @@ -140,11 +147,12 @@ 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 != nil && field.Type.ID() == arrow.TIMESTAMP { + return timestampColumnTypeInfo(field) + } name, geospatial := LogicalGeospatialType(field) if !geospatial { return ColumnTypeInfoFor(field.Type) @@ -160,6 +168,53 @@ func ColumnTypeInfoForField(field arrow.Field) ColumnTypeInfo { return info } +// 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 + } + 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 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 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 new file mode 100644 index 00000000..f3336829 --- /dev/null +++ b/internal/arrowscan/coltype_ntz_test.go @@ -0,0 +1,97 @@ +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") + } + }) +} + +// 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 + 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 00000000..2115c971 Binary files /dev/null and b/internal/arrowscan/testdata/timestamp_ntz_kernel_schema.arrow differ 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/arrowscan/timestamp_ntz_nested_test.go b/internal/arrowscan/timestamp_ntz_nested_test.go new file mode 100644 index 00000000..c8ac6f55 --- /dev/null +++ b/internal/arrowscan/timestamp_ntz_nested_test.go @@ -0,0 +1,101 @@ +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) + } + }) + + // 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() + 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/arrowscan/timestamp_ntz_test.go b/internal/arrowscan/timestamp_ntz_test.go new file mode 100644 index 00000000..cb71b1d5 --- /dev/null +++ b/internal/arrowscan/timestamp_ntz_test.go @@ -0,0 +1,92 @@ +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) + } + 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) + } + }) + + // 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() + 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) + } + 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 4eb22c8f..6429e8e1 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,39 @@ 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), 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) + 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..697dd71f 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,6 +121,7 @@ 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 = timestampNTZColumns(fields) for i, f := range fields { r.cols[i] = f.Name r.colTypes[i] = arrowscan.ColumnTypeInfoForField(f) @@ -132,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 } @@ -235,7 +251,11 @@ 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 + if r.ntz[c] { + 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 58970d85..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,7 +659,11 @@ func (vcm *arrowValueContainerMaker) makeColumnValueContainer(t arrow.DataType, return &dateValueContainer{location: location}, nil case *arrow.TimestampType: - return ×tampValueContainer{location: location, toTimestampFn: toTimestampFn}, 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 @@ -657,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 } @@ -670,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 } @@ -697,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 713ab314..77caff7c 100644 --- a/internal/rows/arrowbased/arrowRows_test.go +++ b/internal/rows/arrowbased/arrowRows_test.go @@ -2398,3 +2398,107 @@ func TestDecimalInComplexTypes(t *testing.T) { assert.Equal(t, `{"col2":null}`, v) }) } + +// 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 +// 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") + }) + } +} 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..f6d877d0 100644 --- a/internal/rows/arrowbased/columnValues.go +++ b/internal/rows/arrowbased/columnValues.go @@ -477,13 +477,24 @@ type timestampValueContainer struct { timeArray *array.Timestamp location *time.Location toTimestampFn func(arrow.Timestamp) time.Time + // 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). + 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..938b41cd --- /dev/null +++ b/internal/rows/arrowbased/timestamp_ntz_test.go @@ -0,0 +1,62 @@ +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) + } + 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) { + 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) + } + if got.Location() != loc { + t.Errorf("zoned location = %q, want %q", got.Location(), loc) + } + }) +} diff --git a/internal/rows/rows.go b/internal/rows/rows.go index 1ba3bb7c..f7a7a0f4 100644 --- a/internal/rows/rows.go +++ b/internal/rows/rows.go @@ -1,16 +1,20 @@ package rows import ( + "bytes" "context" "database/sql" "database/sql/driver" "io" "math" "reflect" + "sync" "time" + "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" @@ -43,8 +47,12 @@ 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 + // databaseTypeNameOverridesOnce guards the lazy, one-time build of + // databaseTypeNameOverrides (ColumnTypes may be called concurrently). + databaseTypeNameOverridesOnce sync.Once config *config.Config @@ -364,11 +372,73 @@ func (r *rows) ColumnTypeDatabaseTypeName(index int) string { return "" } + r.cacheDatabaseTypeNameOverrides() + if name, ok := r.databaseTypeNameOverrides[index]; ok { + return name + } + dbtype := rowscanner.GetDBTypeName(column) return dbtype } +// 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 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. +// +// 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() { + 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 overrides + } + if r.config != nil && !r.config.UseArrowBatches { + 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 overrides + } + } + + 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 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) { + overrides[i] = "TIMESTAMP_NTZ" + } + } + return overrides +} + // 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 e04f9ef5..95688399 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) { @@ -708,6 +712,152 @@ func TestColumnTypeDatabaseTypeName(t *testing.T) { assert.Equal(t, expectedScanTypes, scanTypes) } +func TestColumnTypeDatabaseTypeNameArrowTimestamps(t *testing.T) { + 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{ + 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}, + }}}, + }) + } + 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) + }) + } +} + +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 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)) + }) + } +} + +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{ 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 00000000..3cd1798e Binary files /dev/null and b/internal/rows/testdata/timestamp_ntz_arrow_schema.arrow differ diff --git a/internal/rows/testdata/timestamp_ntz_arrow_schema.md b/internal/rows/testdata/timestamp_ntz_arrow_schema.md new file mode 100644 index 00000000..e74d47b0 --- /dev/null +++ b/internal/rows/testdata/timestamp_ntz_arrow_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, CAST(7 AS BIGINT) AS number +``` diff --git a/kernel_e2e_test.go b/kernel_e2e_test.go index 18eea660..55feacae 100644 --- a/kernel_e2e_test.go +++ b/kernel_e2e_test.go @@ -123,14 +123,8 @@ 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 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, @@ -146,28 +140,41 @@ func TestKernelE2ETimestampNTZ(t *testing.T) { } ctx := context.Background() - var ltz time.Time - if err := db.QueryRowContext(ctx, - "SELECT CAST('2026-07-09 12:00:00' AS TIMESTAMP)").Scan(<z); err != nil { - t.Fatalf("TIMESTAMP query: %v", err) + result, err := db.QueryContext(ctx, + "SELECT TIMESTAMP '2026-07-09 12:00:00' AS zoned, TIMESTAMP_NTZ '2026-07-09 12:00:00' AS ntz") + if err != nil { + t.Fatalf("metadata query: %v", err) + } + defer result.Close() //nolint:errcheck + columnTypes, err := result.ColumnTypes() + if err != nil { + t.Fatalf("ColumnTypes: %v", err) + } + if len(columnTypes) != 2 { + t.Fatalf("column count = %d, want 2", len(columnTypes)) + } + for i, want := range []string{"TIMESTAMP", "TIMESTAMP_NTZ"} { + if got := columnTypes[i].DatabaseTypeName(); got != want { + t.Errorf("column %d type = %q, want %q", i, got, want) + } + } + if !result.Next() { + t.Fatalf("expected a timestamp row: %v", result.Err()) + } + var ltz, ntz time.Time + if err := result.Scan(<z, &ntz); err != nil { + t.Fatalf("Scan: %v", err) } // A zoned TIMESTAMP literal is the wall-clock in the session tz. if want := time.Date(2026, 7, 9, 12, 0, 0, 0, loc); !ltz.Equal(want) { t.Errorf("TIMESTAMP = %s, want %s", ltz, want) } - var ntz time.Time - if err := db.QueryRowContext(ctx, - "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()) + 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().String() != tz { - t.Errorf("TIMESTAMP_NTZ location = %q, want %q", ntz.Location(), tz) + if ntz.Location() != time.UTC { + t.Errorf("TIMESTAMP_NTZ location = %q, want UTC", ntz.Location()) } }