fix(arrow/array): clamp run ends when concatenating a sliced RunEndEncoded array

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
 }

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