Merge 4d5e821 into e1ce008

ydb/library/yql/udfs/common/roaring/roaring.cpp

@@ -7,6 +7,7 @@
 #include <contrib/libs/croaring/include/roaring/memory.h>
 #include <contrib/libs/croaring/include/roaring/roaring.h>
 
+#include <util/generic/array_ref.h>
 #include <util/generic/vector.h>
 #include <util/string/builder.h>
 #include <util/system/yassert.h>
@@ -30,6 +31,11 @@ namespace {
         {
         }
 
+        TRoaringWrapper(roaring_bitmap_t* bitmap)
+            : Roaring(bitmap)
+        {
+        }
+
         ~TRoaringWrapper() {
             roaring_bitmap_free(Roaring);
         }
@@ -105,6 +111,47 @@ namespace {
         }
     };
 
+    class TRoaringAndNotWithBinary: public TBoxedValue {
+    public:
+        TRoaringAndNotWithBinary() {
+        }
+
+        static TStringRef Name() {
+            return TStringRef::Of("AndNotWithBinary");
+        }
+
+    private:
+        TUnboxedValue Run(const IValueBuilder* valueBuilder,
+                          const TUnboxedValuePod* args) const override {
+            Y_UNUSED(valueBuilder);
+            auto binaryString = args[1].AsStringRef();
+            auto bitmap = DeserializePortable(binaryString);
+
+            roaring_bitmap_andnot_inplace(GetBitmapFromArg(args[0]), bitmap);
+            roaring_bitmap_free(bitmap);
+
+            return args[0];
+        }
+    };
+
+    class TRoaringAndNot: public TBoxedValue {
+    public:
+        TRoaringAndNot() {
+        }
+
+        static TStringRef Name() {
+            return TStringRef::Of("AndNot");
+        }
+
+    private:
+        TUnboxedValue Run(const IValueBuilder* valueBuilder,
+                          const TUnboxedValuePod* args) const override {
+            Y_UNUSED(valueBuilder);
+            roaring_bitmap_andnot_inplace(GetBitmapFromArg(args[0]), GetBitmapFromArg(args[1]));
+            return args[0];
+        }
+    };
+
     class TRoaringOr: public TBoxedValue {
     public:
         TRoaringOr() {
@@ -223,6 +270,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<ui32>());
+                    }
+                } else {
+                    TUnboxedValue value;
+                    const auto it = vector.GetListIterator();
+                    while (it.Next(value)) {
+                        roaring_bitmap_add(b, value.Get<ui32>());
+                    }
+                }
+
+                return TUnboxedValuePod(new TRoaringWrapper(b));
+            } catch (const std::exception& e) {
+                UdfTerminate((TStringBuilder() << Pos_ << " " << e.what()).data());
+            }
+        }
+        TSourcePosition Pos_;
+    };
+
     class TRoaringSerialize: public TBoxedValue {
     public:
         TRoaringSerialize() {
@@ -266,6 +353,25 @@ namespace {
         }
     };
 
+    class TRoaringRunOptimize: public TBoxedValue {
+    public:
+        TRoaringRunOptimize() {
+        }
+
+        static TStringRef Name() {
+            return TStringRef::Of("RunOptimize");
+        }
+
+    private:
+        TUnboxedValue Run(const IValueBuilder* valueBuilder,
+                          const TUnboxedValuePod* args) const override {
+            Y_UNUSED(valueBuilder);
+            auto bitmap = GetBitmapFromArg(args[0]);
+            roaring_bitmap_run_optimize(bitmap);
+            return args[0];
+        }
+    };
+
     class TRoaringModule: public IUdfModule {
     public:
         TRoaringModule() {
@@ -282,6 +388,7 @@ namespace {
         void GetAllFunctions(IFunctionsSink& sink) const final {
             sink.Add(TRoaringSerialize::Name());
             sink.Add(TRoaringDeserialize::Name());
+            sink.Add(TRoaringFromUint32List::Name());
 
             sink.Add(TRoaringCardinality::Name());
 
@@ -292,6 +399,11 @@ namespace {
 
             sink.Add(TRoaringAndWithBinary::Name());
             sink.Add(TRoaringAnd::Name());
+
+            sink.Add(TRoaringAndNotWithBinary::Name());
+            sink.Add(TRoaringAndNot::Name());
+
+            sink.Add(TRoaringRunOptimize::Name());
         }
 
         void CleanupOnTerminate() const final {
@@ -312,6 +424,12 @@ namespace {
                     if (!typesOnly) {
                         builder.Implementation(new TRoaringDeserialize(builder.GetSourcePosition()));
                     }
+                } else if (TRoaringFromUint32List::Name() == name) {
+                    builder.Returns<TResource<RoaringResourceName>>().Args()->Add<TListType<ui32>>();
+
+                    if (!typesOnly) {
+                        builder.Implementation(new TRoaringFromUint32List(builder.GetSourcePosition()));
+                    }
                 } else if (TRoaringSerialize::Name() == name) {
                     builder.Returns(builder.SimpleType<char*>())
                         .Args()
@@ -372,6 +490,32 @@ namespace {
                     if (!typesOnly) {
                         builder.Implementation(new TRoaringAnd());
                     }
+                } else if (TRoaringAndNotWithBinary::Name() == name) {
+                    builder.Returns<TResource<RoaringResourceName>>()
+                        .Args()
+                        ->Add<TAutoMap<TResource<RoaringResourceName>>>()
+                        .Add<TAutoMap<char*>>();
+
+                    if (!typesOnly) {
+                        builder.Implementation(new TRoaringAndNotWithBinary());
+                    }
+                } else if (TRoaringAndNot::Name() == name) {
+                    builder.Returns<TResource<RoaringResourceName>>()
+                        .Args()
+                        ->Add<TAutoMap<TResource<RoaringResourceName>>>()
+                        .Add<TAutoMap<TResource<RoaringResourceName>>>();
+
+                    if (!typesOnly) {
+                        builder.Implementation(new TRoaringAndNot());
+                    }
+                } else if (TRoaringRunOptimize::Name() == name) {
+                    builder.Returns<TResource<RoaringResourceName>>()
+                        .Args()
+                        ->Add<TAutoMap<TResource<RoaringResourceName>>>();
+
+                    if (!typesOnly) {
+                        builder.Implementation(new TRoaringRunOptimize());
+                    }
                 } else {
                     TStringBuilder sb;
                     sb << "Unknown function: " << name.Data();

ydb/library/yql/udfs/common/roaring/test/canondata/test.test_intersect_/results.txt

@@ -102,5 +102,41 @@
                 ]
             }
         ]
+    };
+    {
+        "Write" = [
+            {
+                "Type" = [
+                    "ListType";
+                    [
+                        "StructType";
+                        [
+                            [
+                                "AndNotList";
+                                [
+                                    "OptionalType";
+                                    [
+                                        "ListType";
+                                        [
+                                            "DataType";
+                                            "Uint32"
+                                        ]
+                                    ]
+                                ]
+                            ]
+                        ]
+                    ]
+                ];
+                "Data" = [
+                    [
+                        [
+                            [
+                                "42"
+                            ]
+                        ]
+                    ]
+                ]
+            }
+        ]
     }
 ]

ydb/library/yql/udfs/common/roaring/test/canondata/test.test_run_optimize_/results.txt

@@ -0,0 +1,35 @@
+[
+    {
+        "Write" = [
+            {
+                "Type" = [
+                    "ListType";
+                    [
+                        "StructType";
+                        [
+                            [
+                                "RunOptimizeList";
+                                [
+                                    "ListType";
+                                    [
+                                        "DataType";
+                                        "Uint32"
+                                    ]
+                                ]
+                            ]
+                        ]
+                    ]
+                ];
+                "Data" = [
+                    [
+                        [
+                            "10";
+                            "42";
+                            "567"
+                        ]
+                    ]
+                ]
+            }
+        ]
+    }
+]

ydb/library/yql/udfs/common/roaring/test/canondata/test.test_serialize_deserialize_/results.txt

@@ -172,5 +172,38 @@
                 ]
             }
         ]
