diff --git a/velox/experimental/cudf/exec/DecimalAggregationDevice.cu b/velox/experimental/cudf/exec/DecimalAggregationDevice.cu index 1c2a89d2ebf..87a31f9badb 100644 --- a/velox/experimental/cudf/exec/DecimalAggregationDevice.cu +++ b/velox/experimental/cudf/exec/DecimalAggregationDevice.cu @@ -98,16 +98,18 @@ struct UnpackStateFunctor { const uint8_t* chars; cuda::std::span<__int128_t> sums; cuda::std::span counts; + cudf::size_type rowOffset; cudf::bitmask_type const* nullMask; __device__ void operator()(cudf::size_type idx) const { - if (nullMask && !cudf::bit_is_set(nullMask, idx)) { + auto const inputIdx = idx + rowOffset; + if (nullMask && !cudf::bit_is_set(nullMask, inputIdx)) { return; } assert( - offsets[idx + 1] - offsets[idx] == + offsets[inputIdx + 1] - offsets[inputIdx] == static_cast(detail::kDecimalSumStateSize)); - int64_t offset = static_cast(offsets[idx]); + int64_t offset = static_cast(offsets[inputIdx]); auto* state = reinterpret_cast(chars + offset); counts[idx] = state->count; sums[idx] = (static_cast<__int128_t>(state->upper) << 64) | state->lower; @@ -239,6 +241,7 @@ struct unpackDecimalSumStateKernel { cudf::mutable_column_view sumView; cudf::mutable_column_view countView; cudf::size_type numRows; + cudf::size_type rowOffset; cudf::bitmask_type const* nullMask; rmm::cuda_stream_view stream; @@ -246,15 +249,17 @@ struct unpackDecimalSumStateKernel { requires OffsetStorageType void operator()() const { auto const n = static_cast(numRows); + auto const inputSize = static_cast(rowOffset) + n; launchDeviceFor( numRows, [&] { return UnpackStateFunctor{ cuda::std::span{ - offsetsView.data(), n + 1}, + offsetsView.data(), inputSize + 1}, chars, cuda::std::span<__int128_t>{sumView.data<__int128_t>(), n}, cuda::std::span{countView.data(), n}, + rowOffset, nullMask}; }, stream); @@ -345,12 +350,20 @@ void unpackDecimalSumState( cudf::mutable_column_view sumView, cudf::mutable_column_view countView, cudf::size_type numRows, + cudf::size_type rowOffset, cudf::bitmask_type const* nullMask, rmm::cuda_stream_view stream) { cudf::type_dispatcher( cudf::data_type{offsetType}, unpackDecimalSumStateKernel{ - offsetsView, chars, sumView, countView, numRows, nullMask, stream}); + offsetsView, + chars, + sumView, + countView, + numRows, + rowOffset, + nullMask, + stream}); } void averageRoundDecimalSum( diff --git a/velox/experimental/cudf/exec/DecimalAggregationDevice.h b/velox/experimental/cudf/exec/DecimalAggregationDevice.h index 63bc3a0fbad..afd96d07b27 100644 --- a/velox/experimental/cudf/exec/DecimalAggregationDevice.h +++ b/velox/experimental/cudf/exec/DecimalAggregationDevice.h @@ -82,6 +82,8 @@ void packDecimalSumState( * @param sumView output per-row DECIMAL128 sums. * @param countView output per-row counts. * @param numRows number of rows. + * @param rowOffset row offset of a sliced parent STRING column. Applied to + * both offsets and null-mask indexing. * @param nullMask device null-mask bitmap; null rows are skipped to avoid * out-of-bounds reads when Arrow compacts null payloads. Pass nullptr * when no mask is present. @@ -94,6 +96,7 @@ void unpackDecimalSumState( cudf::mutable_column_view sumView, cudf::mutable_column_view countView, cudf::size_type numRows, + cudf::size_type rowOffset, cudf::bitmask_type const* nullMask, rmm::cuda_stream_view stream); diff --git a/velox/experimental/cudf/exec/DecimalAggregationState.cpp b/velox/experimental/cudf/exec/DecimalAggregationState.cpp index b3d15d24df7..3e0ca016496 100644 --- a/velox/experimental/cudf/exec/DecimalAggregationState.cpp +++ b/velox/experimental/cudf/exec/DecimalAggregationState.cpp @@ -23,6 +23,7 @@ #include #include +#include #include #include @@ -81,7 +82,9 @@ DecimalSumStateColumns deserializeDecimalSumState( cudf::strings_column_view strings(stateCol); auto const nullCount = stateCol.nullable() ? stateCol.null_count() : 0; - auto const payloadSize = strings.chars_size(stream); + auto const [payloadBegin, payloadEnd] = + cudf::strings::detail::get_first_and_last_offset(strings, stream); + auto const payloadSize = payloadEnd - payloadBegin; // serializeDecimalSumState writes 32 bytes for every row (including nulls), // but an Arrow round-trip compacts null rows to 0 bytes. Accept both. auto const fullPayloadSize = @@ -128,6 +131,7 @@ DecimalSumStateColumns deserializeDecimalSumState( sumView, countView, numRows, + stateCol.offset(), stateCol.null_mask(), stream); diff --git a/velox/experimental/cudf/tests/DecimalAggregationTest.cpp b/velox/experimental/cudf/tests/DecimalAggregationTest.cpp index 6c9f2ea7786..a23ee9bdad8 100644 --- a/velox/experimental/cudf/tests/DecimalAggregationTest.cpp +++ b/velox/experimental/cudf/tests/DecimalAggregationTest.cpp @@ -30,6 +30,7 @@ #include "velox/type/DecimalUtil.h" #include +#include #include #include #include @@ -1281,6 +1282,39 @@ TEST_F(CudfDecimalTest, decimalDeserializeSumStatePartialNullCompact) { EXPECT_EQ(outCount[2], 2); } +TEST_F(CudfDecimalTest, decimalDeserializeSumStateSlice) { + auto stream = cudf::get_default_stream(); + auto mr = cudf::get_current_device_resource_ref(); + + // Slice [1, 4) so the deserializer must apply a non-zero parent offset to + // both the strings offsets child and the parent validity mask. + std::vector sums = {10, 20, 0, 40, 50}; + std::vector counts = {1, 2, 0, 4, 5}; + std::vector sumValid = {true, true, false, true, true}; + auto sumCol = makeDecimalColumn(sums, 2, &sumValid, stream); + auto countCol = makeInt64Column(counts, nullptr, stream); + auto stateCol = + serializeDecimalSumState(sumCol->view(), countCol->view(), stream, mr); + + auto slices = cudf::slice(stateCol->view(), {1, 4}); + ASSERT_EQ(slices.size(), 1); + ASSERT_EQ(slices.front().offset(), 1); + + auto result = deserializeDecimalSumState(slices.front(), 2, stream); + auto outSum = copyColumnData<__int128_t>(result.sum->view(), stream); + auto outCount = copyColumnData(result.count->view(), stream); + auto outMask = copyNullMask(result.sum->view(), stream); + + ASSERT_EQ(outSum.size(), 3); + EXPECT_TRUE(isValidAt(outMask, 0)); + EXPECT_FALSE(isValidAt(outMask, 1)); + EXPECT_TRUE(isValidAt(outMask, 2)); + EXPECT_EQ(outSum[0], static_cast<__int128_t>(20)); + EXPECT_EQ(outCount[0], 2); + EXPECT_EQ(outSum[2], static_cast<__int128_t>(40)); + EXPECT_EQ(outCount[2], 4); +} + // Trailing null: the offset for the last row equals chars_size, so the kernel // would read 32 bytes past the buffer end without the null-mask guard. TEST_F(CudfDecimalTest, decimalDeserializeSumStateTrailingNullCompact) {