You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by ze...@apache.org on 2023/01/25 22:54:53 UTC
[arrow] branch master updated: GH-32947: [Go] Implement Concatenate for REE Array (#14126)
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 3b21e0c4b2 GH-32947: [Go] Implement Concatenate for REE Array (#14126)
3b21e0c4b2 is described below
commit 3b21e0c4b23fd69aceae66534e581620f149283d
Author: Matt Topol <zo...@gmail.com>
AuthorDate: Wed Jan 25 17:54:46 2023 -0500
GH-32947: [Go] Implement Concatenate for REE Array (#14126)
* Closes: #32947
Authored-by: Matt Topol <zo...@gmail.com>
Signed-off-by: Matt Topol <zo...@gmail.com>
---
go/arrow/array/concat.go | 172 ++++++++++++++++++++++++++++++++++++++++-
go/arrow/array/concat_test.go | 175 ++++++++++++++++++++++++++++++++++++++++++
go/arrow/type_string.go | 2 +-
3 files changed, 347 insertions(+), 2 deletions(-)
diff --git a/go/arrow/array/concat.go b/go/arrow/array/concat.go
index 8ba1973d9c..2fc73a5d59 100644
--- a/go/arrow/array/concat.go
+++ b/go/arrow/array/concat.go
@@ -24,6 +24,7 @@ import (
"github.com/apache/arrow/go/v11/arrow"
"github.com/apache/arrow/go/v11/arrow/bitutil"
+ "github.com/apache/arrow/go/v11/arrow/encoded"
"github.com/apache/arrow/go/v11/arrow/internal/debug"
"github.com/apache/arrow/go/v11/arrow/memory"
"github.com/apache/arrow/go/v11/internal/bitutils"
@@ -42,7 +43,12 @@ func Concatenate(arrs []arrow.Array, mem memory.Allocator) (result arrow.Array,
defer func() {
if pErr := recover(); pErr != nil {
- err = fmt.Errorf("arrow/concat: unknown error: %v", pErr)
+ switch e := pErr.(type) {
+ case error:
+ err = fmt.Errorf("arrow/concat: %w", e)
+ default:
+ err = fmt.Errorf("arrow/concat: unknown error: %v", pErr)
+ }
}
}()
@@ -517,6 +523,50 @@ func concat(data []arrow.ArrayData, mem memory.Allocator) (arrow.ArrayData, erro
if err != nil {
return nil, err
}
+ case *arrow.RunEndEncodedType:
+ physicalLength, overflow := int(0), false
+ // we can't use gatherChildren because the Offset and Len of
+ // data doesn't correspond to the physical length or offset
+ runs := make([]arrow.ArrayData, len(data))
+ values := make([]arrow.ArrayData, len(data))
+ for i, d := range data {
+ plen := encoded.GetPhysicalLength(d)
+ off := encoded.FindPhysicalOffset(d)
+
+ runs[i] = NewSliceData(d.Children()[0], int64(off), int64(off+plen))
+ defer runs[i].Release()
+ values[i] = NewSliceData(d.Children()[1], int64(off), int64(off+plen))
+ defer values[i].Release()
+
+ physicalLength, overflow = addOvf(physicalLength, plen)
+ if overflow {
+ return nil, fmt.Errorf("%w: run end encoded array length must fit into a 32-bit signed integer",
+ arrow.ErrInvalid)
+ }
+ }
+
+ runEndsByteWidth := runs[0].DataType().(arrow.FixedWidthDataType).Bytes()
+ runEndsBuffers := gatherFixedBuffers(runs, 1, runEndsByteWidth)
+ outRunEndsLen := physicalLength * runEndsByteWidth
+ outRunEndsBuf := memory.NewResizableBuffer(mem)
+ outRunEndsBuf.Resize(outRunEndsLen)
+ defer outRunEndsBuf.Release()
+
+ if err := updateRunEnds(runEndsByteWidth, data, runEndsBuffers, outRunEndsBuf); err != nil {
+ return nil, err
+ }
+
+ out.childData = make([]arrow.ArrayData, 2)
+ out.childData[0] = NewData(data[0].Children()[0].DataType(), int(physicalLength),
+ []*memory.Buffer{nil, outRunEndsBuf}, nil, 0, 0)
+
+ var err error
+ out.childData[1], err = concat(values, mem)
+ if err != nil {
+ out.childData[0].Release()
+ return nil, err
+ }
+
default:
return nil, fmt.Errorf("concatenate not implemented for type %s", dt)
}
@@ -565,3 +615,123 @@ func concatBitmaps(bitmaps []bitmap, mem memory.Allocator) (*memory.Buffer, erro
}
return out, nil
}
+
+func updateRunEnds(byteWidth int, inputData []arrow.ArrayData, inputBuffers []*memory.Buffer, outputBuffer *memory.Buffer) error {
+ switch byteWidth {
+ case 2:
+ out := arrow.Int16Traits.CastFromBytes(outputBuffer.Bytes())
+ return updateRunsInt16(inputData, inputBuffers, out)
+ case 4:
+ out := arrow.Int32Traits.CastFromBytes(outputBuffer.Bytes())
+ return updateRunsInt32(inputData, inputBuffers, out)
+ case 8:
+ out := arrow.Int64Traits.CastFromBytes(outputBuffer.Bytes())
+ return updateRunsInt64(inputData, inputBuffers, out)
+ }
+ return fmt.Errorf("%w: invalid dataType for RLE runEnds", arrow.ErrInvalid)
+}
+
+func updateRunsInt16(inputData []arrow.ArrayData, inputBuffers []*memory.Buffer, output []int16) error {
+ // for now we will not attempt to optimize by checking if we
+ // can fold the end and beginning of each array we're concatenating
+ // into a single run
+ pos := 0
+ for i, buf := range inputBuffers {
+ if buf.Len() == 0 {
+ continue
+ }
+ src := arrow.Int16Traits.CastFromBytes(buf.Bytes())
+ if pos == 0 {
+ pos += copy(output, src)
+ continue
+ }
+
+ lastEnd := output[pos-1]
+ // we can check the last runEnd in the src and add it to the
+ // last value that we're adjusting them all by to see if we
+ // are going to overflow
+ if int64(lastEnd)+int64(int(src[len(src)-1])-inputData[i].Offset()) > math.MaxInt16 {
+ return fmt.Errorf("%w: overflow in run-length-encoded run ends concat", arrow.ErrInvalid)
+ }
+
+ // adjust all of the run ends by first normalizing them (e - data[i].offset)
+ // then adding the previous value we ended on. Since the offset
+ // is a logical length offset it should be accurate to just subtract
+ // it from each value.
+ for j, e := range src {
+ output[pos+j] = lastEnd + int16(int(e)-inputData[i].Offset())
+ }
+ pos += len(src)
+ }
+ return nil
+}
+
+func updateRunsInt32(inputData []arrow.ArrayData, inputBuffers []*memory.Buffer, output []int32) error {
+ // for now we will not attempt to optimize by checking if we
+ // can fold the end and beginning of each array we're concatenating
+ // into a single run
+ pos := 0
+ for i, buf := range inputBuffers {
+ if buf.Len() == 0 {
+ continue
+ }
+ src := arrow.Int32Traits.CastFromBytes(buf.Bytes())
+ if pos == 0 {
+ pos += copy(output, src)
+ continue
+ }
+
+ lastEnd := output[pos-1]
+ // we can check the last runEnd in the src and add it to the
+ // last value that we're adjusting them all by to see if we
+ // are going to overflow
+ if int64(lastEnd)+int64(int(src[len(src)-1])-inputData[i].Offset()) > math.MaxInt32 {
+ return fmt.Errorf("%w: overflow in run-length-encoded run ends concat", arrow.ErrInvalid)
+ }
+
+ // adjust all of the run ends by first normalizing them (e - data[i].offset)
+ // then adding the previous value we ended on. Since the offset
+ // is a logical length offset it should be accurate to just subtract
+ // it from each value.
+ for j, e := range src {
+ output[pos+j] = lastEnd + int32(int(e)-inputData[i].Offset())
+ }
+ pos += len(src)
+ }
+ return nil
+}
+
+func updateRunsInt64(inputData []arrow.ArrayData, inputBuffers []*memory.Buffer, output []int64) error {
+ // for now we will not attempt to optimize by checking if we
+ // can fold the end and beginning of each array we're concatenating
+ // into a single run
+ pos := 0
+ for i, buf := range inputBuffers {
+ if buf.Len() == 0 {
+ continue
+ }
+ src := arrow.Int64Traits.CastFromBytes(buf.Bytes())
+ if pos == 0 {
+ pos += copy(output, src)
+ continue
+ }
+
+ lastEnd := output[pos-1]
+ // we can check the last runEnd in the src and add it to the
+ // last value that we're adjusting them all by to see if we
+ // are going to overflow
+ if uint64(lastEnd)+uint64(int(src[len(src)-1])-inputData[i].Offset()) > math.MaxInt64 {
+ return fmt.Errorf("%w: overflow in run-length-encoded run ends concat", arrow.ErrInvalid)
+ }
+
+ // adjust all of the run ends by first normalizing them (e - data[i].offset)
+ // then adding the previous value we ended on. Since the offset
+ // is a logical length offset it should be accurate to just subtract
+ // it from each value.
+ for j, e := range src {
+ output[pos+j] = lastEnd + e - int64(inputData[i].Offset())
+ }
+ pos += len(src)
+ }
+ return nil
+}
diff --git a/go/arrow/array/concat_test.go b/go/arrow/array/concat_test.go
index 8da274319f..46430bd795 100644
--- a/go/arrow/array/concat_test.go
+++ b/go/arrow/array/concat_test.go
@@ -20,6 +20,7 @@ import (
"fmt"
"math"
"sort"
+ "strings"
"testing"
"github.com/apache/arrow/go/v11/arrow"
@@ -568,3 +569,177 @@ func TestConcatDictionaryNullSlots(t *testing.T) {
assert.Truef(t, array.Equal(actual, expected), "got: %s, expected: %s", actual, expected)
}
+
+func TestConcatRunEndEncoded(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ tests := []struct {
+ offsetType arrow.DataType
+ expected interface{}
+ }{
+ {arrow.PrimitiveTypes.Int16, []int16{1, 11, 111, 211, 311, 411, 500, 600}},
+ {arrow.PrimitiveTypes.Int32, []int32{1, 11, 111, 211, 311, 411, 500, 600}},
+ {arrow.PrimitiveTypes.Int64, []int64{1, 11, 111, 211, 311, 411, 500, 600}},
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.offsetType.String(), func(t *testing.T) {
+
+ arrs := make([]arrow.Array, 0)
+ bldr := array.NewRunEndEncodedBuilder(mem, tt.offsetType, arrow.BinaryTypes.String)
+ defer bldr.Release()
+ valBldr := bldr.ValueBuilder().(*array.StringBuilder)
+
+ bldr.Append(1)
+ valBldr.Append("Hello")
+ bldr.AppendNull()
+ bldr.ContinueRun(9)
+
+ bldr.Append(100)
+ valBldr.Append("World")
+ arrs = append(arrs, bldr.NewArray())
+
+ bldr.Append(100)
+ valBldr.Append("Goku")
+ bldr.Append(100)
+ valBldr.Append("Gohan")
+ bldr.Append(100)
+ valBldr.Append("Goten")
+ arrs = append(arrs, bldr.NewArray())
+
+ bldr.AppendNull()
+ bldr.ContinueRun(99)
+ bldr.Append(100)
+ valBldr.Append("Vegeta")
+ bldr.Append(100)
+ valBldr.Append("Trunks")
+ next := bldr.NewArray()
+ defer next.Release()
+ // remove the initial null with an offset and dig into the next run
+ arrs = append(arrs, array.NewSlice(next, 111, int64(next.Len())))
+
+ for _, a := range arrs {
+ defer a.Release()
+ }
+
+ result, err := array.Concatenate(arrs, mem)
+ assert.NoError(t, err)
+ defer result.Release()
+
+ rle := result.(*array.RunEndEncoded)
+ assert.EqualValues(t, 8, rle.GetPhysicalLength())
+ assert.EqualValues(t, 0, rle.GetPhysicalOffset())
+
+ var values interface{}
+ switch endsArr := rle.RunEndsArr().(type) {
+ case *array.Int16:
+ values = endsArr.Int16Values()
+ case *array.Int32:
+ values = endsArr.Int32Values()
+ case *array.Int64:
+ values = endsArr.Int64Values()
+ }
+ assert.Equal(t, tt.expected, values)
+
+ expectedValues, _, _ := array.FromJSON(mem, arrow.BinaryTypes.String,
+ strings.NewReader(`["Hello", null, "World", "Goku", "Gohan", "Goten", "Vegeta", "Trunks"]`))
+ defer expectedValues.Release()
+ assert.Truef(t, array.Equal(expectedValues, rle.Values()), "expected: %s\ngot: %s", expectedValues, rle.Values())
+ })
+ }
+}
+
+func TestConcatAlmostOverflowRunEndEncoding(t *testing.T) {
+ tests := []struct {
+ offsetType arrow.DataType
+ max uint64
+ }{
+ {arrow.PrimitiveTypes.Int16, math.MaxInt16},
+ {arrow.PrimitiveTypes.Int32, math.MaxInt32},
+ {arrow.PrimitiveTypes.Int64, math.MaxInt64},
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.offsetType.String(), func(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ arrs := make([]arrow.Array, 0)
+ bldr := array.NewRunEndEncodedBuilder(mem, tt.offsetType, arrow.BinaryTypes.String)
+ defer bldr.Release()
+ valBldr := bldr.ValueBuilder().(*array.StringBuilder)
+
+ // max is not evently divisible by 4, so we add one to each
+ // to account for that so our final concatenate will overflow
+ bldr.Append((tt.max / 4) + 1)
+ valBldr.Append("foo")
+ bldr.Append((tt.max / 4) + 1)
+ valBldr.Append("bar")
+ arrs = append(arrs, bldr.NewArray())
+
+ bldr.Append((tt.max / 4) + 1)
+ valBldr.Append("baz")
+ bldr.Append((tt.max / 4))
+ valBldr.Append("bop")
+ arrs = append(arrs, bldr.NewArray())
+
+ defer func() {
+ for _, a := range arrs {
+ a.Release()
+ }
+ }()
+
+ arr, err := array.Concatenate(arrs, mem)
+ assert.NoError(t, err)
+ defer arr.Release()
+ })
+ }
+}
+
+func TestConcatOverflowRunEndEncoding(t *testing.T) {
+ tests := []struct {
+ offsetType arrow.DataType
+ max uint64
+ }{
+ {arrow.PrimitiveTypes.Int16, math.MaxInt16},
+ {arrow.PrimitiveTypes.Int32, math.MaxInt32},
+ {arrow.PrimitiveTypes.Int64, math.MaxInt64},
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.offsetType.String(), func(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ arrs := make([]arrow.Array, 0)
+ bldr := array.NewRunEndEncodedBuilder(mem, tt.offsetType, arrow.BinaryTypes.String)
+ defer bldr.Release()
+ valBldr := bldr.ValueBuilder().(*array.StringBuilder)
+
+ // max is not evently divisible by 4, so we add one to each
+ // to account for that so our final concatenate will overflow
+ bldr.Append((tt.max / 4) + 1)
+ valBldr.Append("foo")
+ bldr.Append((tt.max / 4) + 1)
+ valBldr.Append("bar")
+ arrs = append(arrs, bldr.NewArray())
+
+ bldr.Append((tt.max / 4) + 1)
+ valBldr.Append("baz")
+ bldr.Append((tt.max / 4) + 1)
+ valBldr.Append("bop")
+ arrs = append(arrs, bldr.NewArray())
+
+ defer func() {
+ for _, a := range arrs {
+ a.Release()
+ }
+ }()
+
+ arr, err := array.Concatenate(arrs, mem)
+ assert.Nil(t, arr)
+ assert.ErrorIs(t, err, arrow.ErrInvalid)
+ })
+ }
+}
diff --git a/go/arrow/type_string.go b/go/arrow/type_string.go
index 41a4073863..60c23514d1 100644
--- a/go/arrow/type_string.go
+++ b/go/arrow/type_string.go
@@ -1,4 +1,4 @@
-// Code generated by "stringer -type=Type"; DO NOT EDIT.
+// Code generated by "stringer -type Type"; DO NOT EDIT.
package arrow