+    };
+    {
+        "Write" = [
+            {
+                "Type" = [
+                    "ListType";
+                    [
+                        "StructType";
+                        [
+                            [
+                                "DeserializedList";
+                                [
+                                    "ListType";
+                                    [
+                                        "DataType";
+                                        "Uint32"
+                                    ]
+                                ]
+                            ]
+                        ]
+                    ]
+                ];
+                "Data" = [
+                    [
+                        [
+                            "10";
+                            "42";
+                            "567"
+                        ]
+                    ]
+                ]
+            }
+        ]
     }
 ]

ydb/library/yql/udfs/common/roaring/test/cases/intersect.sql

@@ -1,3 +1,5 @@
 SELECT Roaring::Uint32List(Roaring::And(Roaring::Deserialize(left), Roaring::Deserialize(right))) AS AndList FROM Input;
 SELECT Roaring::Uint32List(Roaring::AndWithBinary(Roaring::Deserialize(right), left)) AS AndWithBinaryList FROM Input;
 SELECT Roaring::Uint32List(Roaring::AndWithBinary(Roaring::Deserialize(right), NULL)) AS AndWithBinaryListEmpty FROM Input;
+
+SELECT Roaring::Uint32List(Roaring::AndNot(Roaring::FromUint32List(AsList(1, 10, 42)), Roaring::FromUint32List(AsList(10, 1)))) AS AndNotList;

ydb/library/yql/udfs/common/roaring/test/cases/run_optimize.sql

@@ -0,0 +1 @@
+SELECT Roaring::Uint32List(Roaring::RunOptimize(Roaring::FromUint32List(AsList(10, 567, 42)))) AS RunOptimizeList;

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, 42))) AS DeserializedList
+FROM Input;