diff options
| author | xyligansereja <[email protected]> | 2026-07-10 20:02:07 +0300 |
|---|---|---|
| committer | xyligansereja <[email protected]> | 2026-07-10 20:40:36 +0300 |
| commit | fb8afa3946bca4bc2e84bc59cba90c58774d84d4 (patch) | |
| tree | 19a891eecb477bcd0ba03b7148235841f8c40d82 | |
| parent | c8b89209e7704a4fab0af5e9060c573bd061283e (diff) | |
Support Uuid and DyNumber in minikql blocks and Arrow
Изменения в рамках этого тикета <https://github.com/ydb-platform/ydb/issues/44190>
#### Support Uuid and DyNumber in minikql blocks and Arrow ✎
- Added support for Uuid and DyNumber data types within MiniKQL blocks and Arrow computations, enabling their use in block-based operations and serialization
- Implemented proper conversion, serialization, and deserialization logic for both data types in the Arrow format
- Extended block item handling to support Uuid and DyNumber, including comparator and hasher implementations
- Added comprehensive unit tests covering roundtrip operations, serialization, comparison, and conversion for both data types
- Integrated Uuid and DyNumber support into built-in functions like Min/Max and their aggregate variants
- Updated type system and block builder components to recognize and handle the new data types correctly
- Added necessary includes and utility functions to work with Uuid and DyNumber in MiniKQL contexts
<a href="https://nda.ya.ru/t/qa0kX64r7DqvtN"><font size="2">Autodescription by Yandex Code Assistant</font></a>
commit_hash:d1ab82c060cfa90b7cd13a1cb3efbcb9615e6291
23 files changed, 878 insertions, 59 deletions
diff --git a/yql/essentials/minikql/arrow/arrow_util.cpp b/yql/essentials/minikql/arrow/arrow_util.cpp index b59c58465f0..649c2c5eeaa 100644 --- a/yql/essentials/minikql/arrow/arrow_util.cpp +++ b/yql/essentials/minikql/arrow/arrow_util.cpp @@ -70,4 +70,9 @@ void UntrackDatum(const arrow::Datum& datum) { } } +std::shared_ptr<arrow::DataType> GetUuidArrowType() { + static const std::shared_ptr<arrow::DataType> Type = arrow::fixed_size_binary(UuidBinarySize); + return Type; +} + } // namespace NKikimr::NMiniKQL diff --git a/yql/essentials/minikql/arrow/arrow_util.h b/yql/essentials/minikql/arrow/arrow_util.h index 8a2e80009c7..5cc9aebd434 100644 --- a/yql/essentials/minikql/arrow/arrow_util.h +++ b/yql/essentials/minikql/arrow/arrow_util.h @@ -11,6 +11,7 @@ #include <yql/essentials/minikql/mkql_node.h> #include <yql/essentials/minikql/arrow/mkql_bit_utils.h> #include <yql/essentials/public/udf/arrow/util.h> +#include <yql/essentials/public/udf/udf_data_type.h> namespace NKikimr::NMiniKQL { @@ -202,6 +203,10 @@ struct TPrimitiveDataType<NYql::NDecimal::TInt128> { }; }; +inline constexpr size_t UuidBinarySize = NYql::NUdf::UUID_SIZE; + +std::shared_ptr<arrow::DataType> GetUuidArrowType(); + template <typename T, typename = typename std::enable_if<std::is_arithmetic<T>::value>::type> inline arrow::Datum MakeScalarDatum(T value) { return arrow::Datum(std::make_shared<typename TPrimitiveDataType<T>::TScalarResult>(value)); diff --git a/yql/essentials/minikql/comp_nodes/ut/mkql_block_dynumber_ut.cpp b/yql/essentials/minikql/comp_nodes/ut/mkql_block_dynumber_ut.cpp new file mode 100644 index 00000000000..5a1d8f78e6a --- /dev/null +++ b/yql/essentials/minikql/comp_nodes/ut/mkql_block_dynumber_ut.cpp @@ -0,0 +1,160 @@ +#include "mkql_block_serializer_test_utils.h" +#include "mkql_block_test_helper.h" +#include "mkql_computation_node_ut.h" +#include "mkql_program_builder_test_utils.h" + +#include <yql/essentials/minikql/computation/mkql_block_reader.h> +#include <yql/essentials/minikql/computation/mkql_value_builder.h> +#include <yql/essentials/minikql/mkql_type_builder.h> +#include <yql/essentials/minikql/udf_value_test_support/udf_value_comparator_utils.h> +#include <yql/essentials/public/udf/arrow/block_builder.h> + +namespace NKikimr::NMiniKQL { + +namespace { + +constexpr TStringBuf DyNumber1 = "0"; +constexpr TStringBuf DyNumber2 = "-123.45e3"; +constexpr TStringBuf DyNumber3 = "150e2"; + +template <typename T> +void TestIdentityKernel(T input, T expected) { + TBlockHelper().TestKernelFuzzied(input, expected, + [](TSetup<false>&, TRuntimeNode node) { return node; }); +} + +} // namespace + +Y_UNIT_TEST_SUITE(TMiniKQLBlockDyNumberTest) { + +Y_UNIT_TEST(DyNumberKernel) { + TestIdentityKernel( + TVector<NTest::TTestDyNumber>{ + NTest::TTestDyNumber{DyNumber1}, + NTest::TTestDyNumber{DyNumber2}, + NTest::TTestDyNumber{DyNumber3}, + }, + TVector<NTest::TTestDyNumber>{ + NTest::TTestDyNumber{DyNumber1}, + NTest::TTestDyNumber{DyNumber2}, + NTest::TTestDyNumber{DyNumber3}, + }); +} + +Y_UNIT_TEST(DyNumberKernelWithNulls) { + TestIdentityKernel( + TVector<TMaybe<NTest::TTestDyNumber>>{ + Nothing(), + NTest::TTestDyNumber{DyNumber1}, + NTest::TTestDyNumber{DyNumber2}, + Nothing(), + }, + TVector<TMaybe<NTest::TTestDyNumber>>{ + Nothing(), + NTest::TTestDyNumber{DyNumber1}, + NTest::TTestDyNumber{DyNumber2}, + Nothing(), + }); +} + +Y_UNIT_TEST(DyNumberSerializerRoundtrip) { + TBlockHelper helper; + + const TVector<NTest::TTestDyNumber> data = { + NTest::TTestDyNumber{DyNumber1}, + NTest::TTestDyNumber{DyNumber2}, + NTest::TTestDyNumber{DyNumber3}, + }; + auto [graph, value, itemType, blockType] = helper.GetArrowBlock(data); + Y_UNUSED(graph); + auto arrayData = TArrowBlock::From(value).GetDatum().array(); + auto restored = DoSerializerRoundtrip(arrayData, itemType, blockType); + + auto reader = NYql::NUdf::MakeBlockReader(TTypeInfoHelper(), itemType); + const TVector<NTest::TTestDyNumber> expected = { + NTest::TTestDyNumber{DyNumber1}, + NTest::TTestDyNumber{DyNumber2}, + NTest::TTestDyNumber{DyNumber3}, + }; + + for (size_t i = 0; i < expected.size(); ++i) { + NYql::NUdf::AssertUnboxedValueElementEqual(reader->GetItem(*restored, i), expected[i]); + } +} + +Y_UNIT_TEST(DyNumberComparator) { + TBlockHelper helper; + + const TVector<NTest::TTestDyNumber> data = { + NTest::TTestDyNumber{DyNumber2}, + NTest::TTestDyNumber{DyNumber3}, + NTest::TTestDyNumber{DyNumber2}, + }; + auto [graph, value, itemType, blockType] = helper.GetArrowBlock(data); + Y_UNUSED(graph, blockType); + auto arrayData = TArrowBlock::From(value).GetDatum().array(); + + auto reader0 = NYql::NUdf::MakeBlockReader(TTypeInfoHelper(), itemType); + auto reader1 = NYql::NUdf::MakeBlockReader(TTypeInfoHelper(), itemType); + auto reader2 = NYql::NUdf::MakeBlockReader(TTypeInfoHelper(), itemType); + const auto itemSmaller = reader0->GetItem(*arrayData, 0); + const auto itemLarger = reader1->GetItem(*arrayData, 1); + const auto itemEqual = reader2->GetItem(*arrayData, 2); + + auto comparator = TBlockTypeHelper().MakeComparator(itemType); + UNIT_ASSERT(comparator->Less(itemSmaller, itemLarger)); + UNIT_ASSERT(!comparator->Equals(itemSmaller, itemLarger)); + UNIT_ASSERT(comparator->Equals(itemSmaller, itemEqual)); +} + +Y_UNIT_TEST(DyNumberHasher) { + TBlockHelper helper; + + const TVector<NTest::TTestDyNumber> data = { + NTest::TTestDyNumber{DyNumber1}, + NTest::TTestDyNumber{DyNumber2}, + }; + auto [graph, value, itemType, blockType] = helper.GetArrowBlock(data); + Y_UNUSED(graph, blockType); + auto arrayData = TArrowBlock::From(value).GetDatum().array(); + + auto reader0 = NYql::NUdf::MakeBlockReader(TTypeInfoHelper(), itemType); + auto reader1 = NYql::NUdf::MakeBlockReader(TTypeInfoHelper(), itemType); + const auto item0 = reader0->GetItem(*arrayData, 0); + const auto item1 = reader1->GetItem(*arrayData, 1); + + auto hasher = TBlockTypeHelper().MakeHasher(itemType); + UNIT_ASSERT_VALUES_EQUAL(hasher->Hash(item0), hasher->Hash(item0)); + UNIT_ASSERT_VALUES_UNEQUAL(hasher->Hash(item0), hasher->Hash(item1)); +} + +Y_UNIT_TEST(DyNumberBlockItemConverter) { + TBlockHelper helper; + + const TVector<NTest::TTestDyNumber> data = { + NTest::TTestDyNumber{DyNumber1}, + NTest::TTestDyNumber{DyNumber2}, + NTest::TTestDyNumber{DyNumber3}, + }; + auto [graph, value, itemType, blockType] = helper.GetArrowBlock(data); + Y_UNUSED(blockType); + auto arrayData = TArrowBlock::From(value).GetDatum().array(); + + const THolderFactory& holderFactory = graph->GetHolderFactory(); + TDefaultValueBuilder valueBuilder(holderFactory); + auto converter = MakeBlockItemConverter(TTypeInfoHelper(), itemType, valueBuilder.GetPgBuilder()); + auto reader = NYql::NUdf::MakeBlockReader(TTypeInfoHelper(), itemType); + + for (size_t i = 0; i < data.size(); ++i) { + const TBlockItem blockItem = reader->GetItem(*arrayData, i); + NUdf::TUnboxedValue fromBlock(converter->MakeValue(blockItem, holderFactory)); + NYql::NUdf::AssertUnboxedValueElementEqual(fromBlock, data[i]); + + NUdf::TUnboxedValue roundTripped(converter->MakeValue(converter->MakeItem(fromBlock), holderFactory)); + NYql::NUdf::AssertUnboxedValueElementEqual(roundTripped, data[i]); + } +} + +} // Y_UNIT_TEST_SUITE(TMiniKQLBlockDyNumberTest) + +} // namespace NKikimr::NMiniKQL diff --git a/yql/essentials/minikql/comp_nodes/ut/mkql_block_serializer_test_utils.cpp b/yql/essentials/minikql/comp_nodes/ut/mkql_block_serializer_test_utils.cpp new file mode 100644 index 00000000000..1fb8135b3ec --- /dev/null +++ b/yql/essentials/minikql/comp_nodes/ut/mkql_block_serializer_test_utils.cpp @@ -0,0 +1,35 @@ +#include "mkql_block_serializer_test_utils.h" + +#include <yql/essentials/minikql/computation/mkql_block_transport.h> +#include <yql/essentials/minikql/computation/mkql_datum_validate.h> +#include <yql/essentials/minikql/computation/mkql_computation_node_holders.h> +#include <yql/essentials/utils/chunked_buffer.h> + +#include <library/cpp/testing/unittest/registar.h> + +namespace NKikimr::NMiniKQL { + +std::shared_ptr<arrow::ArrayData> DoSerializerRoundtrip( + const std::shared_ptr<arrow::ArrayData>& arrayData, TType* itemType, TType* blockType) +{ + const ui64 blockLen = static_cast<ui64>(arrayData->length); + auto* pool = arrow::default_memory_pool(); + TBlockSerializerParams params(pool, Nothing(), /*shouldSerializeOffset=*/true); + auto serializer = MakeBlockSerializer(TTypeInfoHelper(), itemType, params); + auto deserializer = MakeBlockDeserializer(TTypeInfoHelper(), itemType, params); + + TVector<ui64> metadata; + serializer->StoreMetadata(*arrayData, [&](ui64 meta) { metadata.push_back(meta); }); + NYql::TChunkedBuffer buffer; + serializer->StoreArray(*arrayData, buffer); + + size_t metaIdx = 0; + deserializer->LoadMetadata([&]() -> ui64 { return metadata[metaIdx++]; }); + auto restored = deserializer->LoadArray(buffer, blockLen, TMaybe<size_t>(0)); + + ValidateDatum(restored, Nothing(), blockType, NYql::EDatumValidationMode::Expensive); + UNIT_ASSERT_VALUES_EQUAL(restored->length, arrayData->length); + return restored; +} + +} // namespace NKikimr::NMiniKQL diff --git a/yql/essentials/minikql/comp_nodes/ut/mkql_block_serializer_test_utils.h b/yql/essentials/minikql/comp_nodes/ut/mkql_block_serializer_test_utils.h new file mode 100644 index 00000000000..009cf1008ba --- /dev/null +++ b/yql/essentials/minikql/comp_nodes/ut/mkql_block_serializer_test_utils.h @@ -0,0 +1,16 @@ +#pragma once + +#include <yql/essentials/minikql/mkql_node.h> + +#include <memory> + +namespace arrow { +struct ArrayData; +} + +namespace NKikimr::NMiniKQL { + +std::shared_ptr<arrow::ArrayData> DoSerializerRoundtrip( + const std::shared_ptr<arrow::ArrayData>& arrayData, TType* itemType, TType* blockType); + +} // namespace NKikimr::NMiniKQL diff --git a/yql/essentials/minikql/comp_nodes/ut/mkql_block_uuid_ut.cpp b/yql/essentials/minikql/comp_nodes/ut/mkql_block_uuid_ut.cpp new file mode 100644 index 00000000000..241c524dd8f --- /dev/null +++ b/yql/essentials/minikql/comp_nodes/ut/mkql_block_uuid_ut.cpp @@ -0,0 +1,145 @@ +#include "mkql_block_serializer_test_utils.h" +#include "mkql_block_test_helper.h" +#include "mkql_computation_node_ut.h" +#include "mkql_program_builder_test_utils.h" + +#include <yql/essentials/minikql/computation/mkql_block_reader.h> +#include <yql/essentials/minikql/computation/mkql_value_builder.h> +#include <yql/essentials/minikql/mkql_type_builder.h> +#include <yql/essentials/minikql/udf_value_test_support/udf_value_comparator_utils.h> +#include <yql/essentials/public/udf/arrow/block_builder.h> + +#include <util/generic/guid.h> + +namespace NKikimr::NMiniKQL { + +namespace { + +TGUID ParseUuid(TStringBuf str) { + TGUID guid; + Y_ABORT_UNLESS(GetGuid(str, guid), "invalid uuid literal"); + return guid; +} + +const TGUID Uuid1 = ParseUuid("aaaaaaaa-aaaaaaaa-aaaaaaaa-aaaaaaaa"); +const TGUID Uuid2 = ParseUuid("bbbbbbbb-bbbbbbbb-bbbbbbbb-bbbbbbbb"); +const TGUID Uuid3 = ParseUuid("cccccccc-cccccccc-cccccccc-cccccccc"); + +template <typename T> +void TestIdentityKernel(T input, T expected) { + TBlockHelper().TestKernelFuzzied(input, expected, + [](TSetup<false>&, TRuntimeNode node) { return node; }); +} + +} // namespace + +Y_UNIT_TEST_SUITE(TMiniKQLBlockUuidTest) { + +Y_UNIT_TEST(UuidRoundtrip) { + TSetup<false> setup; + TProgramBuilder& pb = *setup.PgmBuilder; + + using TVar = std::variant<TGUID, ui64>; + const TVector<TVar> stored = {Uuid1, ui64(42), Uuid2}; + const auto list = NTest::ConvertValueToLiteralNode(pb, stored); + + const auto graph = setup.BuildGraph(pb.ForwardList(pb.FromBlocks(pb.ToBlocks(pb.ToFlow(list, {}))))); + + const TVector<TVar> expected = {Uuid1, ui64(42), Uuid2}; + NYql::NUdf::AssertUnboxedValueElementEqual(graph->GetValue(), expected); +} + +Y_UNIT_TEST(UuidKernel) { + TestIdentityKernel( + TVector<TGUID>{Uuid1, Uuid2, Uuid3}, + TVector<TGUID>{Uuid1, Uuid2, Uuid3}); +} + +Y_UNIT_TEST(UuidKernelWithNulls) { + TestIdentityKernel( + TVector<TMaybe<TGUID>>{Nothing(), Uuid1, Uuid2, Nothing()}, + TVector<TMaybe<TGUID>>{Nothing(), Uuid1, Uuid2, Nothing()}); +} + +Y_UNIT_TEST(UuidSerializerRoundtrip) { + TBlockHelper helper; + + const TVector<TGUID> data = {Uuid1, Uuid2, Uuid3}; + auto [graph, value, itemType, blockType] = helper.GetArrowBlock(data); + Y_UNUSED(graph); + auto arrayData = TArrowBlock::From(value).GetDatum().array(); + auto restored = DoSerializerRoundtrip(arrayData, itemType, blockType); + + auto reader = NYql::NUdf::MakeBlockReader(TTypeInfoHelper(), itemType); + const TVector<TGUID> expected = {Uuid1, Uuid2, Uuid3}; + + for (size_t i = 0; i < expected.size(); ++i) { + NYql::NUdf::AssertUnboxedValueElementEqual(reader->GetItem(*restored, i), expected[i]); + } +} + +Y_UNIT_TEST(UuidComparator) { + TBlockHelper helper; + + const TVector<TGUID> data = {Uuid1, Uuid2, Uuid1}; + auto [graph, value, itemType, blockType] = helper.GetArrowBlock(data); + Y_UNUSED(graph, blockType); + auto arrayData = TArrowBlock::From(value).GetDatum().array(); + + auto reader0 = NYql::NUdf::MakeBlockReader(TTypeInfoHelper(), itemType); + auto reader1 = NYql::NUdf::MakeBlockReader(TTypeInfoHelper(), itemType); + auto reader2 = NYql::NUdf::MakeBlockReader(TTypeInfoHelper(), itemType); + const auto itemSmaller = reader0->GetItem(*arrayData, 0); + const auto itemLarger = reader1->GetItem(*arrayData, 1); + const auto itemEqual = reader2->GetItem(*arrayData, 2); + + auto comparator = TBlockTypeHelper().MakeComparator(itemType); + UNIT_ASSERT(comparator->Less(itemSmaller, itemLarger)); + UNIT_ASSERT(!comparator->Equals(itemSmaller, itemLarger)); + UNIT_ASSERT(comparator->Equals(itemSmaller, itemEqual)); +} + +Y_UNIT_TEST(UuidHasher) { + TBlockHelper helper; + + const TVector<TGUID> data = {Uuid1, Uuid2}; + auto [graph, value, itemType, blockType] = helper.GetArrowBlock(data); + Y_UNUSED(graph, blockType); + auto arrayData = TArrowBlock::From(value).GetDatum().array(); + + auto reader0 = NYql::NUdf::MakeBlockReader(TTypeInfoHelper(), itemType); + auto reader1 = NYql::NUdf::MakeBlockReader(TTypeInfoHelper(), itemType); + const auto item0 = reader0->GetItem(*arrayData, 0); + const auto item1 = reader1->GetItem(*arrayData, 1); + + auto hasher = TBlockTypeHelper().MakeHasher(itemType); + UNIT_ASSERT_VALUES_EQUAL(hasher->Hash(item0), hasher->Hash(item0)); + UNIT_ASSERT_VALUES_UNEQUAL(hasher->Hash(item0), hasher->Hash(item1)); +} + +Y_UNIT_TEST(UuidBlockItemConverter) { + TBlockHelper helper; + + const TVector<TGUID> data = {Uuid1, Uuid2, Uuid3}; + auto [graph, value, itemType, blockType] = helper.GetArrowBlock(data); + Y_UNUSED(blockType); + auto arrayData = TArrowBlock::From(value).GetDatum().array(); + + const THolderFactory& holderFactory = graph->GetHolderFactory(); + TDefaultValueBuilder valueBuilder(holderFactory); + auto converter = MakeBlockItemConverter(TTypeInfoHelper(), itemType, valueBuilder.GetPgBuilder()); + auto reader = NYql::NUdf::MakeBlockReader(TTypeInfoHelper(), itemType); + + for (size_t i = 0; i < data.size(); ++i) { + const TBlockItem blockItem = reader->GetItem(*arrayData, i); + NUdf::TUnboxedValue fromBlock(converter->MakeValue(blockItem, holderFactory)); + NYql::NUdf::AssertUnboxedValueElementEqual(fromBlock, data[i]); + + NUdf::TUnboxedValue roundTripped(converter->MakeValue(converter->MakeItem(fromBlock), holderFactory)); + NYql::NUdf::AssertUnboxedValueElementEqual(roundTripped, data[i]); + } +} + +} // Y_UNIT_TEST_SUITE(TMiniKQLBlockUuidTest) + +} // namespace NKikimr::NMiniKQL diff --git a/yql/essentials/minikql/comp_nodes/ut/mkql_block_variant_ut.cpp b/yql/essentials/minikql/comp_nodes/ut/mkql_block_variant_ut.cpp index 4fca8025b4a..8965f46654e 100644 --- a/yql/essentials/minikql/comp_nodes/ut/mkql_block_variant_ut.cpp +++ b/yql/essentials/minikql/comp_nodes/ut/mkql_block_variant_ut.cpp @@ -1,8 +1,8 @@ +#include "mkql_block_serializer_test_utils.h" #include "mkql_block_test_helper.h" #include "mkql_computation_node_ut.h" #include <yql/essentials/minikql/computation/mkql_computation_node_holders.h> -#include <yql/essentials/minikql/computation/mkql_datum_validate.h> #include <yql/essentials/minikql/computation/mkql_block_transport.h> #include <yql/essentials/minikql/computation/mkql_block_reader.h> #include <yql/essentials/minikql/computation/mkql_block_trimmer.h> @@ -178,33 +178,6 @@ Y_UNIT_TEST(TypeBuilderRejectsTooManyAlternatives) { UNIT_ASSERT_C(!ok, "Expected ConvertArrowType to fail for >127 alternatives"); } -namespace { - -std::shared_ptr<arrow::ArrayData> DoSerializerRoundtrip( - const std::shared_ptr<arrow::ArrayData>& arrayData, TType* itemType, TType* blockType) -{ - const ui64 blockLen = static_cast<ui64>(arrayData->length); - auto* pool = arrow::default_memory_pool(); - TBlockSerializerParams params(pool, Nothing(), /*shouldSerializeOffset=*/true); - auto serializer = MakeBlockSerializer(TTypeInfoHelper(), itemType, params); - auto deserializer = MakeBlockDeserializer(TTypeInfoHelper(), itemType, params); - - TVector<ui64> metadata; - serializer->StoreMetadata(*arrayData, [&](ui64 meta) { metadata.push_back(meta); }); - NYql::TChunkedBuffer buffer; - serializer->StoreArray(*arrayData, buffer); - - size_t metaIdx = 0; - deserializer->LoadMetadata([&]() -> ui64 { return metadata[metaIdx++]; }); - auto restored = deserializer->LoadArray(buffer, blockLen, TMaybe<size_t>(0)); - - TBlockHelper().ValidateDatum(restored, Nothing(), blockType); - UNIT_ASSERT_VALUES_EQUAL(restored->length, arrayData->length); - return restored; -} - -} // namespace - Y_UNIT_TEST(SerializerRoundtrip) { using TVar = std::variant<ui32, ui64>; TBlockHelper helper; diff --git a/yql/essentials/minikql/comp_nodes/ut/mkql_compare_ut.cpp b/yql/essentials/minikql/comp_nodes/ut/mkql_compare_ut.cpp index 57be2838634..3472b124cf1 100644 --- a/yql/essentials/minikql/comp_nodes/ut/mkql_compare_ut.cpp +++ b/yql/essentials/minikql/comp_nodes/ut/mkql_compare_ut.cpp @@ -13,9 +13,29 @@ #include <cfloat> #include <utility> +#include <util/generic/guid.h> + namespace NKikimr { namespace NMiniKQL { +namespace { + +TGUID ParseUuid(TStringBuf str) { + TGUID guid; + Y_ABORT_UNLESS(GetGuid(str, guid), "invalid uuid literal"); + return guid; +} + +const TGUID UuidLo = ParseUuid("aaaaaaaa-aaaaaaaa-aaaaaaaa-aaaaaaaa"); +const TGUID UuidMid = ParseUuid("bbbbbbbb-bbbbbbbb-bbbbbbbb-bbbbbbbb"); +const TGUID UuidHi = ParseUuid("cccccccc-cccccccc-cccccccc-cccccccc"); + +const NTest::TTestDyNumber DyNumberLo{"-123.45e3"}; +const NTest::TTestDyNumber DyNumberMid{"0"}; +const NTest::TTestDyNumber DyNumberHi{"150e2"}; + +} // namespace + Y_UNIT_TEST_SUITE(TMiniKQLCompareTest) { Y_UNIT_TEST_LLVM(SqlString) { TSetup<LLVM> setup; @@ -626,6 +646,118 @@ Y_UNIT_TEST_LLVM(TestAggrMinMaxFloats) { UNIT_ASSERT(!iterator.Next(item)); UNIT_ASSERT(!iterator.Next(item)); } + +Y_UNIT_TEST_LLVM(TestUuidMinMax) { + TSetup<LLVM> setup; + TProgramBuilder& pb = *setup.PgmBuilder; + + const auto list = NTest::ConvertValueToLiteralNode(pb, TVector<TGUID>{UuidLo, UuidMid, UuidHi}); + const auto pgmReturn = pb.FlatMap(list, + [&](TRuntimeNode left) { + return pb.Map(list, + [&](TRuntimeNode right) { + return pb.NewTuple({pb.Min(left, right), pb.Max(left, right)}); + }); + }); + + const auto graph = setup.BuildGraph(pgmReturn); + using TRow = std::tuple<TGUID, TGUID>; + AssertUnboxedValueElementEqual(graph->GetValue(), TVector<TRow>{ + TRow{UuidLo, UuidLo}, + TRow{UuidLo, UuidMid}, + TRow{UuidLo, UuidHi}, + TRow{UuidLo, UuidMid}, + TRow{UuidMid, UuidMid}, + TRow{UuidMid, UuidHi}, + TRow{UuidLo, UuidHi}, + TRow{UuidMid, UuidHi}, + TRow{UuidHi, UuidHi}, + }); +} + +Y_UNIT_TEST_LLVM(TestDyNumberMinMax) { + TSetup<LLVM> setup; + TProgramBuilder& pb = *setup.PgmBuilder; + + const auto list = NTest::ConvertValueToLiteralNode(pb, TVector<NTest::TTestDyNumber>{DyNumberLo, DyNumberMid, DyNumberHi}); + const auto pgmReturn = pb.FlatMap(list, + [&](TRuntimeNode left) { + return pb.Map(list, + [&](TRuntimeNode right) { + return pb.NewTuple({pb.Min(left, right), pb.Max(left, right)}); + }); + }); + + const auto graph = setup.BuildGraph(pgmReturn); + using TRow = std::tuple<NTest::TTestDyNumber, NTest::TTestDyNumber>; + AssertUnboxedValueElementEqual(graph->GetValue(), TVector<TRow>{ + TRow{DyNumberLo, DyNumberLo}, + TRow{DyNumberLo, DyNumberMid}, + TRow{DyNumberLo, DyNumberHi}, + TRow{DyNumberLo, DyNumberMid}, + TRow{DyNumberMid, DyNumberMid}, + TRow{DyNumberMid, DyNumberHi}, + TRow{DyNumberLo, DyNumberHi}, + TRow{DyNumberMid, DyNumberHi}, + TRow{DyNumberHi, DyNumberHi}, + }); +} + +Y_UNIT_TEST_LLVM(TestAggrUuidMinMax) { + TSetup<LLVM> setup; + TProgramBuilder& pb = *setup.PgmBuilder; + + const auto list = NTest::ConvertValueToLiteralNode(pb, TVector<TGUID>{UuidLo, UuidMid, UuidHi}); + const auto pgmReturn = pb.FlatMap(list, + [&](TRuntimeNode left) { + return pb.Map(list, + [&](TRuntimeNode right) { + return pb.NewTuple({pb.AggrMin(left, right), pb.AggrMax(left, right)}); + }); + }); + + const auto graph = setup.BuildGraph(pgmReturn); + using TRow = std::tuple<TGUID, TGUID>; + AssertUnboxedValueElementEqual(graph->GetValue(), TVector<TRow>{ + TRow{UuidLo, UuidLo}, + TRow{UuidLo, UuidMid}, + TRow{UuidLo, UuidHi}, + TRow{UuidLo, UuidMid}, + TRow{UuidMid, UuidMid}, + TRow{UuidMid, UuidHi}, + TRow{UuidLo, UuidHi}, + TRow{UuidMid, UuidHi}, + TRow{UuidHi, UuidHi}, + }); +} + +Y_UNIT_TEST_LLVM(TestAggrDyNumberMinMax) { + TSetup<LLVM> setup; + TProgramBuilder& pb = *setup.PgmBuilder; + + const auto list = NTest::ConvertValueToLiteralNode(pb, TVector<NTest::TTestDyNumber>{DyNumberLo, DyNumberMid, DyNumberHi}); + const auto pgmReturn = pb.FlatMap(list, + [&](TRuntimeNode left) { + return pb.Map(list, + [&](TRuntimeNode right) { + return pb.NewTuple({pb.AggrMin(left, right), pb.AggrMax(left, right)}); + }); + }); + + const auto graph = setup.BuildGraph(pgmReturn); + using TRow = std::tuple<NTest::TTestDyNumber, NTest::TTestDyNumber>; + AssertUnboxedValueElementEqual(graph->GetValue(), TVector<TRow>{ + TRow{DyNumberLo, DyNumberLo}, + TRow{DyNumberLo, DyNumberMid}, + TRow{DyNumberLo, DyNumberHi}, + TRow{DyNumberLo, DyNumberMid}, + TRow{DyNumberMid, DyNumberMid}, + TRow{DyNumberMid, DyNumberHi}, + TRow{DyNumberLo, DyNumberHi}, + TRow{DyNumberMid, DyNumberHi}, + TRow{DyNumberHi, DyNumberHi}, + }); +} } // Y_UNIT_TEST_SUITE(TMiniKQLCompareTest) } // namespace NMiniKQL diff --git a/yql/essentials/minikql/comp_nodes/ut/mkql_program_builder_test_utils.h b/yql/essentials/minikql/comp_nodes/ut/mkql_program_builder_test_utils.h index d17acb74972..a1a70366a2f 100644 --- a/yql/essentials/minikql/comp_nodes/ut/mkql_program_builder_test_utils.h +++ b/yql/essentials/minikql/comp_nodes/ut/mkql_program_builder_test_utils.h @@ -4,6 +4,10 @@ #include <yql/essentials/minikql/mkql_program_builder.h> #include <yql/essentials/minikql/udf_value_test_support/struct_variant_type.h> #include <yql/essentials/minikql/udf_value_test_support/udf_value_comparator_utils.h> +#include <yql/essentials/types/dynumber/dynumber.h> + +#include <util/generic/guid.h> +#include <util/system/unaligned_mem.h> #include <array> #include <tuple> @@ -91,6 +95,25 @@ public: TSingularNull() = default; }; +struct TTestDyNumber { + TString Value; + + TTestDyNumber() + : Value("0") + { + } + + explicit TTestDyNumber(TStringBuf value) + : Value(TString{value}) + { + } + + explicit TTestDyNumber(const char* value) + : Value(value) + { + } +}; + enum class TTag { A, B, @@ -184,6 +207,17 @@ inline TRuntimeNode ConvertValueToLiteralNode(TProgramBuilder& pb, TStringBuf si return pb.NewDataLiteral<NUdf::EDataSlot::String>(simpleNode); } +inline TRuntimeNode ConvertValueToLiteralNode(TProgramBuilder& pb, const TGUID& uuid) { + return pb.NewDataLiteral<NUdf::EDataSlot::Uuid>( + NUdf::TStringRef(reinterpret_cast<const char*>(&uuid), sizeof(uuid))); +} + +inline TRuntimeNode ConvertValueToLiteralNode(TProgramBuilder& pb, TTestDyNumber value) { + const auto parsed = NDyNumber::ParseDyNumberString(value.Value); + MKQL_ENSURE(parsed, "Invalid DyNumber literal: " << value.Value); + return pb.NewDataLiteral<NUdf::EDataSlot::DyNumber>(*parsed); +} + inline TRuntimeNode ConvertValueToLiteralNode(TProgramBuilder& pb, const TUtf8& utf8Node) { return pb.NewDataLiteral<NUdf::EDataSlot::Utf8>(utf8Node.Value); } @@ -303,4 +337,40 @@ struct TUnboxedValueComparator<NKikimr::NMiniKQL::NTest::TUtf8> { } }; +template <> +struct TUnboxedValueComparator<TGUID> { + template <CComparatorUtilsUdfValue THolder> + static TUnboxedValueComparatorResult IsEqual(const THolder& value, const TGUID& expected) { + const auto ref = value.AsStringRef(); + if (ref.Size() != sizeof(TGUID)) { + return std::unexpected(TStringBuilder() << "Expected Uuid of " << sizeof(TGUID) << " bytes but got " << ref.Size()); + } + const TGUID got = ReadUnaligned<TGUID>(ref.Data()); + if (got != expected) { + return std::unexpected(TStringBuilder() + << "Expected Uuid " << expected.AsGuidString() + << " but got " << got.AsGuidString()); + } + return {}; + } +}; + +template <> +struct TUnboxedValueComparator<NKikimr::NMiniKQL::NTest::TTestDyNumber> { + template <CComparatorUtilsUdfValue THolder> + static TUnboxedValueComparatorResult IsEqual(const THolder& value, const NKikimr::NMiniKQL::NTest::TTestDyNumber& expected) { + const auto parsed = NKikimr::NDyNumber::ParseDyNumberString(expected.Value); + if (!parsed) { + return std::unexpected(TStringBuilder() << "Invalid expected DyNumber: " << expected.Value); + } + if (value.AsStringRef() != NUdf::TStringRef(*parsed)) { + const auto& got = NKikimr::NDyNumber::DyNumberToString(value.AsStringRef()); + return std::unexpected(TStringBuilder() + << "Expected DyNumber \"" << expected.Value << "\" but got \"" + << (got ? *got : TString("<invalid>")) << "\""); + } + return {}; + } +}; + } // namespace NYql::NUdf::NPrivate diff --git a/yql/essentials/minikql/comp_nodes/ut/ya.make.inc b/yql/essentials/minikql/comp_nodes/ut/ya.make.inc index 1572c0c7f13..f5303a0fb1b 100644 --- a/yql/essentials/minikql/comp_nodes/ut/ya.make.inc +++ b/yql/essentials/minikql/comp_nodes/ut/ya.make.inc @@ -33,6 +33,7 @@ SET(ORIG_SOURCES mkql_block_map_join_ut_utils.cpp mkql_block_map_join_ut.cpp mkql_block_test_helper.cpp + mkql_block_serializer_test_utils.cpp mkql_block_fuzzer.cpp mkql_block_getelem_ut.cpp mkql_block_guess_ut.cpp @@ -40,6 +41,8 @@ SET(ORIG_SOURCES mkql_block_just_ut.cpp mkql_block_top_sort_ut.cpp mkql_block_variant_ut.cpp + mkql_block_uuid_ut.cpp + mkql_block_dynumber_ut.cpp mkql_blocks_ut.cpp mkql_callable_ut.cpp mkql_combine_ut.cpp @@ -103,6 +106,7 @@ PEERDIR( yql/essentials/parser/pg_wrapper yql/essentials/parser/pg_wrapper/interface yql/essentials/sql/pg + yql/essentials/types/dynumber yql/essentials/utils/random_data_generator ) diff --git a/yql/essentials/minikql/computation/mkql_block_impl.cpp b/yql/essentials/minikql/computation/mkql_block_impl.cpp index 59940aa2e62..2dd3f62a433 100644 --- a/yql/essentials/minikql/computation/mkql_block_impl.cpp +++ b/yql/essentials/minikql/computation/mkql_block_impl.cpp @@ -131,12 +131,13 @@ arrow::Datum DoConvertScalar(TType* type, const T& value, arrow::MemoryPool& poo case NUdf::EDataSlot::Utf8: case NUdf::EDataSlot::Yson: case NUdf::EDataSlot::Json: - case NUdf::EDataSlot::JsonDocument: { + case NUdf::EDataSlot::JsonDocument: + case NUdf::EDataSlot::DyNumber: { const auto& str = value.AsStringRef(); std::shared_ptr<arrow::Buffer> buffer(ARROW_RESULT(arrow::AllocateBuffer(str.Size(), &pool))); std::memcpy(buffer->mutable_data(), str.Data(), str.Size()); std::shared_ptr<arrow::Scalar> scalar; - if (slot == NUdf::EDataSlot::String || slot == NUdf::EDataSlot::Yson || slot == NUdf::EDataSlot::JsonDocument) { + if (slot == NUdf::EDataSlot::String || slot == NUdf::EDataSlot::Yson || slot == NUdf::EDataSlot::JsonDocument || slot == NUdf::EDataSlot::DyNumber) { scalar = std::make_shared<arrow::BinaryScalar>(buffer, arrow::binary()); } else { // NOTE: Do not use |arrow::BinaryScalar| for utf8 and json types directly. @@ -192,6 +193,13 @@ arrow::Datum DoConvertScalar(TType* type, const T& value, arrow::MemoryPool& poo *reinterpret_cast<NYql::NDecimal::TInt128*>(buffer->mutable_data()) = value.GetInt128(); return arrow::Datum(std::make_shared<TPrimitiveDataType<NYql::NDecimal::TInt128>::TScalarResult>(buffer)); } + case NUdf::EDataSlot::Uuid: { + std::shared_ptr<arrow::Buffer> buffer(ARROW_RESULT(arrow::AllocateBuffer(UuidBinarySize, &pool))); + const auto ref = value.AsStringRef(); + MKQL_ENSURE(ref.Size() == UuidBinarySize, "Wrong Uuid size: " << ref.Size()); + std::memcpy(buffer->mutable_data(), ref.Data(), ref.Size()); + return arrow::Datum(std::make_shared<arrow::FixedSizeBinaryScalar>(std::move(buffer), GetUuidArrowType())); + } default: MKQL_ENSURE(false, "Unsupported data slot " << slot); } diff --git a/yql/essentials/minikql/computation/mkql_block_reader.cpp b/yql/essentials/minikql/computation/mkql_block_reader.cpp index bad2fc17ea3..14a91181325 100644 --- a/yql/essentials/minikql/computation/mkql_block_reader.cpp +++ b/yql/essentials/minikql/computation/mkql_block_reader.cpp @@ -6,6 +6,7 @@ #include <yql/essentials/public/udf/udf_type_inspection.h> #include <yql/essentials/public/udf/udf_value_utils.h> +#include <yql/essentials/public/udf/udf_data_type.h> #include <arrow/array/array_binary.h> #include <arrow/chunked_array.h> @@ -162,6 +163,13 @@ private: i32 TypeLen_ = 0; }; +template <bool Nullable> +using TUuidBlockItemConverter = TStringBlockItemConverter<arrow::BinaryType, Nullable, NUdf::EPgStringType::None>; + +template <bool Nullable> +class TFixedSizeBlockItemConverter<TGUID, Nullable>: public TUuidBlockItemConverter<Nullable> { +}; + template <bool IsNull> class TSingularTypeItemConverter: public IBlockItemConverter { public: diff --git a/yql/essentials/minikql/invoke_builtins/mkql_builtins_max.cpp b/yql/essentials/minikql/invoke_builtins/mkql_builtins_max.cpp index 183c982b27b..56d6623d3d8 100644 --- a/yql/essentials/minikql/invoke_builtins/mkql_builtins_max.cpp +++ b/yql/essentials/minikql/invoke_builtins/mkql_builtins_max.cpp @@ -208,6 +208,8 @@ void RegisterMax(IBuiltinFunctionRegistry& registry) { RegisterCustomSameTypesFunction<NUdf::TDataType<char*>, TCustomMax, TBinaryArgsOpt>(registry, "Max"); RegisterCustomSameTypesFunction<NUdf::TDataType<NUdf::TUtf8>, TCustomMax, TBinaryArgsOpt>(registry, "Max"); + RegisterCustomSameTypesFunction<NUdf::TDataType<NUdf::TDyNumber>, TCustomMax, TBinaryArgsOpt>(registry, "Max"); + RegisterCustomSameTypesFunction<NUdf::TDataType<NUdf::TUuid>, TCustomMax, TBinaryArgsOpt>(registry, "Max"); } void RegisterAggrMax(IBuiltinFunctionRegistry& registry) { @@ -221,6 +223,8 @@ void RegisterAggrMax(IBuiltinFunctionRegistry& registry) { RegisterCustomAggregateFunction<NUdf::TDataType<char*>, TCustomMax, TBinaryArgsSameOpt>(registry, "AggrMax"); RegisterCustomAggregateFunction<NUdf::TDataType<NUdf::TUtf8>, TCustomMax, TBinaryArgsSameOpt>(registry, "AggrMax"); + RegisterCustomAggregateFunction<NUdf::TDataType<NUdf::TDyNumber>, TCustomMax, TBinaryArgsSameOpt>(registry, "AggrMax"); + RegisterCustomAggregateFunction<NUdf::TDataType<NUdf::TUuid>, TCustomMax, TBinaryArgsSameOpt>(registry, "AggrMax"); } } // namespace NMiniKQL diff --git a/yql/essentials/minikql/invoke_builtins/mkql_builtins_min.cpp b/yql/essentials/minikql/invoke_builtins/mkql_builtins_min.cpp index 28cede3d6d8..f2153fe069f 100644 --- a/yql/essentials/minikql/invoke_builtins/mkql_builtins_min.cpp +++ b/yql/essentials/minikql/invoke_builtins/mkql_builtins_min.cpp @@ -207,6 +207,8 @@ void RegisterMin(IBuiltinFunctionRegistry& registry) { RegisterCustomSameTypesFunction<NUdf::TDataType<char*>, TCustomMin, TBinaryArgsOpt>(registry, "Min"); RegisterCustomSameTypesFunction<NUdf::TDataType<NUdf::TUtf8>, TCustomMin, TBinaryArgsOpt>(registry, "Min"); + RegisterCustomSameTypesFunction<NUdf::TDataType<NUdf::TDyNumber>, TCustomMin, TBinaryArgsOpt>(registry, "Min"); + RegisterCustomSameTypesFunction<NUdf::TDataType<NUdf::TUuid>, TCustomMin, TBinaryArgsOpt>(registry, "Min"); } void RegisterAggrMin(IBuiltinFunctionRegistry& registry) { @@ -220,6 +222,8 @@ void RegisterAggrMin(IBuiltinFunctionRegistry& registry) { RegisterCustomAggregateFunction<NUdf::TDataType<char*>, TCustomMin, TBinaryArgsSameOpt>(registry, "AggrMin"); RegisterCustomAggregateFunction<NUdf::TDataType<NUdf::TUtf8>, TCustomMin, TBinaryArgsSameOpt>(registry, "AggrMin"); + RegisterCustomAggregateFunction<NUdf::TDataType<NUdf::TDyNumber>, TCustomMin, TBinaryArgsSameOpt>(registry, "AggrMin"); + RegisterCustomAggregateFunction<NUdf::TDataType<NUdf::TUuid>, TCustomMin, TBinaryArgsSameOpt>(registry, "AggrMin"); } } // namespace NMiniKQL diff --git a/yql/essentials/minikql/mkql_type_builder.cpp b/yql/essentials/minikql/mkql_type_builder.cpp index 5c035dc34ba..4c354629b01 100644 --- a/yql/essentials/minikql/mkql_type_builder.cpp +++ b/yql/essentials/minikql/mkql_type_builder.cpp @@ -12,6 +12,7 @@ #include <library/cpp/containers/stack_vector/stack_vec.h> #include <yql/essentials/minikql/computation/mkql_computation_node_impl.h> +#include <yql/essentials/minikql/arrow/arrow_util.h> #include <yql/essentials/minikql/mkql_runtime_version.h> #include <yql/essentials/minikql/mkql_node_printer.h> #include <yql/essentials/parser/pg_catalog/catalog.h> @@ -1556,14 +1557,16 @@ bool ConvertArrowTypeImpl(NUdf::EDataSlot slot, std::shared_ptr<arrow::DataType> return true; } case NUdf::EDataSlot::Uuid: { - return false; + type = arrow::fixed_size_binary(UuidBinarySize); + return true; } case NUdf::EDataSlot::Decimal: { type = arrow::fixed_size_binary(sizeof(NYql::NUdf::TUnboxedValuePod)); return true; } case NUdf::EDataSlot::DyNumber: { - return false; + type = arrow::binary(); + return true; } } } @@ -2700,14 +2703,13 @@ size_t CalcMaxBlockItemSize(const TType* type) { case NUdf::EDataSlot::TzTimestamp64: return sizeof(typename NUdf::TDataType<NUdf::TTzTimestamp64>::TLayout) + sizeof(NYql::NUdf::TTimezoneId); case NUdf::EDataSlot::Uuid: { - MKQL_ENSURE(false, "Unsupported data slot: " << slot); + return UuidBinarySize; } case NUdf::EDataSlot::Decimal: { return sizeof(NYql::NDecimal::TInt128); } - case NUdf::EDataSlot::DyNumber: { - MKQL_ENSURE(false, "Unsupported data slot: " << slot); - } + case NUdf::EDataSlot::DyNumber: + return sizeof(arrow::BinaryType::offset_type); } } diff --git a/yql/essentials/parser/pg_wrapper/comp_factory.cpp b/yql/essentials/parser/pg_wrapper/comp_factory.cpp index 82c87dbfc0b..78b7e4c59fa 100644 --- a/yql/essentials/parser/pg_wrapper/comp_factory.cpp +++ b/yql/essentials/parser/pg_wrapper/comp_factory.cpp @@ -19,8 +19,10 @@ #include <yql/essentials/minikql/mkql_node_builder.h> #include <yql/essentials/minikql/mkql_string_util.h> #include <yql/essentials/minikql/mkql_type_builder.h> +#include <yql/essentials/minikql/mkql_type_ops.h> #include <yql/essentials/minikql/mkql_safe_arithmetic_ops.h> #include <yql/essentials/types/binary_json/read.h> +#include <yql/essentials/types/dynumber/dynumber.h> #include <yql/essentials/types/uuid/uuid.h> #include <yql/essentials/public/udf/arrow/block_reader.h> #include <yql/essentials/public/udf/arrow/block_builder.cpp> @@ -33,6 +35,7 @@ #include <library/cpp/yson/detail.h> #include <library/cpp/string_utils/base64/base64.h> #include <util/string/split.h> +#include <util/generic/guid.h> #include <util/system/getpid.h> #include "arrow.h" @@ -1624,12 +1627,7 @@ NUdf::TUnboxedValuePod ConvertToPgValue(NUdf::TUnboxedValuePod value, TMaybe<NUd return PointerDatumToPod(PointerGetDatum(res)); } case NUdf::EDataSlot::Uuid: { - TString str; - str.reserve(36); - ui16 dw[8]; - std::memcpy(dw, value.AsStringRef().Data(), sizeof(dw)); - TStringOutput out(str); - NKikimr::NUuid::UuidToString(dw, out); + auto str = NKikimr::NUuid::UuidBytesToString(value.AsStringRef()); auto res = DirectFunctionCall1Coll(uuid_in, DEFAULT_COLLATION_OID, PointerGetDatum(str.c_str())); return PointerDatumToPod(PointerGetDatum((void*)res)); } @@ -2335,9 +2333,39 @@ std::shared_ptr<arrow::compute::ScalarKernel> MakeFromPgKernel(TType* inputType, return kernel; } +namespace { + +template <typename TTzDate, NUdf::EDataSlot Slot> +arrow::Status ConvertTzDateBlockToPg(const arrow::ArrayData& array, size_t length, arrow::compute::KernelContext* ctx, arrow::Datum* res) { + using TLayout = typename NUdf::TDataType<TTzDate>::TLayout; + NUdf::TTzDateBlockReader<TTzDate, true> reader; + NUdf::TStringArrayBuilder<arrow::BinaryType, true> builder( + NKikimr::NMiniKQL::TTypeInfoHelper(), arrow::binary(), *ctx->memory_pool(), length); + for (size_t i = 0; i < length; ++i) { + auto item = reader.GetItem(array, i); + if (!item) { + builder.Add(NUdf::TBlockItem()); + continue; + } + + NUdf::TUnboxedValuePod value{item.template Get<TLayout>()}; + value.SetTimezoneId(item.GetTimezoneId()); + NUdf::TUnboxedValue str = ValueToString(Slot, value); + const auto ref = str.AsStringRef(); + auto ptr = builder.AddPgItem<false, VARHDRSZ>(ref); + UpdateCleanVarSize((text*)(ptr + sizeof(void*)), ref.Size()); + } + + *res = builder.Build(true); + return arrow::Status::OK(); +} + +} // namespace + struct TToPgExec { - TToPgExec(NUdf::EDataSlot sourceDataSlot) + TToPgExec(NUdf::EDataSlot sourceDataSlot, TType* itemType) : SourceDataSlot(sourceDataSlot) + , ItemType(itemType) {} arrow::Status Exec(arrow::compute::KernelContext* ctx, const arrow::compute::ExecBatch& batch, arrow::Datum* res) const { @@ -2581,6 +2609,84 @@ struct TToPgExec { *res = builder.Build(true); break; } + case NUdf::EDataSlot::Uuid: { + NUdf::TFixedSizeBlockReader<TGUID, true> reader; + NUdf::TStringArrayBuilder<arrow::BinaryType, true> builder(NKikimr::NMiniKQL::TTypeInfoHelper(), arrow::binary(), *ctx->memory_pool(), length); + for (size_t i = 0; i < length; ++i) { + auto item = reader.GetItem(array, i); + if (!item) { + builder.Add(NUdf::TBlockItem()); + continue; + } + + auto str = NKikimr::NUuid::UuidBytesToString(item.AsStringRef()); + auto pgDatum = DirectFunctionCall1Coll(uuid_in, DEFAULT_COLLATION_OID, PointerGetDatum(str.c_str())); + const char* pgUuid = static_cast<const char*>(DatumGetPointer(pgDatum)); + auto ref = NUdf::TStringRef(pgUuid, 16); + builder.AddPgItem<false, 0>(ref); + pfree(const_cast<char*>(pgUuid)); + } + + *res = builder.Build(true); + break; + } + case NUdf::EDataSlot::Decimal: { + NUdf::TFixedSizeBlockReader<NYql::NDecimal::TInt128, true> reader; + NUdf::TStringArrayBuilder<arrow::BinaryType, true> builder(NKikimr::NMiniKQL::TTypeInfoHelper(), arrow::binary(), *ctx->memory_pool(), length); + for (size_t i = 0; i < length; ++i) { + auto item = reader.GetItem(array, i); + if (!item) { + builder.Add(NUdf::TBlockItem()); + continue; + } + + NUdf::TUnboxedValuePod value(item.GetInt128()); + auto decimalType = static_cast<TDataDecimalType*>(ItemType); + auto pgRes = DecimalToPgNumeric(value, decimalType->GetParams().first, decimalType->GetParams().second); + auto ref = NUdf::TStringRef((const char*)pgRes, GetFullVarSize((const text*)pgRes)); + auto ptr = builder.AddPgItem<false, 0>(ref); + UpdateCleanVarSize((text*)(ptr + sizeof(void*)), GetCleanVarSize((const text*)pgRes)); + pfree(pgRes); + } + + *res = builder.Build(true); + break; + } + case NUdf::EDataSlot::DyNumber: { + NUdf::TStringBlockReader<arrow::BinaryType, true> reader; + NUdf::TStringArrayBuilder<arrow::BinaryType, true> builder(NKikimr::NMiniKQL::TTypeInfoHelper(), arrow::binary(), *ctx->memory_pool(), length); + for (size_t i = 0; i < length; ++i) { + auto item = reader.GetItem(array, i); + if (!item) { + builder.Add(NUdf::TBlockItem()); + continue; + } + + auto str = NKikimr::NDyNumber::DyNumberToString(item.AsStringRef()); + Y_ENSURE(str); + auto pgRes = (Numeric)DirectFunctionCall3Coll(numeric_in, DEFAULT_COLLATION_OID, + PointerGetDatum(str->c_str()), Int32GetDatum(0), Int32GetDatum(-1)); + auto ref = NUdf::TStringRef((const char*)pgRes, GetFullVarSize((const text*)pgRes)); + auto ptr = builder.AddPgItem<false, 0>(ref); + UpdateCleanVarSize((text*)(ptr + sizeof(void*)), GetCleanVarSize((const text*)pgRes)); + pfree(pgRes); + } + + *res = builder.Build(true); + break; + } + case NUdf::EDataSlot::TzDate: + return ConvertTzDateBlockToPg<NUdf::TTzDate, NUdf::EDataSlot::TzDate>(array, length, ctx, res); + case NUdf::EDataSlot::TzDatetime: + return ConvertTzDateBlockToPg<NUdf::TTzDatetime, NUdf::EDataSlot::TzDatetime>(array, length, ctx, res); + case NUdf::EDataSlot::TzTimestamp: + return ConvertTzDateBlockToPg<NUdf::TTzTimestamp, NUdf::EDataSlot::TzTimestamp>(array, length, ctx, res); + case NUdf::EDataSlot::TzDate32: + return ConvertTzDateBlockToPg<NUdf::TTzDate32, NUdf::EDataSlot::TzDate32>(array, length, ctx, res); + case NUdf::EDataSlot::TzDatetime64: + return ConvertTzDateBlockToPg<NUdf::TTzDatetime64, NUdf::EDataSlot::TzDatetime64>(array, length, ctx, res); + case NUdf::EDataSlot::TzTimestamp64: + return ConvertTzDateBlockToPg<NUdf::TTzTimestamp64, NUdf::EDataSlot::TzTimestamp64>(array, length, ctx, res); default: ythrow yexception() << "Unsupported type: " << NUdf::GetDataTypeInfo(SourceDataSlot).Name; } @@ -2588,6 +2694,7 @@ struct TToPgExec { } const NUdf::EDataSlot SourceDataSlot; + TType* const ItemType; }; std::shared_ptr<arrow::compute::ScalarKernel> MakeToPgKernel(TType* inputType, TType* resultType, NUdf::EDataSlot dataSlot) { @@ -2595,7 +2702,12 @@ std::shared_ptr<arrow::compute::ScalarKernel> MakeToPgKernel(TType* inputType, T std::shared_ptr<arrow::DataType> returnArrowType; MKQL_ENSURE(ConvertArrowType(AS_TYPE(TBlockType, resultType)->GetItemType(), returnArrowType), "Unsupported arrow type"); - auto exec = std::make_shared<TToPgExec>(dataSlot); + std::shared_ptr<TToPgExec> exec; + auto itemType = AS_TYPE(TBlockType, inputType)->GetItemType(); + if (itemType->IsOptional()) { + itemType = AS_TYPE(TOptionalType, itemType)->GetItemType(); + } + exec = std::make_shared<TToPgExec>(dataSlot, itemType); auto kernel = std::make_shared<arrow::compute::ScalarKernel>(ConvertToInputTypes(argTypes), ConvertToOutputType(resultType), [exec](arrow::compute::KernelContext* ctx, const arrow::compute::ExecBatch& batch, arrow::Datum* res) { return exec->Exec(ctx, batch, res); @@ -2627,6 +2739,15 @@ std::shared_ptr<arrow::compute::ScalarKernel> MakeToPgKernel(TType* inputType, T case NUdf::EDataSlot::Yson: case NUdf::EDataSlot::Json: case NUdf::EDataSlot::JsonDocument: + case NUdf::EDataSlot::Uuid: + case NUdf::EDataSlot::DyNumber: + case NUdf::EDataSlot::Decimal: + case NUdf::EDataSlot::TzDate: + case NUdf::EDataSlot::TzDatetime: + case NUdf::EDataSlot::TzTimestamp: + case NUdf::EDataSlot::TzDate32: + case NUdf::EDataSlot::TzDatetime64: + case NUdf::EDataSlot::TzTimestamp64: kernel->null_handling = arrow::compute::NullHandling::COMPUTED_NO_PREALLOCATE; break; default: diff --git a/yql/essentials/public/udf/arrow/block_builder.h b/yql/essentials/public/udf/arrow/block_builder.h index 10beba194d2..b903f35af67 100644 --- a/yql/essentials/public/udf/arrow/block_builder.h +++ b/yql/essentials/public/udf/arrow/block_builder.h @@ -12,6 +12,9 @@ #include <yql/essentials/public/udf/udf_value_builder.h> #include <yql/essentials/public/udf/udf_type_inspection.h> +#include <util/generic/guid.h> +#include <util/system/unaligned_mem.h> + #include <arrow/array/array_base.h> #include <arrow/array/concatenate.h> #include <arrow/array/util.h> @@ -533,6 +536,45 @@ public: }; template <bool Nullable> +class TFixedSizeArrayBuilder<TGUID, Nullable> final: public TFixedSizeArrayBuilderBase<TGUID, Nullable, TFixedSizeArrayBuilder<TGUID, Nullable>> { + using TSelf = TFixedSizeArrayBuilder<TGUID, Nullable>; + using TBase = TFixedSizeArrayBuilderBase<TGUID, Nullable, TSelf>; + using TParams = TArrayBuilderBase::TParams; + +public: + TFixedSizeArrayBuilder(const ITypeInfoHelper& typeInfoHelper, std::shared_ptr<arrow::DataType> arrowType, arrow::MemoryPool& pool, size_t maxLen, const TParams& params = {}) + : TBase(typeInfoHelper, std::move(arrowType), pool, maxLen, params) + { + } + + TFixedSizeArrayBuilder(const TType* type, const ITypeInfoHelper& typeInfoHelper, arrow::MemoryPool& pool, size_t maxLen, const TParams& params = {}) + : TBase(typeInfoHelper, type, pool, maxLen, params) + { + } + + void DoAddNotNullFromStringRef(TStringBuf ref) { + this->PlaceItem(ReadUnaligned<TGUID>(ref.Data())); + } + + void DoAddNotNull(TUnboxedValuePod value) { + DoAddNotNullFromStringRef(value.AsStringRef()); + } + + void DoAddNotNull(TBlockItem value) { + DoAddNotNullFromStringRef(value.AsStringRef()); + } + + void DoAddNotNull(TInputBuffer& input) { + this->PlaceItem(input.PopNumber<TGUID>()); + } + + void DoAddNotNull(TBlockItem value, size_t count) { + const TGUID uuid = ReadUnaligned<TGUID>(value.AsStringRef().Data()); + std::fill(this->DataPtr_ + this->GetCurrLen(), this->DataPtr_ + this->GetCurrLen() + count, uuid); + } +}; + +template <bool Nullable> class TResourceArrayBuilder final: public TFixedSizeArrayBuilderBase<TUnboxedValue, Nullable, TResourceArrayBuilder<Nullable>> { using TBase = TFixedSizeArrayBuilderBase<TUnboxedValue, Nullable, TResourceArrayBuilder<Nullable>>; using TParams = TArrayBuilderBase::TParams; diff --git a/yql/essentials/public/udf/arrow/block_item_comparator.h b/yql/essentials/public/udf/arrow/block_item_comparator.h index 114271fed37..00f4802f985 100644 --- a/yql/essentials/public/udf/arrow/block_item_comparator.h +++ b/yql/essentials/public/udf/arrow/block_item_comparator.h @@ -5,6 +5,9 @@ #include <yql/essentials/public/udf/udf_ptr.h> #include <yql/essentials/public/udf/udf_type_inspection.h> #include <yql/essentials/public/udf/udf_type_size_check.h> +#include <yql/essentials/public/udf/udf_data_type.h> + +#include <util/generic/guid.h> namespace NYql::NUdf { @@ -144,8 +147,8 @@ public: } }; -template <typename TStringType, bool Nullable> -class TStringBlockItemComparator: public TBlockItemComparatorBase<TStringBlockItemComparator<TStringType, Nullable>, Nullable> { +template <bool Nullable> +class TStringBlockItemComparatorBase: public TBlockItemComparatorBase<TStringBlockItemComparatorBase<Nullable>, Nullable> { public: i64 DoCompare(TBlockItem lhs, TBlockItem rhs) const { return lhs.AsStringRef().Compare(rhs.AsStringRef()); @@ -160,6 +163,18 @@ public: } }; +template <typename TStringType, bool Nullable> +class TStringBlockItemComparator: public TStringBlockItemComparatorBase<Nullable> { +}; + +template <bool Nullable> +class TUuidBlockItemComparator: public TStringBlockItemComparatorBase<Nullable> { +}; + +template <bool Nullable> +class TFixedSizeBlockItemComparator<TGUID, Nullable>: public TUuidBlockItemComparator<Nullable> { +}; + class TSingularTypeBlockItemComparator: public TBlockItemComparatorBase<TSingularTypeBlockItemComparator, /*Nullable=*/false> { public: i64 DoCompare(TBlockItem lhs, TBlockItem rhs) const { diff --git a/yql/essentials/public/udf/arrow/block_item_hasher.h b/yql/essentials/public/udf/arrow/block_item_hasher.h index 234d69bba8f..2be0fe7ad2b 100644 --- a/yql/essentials/public/udf/arrow/block_item_hasher.h +++ b/yql/essentials/public/udf/arrow/block_item_hasher.h @@ -6,6 +6,9 @@ #include <yql/essentials/public/udf/udf_type_inspection.h> #include <yql/essentials/public/udf/udf_type_ops.h> #include <yql/essentials/public/udf/udf_type_size_check.h> +#include <yql/essentials/public/udf/udf_data_type.h> + +#include <util/generic/guid.h> namespace NYql::NUdf { @@ -68,14 +71,26 @@ public: } }; -template <typename TStringType, bool Nullable> -class TStringBlockItemHasher: public TBlockItemHasherBase<TStringBlockItemHasher<TStringType, Nullable>, Nullable> { +template <bool Nullable> +class TStringBlockItemHasherBase: public TBlockItemHasherBase<TStringBlockItemHasherBase<Nullable>, Nullable> { public: ui64 DoHash(TBlockItem value) const { return GetStringHash(value.AsStringRef()); } }; +template <typename TStringType, bool Nullable> +class TStringBlockItemHasher: public TStringBlockItemHasherBase<Nullable> { +}; + +template <bool Nullable> +class TUuidBlockItemHasher: public TStringBlockItemHasherBase<Nullable> { +}; + +template <bool Nullable> +class TFixedSizeBlockItemHasher<TGUID, Nullable>: public TUuidBlockItemHasher<Nullable> { +}; + class TSingularTypeBlockItemHaser: public TBlockItemHasherBase<TSingularTypeBlockItemHaser, /*Nullable=*/false> { public: ui64 DoHash(TBlockItem value) const { diff --git a/yql/essentials/public/udf/arrow/block_reader.h b/yql/essentials/public/udf/arrow/block_reader.h index 8b8285f5d3f..b5420c646da 100644 --- a/yql/essentials/public/udf/arrow/block_reader.h +++ b/yql/essentials/public/udf/arrow/block_reader.h @@ -13,8 +13,12 @@ #include <arrow/type.h> #include <yql/essentials/public/decimal/yql_decimal.h> +#include <yql/essentials/public/udf/udf_data_type.h> #include <yql/essentials/public/udf/udf_value_utils.h> +#include <util/generic/guid.h> +#include <util/system/unaligned_mem.h> + namespace NYql::NUdf { class IBlockReader: private TNonCopyable { @@ -61,6 +65,10 @@ protected: template <typename T, bool Nullable, typename TDerived> class TFixedSizeBlockReaderBase: public TBlockReaderBase { + static constexpr bool IsUuid = std::is_same_v<T, TGUID>; + static constexpr bool IsStoredAsStringRef = IsUuid; + static constexpr size_t FixedBinarySize = sizeof(T); + public: TBlockItem GetItem(const arrow::ArrayData& data, size_t index) final { if constexpr (Nullable) { @@ -68,7 +76,13 @@ public: return {}; } } - return static_cast<TDerived*>(this)->MakeBlockItem(data.GetValues<T>(1)[index]); + if constexpr (IsStoredAsStringRef) { + const auto byteOffset = (data.offset + static_cast<int64_t>(index)) * static_cast<int64_t>(FixedBinarySize); + const auto* ptr = reinterpret_cast<const char*>(data.GetValues<uint8_t>(1, byteOffset)); + return TBlockItem(TStringRef(ptr, FixedBinarySize)); + } else { + return static_cast<TDerived*>(this)->MakeBlockItem(data.GetValues<T>(1)[index]); + } } TBlockItem GetScalarItem(const arrow::Scalar& scalar) final { @@ -80,7 +94,12 @@ public: } } - if constexpr (std::is_same_v<T, NYql::NDecimal::TInt128>) { + if constexpr (IsStoredAsStringRef) { + auto& fixedScalar = checked_cast<const arrow::FixedSizeBinaryScalar&>(scalar); + return TBlockItem(TStringRef( + reinterpret_cast<const char*>(fixedScalar.value->data()), + FixedBinarySize)); + } else if constexpr (std::is_same_v<T, NYql::NDecimal::TInt128>) { auto& fixedScalar = checked_cast<const arrow::FixedSizeBinaryScalar&>(scalar); T value; memcpy((void*)&value, fixedScalar.value->data(), sizeof(T)); @@ -107,9 +126,9 @@ public: ui64 GetDefaultValueWeight() const final { if constexpr (Nullable) { - return 1 + sizeof(T); + return 1 + FixedBinarySize; } - return sizeof(T); + return FixedBinarySize; } void SaveItem(const arrow::ArrayData& data, size_t index, TOutputBuffer& out) const final { @@ -120,7 +139,17 @@ public: out.PushChar(1); } - out.PushNumber(data.GetValues<T>(1)[index]); + if constexpr (IsUuid) { + const auto byteOffset = (data.offset + static_cast<int64_t>(index)) * static_cast<int64_t>(FixedBinarySize); + out.PushNumber(ReadUnaligned<TGUID>(data.GetValues<uint8_t>(1, byteOffset))); + } else if constexpr (std::is_same_v<T, NYql::NDecimal::TInt128>) { + T value; + const auto byteOffset = (data.offset + static_cast<int64_t>(index)) * static_cast<int64_t>(sizeof(T)); + std::memcpy(&value, data.GetValues<uint8_t>(1, byteOffset), sizeof(T)); + out.PushNumber(value); + } else { + out.PushNumber(data.GetValues<T>(1)[index]); + } } void SaveScalarItem(const arrow::Scalar& scalar, TOutputBuffer& out) const final { @@ -131,7 +160,10 @@ public: out.PushChar(1); } - if constexpr (std::is_same_v<T, NYql::NDecimal::TInt128>) { + if constexpr (IsUuid) { + auto& fixedScalar = arrow::internal::checked_cast<const arrow::FixedSizeBinaryScalar&>(scalar); + out.PushNumber(ReadUnaligned<TGUID>(fixedScalar.value->data())); + } else if constexpr (std::is_same_v<T, NYql::NDecimal::TInt128>) { auto& fixedScalar = arrow::internal::checked_cast<const arrow::FixedSizeBinaryScalar&>(scalar); T value; memcpy((void*)&value, fixedScalar.value->data(), sizeof(T)); @@ -143,7 +175,7 @@ public: private: ui64 GetDataWeightImpl(i64 dataLength) const { - ui64 size = sizeof(T) * dataLength; + ui64 size = FixedBinarySize * dataLength; if constexpr (Nullable) { size += GetBitmaskDataWeight(dataLength); } diff --git a/yql/essentials/public/udf/arrow/dispatch_traits.h b/yql/essentials/public/udf/arrow/dispatch_traits.h index 975fa1ffaf0..50e9a526e76 100644 --- a/yql/essentials/public/udf/arrow/dispatch_traits.h +++ b/yql/essentials/public/udf/arrow/dispatch_traits.h @@ -3,9 +3,12 @@ #include <yql/essentials/public/udf/arrow/util.h> #include <yql/essentials/public/udf/udf_type_inspection.h> #include <yql/essentials/public/udf/udf_value_builder.h> +#include <yql/essentials/public/udf/udf_data_type.h> #include <arrow/type.h> +#include <util/generic/guid.h> + namespace NYql::NUdf { template <typename TTraits, typename... TArgs> @@ -237,9 +240,15 @@ std::unique_ptr<typename TTraits::TResult> DispatchByArrowTraits(const ITypeInfo Y_ENSURE(false, "Unsupported data slot"); } } - case NUdf::EDataSlot::Uuid: + case NUdf::EDataSlot::Uuid: { + if constexpr (requires { typename TTraits::template TFixedSize<TGUID, true>; }) { + return MakeFixedSizeArrowTraitsImpl<TTraits, TGUID>(isOptional, type, std::forward<TArgs>(args)...); + } else { + Y_ENSURE(false, "Unsupported data slot"); + } + } case NUdf::EDataSlot::DyNumber: - Y_ENSURE(false, "Unsupported data slot"); + return MakeStringArrowTraitsImpl<TTraits, arrow::BinaryType, NUdf::EDataSlot::DyNumber>(isOptional, type, std::forward<TArgs>(args)...); } } diff --git a/yql/essentials/types/uuid/uuid.cpp b/yql/essentials/types/uuid/uuid.cpp index 0fcc7e0cea1..ee723f11752 100644 --- a/yql/essentials/types/uuid/uuid.cpp +++ b/yql/essentials/types/uuid/uuid.cpp @@ -1,5 +1,6 @@ #include "uuid.h" +#include <util/generic/yexception.h> #include <util/stream/str.h> #include <array> @@ -32,7 +33,7 @@ void WriteHex(ui16 bytes, IOutputStream& out, bool reverseBytes = false) { } // namespace -TString UuidBytesToString(const TString& in) { +TString UuidBytesToString(TStringBuf in) { TStringStream ss; UuidBytesToString(in, ss); @@ -40,12 +41,21 @@ TString UuidBytesToString(const TString& in) { return ss.Str(); } -void UuidBytesToString(const TString& in, IOutputStream& out) { +void UuidBytesToString(TStringBuf in, IOutputStream& out) { + Y_ENSURE(in.size() == UUID_LEN, "Invalid uuid bytes size"); std::array<ui16, 8> dw; std::memcpy(dw.data(), in.data(), sizeof(dw)); NUuid::UuidToString(dw.data(), out); } +TString UuidBytesToString(const TString& in) { + return UuidBytesToString(TStringBuf(in)); +} + +void UuidBytesToString(const TString& in, IOutputStream& out) { + UuidBytesToString(TStringBuf(in), out); +} + void UuidHalfsToString(ui64 low, ui64 hi, IOutputStream& out) { union { ui16 Dw[8]; // NOLINT(modernize-avoid-c-arrays) diff --git a/yql/essentials/types/uuid/uuid.h b/yql/essentials/types/uuid/uuid.h index 2eca102a6f8..45de33f6405 100644 --- a/yql/essentials/types/uuid/uuid.h +++ b/yql/essentials/types/uuid/uuid.h @@ -6,12 +6,16 @@ #include <cstring> #include <utility> +#include <util/generic/strbuf.h> + class IOutputStream; namespace NKikimr::NUuid { static constexpr ui32 UUID_LEN = 16; +TString UuidBytesToString(TStringBuf in); +void UuidBytesToString(TStringBuf in, IOutputStream& out); TString UuidBytesToString(const TString& in); void UuidBytesToString(const TString& in, IOutputStream& out); void UuidHalfsToString(ui64 low, ui64 hi, IOutputStream& out); |
