diff --git a/ydb/library/yql/udfs/common/roaring/roaring.cpp b/ydb/library/yql/udfs/common/roaring/roaring.cpp index 9699f8c7a601..00ccf7e01d1a 100644 --- a/ydb/library/yql/udfs/common/roaring/roaring.cpp +++ b/ydb/library/yql/udfs/common/roaring/roaring.cpp @@ -7,6 +7,7 @@ #include #include +#include #include #include #include @@ -30,6 +31,11 @@ namespace { { } + TRoaringWrapper(roaring_bitmap_t* bitmap) + : Roaring(bitmap) + { + } + ~TRoaringWrapper() { roaring_bitmap_free(Roaring); } @@ -223,6 +229,46 @@ namespace { TSourcePosition Pos_; }; + class TRoaringFromUint32List: public TBoxedValue { + public: + TRoaringFromUint32List(TSourcePosition pos) + : Pos_(pos) + { + } + + static TStringRef Name() { + return TStringRef::Of("FromUint32List"); + } + + private: + TUnboxedValue Run(const IValueBuilder* valueBuilder, + const TUnboxedValuePod* args) const override { + Y_UNUSED(valueBuilder); + try { + auto *b = roaring_bitmap_create(); + + const auto vector = args[0]; + const auto* elements = vector.GetElements(); + if (elements) { + for (auto& value : TArrayRef{elements, vector.GetListLength()}) { + roaring_bitmap_add(b, value.Get()); + } + } else { + TUnboxedValue value; + const auto it = vector.GetListIterator(); + while (it.Next(value)) { + roaring_bitmap_add(b, value.Get()); + } + } + + return TUnboxedValuePod(new TRoaringWrapper(b)); + } catch (const std::exception& e) { + UdfTerminate((TStringBuilder() << Pos_ << " " << e.what()).data()); + } + } + TSourcePosition Pos_; + }; + class TRoaringSerialize: public TBoxedValue { public: TRoaringSerialize() { @@ -282,6 +328,7 @@ namespace { void GetAllFunctions(IFunctionsSink& sink) const final { sink.Add(TRoaringSerialize::Name()); sink.Add(TRoaringDeserialize::Name()); + sink.Add(TRoaringFromUint32List::Name()); sink.Add(TRoaringCardinality::Name()); @@ -312,6 +359,12 @@ namespace { if (!typesOnly) { builder.Implementation(new TRoaringDeserialize(builder.GetSourcePosition())); } + } else if (TRoaringFromUint32List::Name() == name) { + builder.Returns>().Args()->Add>(); + + if (!typesOnly) { + builder.Implementation(new TRoaringFromUint32List(builder.GetSourcePosition())); + } } else if (TRoaringSerialize::Name() == name) { builder.Returns(builder.SimpleType()) .Args() diff --git a/ydb/library/yql/udfs/common/roaring/test/canondata/test.test_serialize_deserialize_/results.txt b/ydb/library/yql/udfs/common/roaring/test/canondata/test.test_serialize_deserialize_/results.txt index 05f944c62233..1bebeec56030 100644 --- a/ydb/library/yql/udfs/common/roaring/test/canondata/test.test_serialize_deserialize_/results.txt +++ b/ydb/library/yql/udfs/common/roaring/test/canondata/test.test_serialize_deserialize_/results.txt @@ -172,5 +172,42 @@ ] } ] - } + }; + { + "Write" = [ + { + "Type" = [ + "ListType"; + [ + "StructType"; + [ + [ + "DeserializedList"; + [ + "OptionalType"; + [ + "ListType"; + [ + "DataType"; + "Uint32" + ] + ] + ] + ] + ] + ] + ]; + "Data" = [ + [ + [ + [ + "10"; + "567" + ] + ] + ] + ] + } + ] + }; ] \ No newline at end of file diff --git a/ydb/library/yql/udfs/common/roaring/test/cases/serialize_deserialize.sql b/ydb/library/yql/udfs/common/roaring/test/cases/serialize_deserialize.sql index 9735432181c9..e3151c71dff6 100644 --- a/ydb/library/yql/udfs/common/roaring/test/cases/serialize_deserialize.sql +++ b/ydb/library/yql/udfs/common/roaring/test/cases/serialize_deserialize.sql @@ -13,3 +13,6 @@ FROM Input; SELECT ListTake(ListSkip(Roaring::Uint32List(Roaring::Deserialize(binaryString)), 10), 1) AS EmptyList FROM Input; + +SELECT Roaring::Uint32List(Roaring::FromUint32List(AsList(10, 567))) AS DeserializedList +FROM Input; \ No newline at end of file