github.com/apache/arrow/go/v16@v16.1.0/arrow/array/timestamp_test.go (about)

     1  // Licensed to the Apache Software Foundation (ASF) under one
     2  // or more contributor license agreements.  See the NOTICE file
     3  // distributed with this work for additional information
     4  // regarding copyright ownership.  The ASF licenses this file
     5  // to you under the Apache License, Version 2.0 (the
     6  // "License"); you may not use this file except in compliance
     7  // with the License.  You may obtain a copy of the License at
     8  //
     9  // http://www.apache.org/licenses/LICENSE-2.0
    10  //
    11  // Unless required by applicable law or agreed to in writing, software
    12  // distributed under the License is distributed on an "AS IS" BASIS,
    13  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14  // See the License for the specific language governing permissions and
    15  // limitations under the License.
    16  
    17  package array_test
    18  
    19  import (
    20  	"testing"
    21  	"time"
    22  
    23  	"github.com/apache/arrow/go/v16/arrow"
    24  	"github.com/apache/arrow/go/v16/arrow/array"
    25  	"github.com/apache/arrow/go/v16/arrow/memory"
    26  	"github.com/stretchr/testify/assert"
    27  )
    28  
    29  func TestTimestampStringRoundTrip(t *testing.T) {
    30  	// 1. create array
    31  	mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
    32  	defer mem.AssertSize(t, 0)
    33  
    34  	dt := &arrow.TimestampType{Unit: arrow.Second}
    35  	b := array.NewTimestampBuilder(mem, dt)
    36  	defer b.Release()
    37  
    38  	b.Append(1)
    39  	b.Append(2)
    40  	b.Append(3)
    41  	b.AppendNull()
    42  	b.Append(5)
    43  	b.Append(6)
    44  	b.AppendNull()
    45  	b.Append(8)
    46  	b.Append(9)
    47  	b.Append(10)
    48  
    49  	arr := b.NewArray().(*array.Timestamp)
    50  	defer arr.Release()
    51  
    52  	// 2. create array via AppendValueFromString
    53  	b1 := array.NewTimestampBuilder(mem, dt)
    54  	defer b1.Release()
    55  
    56  	for i := 0; i < arr.Len(); i++ {
    57  		assert.NoError(t, b1.AppendValueFromString(arr.ValueStr(i)))
    58  	}
    59  
    60  	arr1 := b1.NewArray().(*array.Timestamp)
    61  	defer arr1.Release()
    62  
    63  	assert.True(t, array.Equal(arr, arr1))
    64  }
    65  
    66  func TestNewTimestampBuilder(t *testing.T) {
    67  	mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
    68  	defer mem.AssertSize(t, 0)
    69  	timestamp := time.Now()
    70  	dtype := &arrow.TimestampType{Unit: arrow.Second}
    71  	ab := array.NewTimestampBuilder(mem, dtype)
    72  	defer ab.Release()
    73  
    74  	ab.Retain()
    75  	ab.Release()
    76  
    77  	ab.Append(1)
    78  	ab.Append(2)
    79  	ab.Append(3)
    80  	ab.AppendNull()
    81  	ab.Append(5)
    82  	ab.Append(6)
    83  	ab.AppendNull()
    84  	ab.Append(8)
    85  	ab.Append(9)
    86  	ab.Append(10)
    87  	ab.AppendTime(timestamp)
    88  
    89  	// check state of builder before NewTimestampArray
    90  	assert.Equal(t, 11, ab.Len(), "unexpected Len()")
    91  	assert.Equal(t, 2, ab.NullN(), "unexpected NullN()")
    92  
    93  	a := ab.NewTimestampArray()
    94  
    95  	// check state of builder after NewTimestampArray
    96  	assert.Zero(t, ab.Len(), "unexpected ArrayBuilder.Len(), NewTimestampArray did not reset state")
    97  	assert.Zero(t, ab.Cap(), "unexpected ArrayBuilder.Cap(), NewTimestampArray did not reset state")
    98  	assert.Zero(t, ab.NullN(), "unexpected ArrayBuilder.NullN(), NewTimestampArray did not reset state")
    99  
   100  	// check state of array
   101  	assert.Equal(t, 2, a.NullN(), "unexpected null count")
   102  	assert.Equal(t, []arrow.Timestamp{1, 2, 3, 0, 5, 6, 0, 8, 9, 10, arrow.Timestamp(timestamp.Unix())}, a.TimestampValues(), "unexpected TimestampValues")
   103  	assert.Equal(t, []byte{0xb7}, a.NullBitmapBytes()[:1]) // 4 bytes due to minBuilderCapacity
   104  	assert.Len(t, a.TimestampValues(), 11, "unexpected length of TimestampValues")
   105  
   106  	a.Release()
   107  
   108  	ab.Append(7)
   109  	ab.Append(8)
   110  
   111  	a = ab.NewTimestampArray()
   112  
   113  	assert.Equal(t, 0, a.NullN())
   114  	assert.Equal(t, []arrow.Timestamp{7, 8}, a.TimestampValues())
   115  	assert.Len(t, a.TimestampValues(), 2)
   116  
   117  	a.Release()
   118  
   119  	var (
   120  		want   = []arrow.Timestamp{1, 2, 3, 4}
   121  		valids = []bool{true, true, false, true}
   122  	)
   123  
   124  	ab.AppendValues(want, valids)
   125  	a = ab.NewTimestampArray()
   126  
   127  	sub := array.MakeFromData(a.Data())
   128  	defer sub.Release()
   129  
   130  	if got, want := sub.DataType().ID(), a.DataType().ID(); got != want {
   131  		t.Fatalf("invalid type: got=%q, want=%q", got, want)
   132  	}
   133  
   134  	if _, ok := sub.(*array.Timestamp); !ok {
   135  		t.Fatalf("could not type-assert to array.Timestamp")
   136  	}
   137  
   138  	if got, want := a.String(), `[1 2 (null) 4]`; got != want {
   139  		t.Fatalf("got=%q, want=%q", got, want)
   140  	}
   141  
   142  	slice := array.NewSliceData(a.Data(), 2, 4)
   143  	defer slice.Release()
   144  
   145  	sub1 := array.MakeFromData(slice)
   146  	defer sub1.Release()
   147  
   148  	v, ok := sub1.(*array.Timestamp)
   149  	if !ok {
   150  		t.Fatalf("could not type-assert to array.Timestamp")
   151  	}
   152  
   153  	if got, want := v.String(), `[(null) 4]`; got != want {
   154  		t.Fatalf("got=%q, want=%q", got, want)
   155  	}
   156  
   157  	a.Release()
   158  }
   159  
   160  func TestTimestampBuilder_AppendValues(t *testing.T) {
   161  	mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
   162  	defer mem.AssertSize(t, 0)
   163  
   164  	dtype := &arrow.TimestampType{Unit: arrow.Second}
   165  	ab := array.NewTimestampBuilder(mem, dtype)
   166  	defer ab.Release()
   167  
   168  	exp := []arrow.Timestamp{0, 1, 2, 3}
   169  	ab.AppendValues(exp, nil)
   170  	a := ab.NewTimestampArray()
   171  	assert.Equal(t, exp, a.TimestampValues())
   172  
   173  	a.Release()
   174  }
   175  
   176  func TestTimestampBuilder_Empty(t *testing.T) {
   177  	mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
   178  	defer mem.AssertSize(t, 0)
   179  
   180  	dtype := &arrow.TimestampType{Unit: arrow.Second}
   181  	ab := array.NewTimestampBuilder(mem, dtype)
   182  	defer ab.Release()
   183  
   184  	exp := []arrow.Timestamp{0, 1, 2, 3}
   185  
   186  	ab.AppendValues([]arrow.Timestamp{}, nil)
   187  	a := ab.NewTimestampArray()
   188  	assert.Zero(t, a.Len())
   189  	a.Release()
   190  
   191  	ab.AppendValues(nil, nil)
   192  	a = ab.NewTimestampArray()
   193  	assert.Zero(t, a.Len())
   194  	a.Release()
   195  
   196  	ab.AppendValues([]arrow.Timestamp{}, nil)
   197  	ab.AppendValues(exp, nil)
   198  	a = ab.NewTimestampArray()
   199  	assert.Equal(t, exp, a.TimestampValues())
   200  	a.Release()
   201  
   202  	ab.AppendValues(exp, nil)
   203  	ab.AppendValues([]arrow.Timestamp{}, nil)
   204  	a = ab.NewTimestampArray()
   205  	assert.Equal(t, exp, a.TimestampValues())
   206  	a.Release()
   207  }
   208  
   209  func TestTimestampBuilder_Resize(t *testing.T) {
   210  	mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
   211  	defer mem.AssertSize(t, 0)
   212  
   213  	dtype := &arrow.TimestampType{Unit: arrow.Second}
   214  	ab := array.NewTimestampBuilder(mem, dtype)
   215  	defer ab.Release()
   216  
   217  	assert.Equal(t, 0, ab.Cap())
   218  	assert.Equal(t, 0, ab.Len())
   219  
   220  	ab.Reserve(63)
   221  	assert.Equal(t, 64, ab.Cap())
   222  	assert.Equal(t, 0, ab.Len())
   223  
   224  	for i := 0; i < 63; i++ {
   225  		ab.Append(0)
   226  	}
   227  	assert.Equal(t, 64, ab.Cap())
   228  	assert.Equal(t, 63, ab.Len())
   229  
   230  	ab.Resize(5)
   231  	assert.Equal(t, 5, ab.Len())
   232  
   233  	ab.Resize(32)
   234  	assert.Equal(t, 5, ab.Len())
   235  }
   236  
   237  func TestTimestampValueStr(t *testing.T) {
   238  	mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
   239  	defer mem.AssertSize(t, 0)
   240  
   241  	dt := &arrow.TimestampType{Unit: arrow.Second, TimeZone: "America/Phoenix"}
   242  	b := array.NewTimestampBuilder(mem, dt)
   243  	defer b.Release()
   244  
   245  	b.Append(-34226955)
   246  	b.Append(1456767743)
   247  
   248  	arr := b.NewArray()
   249  	defer arr.Release()
   250  
   251  	assert.Equal(t, "1968-11-30 13:30:45-0700", arr.ValueStr(0))
   252  	assert.Equal(t, "2016-02-29 10:42:23-0700", arr.ValueStr(1))
   253  }
   254  
   255  func TestTimestampEquality(t *testing.T) {
   256  	mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
   257  	defer mem.AssertSize(t, 0)
   258  
   259  	tsDatatypes := []*arrow.TimestampType{
   260  		{Unit: arrow.Second},
   261  		{Unit: arrow.Second, TimeZone: "UTC"},
   262  		{Unit: arrow.Second, TimeZone: "America/Phoenix"},
   263  	}
   264  
   265  	arrs := make([]*array.Timestamp, 0, len(tsDatatypes))
   266  	for _, dt := range tsDatatypes {
   267  		bldr := array.NewTimestampBuilder(mem, dt)
   268  		defer bldr.Release()
   269  
   270  		bldr.Append(-34226955)
   271  		bldr.Append(1456767743)
   272  
   273  		arr := bldr.NewTimestampArray()
   274  		defer arr.Release()
   275  
   276  		arrs = append(arrs, arr)
   277  	}
   278  
   279  	// No timezone, "wall clock" semantics
   280  	// These timestamps have no actual timezone, but we still represent as UTC per Go conventions
   281  	assert.Equal(t, "1968-11-30 20:30:45Z", arrs[0].ValueStr(0))
   282  	assert.Equal(t, "2016-02-29 17:42:23Z", arrs[0].ValueStr(1))
   283  
   284  	// UTC timezone, "instant" semantics
   285  	assert.Equal(t, "1968-11-30 20:30:45Z", arrs[1].ValueStr(0))
   286  	assert.Equal(t, "2016-02-29 17:42:23Z", arrs[1].ValueStr(1))
   287  
   288  	// America/Phoenix timezone, "instant" semantics
   289  	assert.Equal(t, "1968-11-30 13:30:45-0700", arrs[2].ValueStr(0))
   290  	assert.Equal(t, "2016-02-29 10:42:23-0700", arrs[2].ValueStr(1))
   291  
   292  	// Despite timezone and semantics, the physical values are equivalent
   293  	assert.Equal(t, arrs[0].Value(0), arrs[1].Value(0))
   294  	assert.Equal(t, arrs[0].Value(0), arrs[2].Value(0))
   295  	assert.Equal(t, arrs[1].Value(0), arrs[2].Value(0))
   296  
   297  	assert.Equal(t, arrs[0].Value(1), arrs[1].Value(1))
   298  	assert.Equal(t, arrs[0].Value(1), arrs[2].Value(1))
   299  	assert.Equal(t, arrs[1].Value(1), arrs[2].Value(1))
   300  }