This is an automated email from the ASF dual-hosted git repository.

zeroshade pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow.git


The following commit(s) were added to refs/heads/master by this push:
     new d829ab1d5f ARROW-16551: [Go] Improve Temporal Types
d829ab1d5f is described below

commit d829ab1d5fca7b38517b54983464b595eb0d64d2
Author: Matthew Topol <[email protected]>
AuthorDate: Mon May 23 14:49:04 2022 -0400

    ARROW-16551: [Go] Improve Temporal Types
    
    Fix JSON reading of Temporal types to allow passing numbers in addition to 
strings that will be parsed. Enhance timestamps to allow specifying the time 
zone when parsing the string. Provide function to retrieve a time.Location from 
a TimestampType
    
    Closes #13132 from zeroshade/arrow-16551-temporal-types
    
    Authored-by: Matthew Topol <[email protected]>
    Signed-off-by: Matthew Topol <[email protected]>
---
 go/arrow/array/numericbuilder.gen.go      |  78 +++++++++-
 go/arrow/array/numericbuilder.gen.go.tmpl |  41 ++++++
 go/arrow/array/util.go                    |  15 ++
 go/arrow/array/util_test.go               |  29 ++--
 go/arrow/datatype_fixedwidth.go           | 227 ++++++++++++++++++++++++++----
 go/arrow/datatype_fixedwidth_test.go      |   2 +-
 6 files changed, 350 insertions(+), 42 deletions(-)

