summaryrefslogtreecommitdiffstats
path: root/yql/essentials
diff options
context:
space:
mode:
authoratarasov5 <[email protected]>2025-06-18 12:30:17 +0300
committeratarasov5 <[email protected]>2025-06-18 14:11:07 +0300
commit1993defec2b80bfe5c878d6cbda5b0f9ca89418a (patch)
treecb87bdbbb4f53d0febf9f0adfa9049d49c9bc5a0 /yql/essentials
parent38cf8b026f3b0377d34b7820d8a51f53a4337403 (diff)
Fix coalesce when bitmap is empty
Раньше наш код считал, что у опциональных MKQL-типов битовая маска в ArrayData — это всегда не nullptr. Однако такой гарантии нет. Теперь, чтобы определить, есть ли в ArrayData опциональные значения, мы больше не смотрим на тип minikql. Вместо этого ориентируемся только на значение `null_count`. Более того, я немного подрезал длину массивов в тестах, так как тесты выполнялись слишком долго. commit_hash:e538ec4f2251d9411f67c195d666401c2d950527
Diffstat (limited to 'yql/essentials')
-rw-r--r--yql/essentials/minikql/comp_nodes/mkql_block_coalesce.cpp35
-rw-r--r--yql/essentials/minikql/comp_nodes/mkql_block_coalesce_blending_helper.h22
-rw-r--r--yql/essentials/minikql/comp_nodes/ut/mkql_block_coalesce_ut.cpp50
3 files changed, 55 insertions, 52 deletions
diff --git a/yql/essentials/minikql/comp_nodes/mkql_block_coalesce.cpp b/yql/essentials/minikql/comp_nodes/mkql_block_coalesce.cpp
index 6185b7688b1..093fd1e5337 100644
--- a/yql/essentials/minikql/comp_nodes/mkql_block_coalesce.cpp
+++ b/yql/essentials/minikql/comp_nodes/mkql_block_coalesce.cpp
@@ -20,8 +20,9 @@ namespace NKikimr::NMiniKQL {
namespace {
template <typename TType>
-void DispatchCoalesceImpl(const arrow::Datum& left, const arrow::Datum& right, arrow::Datum& out, bool outIsOptional, arrow::MemoryPool& pool) {
- auto bitmap = outIsOptional ? ARROW_RESULT(arrow::AllocateBitmap((left.array()->length + left.array()->offset % 8) * sizeof(ui8), &pool)) : nullptr;
+void DispatchCoalesceImpl(const arrow::Datum& left, const arrow::Datum& right, arrow::Datum& out, arrow::MemoryPool& pool) {
+ bool outHasBitmask = (right.is_array() && right.null_count() > 0) || (right.is_scalar() && !right.scalar()->is_valid);
+ auto bitmap = outHasBitmask ? ARROW_RESULT(arrow::AllocateBitmap((left.array()->length + left.array()->offset % 8) * sizeof(ui8), &pool)) : nullptr;
if (bitmap && bitmap->size() > 0) {
// Fill first byte with zero to prevent further uninitialized memory access.
bitmap->mutable_data()[0] = 0;
@@ -30,20 +31,12 @@ void DispatchCoalesceImpl(const arrow::Datum& left, const arrow::Datum& right, a
{std::move(bitmap),
ARROW_RESULT(arrow::AllocateBuffer((left.array()->length + left.array()->offset % 8) * sizeof(TType), &pool))},
arrow::kUnknownNullCount, left.array()->offset % 8);
- if (outIsOptional) {
+ if (outHasBitmask) {
if (right.is_scalar()) {
- if (right.scalar()->is_valid) {
- BlendCoalesce<TType, /*isScalar=*/true, /*rightIsOptional=*/true>(
- TDatumStorageView<TType>(left),
- TDatumStorageView<TType>(right),
- TDatumStorageView<TType>(out),
- left.array()->length);
- } else {
- out = left;
- }
+ out = left;
} else {
MKQL_ENSURE(TDatumStorageView<TType>(right).bitMask(), "Right array must have a null mask");
- BlendCoalesce<TType, /*isScalar=*/false, /*rightIsOptional=*/true>(
+ BlendCoalesce<TType, /*isScalar=*/false, /*rightHasBitmask=*/true>(
TDatumStorageView<TType>(left),
TDatumStorageView<TType>(right),
TDatumStorageView<TType>(out),
@@ -51,13 +44,13 @@ void DispatchCoalesceImpl(const arrow::Datum& left, const arrow::Datum& right, a
}
} else {
if (right.is_scalar()) {
- BlendCoalesce<TType, /*isScalar=*/true, /*rightIsOptional=*/false>(
+ BlendCoalesce<TType, /*isScalar=*/true, /*rightHasBitmask=*/false>(
TDatumStorageView<TType>(left),
TDatumStorageView<TType>(right),
TDatumStorageView<TType>(out),
left.array()->length);
} else {
- BlendCoalesce<TType, /*isScalar=*/false, /*rightIsOptional=*/false>(
+ BlendCoalesce<TType, /*isScalar=*/false, /*rightHasBitmask=*/false>(
TDatumStorageView<TType>(left),
TDatumStorageView<TType>(right),
TDatumStorageView<TType>(out),
@@ -83,18 +76,18 @@ bool DispatchBlendingCoalesce(const arrow::Datum& left, const arrow::Datum& righ
case NYql::NUdf::EDataSlot::Bool:
case NYql::NUdf::EDataSlot::Int8:
case NYql::NUdf::EDataSlot::Uint8:
- DispatchCoalesceImpl<ui8>(left, right, out, /*outIsOptional=*/!needUnwrapFirst, pool);
+ DispatchCoalesceImpl<ui8>(left, right, out, pool);
return true;
case NYql::NUdf::EDataSlot::Int16:
case NYql::NUdf::EDataSlot::Uint16:
case NYql::NUdf::EDataSlot::Date:
- DispatchCoalesceImpl<ui16>(left, right, out, /*outIsOptional=*/!needUnwrapFirst, pool);
+ DispatchCoalesceImpl<ui16>(left, right, out, pool);
return true;
case NYql::NUdf::EDataSlot::Int32:
case NYql::NUdf::EDataSlot::Uint32:
case NYql::NUdf::EDataSlot::Date32:
case NYql::NUdf::EDataSlot::Datetime:
- DispatchCoalesceImpl<ui32>(left, right, out, /*outIsOptional=*/!needUnwrapFirst, pool);
+ DispatchCoalesceImpl<ui32>(left, right, out, pool);
return true;
case NYql::NUdf::EDataSlot::Int64:
case NYql::NUdf::EDataSlot::Uint64:
@@ -103,15 +96,15 @@ bool DispatchBlendingCoalesce(const arrow::Datum& left, const arrow::Datum& righ
case NYql::NUdf::EDataSlot::Interval64:
case NYql::NUdf::EDataSlot::Interval:
case NYql::NUdf::EDataSlot::Timestamp:
- DispatchCoalesceImpl<ui64>(left, right, out, /*outIsOptional=*/!needUnwrapFirst, pool);
+ DispatchCoalesceImpl<ui64>(left, right, out, pool);
return true;
case NYql::NUdf::EDataSlot::Double:
static_assert(sizeof(NUdf::TDataType<double>::TLayout) == sizeof(NUdf::TDataType<ui64>::TLayout));
- DispatchCoalesceImpl<ui64>(left, right, out, /*outIsOptional=*/!needUnwrapFirst, pool);
+ DispatchCoalesceImpl<ui64>(left, right, out, pool);
return true;
case NYql::NUdf::EDataSlot::Float:
static_assert(sizeof(NUdf::TDataType<float>::TLayout) == sizeof(NUdf::TDataType<ui32>::TLayout));
- DispatchCoalesceImpl<ui32>(left, right, out, /*outIsOptional=*/!needUnwrapFirst, pool);
+ DispatchCoalesceImpl<ui32>(left, right, out, pool);
return true;
default:
// Fallback to general builder/reader pipeline.
diff --git a/yql/essentials/minikql/comp_nodes/mkql_block_coalesce_blending_helper.h b/yql/essentials/minikql/comp_nodes/mkql_block_coalesce_blending_helper.h
index 148c2f48472..3af98b7ff04 100644
--- a/yql/essentials/minikql/comp_nodes/mkql_block_coalesce_blending_helper.h
+++ b/yql/essentials/minikql/comp_nodes/mkql_block_coalesce_blending_helper.h
@@ -84,7 +84,7 @@ private:
const arrow::Datum& Datum_;
};
-template <typename TType, bool rightIsScalar, bool rightIsOptional>
+template <typename TType, bool rightIsScalar, bool rightHasBitmask>
void CoalesceByOneElement(size_t elements,
const ui8* leftBitMask,
const ui8* rightBitMask,
@@ -98,7 +98,7 @@ void CoalesceByOneElement(size_t elements,
for (size_t i = 0; i < elements; i++) {
if (arrow::BitUtil::GetBit(leftBitMask, i + leftOffset)) {
out[i + outOffset] = left[i + leftOffset];
- SetBitTo<rightIsOptional>(outBitMask, i + outOffset, true);
+ SetBitTo<rightHasBitmask>(outBitMask, i + outOffset, true);
} else {
if constexpr (rightIsScalar) {
out[i + outOffset] = right[0];
@@ -106,9 +106,9 @@ void CoalesceByOneElement(size_t elements,
out[i + outOffset] = right[i + rightOffset];
}
- SetBitTo<rightIsOptional>(outBitMask,
+ SetBitTo<rightHasBitmask>(outBitMask,
i + outOffset,
- GetBit<rightIsScalar, rightIsOptional>(rightBitMask, i + rightOffset));
+ GetBit<rightIsScalar, rightHasBitmask>(rightBitMask, i + rightOffset));
}
}
}
@@ -126,7 +126,7 @@ Y_FORCE_INLINE ui8 GetMaskValue(const ui8* mask, size_t offset, size_t bitMaskPo
}
}
-template <typename TType, bool rightIsScalar, bool rightIsOptional, bool rightIsAlignedAsLeft>
+template <typename TType, bool rightIsScalar, bool rightHasBitmask, bool rightIsAlignedAsLeft>
void VectorizedCoalesce(const ui8* __restrict leftBitMask,
const ui8* __restrict rightBitMask,
size_t leftOffset,
@@ -146,7 +146,7 @@ void VectorizedCoalesce(const ui8* __restrict leftBitMask,
bytesProcessed < truncatedLengthInBytes;
bytesProcessed += sizeInBytes, elemShift += 8, step += 1) {
auto currentLeftMask = leftBitMask[leftOffset / 8 + step];
- if constexpr (rightIsOptional) {
+ if constexpr (rightHasBitmask) {
// If right is optional, updates the output bit mask.
// Otherwise, the output bit mask doesn't exist.
if constexpr (rightIsScalar) {
@@ -172,7 +172,7 @@ void VectorizedCoalesce(const ui8* __restrict leftBitMask,
// Coalesces data from two inputs based on bit masks, handling either (array, array) or (array, scalar).
// This function efficiently merges data from 'left' and 'right' into 'out' array using 'left.bitMask()' and 'right.bitMask()'.
// The function is vectorization friendly, which can significantly improve performance.
-template <typename TType, bool rightIsScalar, bool rightIsOptional>
+template <typename TType, bool rightIsScalar, bool rightHasBitmask>
void BlendCoalesce(TDatumStorageView<TType> left,
TDatumStorageView<TType> right,
TDatumStorageView<TType> out,
@@ -180,7 +180,7 @@ void BlendCoalesce(TDatumStorageView<TType> left,
Y_ENSURE(left.offset() % 8 == out.offset() % 8);
auto firstElementsToProcess = std::min((8 - left.offset() % 8) % 8, lengthInElements);
// Process one by one until left mask is aligned by byte.
- CoalesceByOneElement<TType, rightIsScalar, rightIsOptional>(firstElementsToProcess,
+ CoalesceByOneElement<TType, rightIsScalar, rightHasBitmask>(firstElementsToProcess,
left.bitMask(),
right.bitMask(),
left.offset(),
@@ -194,7 +194,7 @@ void BlendCoalesce(TDatumStorageView<TType> left,
// Process vectorized.
if (left.offset() % 8 != right.offset() % 8) {
- VectorizedCoalesce<TType, rightIsScalar, rightIsOptional, false>(left.bitMask(),
+ VectorizedCoalesce<TType, rightIsScalar, rightHasBitmask, false>(left.bitMask(),
right.bitMask(),
left.offset() + firstElementsToProcess,
right.offset() + firstElementsToProcess,
@@ -205,7 +205,7 @@ void BlendCoalesce(TDatumStorageView<TType> left,
out.bitMask(),
lengthInElements);
} else {
- VectorizedCoalesce<TType, rightIsScalar, rightIsOptional, true>(left.bitMask(),
+ VectorizedCoalesce<TType, rightIsScalar, rightHasBitmask, true>(left.bitMask(),
right.bitMask(),
left.offset() + firstElementsToProcess,
right.offset() + firstElementsToProcess,
@@ -218,7 +218,7 @@ void BlendCoalesce(TDatumStorageView<TType> left,
}
// Process remaining bits that take less memory than one byte.
size_t remainingBits = (lengthInElements) % 8;
- CoalesceByOneElement<TType, rightIsScalar, rightIsOptional>(remainingBits,
+ CoalesceByOneElement<TType, rightIsScalar, rightHasBitmask>(remainingBits,
left.bitMask(),
right.bitMask(),
left.offset() + firstElementsToProcess + lengthInElements - remainingBits,
diff --git a/yql/essentials/minikql/comp_nodes/ut/mkql_block_coalesce_ut.cpp b/yql/essentials/minikql/comp_nodes/ut/mkql_block_coalesce_ut.cpp
index 57ad9137cdf..d96346fa061 100644
--- a/yql/essentials/minikql/comp_nodes/ut/mkql_block_coalesce_ut.cpp
+++ b/yql/essentials/minikql/comp_nodes/ut/mkql_block_coalesce_ut.cpp
@@ -144,7 +144,8 @@ void TestBlockCoalesceForVector(InputOptionalVector<T> left,
InputOptionalVector<T> right,
InputOptionalVector<T> expected,
size_t leftOffset,
- size_t rightOffset) {
+ size_t rightOffset,
+ bool resetNullBitmapWhenAllNotNull = false) {
using TLayout = typename NUdf::TDataType<T>::TLayout;
TSetup<false> setup;
NYql::TExprContext exprCtx;
@@ -173,6 +174,14 @@ void TestBlockCoalesceForVector(InputOptionalVector<T> left,
} else {
rightOperand = GenerateArray(typeInfoHelper, rightType, right, rightOffset);
}
+ // Reset bitmap that responses for nullability of arrow::ArrayData.
+ // If all elements are not null then we have two options:
+ // 1. All bitmask elements are set to 1.
+ // 2. There is no bitmask at all.
+ // So we want to test both variants via |resetNullBitmapWhenAllNotNull| flag.
+ if (rightOperand.is_array() && resetNullBitmapWhenAllNotNull && rightOperand.array()->GetNullCount() == 0) {
+ rightOperand.array()->buffers[0] = nullptr;
+ }
auto bi = arrow::compute::detail::ExecBatchIterator::Make({leftOperand, rightOperand}, 1000).ValueOrDie();
arrow::compute::ExecBatch batch;
UNIT_ASSERT(bi->Next(&batch));
@@ -203,11 +212,6 @@ void TestBlockCoalesce(InputOptionalVector<T> left,
// Second test different sizes.
// Also test only small subset of offsets to prevent a combinatorial explosion.
while (left.size() > 1 || right.size() > 1 || expected.size() > 1) {
- for (size_t leftOffset = 0; leftOffset < 2; leftOffset++) {
- for (size_t rightOffset = 0; rightOffset < 2; rightOffset++) {
- TestBlockCoalesceForVector<T, rightType>(left, right, expected, leftOffset, rightOffset);
- }
- }
if (left.size() > 1) {
left.pop_back();
}
@@ -217,6 +221,12 @@ void TestBlockCoalesce(InputOptionalVector<T> left,
if (expected.size() > 1) {
expected.pop_back();
}
+ for (size_t leftOffset = 0; leftOffset < 2; leftOffset++) {
+ for (size_t rightOffset = 0; rightOffset < 2; rightOffset++) {
+ TestBlockCoalesceForVector<T, rightType>(left, right, expected, leftOffset, rightOffset);
+ TestBlockCoalesceForVector<T, rightType>(left, right, expected, leftOffset, rightOffset, /*resetNullBitmapWhenAllNotNull=*/true);
+ }
+ }
}
}
@@ -300,53 +310,53 @@ void BlockCoalesceGraphTest(size_t length, size_t offset) {
Y_UNIT_TEST_SUITE(TMiniKQLBlockCoalesceTest) {
Y_UNIT_TEST(CoalesceGraphTest) {
- for (auto offset : {0, 1, 2, 3, 5, 7, 8, 11, 14, 16}) {
- BlockCoalesceGraphTest(1000, offset);
+ for (auto offset : {0, 7, 8, 11,6}) {
+ BlockCoalesceGraphTest(32, offset);
}
}
UNIT_TEST_WITH_INTEGER(KernelRightIsNotNullArray) {
auto max = std::numeric_limits<typename NUdf::TDataType<TTestType>::TLayout>::max();
auto min = std::numeric_limits<typename NUdf::TDataType<TTestType>::TLayout>::min();
- TestBlockCoalesce<TTestType, ERightOperandType::ARRAY>({Nothing(), 2, 3, Nothing(), 5, 6, 7, max, 9, Nothing(), 11, 12, 13, Nothing(), Nothing(), Nothing(), min, Nothing(), 19, 20},
- {101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120},
- {101, 2, 3, 104, 5, 6, 7, max, 9, 110, 11, 12, 13, 114, 115, 116, min, 118, 19, 20});
+ TestBlockCoalesce<TTestType, ERightOperandType::ARRAY>({Nothing(), 2, 3, Nothing(), 5, 6, 7, max, 9, Nothing(), 11, 12, 13, Nothing(), Nothing(), Nothing(), min, Nothing()},
+ {101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118},
+ {101, 2, 3, 104, 5, 6, 7, max, 9, 110, 11, 12, 13, 114, 115, 116, min, 118});
}
UNIT_TEST_WITH_INTEGER(KernelRightIsScalar) {
auto max = std::numeric_limits<typename NUdf::TDataType<TTestType>::TLayout>::max();
auto min = std::numeric_limits<typename NUdf::TDataType<TTestType>::TLayout>::min();
- TestBlockCoalesce<TTestType, ERightOperandType::SCALAR>({Nothing(), 2, 3, Nothing(), 5, 6, 7, max, 9, Nothing(), 11, 12, 13, Nothing(), Nothing(), Nothing(), min, Nothing(), 19, 20},
+ TestBlockCoalesce<TTestType, ERightOperandType::SCALAR>({Nothing(), 2, 3, Nothing(), 5, 6, 7, max, 9, Nothing(), 11, 12, 13, Nothing(), Nothing(), Nothing(), min, Nothing()},
{77},
- {77, 2, 3, 77, 5, 6, 7, max, 9, 77, 11, 12, 13, 77, 77, 77, min, 77, 19, 20});
+ {77, 2, 3, 77, 5, 6, 7, max, 9, 77, 11, 12, 13, 77, 77, 77, min, 77});
}
UNIT_TEST_WITH_INTEGER(KernelRightIsOptionalArray) {
auto max = std::numeric_limits<typename NUdf::TDataType<TTestType>::TLayout>::max();
auto min = std::numeric_limits<typename NUdf::TDataType<TTestType>::TLayout>::min();
- TestBlockCoalesce<TTestType, ERightOperandType::OPTIONAL_ARRAY>({Nothing(), 2, 3, Nothing(), 5, 6, 7, max, 9, Nothing(), 11, 12, 13, Nothing(), Nothing(), Nothing(), min, Nothing(), 19, 20},
- {Nothing(), 102, Nothing(), 104, Nothing(), 106, 107, 108, 109, 110, 111, 112, 113, 114, Nothing(), 116, 117, 118, Nothing(), 120},
- {Nothing(), 2, 3, 104, 5, 6, 7, max, 9, 110, 11, 12, 13, 114, Nothing(), 116, min, 118, 19, 20});
+ TestBlockCoalesce<TTestType, ERightOperandType::OPTIONAL_ARRAY>({Nothing(), 2, 3, Nothing(), 5, 6, 7, max, 9, Nothing(), 11, 12, 13, Nothing(), Nothing(), Nothing(), min, Nothing()},
+ {101, 102, Nothing(), 104, Nothing(), 106, 107, 108, 109, 110, 111, 112, 113, 114, Nothing(), 116, 117, 118},
+ {101, 2, 3, 104, 5, 6, 7, max, 9, 110, 11, 12, 13, 114, Nothing(), 116, min, 118});
}
UNIT_TEST_WITH_INTEGER(KernelRightIsOptionalInvalidScalar) {
auto max = std::numeric_limits<typename NUdf::TDataType<TTestType>::TLayout>::max();
auto min = std::numeric_limits<typename NUdf::TDataType<TTestType>::TLayout>::min();
- TestBlockCoalesce<TTestType, ERightOperandType::OPTIONAL_SCALAR>({Nothing(), 2, 3, Nothing(), 5, 6, 7, max, 9, Nothing(), 11, 12, 13, Nothing(), Nothing(), Nothing(), min, Nothing(), 19, 20},
+ TestBlockCoalesce<TTestType, ERightOperandType::OPTIONAL_SCALAR>({Nothing(), 2, 3, Nothing(), 5, 6, 7, max, 9, Nothing(), 11, 12, 13, Nothing(), Nothing(), Nothing(), min, Nothing()},
{Nothing()},
- {Nothing(), 2, 3, Nothing(), 5, 6, 7, max, 9, Nothing(), 11, 12, 13, Nothing(), Nothing(), Nothing(), min, Nothing(), 19, 20});
+ {Nothing(), 2, 3, Nothing(), 5, 6, 7, max, 9, Nothing(), 11, 12, 13, Nothing(), Nothing(), Nothing(), min, Nothing()});
}
UNIT_TEST_WITH_INTEGER(KernelRightIsOptionalValidScalar) {
auto max = std::numeric_limits<typename NUdf::TDataType<TTestType>::TLayout>::max();
auto min = std::numeric_limits<typename NUdf::TDataType<TTestType>::TLayout>::min();
- TestBlockCoalesce<TTestType, ERightOperandType::OPTIONAL_SCALAR>({Nothing(), 2, 3, Nothing(), 5, 6, 7, max, 9, Nothing(), 11, 12, 13, Nothing(), Nothing(), Nothing(), min, Nothing(), 19, 20},
+ TestBlockCoalesce<TTestType, ERightOperandType::OPTIONAL_SCALAR>({Nothing(), 2, 3, Nothing(), 5, 6, 7, max, 9, Nothing(), 11, 12, 13, Nothing(), Nothing(), Nothing(), min, Nothing()},
{77},
- {77, 2, 3, 77, 5, 6, 7, max, 9, 77, 11, 12, 13, 77, 77, 77, min, 77, 19, 20});
+ {77, 2, 3, 77, 5, 6, 7, max, 9, 77, 11, 12, 13, 77, 77, 77, min, 77});
}
Y_UNIT_TEST(OptionalScalar) {