Skip to content

Commit fe3b0ab

Browse files
winklemadwinklemad
committed
fix(arrow/array): keep REE concat final run-end arithmetic in the run-end type
Converting the final run end through int overflowed on 32-bit targets: a valid run-end-encoded array with final run end math.MaxInt64 and logical length 1 became -1 during concatenation. Compute and clamp the final run end (and the per-run normalization) in the run-end type T, clamping against T(inputData[i].Len()) rather than going through int. Adds a MaxInt64 regression test. Signed-off-by: Madan Kumar <[email protected]>
1 parent 300dbe3 commit fe3b0ab

2 files changed

Lines changed: 36 additions & 6 deletions

File tree

‎arrow/array/concat.go‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -873,9 +873,9 @@ func updateRuns[T int16 | int32 | int64](inputData []arrow.ArrayData, inputBuffe
873873
// input's logical length before both the overflow check and the output write:
874874
// otherwise a valid near-limit slice trips a false overflow, and the written
875875
// run end overshoots (shifting every following array's run ends).
876-
finalEnd := int(src[len(src)-1]) - offset
877-
if finalEnd > inputData[i].Len() {
878-
finalEnd = inputData[i].Len()
876+
finalEnd := src[len(src)-1] - T(offset)
877+
if finalEnd > T(inputData[i].Len()) {
878+
finalEnd = T(inputData[i].Len())
879879
}
880880

881881
if pos == 0 {
@@ -884,7 +884,7 @@ func updateRuns[T int16 | int32 | int64](inputData []arrow.ArrayData, inputBuffe
884884
for j := 0; j < pos; j++ {
885885
output[j] -= T(offset)
886886
}
887-
output[pos-1] = T(finalEnd)
887+
output[pos-1] = finalEnd
888888
continue
889889
}
890890

@@ -900,12 +900,12 @@ func updateRuns[T int16 | int32 | int64](inputData []arrow.ArrayData, inputBuffe
900900
// is a logical length offset it should be accurate to just subtract
901901
// it from each value.
902902
for j, e := range src {
903-
output[pos+j] = lastEnd + T(int(e)-offset)
903+
output[pos+j] = lastEnd + e - T(offset)
904904
}
905905
pos += len(src)
906906
// the write above uses the unclamped physical end for the final run; set it
907907
// to the clamped logical end.
908-
output[pos-1] = lastEnd + T(finalEnd)
908+
output[pos-1] = lastEnd + finalEnd
909909
}
910910
return nil
911911
}

‎arrow/array/concat_test.go‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -892,6 +892,36 @@ func TestConcatRunEndEncodedNearTypeLimitSlice(t *testing.T) {
892892
assert.EqualValues(t, 32761, result.Len())
893893
}
894894

895+
func TestConcatRunEndEncodedInt64FinalRunEndClamp(t *testing.T) {
896+
// A run-end-encoded array whose final physical run end is math.MaxInt64,
897+
// sliced to a small logical length, must clamp the final run end with
898+
// run-end-typed arithmetic. Converting the value through int overflows on
899+
// 32-bit targets, turning a valid run end into a negative output value.
900+
mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
901+
defer mem.AssertSize(t, 0)
902+
903+
prefixBldr := array.NewRunEndEncodedBuilder(mem, arrow.PrimitiveTypes.Int64, arrow.PrimitiveTypes.Int64)
904+
defer prefixBldr.Release()
905+
prefixBldr.Append(5)
906+
prefixBldr.ValueBuilder().(*array.Int64Builder).Append(1)
907+
prefix := prefixBldr.NewArray()
908+
defer prefix.Release()
909+
910+
bigBldr := array.NewRunEndEncodedBuilder(mem, arrow.PrimitiveTypes.Int64, arrow.PrimitiveTypes.Int64)
911+
defer bigBldr.Release()
912+
bigBldr.Append(math.MaxInt64)
913+
bigBldr.ValueBuilder().(*array.Int64Builder).Append(2)
914+
big := bigBldr.NewArray()
915+
defer big.Release()
916+
oneElem := array.NewSlice(big, 0, 1) // first element of the MaxInt64-length run
917+
defer oneElem.Release()
918+
919+
result, err := array.Concatenate([]arrow.Array{prefix, oneElem}, mem)
920+
require.NoError(t, err)
921+
defer result.Release()
922+
assert.EqualValues(t, 6, result.Len())
923+
}
924+
895925
func TestConcatAlmostOverflowRunEndEncoding(t *testing.T) {
896926
tests := []struct {
897927
offsetType arrow.DataType

0 commit comments

Comments
 (0)