diff --git a/go/arrow/array/numericbuilder.gen.go 
b/go/arrow/array/numericbuilder.gen.go
index d3c61c722a..da944aad3e 100644
--- a/go/arrow/array/numericbuilder.gen.go
+++ b/go/arrow/array/numericbuilder.gen.go
@@ -2211,7 +2211,9 @@ func (b *TimestampBuilder) unmarshalOne(dec 
*json.Decoder) error {
        case nil:
                b.AppendNull()
        case string:
-               tm, err := arrow.TimestampFromString(v, b.dtype.Unit)
+               loc, _ := b.dtype.GetZone()
+               tm, err := arrow.TimestampFromStringInLocation(v, b.dtype.Unit, 
loc)
+
                if err != nil {
                        return &json.UnmarshalTypeError{
                                Value:  v,
@@ -2221,6 +2223,18 @@ func (b *TimestampBuilder) unmarshalOne(dec 
*json.Decoder) error {
                }
 
                b.Append(tm)
+       case json.Number:
+               n, err := v.Int64()
+               if err != nil {
+                       return &json.UnmarshalTypeError{
+                               Value:  v.String(),
+                               Type:   reflect.TypeOf(arrow.Timestamp(0)),
+                               Offset: dec.InputOffset(),
+                       }
+               }
+               b.Append(arrow.Timestamp(n))
+       case float64:
+               b.Append(arrow.Timestamp(v))
 
        default:
                return &json.UnmarshalTypeError{
@@ -2404,6 +2418,7 @@ func (b *Time32Builder) unmarshalOne(dec *json.Decoder) 
error {
                b.AppendNull()
        case string:
                tm, err := arrow.Time32FromString(v, b.dtype.Unit)
+
                if err != nil {
                        return &json.UnmarshalTypeError{
                                Value:  v,
@@ -2413,6 +2428,18 @@ func (b *Time32Builder) unmarshalOne(dec *json.Decoder) 
error {
                }
 
                b.Append(tm)
+       case json.Number:
+               n, err := v.Int64()
+               if err != nil {
+                       return &json.UnmarshalTypeError{
+                               Value:  v.String(),
+                               Type:   reflect.TypeOf(arrow.Time32(0)),
+                               Offset: dec.InputOffset(),
+                       }
+               }
+               b.Append(arrow.Time32(n))
+       case float64:
+               b.Append(arrow.Time32(v))
 
        default:
                return &json.UnmarshalTypeError{
@@ -2596,6 +2623,7 @@ func (b *Time64Builder) unmarshalOne(dec *json.Decoder) 
error {
                b.AppendNull()
        case string:
                tm, err := arrow.Time64FromString(v, b.dtype.Unit)
+
                if err != nil {
                        return &json.UnmarshalTypeError{
                                Value:  v,
@@ -2605,6 +2633,18 @@ func (b *Time64Builder) unmarshalOne(dec *json.Decoder) 
error {
                }
 
                b.Append(tm)
+       case json.Number:
+               n, err := v.Int64()
+               if err != nil {
+                       return &json.UnmarshalTypeError{
+                               Value:  v.String(),
+                               Type:   reflect.TypeOf(arrow.Time64(0)),
+                               Offset: dec.InputOffset(),
+                       }
+               }
+               b.Append(arrow.Time64(n))
+       case float64:
+               b.Append(arrow.Time64(v))
 
        default:
                return &json.UnmarshalTypeError{
@@ -2796,6 +2836,18 @@ func (b *Date32Builder) unmarshalOne(dec *json.Decoder) 
error {
                }
 
                b.Append(arrow.Date32FromTime(tm))
+       case json.Number:
+               n, err := v.Int64()
+               if err != nil {
+                       return &json.UnmarshalTypeError{
+                               Value:  v.String(),
+                               Type:   reflect.TypeOf(arrow.Date32(0)),
+                               Offset: dec.InputOffset(),
+                       }
+               }
+               b.Append(arrow.Date32(n))
+       case float64:
+               b.Append(arrow.Date32(v))
 
        default:
                return &json.UnmarshalTypeError{
@@ -2987,6 +3039,18 @@ func (b *Date64Builder) unmarshalOne(dec *json.Decoder) 
error {
                }
 
                b.Append(arrow.Date64FromTime(tm))
+       case json.Number:
+               n, err := v.Int64()
+               if err != nil {
+                       return &json.UnmarshalTypeError{
+                               Value:  v.String(),
+                               Type:   reflect.TypeOf(arrow.Date64(0)),
+                               Offset: dec.InputOffset(),
+                       }
+               }
+               b.Append(arrow.Date64(n))
+       case float64:
+               b.Append(arrow.Date64(v))
 
        default:
                return &json.UnmarshalTypeError{
@@ -3168,6 +3232,18 @@ func (b *DurationBuilder) unmarshalOne(dec 
*json.Decoder) error {
        switch v := t.(type) {
        case nil:
                b.AppendNull()
+       case json.Number:
+               n, err := v.Int64()
+               if err != nil {
+                       return &json.UnmarshalTypeError{
+                               Value:  v.String(),
+                               Type:   reflect.TypeOf(arrow.Duration(0)),
+                               Offset: dec.InputOffset(),
+                       }
+               }
+               b.Append(arrow.Duration(n))
+       case float64:
+               b.Append(arrow.Duration(v))
        case string:
                // be flexible for specifying durations by accepting forms like
                // 3h2m0.5s regardless of the unit and converting it to the 
proper
diff --git a/go/arrow/array/numericbuilder.gen.go.tmpl 
b/go/arrow/array/numericbuilder.gen.go.tmpl
index a14463bdd4..1f67a1b6bc 100644
--- a/go/arrow/array/numericbuilder.gen.go.tmpl
+++ b/go/arrow/array/numericbuilder.gen.go.tmpl
@@ -196,9 +196,26 @@ func (b *{{.Name}}Builder) unmarshalOne(dec *json.Decoder) 
error {
                }
 
                b.Append({{.QualifiedType}}FromTime(tm))
+       case json.Number:
+               n, err := v.Int64()
+               if err != nil {
+                       return &json.UnmarshalTypeError{
+                               Value: v.String(),
+                               Type: reflect.TypeOf({{.QualifiedType}}(0)),
+                               Offset: dec.InputOffset(),
+                       }
+               }
+               b.Append({{.QualifiedType}}(n))
+       case float64:
+               b.Append({{.QualifiedType}}(v))
 {{else if or (eq .Name "Time32") (eq .Name "Time64") (eq .Name "Timestamp") -}}
        case string:
+{{if (eq .Name "Timestamp") -}}
+               loc, _ := b.dtype.GetZone()
+               tm, err := arrow.TimestampFromStringInLocation(v, b.dtype.Unit, 
loc)
+{{else -}}
                tm, err := {{.QualifiedType}}FromString(v, b.dtype.Unit)
+{{end}}
                if err != nil {
                        return &json.UnmarshalTypeError{
                                Value: v,
@@ -208,7 +225,31 @@ func (b *{{.Name}}Builder) unmarshalOne(dec *json.Decoder) 
error {
                }
 
                b.Append(tm)
+       case json.Number:
+               n, err := v.Int64()
+               if err != nil {
+                       return &json.UnmarshalTypeError{
+                               Value: v.String(),
+                               Type: reflect.TypeOf({{.QualifiedType}}(0)),
+                               Offset: dec.InputOffset(),
+                       }
+               }
+               b.Append({{.QualifiedType}}(n))
+       case float64:
+               b.Append({{.QualifiedType}}(v))
 {{else if eq .Name "Duration" -}}
+       case json.Number:
+               n, err := v.Int64()
+               if err != nil {
+                       return &json.UnmarshalTypeError{
+                               Value: v.String(),
+                               Type: reflect.TypeOf({{.QualifiedType}}(0)),
+                               Offset: dec.InputOffset(),
+                       }
+               }
+               b.Append({{.QualifiedType}}(n))
+       case float64:
+               b.Append({{.QualifiedType}}(v))
        case string:
                // be flexible for specifying durations by accepting forms like
                // 3h2m0.5s regardless of the unit and converting it to the 
proper
diff --git a/go/arrow/array/util.go b/go/arrow/array/util.go
index b09c82c5ab..64c8912899 100644
--- a/go/arrow/array/util.go
+++ b/go/arrow/array/util.go
@@ -36,6 +36,7 @@ func min(a, b int) int {
 type fromJSONCfg struct {
        multiDocument bool
        startOffset   int64
+       useNumber     bool
 }
 
 type FromJSONOption func(*fromJSONCfg)
@@ -57,6 +58,16 @@ func WithStartOffset(off int64) FromJSONOption {
        }
 }
 
+// WithUseNumber enables the 'UseNumber' option on the json decoder, using
+// the json.Number type instead of assuming float64 for numbers. This is 
critical
+// if you have numbers that are larger than what can fit into the 53 bits of
+// an IEEE float64 mantissa and want to preserve its value.
+func WithUseNumber() FromJSONOption {
+       return func(c *fromJSONCfg) {
+               c.useNumber = true
+       }
+}
+
 // FromJSON creates an arrow.Array from a corresponding JSON stream and 
defined data type. If the types in the
 // json do not match the type provided, it will return errors. This is *not* 
the integration test format
 // and should not be used as such. This intended to be used by consumers more 
similarly to the current exposing of
@@ -132,6 +143,10 @@ func FromJSON(mem memory.Allocator, dt arrow.DataType, r 
io.Reader, opts ...From
                }
        }()
 
+       if cfg.useNumber {
+               dec.UseNumber()
+       }
+
        if !cfg.multiDocument {
                t, err := dec.Token()
                if err != nil {
diff --git a/go/arrow/array/util_test.go b/go/arrow/array/util_test.go
index 3dea9b4b5e..36b03a94d8 100644
--- a/go/arrow/array/util_test.go
+++ b/go/arrow/array/util_test.go
@@ -288,9 +288,10 @@ func TestDateJSON(t *testing.T) {
                bldr := array.NewDate32Builder(memory.DefaultAllocator)
                defer bldr.Release()
 
-               jsonstr := `["1970-01-06", null, "1970-02-12"]`
+               jsonstr := `["1970-01-06", null, "1970-02-12", 0]`
+               jsonExp := `["1970-01-06", null, "1970-02-12", "1970-01-01"]`
 
-               bldr.AppendValues([]arrow.Date32{5, 0, 42}, []bool{true, false, 
true})
+               bldr.AppendValues([]arrow.Date32{5, 0, 42, 0}, []bool{true, 
false, true, true})
                expected := bldr.NewArray()
                defer expected.Release()
 
@@ -302,15 +303,16 @@ func TestDateJSON(t *testing.T) {
 
                data, err := json.Marshal(arr)
                assert.NoError(t, err)
-               assert.JSONEq(t, jsonstr, string(data))
+               assert.JSONEq(t, jsonExp, string(data))
        })
        t.Run("date64", func(t *testing.T) {
                bldr := array.NewDate64Builder(memory.DefaultAllocator)
                defer bldr.Release()
 
-               jsonstr := `["1970-01-02", null, "2286-11-20"]`
+               jsonstr := `["1970-01-02", null, "2286-11-20", 86400000]`
+               jsonExp := `["1970-01-02", null, "2286-11-20", "1970-01-02"]`
 
-               bldr.AppendValues([]arrow.Date64{86400000, 0, 9999936000000}, 
[]bool{true, false, true})
+               bldr.AppendValues([]arrow.Date64{86400000, 0, 9999936000000, 
86400000}, []bool{true, false, true, true})
                expected := bldr.NewArray()
                defer expected.Release()
 
@@ -322,7 +324,7 @@ func TestDateJSON(t *testing.T) {
 
                data, err := json.Marshal(arr)
                assert.NoError(t, err)
-               assert.JSONEq(t, jsonstr, string(data))
+               assert.JSONEq(t, jsonExp, string(data))
        })
 }
 
@@ -331,12 +333,13 @@ func TestTimeJSON(t *testing.T) {
        tests := []struct {
                dt       arrow.DataType
                jsonstr  string
+               jsonexp  string
                valueadd int
        }{
-               {arrow.FixedWidthTypes.Time32s, `[null, "10:10:10"]`, 123},
-               {arrow.FixedWidthTypes.Time32ms, `[null, "10:10:10.123"]`, 456},
-               {arrow.FixedWidthTypes.Time64us, `[null, "10:10:10.123456"]`, 
789},
-               {arrow.FixedWidthTypes.Time64ns, `[null, 
"10:10:10.123456789"]`, 0},
+               {arrow.FixedWidthTypes.Time32s, `[null, "10:10:10", 36610]`, 
`[null, "10:10:10", "10:10:10"]`, 123},
+               {arrow.FixedWidthTypes.Time32ms, `[null, "10:10:10.123", 
36610123]`, `[null, "10:10:10.123", "10:10:10.123"]`, 456},
+               {arrow.FixedWidthTypes.Time64us, `[null, "10:10:10.123456", 
36610123456]`, `[null, "10:10:10.123456", "10:10:10.123456"]`, 789},
+               {arrow.FixedWidthTypes.Time64ns, `[null, "10:10:10.123456789", 
36610123456789]`, `[null, "10:10:10.123456789", "10:10:10.123456789"]`, 0},
        }
 
        for _, tt := range tests {
@@ -350,9 +353,9 @@ func TestTimeJSON(t *testing.T) {
 
                        switch tt.dt.ID() {
                        case arrow.TIME32:
-                               
bldr.(*array.Time32Builder).AppendValues([]arrow.Time32{0, 
arrow.Time32(tententen)}, []bool{false, true})
+                               
bldr.(*array.Time32Builder).AppendValues([]arrow.Time32{0, 
arrow.Time32(tententen), arrow.Time32(tententen)}, []bool{false, true, true})
                        case arrow.TIME64:
-                               
bldr.(*array.Time64Builder).AppendValues([]arrow.Time64{0, 
arrow.Time64(tententen)}, []bool{false, true})
+                               
bldr.(*array.Time64Builder).AppendValues([]arrow.Time64{0, 
arrow.Time64(tententen), arrow.Time64(tententen)}, []bool{false, true, true})
                        }
 
                        expected := bldr.NewArray()
@@ -366,7 +369,7 @@ func TestTimeJSON(t *testing.T) {
 
                        data, err := json.Marshal(arr)
                        assert.NoError(t, err)
-                       assert.JSONEq(t, tt.jsonstr, string(data))
+                       assert.JSONEq(t, tt.jsonexp, string(data))
                })
        }
 }
diff --git a/go/arrow/datatype_fixedwidth.go b/go/arrow/datatype_fixedwidth.go
index 84170e5631..c8c9b6d44d 100644
--- a/go/arrow/datatype_fixedwidth.go
+++ b/go/arrow/datatype_fixedwidth.go
@@ -59,15 +59,31 @@ type (
 
 // Date32FromTime returns a Date32 value from a time object
 func Date32FromTime(t time.Time) Date32 {
-       return Date32(t.Unix() / int64((time.Hour * 24).Seconds()))
+       if _, offset := t.Zone(); offset != 0 {
+               // properly account for timezone adjustments before we calculate
+               // the number of days by adjusting the time and converting to 
UTC
+               t = t.Add(time.Duration(offset) * time.Second).UTC()
+       }
+       return Date32(t.Truncate(24*time.Hour).Unix() / int64((time.Hour * 
24).Seconds()))
 }
 
 func (d Date32) ToTime() time.Time {
        return time.Unix(0, 0).UTC().AddDate(0, 0, int(d))
 }
 
+func (d Date32) FormattedString() string {
+       return d.ToTime().Format("2006-01-02")
+}
+
 // Date64FromTime returns a Date64 value from a time object
 func Date64FromTime(t time.Time) Date64 {
+       if _, offset := t.Zone(); offset != 0 {
+               // properly account for timezone adjustments before we calculate
+               // the actual value by adjusting the time and converting to UTC
+               t = t.Add(time.Duration(offset) * time.Second).UTC()
+       }
+       // truncate to the start of the day to get the correct value
+       t = t.Truncate(24 * time.Hour)
        return Date64(t.Unix()*1e3 + int64(t.Nanosecond())/1e6)
 }
 
@@ -76,29 +92,46 @@ func (d Date64) ToTime() time.Time {
        return time.Unix(0, 0).UTC().AddDate(0, 0, days)
 }
 
-// TimestampFromString parses a string and returns a timestamp for the given 
unit
-// level.
-//
-// The timestamp should be in one of the following forms, [T] can be either T
-// or a space, and [.zzzzzzzzz] can be either left out or up to 9 digits of
-// fractions of a second.
-//
-//      YYYY-MM-DD
-//      YYYY-MM-DD[T]HH
-//   YYYY-MM-DD[T]HH:MM
-//   YYYY-MM-DD[T]HH:MM:SS[.zzzzzzzz]
-func TimestampFromString(val string, unit TimeUnit) (Timestamp, error) {
-       format := "2006-01-02"
-       if val[len(val)-1] == 'Z' {
-               val = val[:len(val)-1]
+func (d Date64) FormattedString() string {
+       return d.ToTime().Format("2006-01-02")
+}
+
+// TimestampFromStringInLocation is like TimestampFromString, but treats the 
time instant
+// as if it were in the passed timezone before converting to UTC for internal 
representation.
+func TimestampFromStringInLocation(val string, unit TimeUnit, loc 
*time.Location) (Timestamp, error) {
+       if len(val) < 10 {
+               return 0, fmt.Errorf("invalid timestamp string")
+       }
+
+       var (
+               format         = "2006-01-02"
+               zoneFmt        string
+               lenWithoutZone = len(val)
+       )
+
+       if lenWithoutZone > 10 {
+               switch {
+               case val[len(val)-1] == 'Z':
+                       zoneFmt = "Z"
+                       lenWithoutZone--
+               case val[len(val)-3] == '+' || val[len(val)-3] == '-':
+                       zoneFmt = "-07"
+                       lenWithoutZone -= 3
+               case val[len(val)-5] == '+' || val[len(val)-5] == '-':
+                       zoneFmt = "-0700"
+                       lenWithoutZone -= 5
+               case val[len(val)-6] == '+' || val[len(val)-6] == '-':
+                       zoneFmt = "-07:00"
+                       lenWithoutZone -= 6
+               }
        }
 
        switch {
-       case len(val) == 13:
+       case lenWithoutZone == 13:
                format += string(val[10]) + "15"
-       case len(val) == 16:
+       case lenWithoutZone == 16:
                format += string(val[10]) + "15:04"
-       case len(val) >= 19:
+       case lenWithoutZone >= 19:
                format += string(val[10]) + "15:04:05.999999999"
        }
 
@@ -106,18 +139,24 @@ func TimestampFromString(val string, unit TimeUnit) 
(Timestamp, error) {
        // don't need a case for nano as time.Parse will already error if
        // more than nanosecond precision is provided
        switch {
-       case unit == Second && len(val) > 19:
+       case unit == Second && lenWithoutZone > 19:
                return 0, xerrors.New("provided more than second precision for 
timestamp[s]")
-       case unit == Millisecond && len(val) > 23:
+       case unit == Millisecond && lenWithoutZone > 23:
                return 0, xerrors.New("provided more than millisecond precision 
for timestamp[ms]")
-       case unit == Microsecond && len(val) > 26:
+       case unit == Microsecond && lenWithoutZone > 26:
                return 0, xerrors.New("provided more than microsecond precision 
for timestamp[us]")
        }
 
-       out, err := time.ParseInLocation(format, val, time.UTC)
+       format += zoneFmt
+       out, err := time.Parse(format, val)
        if err != nil {
                return 0, err
        }
+       if loc != time.UTC {
+               // convert to UTC by putting the same time instant in the 
desired location
+               // before converting to UTC
+               out = out.In(loc).UTC()
+       }
 
        switch unit {
        case Second:
@@ -132,6 +171,24 @@ func TimestampFromString(val string, unit TimeUnit) 
(Timestamp, error) {
        return 0, fmt.Errorf("unexpected timestamp unit: %s", unit)
 }
 
+// TimestampFromString parses a string and returns a timestamp for the given 
unit
+// level.
+//
+// The timestamp should be in one of the following forms, [T] can be either T
+// or a space, and [.zzzzzzzzz] can be either left out or up to 9 digits of
+// fractions of a second.
+//
+//      YYYY-MM-DD
+//      YYYY-MM-DD[T]HH
+//   YYYY-MM-DD[T]HH:MM
+//   YYYY-MM-DD[T]HH:MM:SS[.zzzzzzzz]
+//
+// You can also optionally have an ending Z to indicate UTC or indicate a 
specific
+// timezone using ±HH, ±HHMM or ±HH:MM at the end of the string.
+func TimestampFromString(val string, unit TimeUnit) (Timestamp, error) {
+       return TimestampFromStringInLocation(val, unit, time.UTC)
+}
+
 func (t Timestamp) ToTime(unit TimeUnit) time.Time {
        if unit == Second {
                return time.Unix(int64(t), 0).UTC()
@@ -162,9 +219,9 @@ func Time32FromString(val string, unit TimeUnit) (Time32, 
error) {
        )
        switch {
        case len(val) == 5:
-               out, err = time.ParseInLocation("15:04", val, time.UTC)
+               out, err = time.Parse("15:04", val)
        default:
-               out, err = time.ParseInLocation("15:04:05.999", val, time.UTC)
+               out, err = time.Parse("15:04:05.999", val)
        }
        if err != nil {
                return 0, err
@@ -180,6 +237,18 @@ func (t Time32) ToTime(unit TimeUnit) time.Time {
        return time.Unix(0, int64(t)*int64(unit.Multiplier())).UTC()
 }
 
+func (t Time32) FormattedString(unit TimeUnit) string {
+       const baseFmt = "15:04:05"
+       tm := t.ToTime(unit)
+       switch unit {
+       case Second:
+               return tm.Format(baseFmt)
+       case Millisecond:
+               return tm.Format(baseFmt + ".000")
+       }
+       return ""
+}
+
 // Time64FromString parses a string to return a Time64 value in the given unit,
 // unit needs to be only microseconds or nanoseconds and the string should be 
in the
 // form of HH:MM or HH:MM:SS[.zzzzzzzzz] where the fractions of a second are 
optional.
@@ -201,9 +270,9 @@ func Time64FromString(val string, unit TimeUnit) (Time64, 
error) {
        )
        switch {
        case len(val) == 5:
-               out, err = time.ParseInLocation("15:04", val, time.UTC)
+               out, err = time.Parse("15:04", val)
        default:
-               out, err = time.ParseInLocation("15:04:05.999999999", val, 
time.UTC)
+               out, err = time.Parse("15:04:05.999999999", val)
        }
        if err != nil {
                return 0, err
@@ -219,6 +288,18 @@ func (t Time64) ToTime(unit TimeUnit) time.Time {
        return time.Unix(0, int64(t)*int64(unit.Multiplier())).UTC()
 }
 
+func (t Time64) FormattedString(unit TimeUnit) string {
+       const baseFmt = "15:04:05.000000"
+       tm := t.ToTime(unit)
+       switch unit {
+       case Microsecond:
+               return tm.Format(baseFmt)
+       case Nanosecond:
+               return tm.Format(baseFmt + "000")
+       }
+       return ""
+}
+
 const (
        Second TimeUnit = iota
        Millisecond
@@ -232,12 +313,19 @@ func (u TimeUnit) Multiplier() time.Duration {
 
 func (u TimeUnit) String() string { return [...]string{"s", "ms", "us", 
"ns"}[uint(u)&3] }
 
+type TemporalWithUnit interface {
+       FixedWidthDataType
+       TimeUnit() TimeUnit
+}
+
 // TimestampType is encoded as a 64-bit signed integer since the UNIX epoch 
(2017-01-01T00:00:00Z).
 // The zero-value is a nanosecond and time zone neutral. Time zone neutral can 
be
 // considered UTC without having "UTC" as a time zone.
 type TimestampType struct {
        Unit     TimeUnit
        TimeZone string
+
+       loc *time.Location
 }
 
 func (*TimestampType) ID() Type     { return TIMESTAMP }
@@ -258,6 +346,85 @@ func (t *TimestampType) Fingerprint() string {
 // BitWidth returns the number of bits required to store a single element of 
this data type in memory.
 func (*TimestampType) BitWidth() int { return 64 }
 
+func (t *TimestampType) TimeUnit() TimeUnit { return t.Unit }
+
+// ClearCachedLocation clears the cached time.Location object in the type.
+// This should be called if you change the value of the TimeZone after having
+// potentially called GetZone.
+func (t *TimestampType) ClearCachedLocation() {
+       t.loc = nil
+}
+
+// GetZone returns a *time.Location that represents the current TimeZone member
+// of the TimestampType. If it is "", "UTC", or "utc", you'll get time.UTC.
+// Otherwise it must either be a valid tzdata string such as "America/New_York"
+// or of the format +HH:MM or -HH:MM indicating an absolute offset.
+//
+// The location object will be cached in the TimestampType for subsequent calls
+// so if you change the value of TimeZone after calling this, make sure to call
+// ClearCachedLocation.
+func (t *TimestampType) GetZone() (*time.Location, error) {
+       if t.loc != nil {
+               return t.loc, nil
+       }
+
+       // the TimeZone string is allowed to be either a valid tzdata string
+       // such as "America/New_York" or an absolute offset of the form -XX:XX
+       // or +XX:XX
+       //
+       // As such we have two methods we can try, first we'll try LoadLocation
+       // and if that fails, we'll test for an absolute offset.
+       if t.TimeZone == "" || t.TimeZone == "UTC" || t.TimeZone == "utc" {
+               t.loc = time.UTC
+               return time.UTC, nil
+       }
+
+       if loc, err := time.LoadLocation(t.TimeZone); err == nil {
+               t.loc = loc
+               return t.loc, err
+       }
+
+       // at this point we know that the timezone isn't empty, and didn't match
+       // anything in the tzdata names. So either it's an absolute offset
+       // or it's invalid.
+       timetz, err := time.Parse("-07:00", t.TimeZone)
+       if err != nil {
+               return time.UTC, fmt.Errorf("could not find timezone location 
for '%s'", t.TimeZone)
+       }
+
+       _, offset := timetz.Zone()
+       t.loc = time.FixedZone(t.TimeZone, offset)
+       return t.loc, nil
+}
+
+// GetToTimeFunc returns a function for converting an arrow.Timestamp value 
into a
+// time.Time object with proper TimeZone and precision. If the TimeZone is 
invalid
+// this will return an error. It calls GetZone to get the timezone for 
consistency.
+func (t *TimestampType) GetToTimeFunc() (func(Timestamp) time.Time, error) {
+       tz, err := t.GetZone()
+       if err != nil {
+               return nil, err
+       }
+
+       switch t.Unit {
+       case Second:
+               return func(v Timestamp) time.Time { return time.Unix(int64(v), 
0).In(tz) }, nil
+       case Millisecond:
+               factor := int64(time.Second / time.Millisecond)
+               return func(v Timestamp) time.Time {
+                       return time.Unix(int64(v)/factor, 
(int64(v)%factor)*int64(time.Millisecond)).In(tz)
+               }, nil
+       case Microsecond:
+               factor := int64(time.Second / time.Microsecond)
+               return func(v Timestamp) time.Time {
+                       return time.Unix(int64(v)/factor, 
(int64(v)%factor)*int64(time.Microsecond)).In(tz)
+               }, nil
+       case Nanosecond:
+               return func(v Timestamp) time.Time { return time.Unix(0, 
int64(v)).In(tz) }, nil
+       }
+       return nil, fmt.Errorf("invalid timestamp unit: %s", t.Unit)
+}
+
 // Time32Type is encoded as a 32-bit signed integer, representing either 
seconds or milliseconds since midnight.
 type Time32Type struct {
        Unit TimeUnit
@@ -271,6 +438,8 @@ func (t *Time32Type) Fingerprint() string {
        return typeFingerprint(t) + string(timeUnitFingerprint(t.Unit))
 }
 
+func (t *Time32Type) TimeUnit() TimeUnit { return t.Unit }
+
 // Time64Type is encoded as a 64-bit signed integer, representing either 
microseconds or nanoseconds since midnight.
 type Time64Type struct {
        Unit TimeUnit
@@ -284,6 +453,8 @@ func (t *Time64Type) Fingerprint() string {
        return typeFingerprint(t) + string(timeUnitFingerprint(t.Unit))
 }
 
+func (t *Time64Type) TimeUnit() TimeUnit { return t.Unit }
+
 // DurationType is encoded as a 64-bit signed integer, representing an amount
 // of elapsed time without any relation to a calendar artifact.
 type DurationType struct {
@@ -298,6 +469,8 @@ func (t *DurationType) Fingerprint() string {
        return typeFingerprint(t) + string(timeUnitFingerprint(t.Unit))
 }
 
+func (t *DurationType) TimeUnit() TimeUnit { return t.Unit }
+
 // Float16Type represents a floating point value encoded with a 16-bit 
precision.
 type Float16Type struct{}
 
diff --git a/go/arrow/datatype_fixedwidth_test.go 
b/go/arrow/datatype_fixedwidth_test.go
index 3e53f0c5e6..9273836c2c 100644
--- a/go/arrow/datatype_fixedwidth_test.go
+++ b/go/arrow/datatype_fixedwidth_test.go
@@ -112,7 +112,7 @@ func TestTimestampType(t *testing.T) {
                {arrow.Second, "", "timestamp[s]"},
        } {
                t.Run(tc.want, func(t *testing.T) {
-                       dt := arrow.TimestampType{tc.unit, tc.timeZone}
+                       dt := arrow.TimestampType{Unit: tc.unit, TimeZone: 
tc.timeZone}
                        if got, want := dt.BitWidth(), 64; got != want {
                                t.Fatalf("invalid bitwidth: got=%d, want=%d", 
got, want)
                        }

Reply via email to