diff --git a/arrow/array/concat.go b/arrow/array/concat.go index 3a9efe9e8..f0546b5cb 100644 --- a/arrow/array/concat.go +++ b/arrow/array/concat.go @@ -866,21 +866,32 @@ func updateRuns[T int16 | int32 | int64](inputData []arrow.ArrayData, inputBuffe continue } src := arrow.GetData[T](buf.Bytes()) + offset := inputData[i].Offset() + + // A slice can end in the middle of a run, leaving this input's final physical + // run end past its logical length. Clamp the normalized final run end to the + // input's logical length before both the overflow check and the output write: + // otherwise a valid near-limit slice trips a false overflow, and the written + // run end overshoots (shifting every following array's run ends). + finalEnd := src[len(src)-1] - T(offset) + if finalEnd > T(inputData[i].Len()) { + finalEnd = T(inputData[i].Len()) + } + if pos == 0 { pos += copy(output, src) // normalize the first run ends by subtracting the offset for j := 0; j < pos; j++ { - output[j] -= T(inputData[i].Offset()) + output[j] -= T(offset) } - + output[pos-1] = finalEnd 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()) > uint64(maxOf[T]()) { + // check whether adding this input's clamped final run end to the previous + // end will overflow the run-end type + if uint64(lastEnd)+uint64(finalEnd) > uint64(maxOf[T]()) { return fmt.Errorf("%w: overflow in run-length-encoded run ends concat", arrow.ErrInvalid) } @@ -889,9 +900,12 @@ func updateRuns[T int16 | int32 | int64](inputData []arrow.ArrayData, inputBuffe // 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 + T(int(e)-inputData[i].Offset()) + output[pos+j] = lastEnd + e - T(offset) } pos += len(src) + // the write above uses the unclamped physical end for the final run; set it + // to the clamped logical end. + output[pos-1] = lastEnd + finalEnd } return nil } diff --git a/arrow/array/concat_test.go b/arrow/array/concat_test.go index df906bd92..a70951344 100644 --- a/arrow/array/concat_test.go +++ b/arrow/array/concat_test.go @@ -819,6 +819,109 @@ func TestConcatRunEndEncoded(t *testing.T) { } } +func TestConcatRunEndEncodedMidRunSlice(t *testing.T) { + // A run-end encoded array sliced in the middle of a run keeps that run's physical end, which + // reaches past the slice's logical length. Concatenating it must clamp that final run end, + // otherwise the overshoot shifts every following array's run ends and silently corrupts values. + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + + bldr := array.NewRunEndEncodedBuilder(mem, arrow.PrimitiveTypes.Int32, arrow.PrimitiveTypes.Int64) + defer bldr.Release() + valBldr := bldr.ValueBuilder().(*array.Int64Builder) + + // runs: 100x3, 200x2, 300x3 -> run ends [3, 5, 8], logical length 8 + bldr.Append(3) + valBldr.Append(100) + bldr.Append(2) + valBldr.Append(200) + bldr.Append(3) + valBldr.Append(300) + full := bldr.NewArray() + defer full.Release() + + // slice [1, 4): logical [100, 100, 200], length 3, ending in the middle of the "200" run + sliced := array.NewSlice(full, 1, 4) + defer sliced.Release() + + bldr.Append(2) + valBldr.Append(700) + tail := bldr.NewArray() + defer tail.Release() + + result, err := array.Concatenate([]arrow.Array{sliced, tail}, mem) + require.NoError(t, err) + defer result.Release() + + rle := result.(*array.RunEndEncoded) + values := rle.Values().(*array.Int64) + got := make([]int64, rle.Len()) + for i := range got { + got[i] = values.Value(rle.GetPhysicalIndex(i)) + } + assert.Equal(t, []int64{100, 100, 200, 700, 700}, got) +} + +func TestConcatRunEndEncodedNearTypeLimitSlice(t *testing.T) { + // A sliced input whose physical final run end is near the run-end type limit must + // not trip a false overflow: the overflow check has to use the clamped logical end, + // not the physical one. int16 prefix of 32760 + a 1-element slice of a physical + // 32767-length run should yield 32761, not an overflow error. + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + + prefixBldr := array.NewRunEndEncodedBuilder(mem, arrow.PrimitiveTypes.Int16, arrow.PrimitiveTypes.Int64) + defer prefixBldr.Release() + prefixBldr.Append(32760) + prefixBldr.ValueBuilder().(*array.Int64Builder).Append(1) + prefix := prefixBldr.NewArray() + defer prefix.Release() + + bigBldr := array.NewRunEndEncodedBuilder(mem, arrow.PrimitiveTypes.Int16, arrow.PrimitiveTypes.Int64) + defer bigBldr.Release() + bigBldr.Append(32767) + bigBldr.ValueBuilder().(*array.Int64Builder).Append(2) + big := bigBldr.NewArray() + defer big.Release() + oneElem := array.NewSlice(big, 0, 1) // first element of the 32767-length run + defer oneElem.Release() + + result, err := array.Concatenate([]arrow.Array{prefix, oneElem}, mem) + require.NoError(t, err) + defer result.Release() + assert.EqualValues(t, 32761, result.Len()) +} + +func TestConcatRunEndEncodedInt64FinalRunEndClamp(t *testing.T) { + // A run-end-encoded array whose final physical run end is math.MaxInt64, + // sliced to a small logical length, must clamp the final run end with + // run-end-typed arithmetic. Converting the value through int overflows on + // 32-bit targets, turning a valid run end into a negative output value. + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + + prefixBldr := array.NewRunEndEncodedBuilder(mem, arrow.PrimitiveTypes.Int64, arrow.PrimitiveTypes.Int64) + defer prefixBldr.Release() + prefixBldr.Append(5) + prefixBldr.ValueBuilder().(*array.Int64Builder).Append(1) + prefix := prefixBldr.NewArray() + defer prefix.Release() + + bigBldr := array.NewRunEndEncodedBuilder(mem, arrow.PrimitiveTypes.Int64, arrow.PrimitiveTypes.Int64) + defer bigBldr.Release() + bigBldr.Append(math.MaxInt64) + bigBldr.ValueBuilder().(*array.Int64Builder).Append(2) + big := bigBldr.NewArray() + defer big.Release() + oneElem := array.NewSlice(big, 0, 1) // first element of the MaxInt64-length run + defer oneElem.Release() + + result, err := array.Concatenate([]arrow.Array{prefix, oneElem}, mem) + require.NoError(t, err) + defer result.Release() + assert.EqualValues(t, 6, result.Len()) +} + func TestConcatAlmostOverflowRunEndEncoding(t *testing.T) { tests := []struct { offsetType arrow.DataType