Allow run-end encoding of TimestampWithOffset

Author: serramatutu

arrow/extensions/timestamp_with_offset.go

@@ -19,6 +19,7 @@ package extensions
 import (
 	"errors"
 	"fmt"
+	"math"
 	"reflect"
 	"strings"
 	"time"
@@ -41,6 +42,13 @@ func isOffsetTypeOk(offsetType arrow.DataType) bool {
 		return true
 	case *arrow.DictionaryType:
 		return arrow.IsInteger(offsetType.IndexType.ID()) && 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
 	}
@@ -140,6 +148,21 @@ func NewTimestampWithOffsetTypeDictionaryEncoded(unit arrow.TimeUnit, index arro
 	return NewTimestampWithOffsetType(unit, &offsetType)
 }
 
+// NewTimestampWithOffsetType 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.
+//
+// The error will be populated if runEnds is not a valid run-end encoding run-ends type.
+func NewTimestampWithOffsetTypeRunEndEncoded(unit arrow.TimeUnit, runEnds arrow.DataType) (*TimestampWithOffsetType, error) {
+	offsetType := arrow.RunEndEncodedOf(runEnds, arrow.PrimitiveTypes.Int16)
+	if !offsetType.ValidRunEndsType(runEnds) {
+		return nil, errors.New(fmt.Sprintf("Invalid run-ends type %s", runEnds))
+	}
+
+	return NewTimestampWithOffsetType(unit, offsetType)
+}
+
+
 func (b *TimestampWithOffsetType) ArrayType() reflect.Type {
 	return reflect.TypeOf(TimestampWithOffsetArray{})
 }
@@ -231,6 +254,9 @@ func fieldValuesFromTime(t time.Time, unit arrow.TimeUnit) (arrow.Timestamp, int
 // Get the raw arrow values at the given index
 //
 // SAFETY: the value at i must not be nil
+//
+// NOTE: the REE implementation of this calls `LogicalValuesArray()`, which can be slow
+// for large arrays.
 func (a *TimestampWithOffsetArray) rawValueUnsafe(i int) (arrow.Timestamp, int16, arrow.TimeUnit) {
 	structs := a.Storage().(*array.Struct)
 
@@ -247,6 +273,10 @@ func (a *TimestampWithOffsetArray) rawValueUnsafe(i int) (arrow.Timestamp, int16
 		offsetMinutes = offsets.Value(i)
 	case *array.Dictionary:
 		offsetMinutes = offsets.Dictionary().(*array.Int16).Value(offsets.GetValueIndex(i))
+	case *array.RunEndEncoded:
+		logicalValues := offsets.LogicalValuesArray().(*array.Int16)
+		defer logicalValues.Release()
+		offsetMinutes = logicalValues.Value(i)
 	}
 
 	return utcTimestamp, offsetMinutes, timeUnit
@@ -262,6 +292,7 @@ func (a *TimestampWithOffsetArray) Value(i int) time.Time {
 
 func (a *TimestampWithOffsetArray) Values() []time.Time {
 	values := make([]time.Time, a.Len())
+	// TODO: optimize for run-end encoding
 	for i := range a.Len() {
 		val := a.Value(i)
 		values[i] = val
@@ -307,6 +338,8 @@ type TimestampWithOffsetBuilder struct {
 	Layout     string
 	unit       arrow.TimeUnit
 	offsetType arrow.DataType
+	// lastOffset is only used to determine when to start new runs with run-end encoded offsets
+	lastOffset int16
 }
 
 // NewTimestampWithOffsetBuilder creates a new TimestampWithOffsetBuilder, exposing a convenient and efficient interface
@@ -320,38 +353,60 @@ func NewTimestampWithOffsetBuilder(mem memory.Allocator, unit arrow.TimeUnit, of
 	return &TimestampWithOffsetBuilder{
 		unit:             unit,
 		offsetType: 	  offsetType,
+		lastOffset:       math.MaxInt16,
 		Layout:           time.RFC3339,
 		ExtensionBuilder: array.NewExtensionBuilder(mem, dataType),
 	}, nil
 }
 
 func (b *TimestampWithOffsetBuilder) Append(v time.Time) {
 	timestamp, offsetMinutes := fieldValuesFromTime(v, b.unit)
+	offsetMinutes16 := int16(offsetMinutes)
 	structBuilder := b.ExtensionBuilder.Builder.(*array.StructBuilder)
 
 	structBuilder.Append(true)
 	structBuilder.FieldBuilder(0).(*array.TimestampBuilder).Append(timestamp)
 
 	switch offsets := structBuilder.FieldBuilder(1).(type) {
 	case *array.Int16Builder:
-		offsets.Append(int16(offsetMinutes))
+		offsets.Append(offsetMinutes16)
 	case *array.Int16DictionaryBuilder:
-		offsets.Append(int16(offsetMinutes))
+		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
 	}
+
 }
 
 func (b *TimestampWithOffsetBuilder) UnsafeAppend(v time.Time) {
 	timestamp, offsetMinutes := fieldValuesFromTime(v, b.unit)
+	offsetMinutes16 := int16(offsetMinutes)
 	structBuilder := b.ExtensionBuilder.Builder.(*array.StructBuilder)
 
 	structBuilder.Append(true)
 	structBuilder.FieldBuilder(0).(*array.TimestampBuilder).UnsafeAppend(timestamp)
 
 	switch offsets := structBuilder.FieldBuilder(1).(type) {
 	case *array.Int16Builder:
-		offsets.UnsafeAppend(int16(offsetMinutes))
+		offsets.UnsafeAppend(offsetMinutes16)
 	case *array.Int16DictionaryBuilder:
-		offsets.Append(int16(offsetMinutes))
+		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
 	}
 }
 
@@ -399,13 +454,31 @@ func (b *TimestampWithOffsetBuilder) AppendValues(values []time.Time, valids []b
 			timestamp, offsetMinutes := fieldValuesFromTime(v, b.unit)
 			if valids[i] {
 				timestamps.UnsafeAppend(timestamp)
-				// TODO: I was here, this needs to be equivalent to UnsafeAppend
 				offsets.Append(offsetMinutes)
 			} else {
 				timestamps.UnsafeAppendBoolToBitmap(false)
 				offsets.UnsafeAppendBoolToBitmap(false)
 			}
 		}
+	case *array.RunEndEncodedBuilder:
+		offsetValuesBuilder := offsets.ValueBuilder().(*array.Int16Builder)
+		for i, v := range values {
+			timestamp, offsetMinutes := fieldValuesFromTime(v, b.unit)
+			if valids[i] {
+				timestamps.UnsafeAppend(timestamp)
+				offsetMinutes16 := int16(offsetMinutes)
+				if offsetMinutes != b.lastOffset {
+					offsets.Append(1)
+					offsetValuesBuilder.Append(offsetMinutes16)
+				} else {
+					offsets.ContinueRun(1)
+				}
+				b.lastOffset = offsetMinutes16
+			} else {
+				timestamps.UnsafeAppendBoolToBitmap(false)
+				offsets.UnsafeAppendBoolToBitmap(false)
+			}
+		}
 	}
 }
 

arrow/extensions/timestamp_with_offset_test.go

@@ -50,6 +50,12 @@ func dict(index arrow.DataType) arrow.DataType {
 	}
 }
 
+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 = []arrow.DataType{
@@ -65,6 +71,11 @@ var allAllowedOffsetTypes = []arrow.DataType{
 	dict(arrow.PrimitiveTypes.Int16),
 	dict(arrow.PrimitiveTypes.Int32),
 	dict(arrow.PrimitiveTypes.Int64),
+
+	// run-end encoded offsetType
+	ree(arrow.PrimitiveTypes.Int16),
+	ree(arrow.PrimitiveTypes.Int32),
+	ree(arrow.PrimitiveTypes.Int64),
 }
 
 func TestTimestampWithOffsetTypePrimitiveBasics(t *testing.T) {
@@ -141,6 +152,47 @@ func TestTimestampWithOffsetTypeDictionaryEncodedBasics(t *testing.T) {
 	}
 }
 
+func TestTimestampWithOffsetTypeRunEndEncodedBasics(t *testing.T) {
+	invalidRunEndsType := arrow.PrimitiveTypes.Float32
+	_, err := extensions.NewTimestampWithOffsetTypeRunEndEncoded(testTimeUnit, invalidRunEndsType)
+	assert.True(t, err != nil, "Err should not be nil if run ends type is invalid")
+
+	runEndsTypes := []arrow.DataType{
+		arrow.PrimitiveTypes.Int16,
+		arrow.PrimitiveTypes.Int32,
+		arrow.PrimitiveTypes.Int64,
+	};
+
+	for _, indexType := range runEndsTypes {
+		typ, err := extensions.NewTimestampWithOffsetTypeRunEndEncoded(testTimeUnit, indexType)
+		assert.True(t, err == nil, "Err should be nil")
+
+		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(indexType),
+					Nullable: false,
+				},
+			),
+			typ.StorageType()))
+
+		assert.Equal(t, "extension<arrow.timestamp_with_offset>", typ.String())
+	}
+}
+
 func TestTimestampWithOffsetExtensionBuilder(t *testing.T) {
 	mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
 	defer mem.AssertSize(t, 0)