diff options
author | atarasov5 <[email protected]> | 2025-06-18 12:30:17 +0300 |
---|---|---|
committer | atarasov5 <[email protected]> | 2025-06-18 14:11:07 +0300 |
commit | 1993defec2b80bfe5c878d6cbda5b0f9ca89418a (patch) | |
tree | cb87bdbbb4f53d0febf9f0adfa9049d49c9bc5a0 /yql/essentials | |
parent | 38cf8b026f3b0377d34b7820d8a51f53a4337403 (diff) |
Fix coalesce when bitmap is empty
Раньше наш код считал, что у опциональных MKQL-типов битовая маска в ArrayData — это всегда не nullptr. Однако такой гарантии нет. Теперь, чтобы определить, есть ли в ArrayData опциональные значения, мы больше не смотрим на тип minikql. Вместо этого ориентируемся только на значение `null_count`.
Более того, я немного подрезал длину массивов в тестах, так как тесты выполнялись слишком долго.
commit_hash:e538ec4f2251d9411f67c195d666401c2d950527
Diffstat (limited to 'yql/essentials')
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) { |