This is an automated email from the ASF dual-hosted git repository.
zeroshade pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-go.git
The following commit(s) were added to refs/heads/main by this push:
new c9791ae4 Add `TimestampWithOffset` canonical extension type (#558)
c9791ae4 is described below
commit c9791ae4041f3b06bda6bb893baa3382474048cd
Author: Lucas Valente <[email protected]>
AuthorDate: Sun Jul 19 03:36:35 2026 +0200
Add `TimestampWithOffset` canonical extension type (#558)
# Which issue does this PR close?
Implements the new `arrow.timestamp_with_offset` canonical extension
type for arrow-go.
# Rationale for this change
Be compatible with the Arrow columnar spec's
`arrow.timestamp_with_offset` canonical extension type.
# What changes are included in this PR?
Adds a `TimestampWithOffset` extension type. This type represents a
timestamp column that stores a potentially different timezone offset per
value: the timestamp is stored in UTC alongside the original timezone
offset in minutes. The offset-in-minutes field can be primitive-,
dictionary-, or run-end-encoded.
**This PR has been decoupled from the inner-nullability parity work it
was previously based on.** Those comparison / JSON-marshaling changes
now live in #918 and can be reviewed and merged independently. This PR
is now scoped to just the extension type and touches only:
- `arrow/extensions/timestamp_with_offset.go` (new)
- `arrow/extensions/timestamp_with_offset_test.go` (new)
- `arrow/extensions/extensions.go` (register the canonical type)
To keep this PR self-contained, the JSON round-trip test asserts the
values the type guarantees (per-row instant, timezone offset, and
validity) rather than `array.RecordEqual`, which would compare inner
non-nullable struct child validity bitmaps — the parity concern tracked
separately in #918.
# Are these changes tested?
Yes — primitive/dictionary/run-end offset encodings, JSON marshal +
round-trip, and IPC round-trip.
# Are there any user-facing changes?
Yes, this is a new canonical extension type.
---------
Co-authored-by: Matt Topol <[email protected]>
Co-authored-by: Matt Topol <[email protected]>
---
arrow/extensions/extensions.go | 7 +-
arrow/extensions/timestamp_with_offset.go | 608 +++++++++++++++++++++++
arrow/extensions/timestamp_with_offset_test.go | 646 +++++++++++++++++++++++++
3 files changed, 1258 insertions(+), 3 deletions(-)
diff --git a/arrow/extensions/extensions.go b/arrow/extensions/extensions.go
index 04566c75..6f13aa64 100644
--- a/arrow/extensions/extensions.go
+++ b/arrow/extensions/extensions.go
@@ -21,11 +21,12 @@ import (
)
var canonicalExtensionTypes = []arrow.ExtensionType{
- NewBool8Type(),
- NewUUIDType(),
- &OpaqueType{},
&JSONType{},
+ &OpaqueType{},
+ &TimestampWithOffsetType{},
&VariantType{},
+ NewBool8Type(),
+ NewUUIDType(),
}
func init() {
diff --git a/arrow/extensions/timestamp_with_offset.go
b/arrow/extensions/timestamp_with_offset.go
new file mode 100644
index 00000000..47f8d241
--- /dev/null
+++ b/arrow/extensions/timestamp_with_offset.go
@@ -0,0 +1,608 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package extensions
+
+import (
+ "fmt"
+ "iter"
+ "math"
+ "reflect"
+ "slices"
+ "strings"
+ "time"
+
+ "github.com/apache/arrow-go/v18/arrow"
+ "github.com/apache/arrow-go/v18/arrow/array"
+ "github.com/apache/arrow-go/v18/arrow/memory"
+ "github.com/apache/arrow-go/v18/internal/json"
+)
+
+// TimestampWithOffsetType represents a timestamp column that stores a
timezone offset per row instead of
+// applying the same timezone offset to the entire column.
+type TimestampWithOffsetType struct {
+ arrow.ExtensionBase
+}
+
+func isOffsetTypeOk(offsetType arrow.DataType) bool {
+ switch offsetType := offsetType.(type) {
+ case *arrow.Int16Type:
+ return true
+ case *arrow.DictionaryType:
+ return arrow.TypeEqual(offsetType.ValueType,
arrow.PrimitiveTypes.Int16)
+ case *arrow.RunEndEncodedType:
+ return offsetType.ValidRunEndsType(offsetType.RunEnds()) &&
+ arrow.TypeEqual(offsetType.Encoded(),
arrow.PrimitiveTypes.Int16)
+ // FIXME: Technically this should be non-nullable, but a Arrow
IPC does not deserialize
+ // ValueNullable properly, so enforcing this here would always
fail when reading from an IPC
+ // stream
+ // !offsetType.ValueNullable
+ default:
+ return false
+ }
+}
+
+// Whether the storageType is compatible with TimestampWithOffset.
+//
+// Returns (time_unit, offset_type, ok). If ok is false, time_unit and
offset_type are garbage.
+func isDataTypeCompatible(storageType arrow.DataType) (unit arrow.TimeUnit,
offsetType arrow.DataType, ok bool) {
+ unit = arrow.Second
+ offsetType = arrow.PrimitiveTypes.Int16
+ ok = false
+
+ st, compat := storageType.(*arrow.StructType)
+ if !compat || st.NumFields() != 2 {
+ return
+ }
+
+ if ts, compat := st.Field(0).Type.(*arrow.TimestampType); compat &&
ts.TimeZone == "UTC" {
+ unit = ts.TimeUnit()
+ } else {
+ return
+ }
+
+ maybeOffset := st.Field(1)
+ offsetType = maybeOffset.Type
+
+ ok = st.Field(0).Name == "timestamp" &&
+ !st.Field(0).Nullable &&
+ maybeOffset.Name == "offset_minutes" &&
+ isOffsetTypeOk(offsetType) &&
+ !maybeOffset.Nullable
+ return
+}
+
+// NewTimestampWithOffsetType creates a new TimestampWithOffsetType with the
underlying storage type set correctly to
+// Struct(timestamp=Timestamp(T, "UTC"), offset_minutes=Int16), where T is any
TimeUnit.
+func NewTimestampWithOffsetType(unit arrow.TimeUnit) *TimestampWithOffsetType {
+ v, _ := NewTimestampWithOffsetTypeCustomOffset(unit,
arrow.PrimitiveTypes.Int16)
+ // SAFETY: This should never error as Int16 is always a valid offset
type
+
+ return v
+}
+
+// NewTimestampWithOffsetTypeCustomOffset creates a new
TimestampWithOffsetType with the underlying storage type set correctly to
+// Struct(timestamp=Timestamp(T, "UTC"), offset_minutes=O), where T is any
TimeUnit and O is a valid offset type.
+//
+// The error will be populated if the data type is not a valid encoding of the
offsets field.
+func NewTimestampWithOffsetTypeCustomOffset(unit arrow.TimeUnit, offsetType
arrow.DataType) (*TimestampWithOffsetType, error) {
+ if !isOffsetTypeOk(offsetType) {
+ return nil, fmt.Errorf("invalid offset type %s", offsetType)
+ }
+
+ return &TimestampWithOffsetType{
+ ExtensionBase: arrow.ExtensionBase{
+ Storage: arrow.StructOf(
+ arrow.Field{
+ Name: "timestamp",
+ Type: &arrow.TimestampType{
+ Unit: unit,
+ TimeZone: "UTC",
+ },
+ Nullable: false,
+ },
+ arrow.Field{
+ Name: "offset_minutes",
+ Type: offsetType,
+ Nullable: false,
+ },
+ ),
+ },
+ }, nil
+}
+
+type DictIndexType interface {
+ *arrow.Int8Type | *arrow.Int16Type | *arrow.Int32Type |
*arrow.Int64Type |
+ *arrow.Uint8Type | *arrow.Uint16Type | *arrow.Uint32Type |
*arrow.Uint64Type
+}
+
+type RunEndsType interface {
+ *arrow.Int16Type | *arrow.Int32Type | *arrow.Int64Type
+}
+
+// NewTimestampWithOffsetTypeDictionaryEncoded creates a new
TimestampWithOffsetType with the underlying storage type set correctly to
+// Struct(timestamp=Timestamp(T, "UTC"), offset_minutes=Dictionary(I, Int16)),
where T is any TimeUnit and I is a
+// valid Dictionary index type.
+func NewTimestampWithOffsetTypeDictionaryEncoded[I DictIndexType](unit
arrow.TimeUnit, index I) *TimestampWithOffsetType {
+ offsetType := arrow.DictionaryType{
+ IndexType: arrow.DataType(index),
+ ValueType: arrow.PrimitiveTypes.Int16,
+ Ordered: false,
+ }
+ v, _ := NewTimestampWithOffsetTypeCustomOffset(unit, &offsetType)
+ // SAFETY: This should never error as DictIndexType is always a valid
index type
+
+ return v
+}
+
+// NewTimestampWithOffsetTypeRunEndEncoded creates a new
TimestampWithOffsetType with the underlying storage type set correctly to
+// Struct(timestamp=Timestamp(T, "UTC"), offset_minutes=RunEndEncoded(E,
Int16)), where T is any TimeUnit and E is a
+// valid run-ends type.
+func NewTimestampWithOffsetTypeRunEndEncoded[E RunEndsType](unit
arrow.TimeUnit, runEnds E) *TimestampWithOffsetType {
+ offsetType := arrow.RunEndEncodedOf(arrow.DataType(runEnds),
arrow.PrimitiveTypes.Int16)
+
+ v, _ := NewTimestampWithOffsetTypeCustomOffset(unit, offsetType)
+ // SAFETY: This should never error as RunEndsType always a valid run
ends type
+
+ return v
+
+}
+
+func (b *TimestampWithOffsetType) ArrayType() reflect.Type {
+ return reflect.TypeOf(TimestampWithOffsetArray{})
+}
+
+func (b *TimestampWithOffsetType) ExtensionName() string { return
"arrow.timestamp_with_offset" }
+
+func (b *TimestampWithOffsetType) String() string {
+ return fmt.Sprintf("extension<%s>", b.ExtensionName())
+}
+
+func (e *TimestampWithOffsetType) MarshalJSON() ([]byte, error) {
+ return []byte(fmt.Sprintf(`{"name":"%s","metadata":%s}`,
e.ExtensionName(), e.Serialize())), nil
+}
+
+func (b *TimestampWithOffsetType) Serialize() string { return "" }
+
+func (b *TimestampWithOffsetType) Deserialize(storageType arrow.DataType, data
string) (arrow.ExtensionType, error) {
+ if data != "" && data != "{}" {
+ return nil, fmt.Errorf("serialized metadata for
TimestampWithOffset extension type must be '' or '{}', found: %s", data)
+ }
+
+ timeUnit, offsetType, ok := isDataTypeCompatible(storageType)
+ if !ok {
+ return nil, fmt.Errorf("invalid storage type for
TimestampWithOffsetType: %s", storageType.Name())
+ }
+
+ return NewTimestampWithOffsetTypeCustomOffset(timeUnit, offsetType)
+}
+
+func (b *TimestampWithOffsetType) ExtensionEquals(other arrow.ExtensionType)
bool {
+ return b.ExtensionName() == other.ExtensionName() &&
+ arrow.TypeEqual(b.StorageType(), other.StorageType())
+}
+
+func (b *TimestampWithOffsetType) OffsetType() arrow.DataType {
+ return b.ExtensionBase.Storage.(*arrow.StructType).Field(1).Type
+}
+
+func (b *TimestampWithOffsetType) TimeUnit() arrow.TimeUnit {
+ return
b.ExtensionBase.Storage.(*arrow.StructType).Field(0).Type.(*arrow.TimestampType).TimeUnit()
+}
+
+func (b *TimestampWithOffsetType) NewBuilder(mem memory.Allocator)
array.Builder {
+ v, _ := NewTimestampWithOffsetBuilder(mem, b.TimeUnit(), b.OffsetType())
+ return v
+}
+
+// TimestampWithOffsetArray is a simple array of struct
+type TimestampWithOffsetArray struct {
+ array.ExtensionArrayBase
+}
+
+func (a *TimestampWithOffsetArray) String() string {
+ var o strings.Builder
+ o.WriteString("[")
+ for i := 0; i < a.Len(); i++ {
+ if i > 0 {
+ o.WriteString(" ")
+ }
+ switch {
+ case a.IsNull(i):
+ o.WriteString(array.NullValueStr)
+ default:
+ fmt.Fprintf(&o, "\"%s\"", a.Value(i))
+ }
+ }
+ o.WriteString("]")
+ return o.String()
+}
+
+func timeFromFieldValues(utcTimestamp arrow.Timestamp, offsetMinutes int16,
unit arrow.TimeUnit) time.Time {
+ // Derive the sign from the whole offset: integer hours is 0 for any
offset
+ // with magnitude below one hour and so cannot carry a negative sign
(e.g.
+ // -30 minutes must format as "UTC-00:30", not "UTC+00:30").
+ sign := "+"
+ abs := int(offsetMinutes)
+ if abs < 0 {
+ sign = "-"
+ abs = -abs
+ }
+
+ name := fmt.Sprintf("UTC%s%02d:%02d", sign, abs/60, abs%60)
+ loc := time.FixedZone(name, int(offsetMinutes)*60)
+ return utcTimestamp.ToTime(unit).In(loc)
+}
+
+func fieldValuesFromTime(t time.Time, unit arrow.TimeUnit) (arrow.Timestamp,
int16) {
+ _, offsetSeconds := t.Zone()
+ offsetMinutes := int16(offsetSeconds / 60)
+ ts, _ := arrow.TimestampFromTime(t.UTC(), unit)
+ return ts, offsetMinutes
+}
+
+// Get the raw arrow values at the given index
+//
+// SAFETY: the value at i must not be nil
+func (a *TimestampWithOffsetArray) rawValueUnsafe(i int) (arrow.Timestamp,
int16, arrow.TimeUnit) {
+ structs := a.Storage().(*array.Struct)
+
+ timestampField := structs.Field(0)
+ timestamps := timestampField.(*array.Timestamp)
+
+ timeUnit := timestampField.DataType().(*arrow.TimestampType).Unit
+ utcTimestamp := timestamps.Value(i)
+
+ var offsetMinutes int16
+
+ switch offsets := structs.Field(1).(type) {
+ case *array.Int16:
+ offsetMinutes = offsets.Value(i)
+ case *array.Dictionary:
+ offsetMinutes =
offsets.Dictionary().(*array.Int16).Value(offsets.GetValueIndex(i))
+ case *array.RunEndEncoded:
+ offsetMinutes =
offsets.Values().(*array.Int16).Value(offsets.GetPhysicalIndex(i))
+ }
+
+ return utcTimestamp, offsetMinutes, timeUnit
+}
+
+func (a *TimestampWithOffsetArray) Value(i int) time.Time {
+ if a.IsNull(i) {
+ return time.Time{}
+ }
+ utcTimestamp, offsetMinutes, timeUnit := a.rawValueUnsafe(i)
+ return timeFromFieldValues(utcTimestamp, offsetMinutes, timeUnit)
+}
+
+// Iterates over the array returning the timestamp at each position.
+//
+// The second parameter indicates whether the timestamp is valid or not.
+//
+// This will iterate using the fastest method given the underlying storage
array
+func (a *TimestampWithOffsetArray) iterValues() iter.Seq2[time.Time, bool] {
+ return func(yield func(time.Time, bool) bool) {
+ structs := a.Storage().(*array.Struct)
+ offsets := structs.Field(1)
+ if reeOffsets, isRee := offsets.(*array.RunEndEncoded); isRee {
+ timestampField := structs.Field(0)
+ timeUnit :=
timestampField.DataType().(*arrow.TimestampType).Unit
+ timestamps := timestampField.(*array.Timestamp)
+
+ offsetValues := reeOffsets.Values().(*array.Int16)
+ // Run-ends are absolute over the unsliced offsets
child, so a sliced
+ // array must begin at the physical run covering its
logical offset and
+ // advance using absolute positions (logicalOffset + i).
+ logicalOffset := reeOffsets.Offset()
+ offsetPhysicalIdx := reeOffsets.GetPhysicalOffset()
+
+ var getRunEnd (func(int) int)
+ switch arr := reeOffsets.RunEndsArr().(type) {
+ case *array.Int16:
+ getRunEnd = func(idx int) int { return
int(arr.Value(idx)) }
+ case *array.Int32:
+ getRunEnd = func(idx int) int { return
int(arr.Value(idx)) }
+ case *array.Int64:
+ getRunEnd = func(idx int) int { return
int(arr.Value(idx)) }
+ }
+
+ for i := 0; i < a.Len(); i++ {
+ if logicalOffset+i >=
getRunEnd(offsetPhysicalIdx) {
+ offsetPhysicalIdx += 1
+ }
+
+ var ts time.Time
+ valid := a.IsValid(i)
+ if valid {
+ utcTimestamp := timestamps.Value(i)
+ offsetMinutes :=
offsetValues.Value(offsetPhysicalIdx)
+ v := timeFromFieldValues(utcTimestamp,
offsetMinutes, timeUnit)
+ ts = v
+ }
+
+ if !yield(ts, valid) {
+ return
+ }
+ }
+ } else {
+ for i := 0; i < a.Len(); i++ {
+ var ts time.Time
+ valid := a.IsValid(i)
+ if valid {
+ utcTimestamp, offsetMinutes, timeUnit
:= a.rawValueUnsafe(i)
+ v := timeFromFieldValues(utcTimestamp,
offsetMinutes, timeUnit)
+ ts = v
+ }
+
+ if !yield(ts, valid) {
+ return
+ }
+ }
+ }
+ }
+}
+
+func (a *TimestampWithOffsetArray) Values() []time.Time {
+ return slices.Collect(func(yield func(time.Time) bool) {
+ for t := range a.iterValues() {
+ if !yield(t) {
+ return
+ }
+ }
+ })
+}
+
+func (a *TimestampWithOffsetArray) ValueStr(i int) string {
+ switch {
+ case a.IsNull(i):
+ return array.NullValueStr
+ default:
+ return a.Value(i).String()
+ }
+}
+
+func (a *TimestampWithOffsetArray) MarshalJSON() ([]byte, error) {
+ values := make([]interface{}, a.Len())
+ i := 0
+ for ts, valid := range a.iterValues() {
+ if !valid {
+ values[i] = nil
+ } else {
+ values[i] = ts
+ }
+ i += 1
+ }
+ return json.Marshal(values)
+}
+
+func (a *TimestampWithOffsetArray) GetOneForMarshal(i int) interface{} {
+ if a.IsNull(i) {
+ return nil
+ }
+ return a.Value(i)
+}
+
+// noLastOffset is the sentinel value for TimestampWithOffsetBuilder.lastOffset
+// indicating that no run-end-encoded run has been started yet. It is
deliberately
+// outside the range of valid timezone offsets in minutes (roughly [-720,
840]) so
+// it can never compare equal to a real offset.
+const noLastOffset int16 = math.MaxInt16
+
+// TimestampWithOffsetBuilder is a convenience builder for the
TimestampWithOffset extension type,
+// allowing arrays to be built with boolean values rather than the underlying
storage type.
+type TimestampWithOffsetBuilder struct {
+ *array.ExtensionBuilder
+
+ // The layout used to parse any timestamps from strings. Defaults to
time.RFC3339
+ Layout string
+ unit arrow.TimeUnit
+ offsetType arrow.DataType
+ // lastOffset tracks the offset of the current run when the offsets are
run-end
+ // encoded, so we know when to start a new run. It is initialized to
noLastOffset
+ // to signal that no run has been started yet.
+ lastOffset int16
+}
+
+// NewTimestampWithOffsetBuilder creates a new TimestampWithOffsetBuilder,
exposing a convenient and efficient interface
+// for writing time.Time values to the underlying storage array.
+func NewTimestampWithOffsetBuilder(mem memory.Allocator, unit arrow.TimeUnit,
offsetType arrow.DataType) (*TimestampWithOffsetBuilder, error) {
+ dataType, err := NewTimestampWithOffsetTypeCustomOffset(unit,
offsetType)
+ if err != nil {
+ return nil, err
+ }
+
+ return &TimestampWithOffsetBuilder{
+ unit: unit,
+ offsetType: offsetType,
+ lastOffset: noLastOffset,
+ Layout: time.RFC3339,
+ ExtensionBuilder: array.NewExtensionBuilder(mem, dataType),
+ }, nil
+}
+
+// NewArray must route through this type's NewExtensionArray (not the embedded
+// ExtensionBuilder's) so that the run-end-encoding tracker is reset on reuse.
+func (b *TimestampWithOffsetBuilder) NewArray() arrow.Array {
+ return b.NewExtensionArray()
+}
+
+// NewExtensionArray finalizes the current array and resets lastOffset so a
+// reused builder starts a fresh run instead of continuing a run that belonged
+// to the array just finalized (the underlying REE builder is reset too).
+func (b *TimestampWithOffsetBuilder) NewExtensionArray() array.ExtensionArray {
+ arr := b.ExtensionBuilder.NewExtensionArray()
+ b.lastOffset = noLastOffset
+ return arr
+}
+
+// AppendNull resets the run-end-encoding tracker: the embedded struct builder
+// appends a null offset run, so a following value must start a new run instead
+// of continuing the null run.
+//
+// NOTE(#918): this writes a null into the non-nullable offset_minutes storage
+// field rather than a default value. This is intentional; comparison and JSON
+// round-trip parity for non-nullable inner struct fields is tracked separately
+// in https://github.com/apache/arrow-go/issues/918. Logical values (instant,
+// offset, validity) are preserved, but a full array.RecordEqual round-trip is
+// not guaranteed here.
+func (b *TimestampWithOffsetBuilder) AppendNull() {
+ b.ExtensionBuilder.AppendNull()
+ b.lastOffset = noLastOffset
+}
+
+func (b *TimestampWithOffsetBuilder) AppendNulls(n int) {
+ b.ExtensionBuilder.AppendNulls(n)
+ b.lastOffset = noLastOffset
+}
+
+func (b *TimestampWithOffsetBuilder) Append(v time.Time) {
+ timestamp, offsetMinutes := fieldValuesFromTime(v, b.unit)
+ offsetMinutes16 := int16(offsetMinutes)
+ structBuilder := b.Builder.(*array.StructBuilder)
+
+ structBuilder.Append(true)
+
structBuilder.FieldBuilder(0).(*array.TimestampBuilder).Append(timestamp)
+
+ switch offsets := structBuilder.FieldBuilder(1).(type) {
+ case *array.Int16Builder:
+ offsets.Append(offsetMinutes16)
+ case *array.Int16DictionaryBuilder:
+ offsets.Append(offsetMinutes16)
+ case *array.RunEndEncodedBuilder:
+ if offsetMinutes != b.lastOffset {
+ offsets.Append(1)
+
offsets.ValueBuilder().(*array.Int16Builder).Append(offsetMinutes16)
+ } else {
+ offsets.ContinueRun(1)
+ }
+
+ b.lastOffset = offsetMinutes16
+ }
+
+}
+
+// By default, this will try to parse the string using the RFC3339 layout.
+//
+// You can change the default layout by using builder.SetLayout()
+func (b *TimestampWithOffsetBuilder) AppendValueFromString(s string) error {
+ if s == array.NullValueStr {
+ b.AppendNull()
+ return nil
+ }
+
+ parsed, err := time.Parse(b.Layout, s)
+ if err != nil {
+ return err
+ }
+
+ b.Append(parsed)
+ return nil
+}
+
+func (b *TimestampWithOffsetBuilder) AppendValues(values []time.Time, valids
[]bool) {
+ if len(valids) != len(values) && len(valids) != 0 {
+ panic("len(values) != len(valids) && len(valids) != 0")
+ }
+ if len(valids) == 0 {
+ valids = make([]bool, len(values))
+ for i := range valids {
+ valids[i] = true
+ }
+ }
+
+ structBuilder := b.Builder.(*array.StructBuilder)
+ timestamps := structBuilder.FieldBuilder(0).(*array.TimestampBuilder)
+
+ structBuilder.AppendValues(valids)
+ // SAFETY: by this point we know all buffers have available space given
the earlier
+ // call to structBuilder.AppendValues which calls Reserve internally,
so it's OK to
+ // call UnsafeAppend on inner builders
+
+ switch offsets := structBuilder.FieldBuilder(1).(type) {
+ case *array.Int16Builder:
+ for _, v := range values {
+ timestamp, offsetMinutes := fieldValuesFromTime(v,
b.unit)
+ timestamps.UnsafeAppend(timestamp)
+ offsets.UnsafeAppend(offsetMinutes)
+ }
+ case *array.Int16DictionaryBuilder:
+ for _, v := range values {
+ timestamp, offsetMinutes := fieldValuesFromTime(v,
b.unit)
+ timestamps.UnsafeAppend(timestamp)
+ offsets.UnsafeAppend(offsetMinutes)
+ }
+ case *array.RunEndEncodedBuilder:
+ offsetValuesBuilder :=
offsets.ValueBuilder().(*array.Int16Builder)
+ for i, v := range values {
+ timestamp, offsetMinutes := fieldValuesFromTime(v,
b.unit)
+ timestamps.UnsafeAppend(timestamp)
+
+ // A null row's offset is masked by the struct validity
bitmap, so we
+ // continue the current run to maximize compression. A
run must still be
+ // started when none exists yet (e.g. leading null
rows), otherwise
+ // ContinueRun would advance the run-ends without a
matching entry in the
+ // values child and produce an invalid run-end encoded
array. lastOffset
+ // is only updated when a new run starts, so a null row
never splits a
+ // contiguous run into two adjacent runs sharing the
same value.
+ valid := valids[i]
+ if b.lastOffset == noLastOffset || (valid &&
offsetMinutes != b.lastOffset) {
+ offsets.Append(1)
+ offsetValuesBuilder.Append(offsetMinutes)
+ b.lastOffset = offsetMinutes
+ } else {
+ offsets.ContinueRun(1)
+ }
+ }
+ }
+}
+
+func (b *TimestampWithOffsetBuilder) UnmarshalOne(dec *json.Decoder) error {
+ tok, err := dec.Token()
+ if err != nil {
+ return fmt.Errorf("failed to decode json: %w", err)
+ }
+
+ switch raw := tok.(type) {
+ case string:
+ t, err := time.Parse(b.Layout, raw)
+ if err != nil {
+ return fmt.Errorf("failed to parse string \"%s\" as
time.Time using layout \"%s\"", raw, b.Layout)
+ }
+ b.Append(t)
+ case nil:
+ b.AppendNull()
+ default:
+ return fmt.Errorf("expected date string")
+ }
+
+ return nil
+}
+
+func (b *TimestampWithOffsetBuilder) Unmarshal(dec *json.Decoder) error {
+ for dec.More() {
+ if err := b.UnmarshalOne(dec); err != nil {
+ return err
+ }
+ }
+ return nil
+}
+
+var (
+ _ arrow.ExtensionType = (*TimestampWithOffsetType)(nil)
+ _ array.CustomExtensionBuilder = (*TimestampWithOffsetType)(nil)
+ _ array.ExtensionArray = (*TimestampWithOffsetArray)(nil)
+ _ array.Builder = (*TimestampWithOffsetBuilder)(nil)
+)
diff --git a/arrow/extensions/timestamp_with_offset_test.go
b/arrow/extensions/timestamp_with_offset_test.go
new file mode 100644
index 00000000..efbdfd28
--- /dev/null
+++ b/arrow/extensions/timestamp_with_offset_test.go
@@ -0,0 +1,646 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package extensions_test
+
+import (
+ "bytes"
+ "fmt"
+ "testing"
+ "time"
+
+ "github.com/apache/arrow-go/v18/arrow"
+ "github.com/apache/arrow-go/v18/arrow/array"
+ "github.com/apache/arrow-go/v18/arrow/extensions"
+ "github.com/apache/arrow-go/v18/arrow/ipc"
+ "github.com/apache/arrow-go/v18/arrow/memory"
+ "github.com/apache/arrow-go/v18/internal/json"
+ "github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
+)
+
+var testTimeUnit = arrow.Microsecond
+
+var epoch = time.Unix(0, 0).In(time.FixedZone("UTC+00:00", 0))
+
+var testDate0 = time.Date(2025, 01, 01, 00, 00, 00, 00,
time.FixedZone("UTC+00:00", 0))
+
+var testZone1 = time.FixedZone("UTC-08:30", -8*60*60-30*60)
+var testDate1 = testDate0.In(testZone1)
+
+var testZone2 = time.FixedZone("UTC+11:00", +11*60*60)
+var testDate2 = testDate0.In(testZone2)
+
+func dict(index arrow.DataType) arrow.DataType {
+ return &arrow.DictionaryType{
+ IndexType: index,
+ ValueType: arrow.PrimitiveTypes.Int16,
+ Ordered: false,
+ }
+}
+
+func ree(runEnds arrow.DataType) arrow.DataType {
+ v := arrow.RunEndEncodedOf(runEnds, arrow.PrimitiveTypes.Int16)
+ v.ValueNullable = false
+ return v
+}
+
+// All tests use this in a for loop to make sure everything works for every
possible
+// encoding of offsets (primitive, dictionary, run-end)
+var allAllowedOffsetTypes = make(map[string]arrow.DataType)
+
+func init() {
+ // primitive offsetType
+ allAllowedOffsetTypes["primitive-int16"] = arrow.PrimitiveTypes.Int16
+
+ // dict-encoded offsetType
+ allAllowedOffsetTypes["dict-Uint8"] = dict(arrow.PrimitiveTypes.Uint8)
+ allAllowedOffsetTypes["dict-Uint16"] = dict(arrow.PrimitiveTypes.Uint16)
+ allAllowedOffsetTypes["dict-Uint32"] = dict(arrow.PrimitiveTypes.Uint32)
+ allAllowedOffsetTypes["dict-Uint64"] = dict(arrow.PrimitiveTypes.Uint64)
+ allAllowedOffsetTypes["dict-Int8"] = dict(arrow.PrimitiveTypes.Int8)
+ allAllowedOffsetTypes["dict-Int16"] = dict(arrow.PrimitiveTypes.Int16)
+ allAllowedOffsetTypes["dict-Int32"] = dict(arrow.PrimitiveTypes.Int32)
+ allAllowedOffsetTypes["dict-Int64"] = dict(arrow.PrimitiveTypes.Int64)
+
+ // run-end encoded offsetType
+ allAllowedOffsetTypes["ree-Int16"] = ree(arrow.PrimitiveTypes.Int16)
+ allAllowedOffsetTypes["ree-Int32"] = ree(arrow.PrimitiveTypes.Int32)
+ allAllowedOffsetTypes["ree-Int64"] = ree(arrow.PrimitiveTypes.Int64)
+}
+
+func TestTimestampWithOffsetTypePrimitiveBasics(t *testing.T) {
+ typ := extensions.NewTimestampWithOffsetType(testTimeUnit)
+
+ assert.Equal(t, "arrow.timestamp_with_offset", typ.ExtensionName())
+ assert.True(t, typ.ExtensionEquals(typ))
+
+ assert.True(t, arrow.TypeEqual(typ, typ))
+ assert.True(t, arrow.TypeEqual(
+ arrow.StructOf(
+ arrow.Field{
+ Name: "timestamp",
+ Type: &arrow.TimestampType{
+ Unit: testTimeUnit,
+ TimeZone: "UTC",
+ },
+ Nullable: false,
+ },
+ arrow.Field{
+ Name: "offset_minutes",
+ Type: arrow.PrimitiveTypes.Int16,
+ Nullable: false,
+ },
+ ),
+ typ.StorageType()))
+
+ assert.Equal(t, "extension<arrow.timestamp_with_offset>", typ.String())
+}
+
+func TestTimestampWithOffsetTypeDeserializeMetadata(t *testing.T) {
+ typ := extensions.NewTimestampWithOffsetType(testTimeUnit)
+ storage := typ.StorageType()
+
+ for _, data := range []string{"", "{}"} {
+ got, err := typ.Deserialize(storage, data)
+ assert.NoError(t, err)
+ assert.True(t, typ.ExtensionEquals(got))
+ }
+
+ _, err := typ.Deserialize(storage, "not-empty")
+ assert.Error(t, err)
+}
+
+func TestTimestampWithOffsetTypeDeserializeInvalidStorage(t *testing.T) {
+ base := extensions.NewTimestampWithOffsetType(testTimeUnit)
+
+ utcTS := &arrow.TimestampType{Unit: testTimeUnit, TimeZone: "UTC"}
+ tsField := arrow.Field{Name: "timestamp", Type: utcTS}
+ offField := arrow.Field{Name: "offset_minutes", Type:
arrow.PrimitiveTypes.Int16}
+
+ badDict := &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int8,
ValueType: arrow.PrimitiveTypes.Int32}
+ badREE := arrow.RunEndEncodedOf(arrow.PrimitiveTypes.Int32,
arrow.PrimitiveTypes.Int32)
+
+ valid, err := base.Deserialize(base.StorageType(), "")
+ require.NoError(t, err)
+ require.NotNil(t, valid)
+
+ invalid := map[string]arrow.DataType{
+ "not a struct": arrow.PrimitiveTypes.Int16,
+ "too few fields": arrow.StructOf(tsField),
+ "too many fields": arrow.StructOf(tsField,
offField, offField),
+ "timestamp wrong name":
arrow.StructOf(arrow.Field{Name: "ts", Type: utcTS}, offField),
+ "timestamp not a timestamp":
arrow.StructOf(arrow.Field{Name: "timestamp", Type:
arrow.PrimitiveTypes.Int64}, offField),
+ "timestamp non-UTC zone":
arrow.StructOf(arrow.Field{Name: "timestamp", Type: &arrow.TimestampType{Unit:
testTimeUnit, TimeZone: "America/New_York"}}, offField),
+ "timestamp empty zone":
arrow.StructOf(arrow.Field{Name: "timestamp", Type: &arrow.TimestampType{Unit:
testTimeUnit}}, offField),
+ "timestamp nullable":
arrow.StructOf(arrow.Field{Name: "timestamp", Type: utcTS, Nullable: true},
offField),
+ "offset wrong name": arrow.StructOf(tsField,
arrow.Field{Name: "offset", Type: arrow.PrimitiveTypes.Int16}),
+ "offset wrong primitive type": arrow.StructOf(tsField,
arrow.Field{Name: "offset_minutes", Type: arrow.PrimitiveTypes.Int32}),
+ "offset nullable": arrow.StructOf(tsField,
arrow.Field{Name: "offset_minutes", Type: arrow.PrimitiveTypes.Int16, Nullable:
true}),
+ "offset dict value not int16": arrow.StructOf(tsField,
arrow.Field{Name: "offset_minutes", Type: badDict}),
+ "offset ree encoded not int16": arrow.StructOf(tsField,
arrow.Field{Name: "offset_minutes", Type: badREE}),
+ "fields swapped": arrow.StructOf(offField,
tsField),
+ }
+
+ for name, storage := range invalid {
+ t.Run(name, func(t *testing.T) {
+ _, err := base.Deserialize(storage, "")
+ assert.Error(t, err)
+ })
+ }
+}
+
+func assertDictBasics[I extensions.DictIndexType](t *testing.T, indexType I) {
+ typ :=
extensions.NewTimestampWithOffsetTypeDictionaryEncoded(testTimeUnit, indexType)
+
+ assert.Equal(t, "arrow.timestamp_with_offset", typ.ExtensionName())
+ assert.True(t, typ.ExtensionEquals(typ))
+
+ assert.True(t, arrow.TypeEqual(typ, typ))
+ assert.True(t, arrow.TypeEqual(
+ arrow.StructOf(
+ arrow.Field{
+ Name: "timestamp",
+ Type: &arrow.TimestampType{
+ Unit: testTimeUnit,
+ TimeZone: "UTC",
+ },
+ Nullable: false,
+ },
+ arrow.Field{
+ Name: "offset_minutes",
+ Type: dict(arrow.DataType(indexType)),
+ Nullable: false,
+ },
+ ),
+ typ.StorageType()))
+
+ assert.Equal(t, "extension<arrow.timestamp_with_offset>", typ.String())
+}
+
+func TestTimestampWithOffsetTypeDictionaryEncodedBasics(t *testing.T) {
+ assertDictBasics(t, &arrow.Uint8Type{})
+ assertDictBasics(t, &arrow.Uint16Type{})
+ assertDictBasics(t, &arrow.Uint32Type{})
+ assertDictBasics(t, &arrow.Uint64Type{})
+ assertDictBasics(t, &arrow.Int8Type{})
+ assertDictBasics(t, &arrow.Int16Type{})
+ assertDictBasics(t, &arrow.Int32Type{})
+ assertDictBasics(t, &arrow.Int64Type{})
+}
+
+func assertReeBasics[E extensions.RunEndsType](t *testing.T, runEndsType E) {
+ typ := extensions.NewTimestampWithOffsetTypeRunEndEncoded(testTimeUnit,
runEndsType)
+
+ assert.Equal(t, "arrow.timestamp_with_offset", typ.ExtensionName())
+ assert.True(t, typ.ExtensionEquals(typ))
+
+ assert.True(t, arrow.TypeEqual(typ, typ))
+ assert.True(t, arrow.TypeEqual(
+ arrow.StructOf(
+ arrow.Field{
+ Name: "timestamp",
+ Type: &arrow.TimestampType{
+ Unit: testTimeUnit,
+ TimeZone: "UTC",
+ },
+ Nullable: false,
+ },
+ arrow.Field{
+ Name: "offset_minutes",
+ Type: ree(arrow.DataType(runEndsType)),
+ Nullable: false,
+ },
+ ),
+ typ.StorageType()))
+
+ assert.Equal(t, "extension<arrow.timestamp_with_offset>", typ.String())
+}
+
+func TestTimestampWithOffsetTypeRunEndEncodedBasics(t *testing.T) {
+ assertReeBasics(t, &arrow.Int16Type{})
+ assertReeBasics(t, &arrow.Int32Type{})
+ assertReeBasics(t, &arrow.Int64Type{})
+}
+
+func TestTimestampWithOffsetEquals(t *testing.T) {
+ // Completely different types are not equal
+ assert.False(t,
extensions.NewTimestampWithOffsetType(arrow.Nanosecond).ExtensionEquals(extensions.NewBool8Type()))
+
+ // Different time units are not equal
+ assert.False(t,
extensions.NewTimestampWithOffsetType(arrow.Nanosecond).ExtensionEquals(extensions.NewTimestampWithOffsetType(arrow.Microsecond)))
+ assert.False(t,
extensions.NewTimestampWithOffsetType(arrow.Nanosecond).ExtensionEquals(extensions.NewTimestampWithOffsetType(arrow.Second)))
+ assert.False(t,
extensions.NewTimestampWithOffsetType(arrow.Microsecond).ExtensionEquals(extensions.NewTimestampWithOffsetType(arrow.Second)))
+
+ // Different underlying storage type is not equal
+ assert.False(t,
extensions.NewTimestampWithOffsetType(arrow.Microsecond).ExtensionEquals(extensions.NewTimestampWithOffsetTypeDictionaryEncoded(arrow.Microsecond,
&arrow.Int16Type{})))
+ assert.False(t,
extensions.NewTimestampWithOffsetType(arrow.Microsecond).ExtensionEquals(extensions.NewTimestampWithOffsetTypeRunEndEncoded(arrow.Microsecond,
&arrow.Int16Type{})))
+ assert.False(t,
extensions.NewTimestampWithOffsetTypeDictionaryEncoded(arrow.Microsecond,
&arrow.Int16Type{}).ExtensionEquals(extensions.NewTimestampWithOffsetTypeRunEndEncoded(arrow.Microsecond,
&arrow.Int16Type{})))
+
+ // Dict-encoding key type is not equal
+ assert.False(t,
extensions.NewTimestampWithOffsetTypeDictionaryEncoded(arrow.Microsecond,
&arrow.Int16Type{}).ExtensionEquals(extensions.NewTimestampWithOffsetTypeDictionaryEncoded(arrow.Microsecond,
&arrow.Uint16Type{})))
+
+ // REE index type is not equal
+ assert.False(t,
extensions.NewTimestampWithOffsetTypeRunEndEncoded(arrow.Microsecond,
&arrow.Int16Type{}).ExtensionEquals(extensions.NewTimestampWithOffsetTypeRunEndEncoded(arrow.Microsecond,
&arrow.Int32Type{})))
+
+ // Equals OK
+ assert.True(t,
extensions.NewTimestampWithOffsetType(arrow.Nanosecond).ExtensionEquals(extensions.NewTimestampWithOffsetType(arrow.Nanosecond)))
+ assert.True(t,
extensions.NewTimestampWithOffsetType(arrow.Microsecond).ExtensionEquals(extensions.NewTimestampWithOffsetType(arrow.Microsecond)))
+ assert.True(t,
extensions.NewTimestampWithOffsetTypeDictionaryEncoded(arrow.Microsecond,
&arrow.Int16Type{}).ExtensionEquals(extensions.NewTimestampWithOffsetTypeDictionaryEncoded(arrow.Microsecond,
&arrow.Int16Type{})))
+ assert.True(t,
extensions.NewTimestampWithOffsetTypeDictionaryEncoded(arrow.Microsecond,
&arrow.Uint16Type{}).ExtensionEquals(extensions.NewTimestampWithOffsetTypeDictionaryEncoded(arrow.Microsecond,
&arrow.Uint16Type{})))
+ assert.True(t,
extensions.NewTimestampWithOffsetTypeRunEndEncoded(arrow.Microsecond,
&arrow.Int16Type{}).ExtensionEquals(extensions.NewTimestampWithOffsetTypeRunEndEncoded(arrow.Microsecond,
&arrow.Int16Type{})))
+ assert.True(t,
extensions.NewTimestampWithOffsetTypeRunEndEncoded(arrow.Microsecond,
&arrow.Int32Type{}).ExtensionEquals(extensions.NewTimestampWithOffsetTypeRunEndEncoded(arrow.Microsecond,
&arrow.Int32Type{})))
+}
+
+func TestTimestampWithOffsetExtensionBuilder(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ for name, offsetType := range allAllowedOffsetTypes {
+ t.Run(name, func(t *testing.T) {
+ builder, _ :=
extensions.NewTimestampWithOffsetBuilder(mem, testTimeUnit, offsetType)
+
+ builder.Append(testDate0)
+ builder.AppendNull()
+ builder.Append(testDate1)
+ builder.Append(testDate2)
+ builder.Append(epoch)
+
+ // it should build the array with the correct size
+ arr := builder.NewArray()
+ typedArr := arr.(*extensions.TimestampWithOffsetArray)
+ assert.Equal(t, 5, arr.Data().Len())
+ defer arr.Release()
+
+ // typedArr.Value(i) should return values adjusted for
their original timezone
+ assert.Equal(t, testDate0, typedArr.Value(0))
+ assert.Equal(t, testDate1, typedArr.Value(2))
+ assert.Equal(t, testDate2, typedArr.Value(3))
+ assert.Equal(t, epoch, typedArr.Value(4))
+
+ // storage TimeUnit should be the same as we pass in to
the builder, and storage timezone should be UTC
+ timestampStructField :=
typedArr.Storage().(*array.Struct).Field(0)
+ timestampStructDataType :=
timestampStructField.DataType().(*arrow.TimestampType)
+ assert.Equal(t, timestampStructDataType.Unit,
testTimeUnit)
+ assert.Equal(t, timestampStructDataType.TimeZone, "UTC")
+
+ // stored values should be equivalent to the raw values
in UTC
+ timestampsArr := timestampStructField.(*array.Timestamp)
+ assert.Equal(t, testDate0.In(time.UTC),
timestampsArr.Value(0).ToTime(testTimeUnit))
+ assert.Equal(t, testDate1.In(time.UTC),
timestampsArr.Value(2).ToTime(testTimeUnit))
+ assert.Equal(t, testDate2.In(time.UTC),
timestampsArr.Value(3).ToTime(testTimeUnit))
+ assert.Equal(t, epoch.In(time.UTC),
timestampsArr.Value(4).ToTime(testTimeUnit))
+
+ // the array should encode itself as JSON and string
+ arrStr := arr.String()
+ assert.Equal(t, fmt.Sprintf(`["%[1]s" (null) "%[2]s"
"%[3]s" "%[4]s"]`, testDate0, testDate1, testDate2, epoch), arrStr)
+ jsonStr, err := json.Marshal(arr)
+ assert.NoError(t, err)
+
+ // roundtripping from JSON with array.FromJSON should
work
+ expectedDataType, _ :=
extensions.NewTimestampWithOffsetTypeCustomOffset(testTimeUnit, offsetType)
+ roundtripped, _, err := array.FromJSON(mem,
expectedDataType, bytes.NewReader(jsonStr))
+ defer roundtripped.Release()
+ assert.NoError(t, err)
+ assert.Truef(t, array.Equal(arr, roundtripped),
"expected %s\n\ngot %s", arr, roundtripped)
+ })
+ }
+}
+
+func TestTimestampWithOffsetBuilderAppendValuesLeadingNull(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ for name, offsetType := range allAllowedOffsetTypes {
+ t.Run(name, func(t *testing.T) {
+ builder, err :=
extensions.NewTimestampWithOffsetBuilder(mem, testTimeUnit, offsetType)
+ require.NoError(t, err)
+
+ // A leading null as the very first operation on a
fresh builder must not
+ // corrupt the offsets. Reading every value below would
panic on a
+ // run-end encoded array whose run-ends and values
children fell out of
+ // sync (which happened when a run was continued before
one existed).
+ builder.AppendValues([]time.Time{epoch, testDate1,
testDate2}, []bool{false, true, true})
+
+ arr := builder.NewArray()
+ defer arr.Release()
+ typedArr := arr.(*extensions.TimestampWithOffsetArray)
+
+ require.Equal(t, 3, arr.Data().Len())
+ assert.True(t, typedArr.IsNull(0))
+ assert.Equal(t, testDate1, typedArr.Value(1))
+ assert.Equal(t, testDate2, typedArr.Value(2))
+ })
+ }
+}
+
+func TestTimestampWithOffsetBuilderRunEndEncodedNullContinuesRun(t *testing.T)
{
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ builder, err := extensions.NewTimestampWithOffsetBuilder(mem,
testTimeUnit, ree(arrow.PrimitiveTypes.Int16))
+ require.NoError(t, err)
+
+ // The same offset surrounds a null. The null must continue the current
run
+ // rather than reset the tracked offset, otherwise one logical run is
split
+ // into two adjacent runs that share the same value.
+ builder.AppendValues([]time.Time{testDate1, epoch, testDate1},
[]bool{true, false, true})
+
+ arr := builder.NewArray()
+ defer arr.Release()
+ typedArr := arr.(*extensions.TimestampWithOffsetArray)
+
+ offsets :=
typedArr.Storage().(*array.Struct).Field(1).(*array.RunEndEncoded)
+ assert.Equal(t, 1, offsets.Values().Len())
+
+ assert.Equal(t, testDate1, typedArr.Value(0))
+ assert.True(t, typedArr.IsNull(1))
+ assert.Equal(t, testDate1, typedArr.Value(2))
+}
+
+func TestTimestampWithOffsetBuilderAppendValuesNilValids(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ // nil and empty validity slices must both be treated as all-valid so
the
+ // parent struct and its child builders stay the same length.
+ validsVariants := map[string][]bool{"nil": nil, "empty": {}}
+
+ for name, offsetType := range allAllowedOffsetTypes {
+ for validsName, valids := range validsVariants {
+ t.Run(name+"/"+validsName, func(t *testing.T) {
+ builder, err :=
extensions.NewTimestampWithOffsetBuilder(mem, testTimeUnit, offsetType)
+ require.NoError(t, err)
+
+ builder.AppendValues([]time.Time{testDate0,
testDate1, testDate2}, valids)
+
+ arr := builder.NewArray()
+ defer arr.Release()
+ typedArr :=
arr.(*extensions.TimestampWithOffsetArray)
+
+ require.Equal(t, 3, arr.Data().Len())
+ assert.Equal(t, testDate0, typedArr.Value(0))
+ assert.Equal(t, testDate1, typedArr.Value(1))
+ assert.Equal(t, testDate2, typedArr.Value(2))
+ })
+ }
+ }
+}
+
+func TestTimestampWithOffsetBuilderAppendValuesMismatchedValidsPanics(t
*testing.T) {
+ builder, err :=
extensions.NewTimestampWithOffsetBuilder(memory.DefaultAllocator, testTimeUnit,
arrow.PrimitiveTypes.Int16)
+ require.NoError(t, err)
+ defer builder.Release()
+
+ assert.Panics(t, func() {
+ builder.AppendValues([]time.Time{testDate0, testDate1},
[]bool{true})
+ })
+}
+
+func TestTimestampWithOffsetBuilderReuseAfterNewArray(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ for name, offsetType := range allAllowedOffsetTypes {
+ t.Run(name, func(t *testing.T) {
+ builder, err :=
extensions.NewTimestampWithOffsetBuilder(mem, testTimeUnit, offsetType)
+ require.NoError(t, err)
+
+ builder.Append(testDate1)
+ arr1 := builder.NewArray()
+ defer arr1.Release()
+
+ // Reusing the builder must start a fresh run even when
the first value
+ // repeats the previous array's offset; otherwise a
run-end-encoded
+ // offset is continued on a freshly reset (empty)
builder.
+ builder.Append(testDate1)
+ builder.Append(testDate2)
+ arr2 := builder.NewArray()
+ defer arr2.Release()
+
+ typedArr := arr2.(*extensions.TimestampWithOffsetArray)
+ require.Equal(t, 2, arr2.Data().Len())
+ assert.Equal(t, testDate1, typedArr.Value(0))
+ assert.Equal(t, testDate2, typedArr.Value(1))
+ })
+ }
+}
+
+func TestTimestampWithOffsetBuilderAppendNullBetweenEqualOffsets(t *testing.T)
{
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ builder, err := extensions.NewTimestampWithOffsetBuilder(mem,
testTimeUnit, ree(arrow.PrimitiveTypes.Int16))
+ require.NoError(t, err)
+
+ // AppendNull appends a null offset run; the following value repeats the
+ // pre-null offset and must start its own run instead of being folded
into
+ // the null run (which would decode the trailing value's offset as
null).
+ builder.Append(testDate1)
+ builder.AppendNull()
+ builder.Append(testDate1)
+
+ arr := builder.NewArray()
+ defer arr.Release()
+ typedArr := arr.(*extensions.TimestampWithOffsetArray)
+
+ require.Equal(t, 3, arr.Data().Len())
+ assert.Equal(t, testDate1, typedArr.Value(0))
+ assert.True(t, typedArr.IsNull(1))
+ assert.Equal(t, testDate1, typedArr.Value(2))
+}
+
+func TestTimestampWithOffsetRunEndEncodedSliced(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ for _, runEnds := range []arrow.DataType{
+ arrow.PrimitiveTypes.Int16, arrow.PrimitiveTypes.Int32,
arrow.PrimitiveTypes.Int64,
+ } {
+ t.Run(runEnds.String(), func(t *testing.T) {
+ builder, err :=
extensions.NewTimestampWithOffsetBuilder(mem, testTimeUnit, ree(runEnds))
+ require.NoError(t, err)
+
+ // Offsets [0,0,-510,-510,-510,660,0] give run-ends
[2,5,6,7], so a
+ // slice starting at index 3 begins inside a later run:
the case where
+ // an offset-unaware physical index decodes the wrong
run.
+ values := []time.Time{testDate0, testDate0, testDate1,
testDate1, testDate1, testDate2, epoch}
+ builder.AppendValues(values, nil)
+ arr :=
builder.NewArray().(*extensions.TimestampWithOffsetArray)
+ defer arr.Release()
+
+ want := arr.Values()
+
+ for start := 0; start < arr.Len(); start++ {
+ for end := start; end <= arr.Len(); end++ {
+ sliced := array.NewSlice(arr,
int64(start), int64(end)).(*extensions.TimestampWithOffsetArray)
+ got := sliced.Values()
+ require.Equalf(t, end-start, len(got),
"slice [%d:%d]", start, end)
+ for k := 0; k < end-start; k++ {
+ require.Truef(t,
want[start+k].Equal(got[k]), "slice [%d:%d] row %d instant: want %s got %s",
start, end, k, want[start+k], got[k])
+ _, wantOff :=
want[start+k].Zone()
+ _, gotOff := got[k].Zone()
+ require.Equalf(t, wantOff,
gotOff, "slice [%d:%d] row %d offset", start, end, k)
+ }
+ sliced.Release()
+ }
+ }
+ })
+ }
+}
+
+func TestTimestampWithOffsetSubHourNegativeZone(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ // A negative offset smaller than one hour must keep its sign in the
zone
+ // name (regression: -30 minutes formatted as "UTC+00:30" instead of
"UTC-00:30").
+ b, err := extensions.NewTimestampWithOffsetBuilder(mem, testTimeUnit,
arrow.PrimitiveTypes.Int16)
+ require.NoError(t, err)
+ b.Append(testDate0.In(time.FixedZone("UTC-00:30", -30*60)))
+ arr := b.NewArray().(*extensions.TimestampWithOffsetArray)
+ defer arr.Release()
+
+ name, off := arr.Value(0).Zone()
+ require.Equal(t, -30*60, off)
+ require.Equal(t, "UTC-00:30", name)
+}
+
+func TestTimestampWithOffsetExtensionRecordBuilder(t *testing.T) {
+ for name, offsetType := range allAllowedOffsetTypes {
+ t.Run(name, func(t *testing.T) {
+ dataType, _ :=
extensions.NewTimestampWithOffsetTypeCustomOffset(testTimeUnit, offsetType)
+ schema := arrow.NewSchema([]arrow.Field{
+ {
+ Name: "timestamp_with_offset",
+ Nullable: true,
+ Type: dataType,
+ },
+ }, nil)
+ builder :=
array.NewRecordBuilder(memory.DefaultAllocator, schema)
+ defer builder.Release()
+
+ fieldBuilder :=
builder.Field(0).(*extensions.TimestampWithOffsetBuilder)
+
+ // append a simple time.Time
+ fieldBuilder.Append(testDate0)
+
+ // append the epoch
+ fieldBuilder.Append(epoch)
+
+ // append a null and 2 time.Time all at once
+ values := []time.Time{
+ time.Unix(0, 0).In(time.UTC),
+ testDate1,
+ testDate2,
+ }
+ valids := []bool{false, true, true}
+ fieldBuilder.AppendValues(values, valids)
+
+ // append a value from RFC3339 string
+
fieldBuilder.AppendValueFromString(testDate0.Format(time.RFC3339))
+
+ // append value formatted in a different string layout
+ fieldBuilder.Layout = time.RFC3339Nano
+
fieldBuilder.AppendValueFromString(testDate1.Format(time.RFC3339Nano))
+
+ record := builder.NewRecordBatch()
+
+ // Record batch should JSON-encode values containing
per-row timezone info
+ json, err := record.MarshalJSON()
+ require.NoError(t, err)
+ expect :=
`[{"timestamp_with_offset":"2025-01-01T00:00:00Z"}
+,{"timestamp_with_offset":"1970-01-01T00:00:00Z"}
+,{"timestamp_with_offset":null}
+,{"timestamp_with_offset":"2024-12-31T15:30:00-08:30"}
+,{"timestamp_with_offset":"2025-01-01T11:00:00+11:00"}
+,{"timestamp_with_offset":"2025-01-01T00:00:00Z"}
+,{"timestamp_with_offset":"2024-12-31T15:30:00-08:30"}
+]`
+ require.Equal(t, expect, string(json))
+
+ // Record batch roundtrip to JSON should work
+ roundtripped, _, err :=
array.RecordFromJSON(memory.DefaultAllocator, schema, bytes.NewReader(json))
+ require.NoError(t, err)
+ defer roundtripped.Release()
+ require.Equal(t, schema, roundtripped.Schema())
+ // FIXME: ideally array.RecordEqual would succeed after
roundtripping.
+ // It currently doesn't because it compares the inner
non-nullable struct
+ // child validity bitmaps, which differ between builder
output and JSON
+ // reader output. That inner-nullability parity is
decoupled from this
+ // type (tracked in apache/arrow-go#918), so for now
avoid RecordEqual and
+ // assert only what this type guarantees: per-row
instant, offset, validity.
+ origArr :=
record.Column(0).(*extensions.TimestampWithOffsetArray)
+ gotArr :=
roundtripped.Column(0).(*extensions.TimestampWithOffsetArray)
+ require.Equal(t, origArr.Len(), gotArr.Len())
+ for i := 0; i < origArr.Len(); i++ {
+ require.Equalf(t, origArr.IsNull(i),
gotArr.IsNull(i), "validity mismatch at row %d", i)
+ if origArr.IsValid(i) {
+ want, got := origArr.Value(i),
gotArr.Value(i)
+ require.Truef(t, want.Equal(got),
"instant mismatch at row %d: %s vs %s", i, want, got)
+ _, wantOff := want.Zone()
+ _, gotOff := got.Zone()
+ require.Equalf(t, wantOff, gotOff,
"offset mismatch at row %d", i)
+ }
+ }
+ })
+ }
+}
+
+func TestTimestampWithOffsetTypeBatchIPCRoundTrip(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ for name, offsetType := range allAllowedOffsetTypes {
+ t.Run(name, func(t *testing.T) {
+ builder, _ :=
extensions.NewTimestampWithOffsetBuilder(mem, testTimeUnit, offsetType)
+ builder.Append(testDate0)
+ builder.AppendNull()
+ builder.Append(testDate1)
+ builder.Append(testDate2)
+ builder.Append(epoch)
+ arr := builder.NewArray()
+ defer arr.Release()
+
+ typ, _ :=
extensions.NewTimestampWithOffsetTypeCustomOffset(testTimeUnit, offsetType)
+
+ batch :=
array.NewRecordBatch(arrow.NewSchema([]arrow.Field{{Name:
"timestamp_with_offset", Type: typ, Nullable: true}}, nil), []arrow.Array{arr},
-1)
+ defer batch.Release()
+
+ var written arrow.RecordBatch
+ {
+ var buf bytes.Buffer
+ wr := ipc.NewWriter(&buf,
ipc.WithSchema(batch.Schema()))
+ require.NoError(t, wr.Write(batch))
+ require.NoError(t, wr.Close())
+
+ rdr, err := ipc.NewReader(&buf)
+ require.NoError(t, err)
+ written, err = rdr.Read()
+ require.NoError(t, err)
+ written.Retain()
+ defer written.Release()
+ rdr.Release()
+ }
+
+ assert.Truef(t, batch.Schema().Equal(written.Schema()),
"expected: %s\n\ngot: %s",
+ batch.Schema(), written.Schema())
+
+ assert.Truef(t, array.RecordEqual(batch, written),
"expected: %s\n\ngot: %s",
+ batch, written)
+ })
+ }
+}
