YQL-20086 udfs

init
commit_hash:f9684778bf1ea956965f2360b80b91edb7d4ffbe

1 files changed, 116 insertions, 81 deletions

diff --git a/yql/essentials/udfs/common/python/bindings/py_lazy_mkql_dict.cpp b/yql/essentials/udfs/common/python/bindings/py_lazy_mkql_dict.cpp
index 1408f185a75..4ffaee2b66c 100644
--- a/yql/essentials/udfs/common/python/bindings/py_lazy_mkql_dict.cpp
+++ b/yql/essentials/udfs/common/python/bindings/py_lazy_mkql_dict.cpp
@@ -18,14 +18,16 @@ namespace {
 //////////////////////////////////////////////////////////////////////////////
 // TLazyDictBase
 //////////////////////////////////////////////////////////////////////////////
-class TLazyDictBase: public NUdf::TBoxedValue
-{
+class TLazyDictBase: public NUdf::TBoxedValue {
 protected:
     class TIterator: public NUdf::TBoxedValue {
     public:
         TIterator(const TPyCastContext::TPtr& ctx, const NUdf::TType* type, TPyObjectPtr&& pyIter)
-            : CastCtx_(ctx), ItemType_(type), PyIter_(std::move(pyIter))
-        {}
+            : CastCtx_(ctx)
+            , ItemType_(type)
+            , PyIter_(std::move(pyIter))
+        {
+        }
 
         ~TIterator() {
             const TPyGilLocker lock;
@@ -80,8 +82,12 @@ protected:
     class TPairIterator: public NUdf::TBoxedValue {
     public:
         TPairIterator(const TPyCastContext::TPtr& ctx, const NUdf::TType* keyType, const NUdf::TType* payType, TPyObjectPtr&& pyIter)
-            : CastCtx_(ctx),  KeyType_(keyType), PayType_(payType), PyIter_(std::move(pyIter))
-        {}
+            : CastCtx_(ctx)
+            , KeyType_(keyType)
+            , PayType_(payType)
+            , PyIter_(std::move(pyIter))
+        {
+        }
 
         ~TPairIterator() {
             const TPyGilLocker lock;
@@ -131,8 +137,11 @@ protected:
     };
 
     TLazyDictBase(const TPyCastContext::TPtr& castCtx, const NUdf::TType* itemType, PyObject* pyObject)
-        : CastCtx_(castCtx), ItemType_(itemType), PyObject_(pyObject, TPyObjectPtr::AddRef())
-    {}
+        : CastCtx_(castCtx)
+        , ItemType_(itemType)
+        , PyObject_(pyObject, TPyObjectPtr::AddRef())
+    {
+    }
 
     ~TLazyDictBase() {
         const TPyGilLocker lock;
@@ -146,8 +155,7 @@ protected:
             UdfTerminate((TStringBuilder() << CastCtx_->PyCtx->Pos << GetLastErrorAsString()).c_str());
         }
         return bool(has);
-    }
-    catch (const yexception& e) {
+    } catch (const yexception& e) {
         UdfTerminate((TStringBuilder() << CastCtx_->PyCtx->Pos << e.what()).c_str());
     }
 
@@ -159,15 +167,18 @@ protected:
 //////////////////////////////////////////////////////////////////////////////
 // TLazyMapping
 //////////////////////////////////////////////////////////////////////////////
-class TLazyMapping: public TLazyDictBase
-{
+class TLazyMapping: public TLazyDictBase {
 public:
     TLazyMapping(const TPyCastContext::TPtr& ctx, const NUdf::TType* keyType, const NUdf::TType* payType, PyObject* dict)
-        : TLazyDictBase(ctx, keyType, dict), PayType_(payType)
-    {}
+        : TLazyDictBase(ctx, keyType, dict)
+        , PayType_(payType)
+    {
+    }
 
 private:
-    bool IsSortedDict() const override { return false; }
+    bool IsSortedDict() const override {
+        return false;
+    }
 
     ui64 GetDictLength() const override try {
         const TPyGilLocker lock;
@@ -238,9 +249,7 @@ private:
         const TPyGilLocker lock;
         if (const TPyObjectPtr pyKey = ToPyObject(CastCtx_, ItemType_, key)) {
             const auto map = PyObject_.Get();
-            const auto has = map->ob_type->tp_as_sequence && map->ob_type->tp_as_sequence->sq_contains ?
-                (map->ob_type->tp_as_sequence->sq_contains)(map, pyKey.Get()) :
-                PyMapping_HasKey(map, pyKey.Get());
+            const auto has = map->ob_type->tp_as_sequence && map->ob_type->tp_as_sequence->sq_contains ? (map->ob_type->tp_as_sequence->sq_contains)(map, pyKey.Get()) : PyMapping_HasKey(map, pyKey.Get());
 
             if (has >= 0) {
                 return bool(has);
@@ -258,15 +267,18 @@ private:
 //////////////////////////////////////////////////////////////////////////////
 // TLazyDict
 //////////////////////////////////////////////////////////////////////////////
-class TLazyDict: public TLazyDictBase
-{
+class TLazyDict: public TLazyDictBase {
 public:
     TLazyDict(const TPyCastContext::TPtr& ctx, const NUdf::TType* keyType, const NUdf::TType* payType, PyObject* dict)
-        : TLazyDictBase(ctx, keyType, dict), PayType_(payType)
-    {}
+        : TLazyDictBase(ctx, keyType, dict)
+        , PayType_(payType)
+    {
+    }
 
 private:
-    bool IsSortedDict() const override { return false; }
+    bool IsSortedDict() const override {
+        return false;
+    }
 
     ui64 GetDictLength() const override try {
         const TPyGilLocker lock;
@@ -349,15 +361,17 @@ private:
 //////////////////////////////////////////////////////////////////////////////
 // TLazySet
 //////////////////////////////////////////////////////////////////////////////
-class TLazySet: public TLazyDictBase
-{
+class TLazySet: public TLazyDictBase {
 public:
     TLazySet(const TPyCastContext::TPtr& ctx, const NUdf::TType* itemType, PyObject* set)
         : TLazyDictBase(ctx, itemType, set)
-    {}
+    {
+    }
 
 private:
-    bool IsSortedDict() const override { return false; }
+    bool IsSortedDict() const override {
+        return false;
+    }
 
     ui64 GetDictLength() const override try {
         const TPyGilLocker lock;
@@ -425,15 +439,17 @@ private:
 //////////////////////////////////////////////////////////////////////////////
 // TLazySequenceAsSet
 //////////////////////////////////////////////////////////////////////////////
-class TLazySequenceAsSet: public TLazyDictBase
-{
+class TLazySequenceAsSet: public TLazyDictBase {
 public:
     TLazySequenceAsSet(const TPyCastContext::TPtr& ctx, const NUdf::TType* keyType, PyObject* sequence)
         : TLazyDictBase(ctx, keyType, sequence)
-    {}
+    {
+    }
 
 private:
-    bool IsSortedDict() const override { return false; }
+    bool IsSortedDict() const override {
+        return false;
+    }
 
     ui64 GetDictLength() const override try {
         const TPyGilLocker lock;
@@ -501,28 +517,31 @@ private:
 //////////////////////////////////////////////////////////////////////////////
 // TLazySequenceAsDict
 //////////////////////////////////////////////////////////////////////////////
-template<typename KeyType>
-class TLazySequenceAsDict: public NUdf::TBoxedValue
-{
+template <typename KeyType>
+class TLazySequenceAsDict: public NUdf::TBoxedValue {
 private:
     class TKeyIterator: public NUdf::TBoxedValue {
     public:
         TKeyIterator(Py_ssize_t size)
-            : Size_(size), Index_(0)
-        {}
+            : Size_(size)
+            , Index_(0)
+        {
+        }
 
     private:
         bool Skip() override {
-            if (Index_ >= Size_)
+            if (Index_ >= Size_) {
                 return false;
+            }
 
             ++Index_;
             return true;
         }
 
         bool Next(NUdf::TUnboxedValue& value) override {
-            if (Index_ >= Size_)
+            if (Index_ >= Size_) {
                 return false;
+            }
 
             value = NUdf::TUnboxedValuePod(KeyType(Index_++));
             return true;
@@ -536,8 +555,13 @@ private:
     class TIterator: public NUdf::TBoxedValue {
     public:
         TIterator(const TPyCastContext::TPtr& ctx, const NUdf::TType* itemType, Py_ssize_t size, const TPyObjectPtr& pySeq)
-            : CastCtx_(ctx),  ItemType_(itemType), PySeq_(pySeq), Size_(size), Index_(0)
-        {}
+            : CastCtx_(ctx)
+            , ItemType_(itemType)
+            , PySeq_(pySeq)
+            , Size_(size)
+            , Index_(0)
+        {
+        }
 
         ~TIterator() {
             const TPyGilLocker lock;
@@ -546,16 +570,18 @@ private:
 
     private:
         bool Skip() override {
-            if (Index_ >= Size_)
+            if (Index_ >= Size_) {
                 return false;
+            }
 
             ++Index_;
             return true;
         }
 
         bool Next(NUdf::TUnboxedValue& value) override try {
-            if (Index_ >= Size_)
+            if (Index_ >= Size_) {
                 return false;
+            }
 
             const TPyGilLocker lock;
             value = FromPyObject(CastCtx_, ItemType_, PySequence_Fast_GET_ITEM(PySeq_.Get(), Index_++));
@@ -565,8 +591,9 @@ private:
         }
 
         bool NextPair(NUdf::TUnboxedValue& key, NUdf::TUnboxedValue& pay) override try {
-            if (Index_ >= Size_)
+            if (Index_ >= Size_) {
                 return false;
+            }
 
             const TPyGilLocker lock;
             key = NUdf::TUnboxedValuePod(KeyType(Index_));
@@ -586,8 +613,12 @@ private:
 
 public:
     TLazySequenceAsDict(const TPyCastContext::TPtr& ctx, const NUdf::TType* itemType, TPyObjectPtr&& sequence, Py_ssize_t size)
-        : CastCtx_(ctx), ItemType_(itemType), Size_(size), PySeq_(std::move(sequence))
-    {}
+        : CastCtx_(ctx)
+        , ItemType_(itemType)
+        , Size_(size)
+        , PySeq_(std::move(sequence))
+    {
+    }
 
     ~TLazySequenceAsDict()
     {
@@ -596,7 +627,9 @@ public:
     }
 
 private:
-    bool IsSortedDict() const override { return true; }
+    bool IsSortedDict() const override {
+        return true;
+    }
 
     bool HasDictItems() const override {
         return Size_ > 0;
@@ -608,15 +641,17 @@ private:
 
     NUdf::TUnboxedValue Lookup(const NUdf::TUnboxedValuePod& key) const override {
         const Py_ssize_t index = key.Get<KeyType>();
-        if (index >= -Size_ && index < Size_) try {
-            const TPyGilLocker lock;
-            if (const auto item = PySequence_Fast_GET_ITEM(PySeq_.Get(), index >= 0 ? index : Size_ + index)) {
-                return FromPyObject(CastCtx_, ItemType_, item).Release().MakeOptional();
-            } else if (PyErr_Occurred()) {
-                UdfTerminate((TStringBuilder() << CastCtx_->PyCtx->Pos << GetLastErrorAsString()).c_str());
+        if (index >= -Size_ && index < Size_) {
+            try {
+                const TPyGilLocker lock;
+                if (const auto item = PySequence_Fast_GET_ITEM(PySeq_.Get(), index >= 0 ? index : Size_ + index)) {
+                    return FromPyObject(CastCtx_, ItemType_, item).Release().MakeOptional();
+                } else if (PyErr_Occurred()) {
+                    UdfTerminate((TStringBuilder() << CastCtx_->PyCtx->Pos << GetLastErrorAsString()).c_str());
+                }
+            } catch (const yexception& e) {
+                UdfTerminate((TStringBuilder() << CastCtx_->PyCtx->Pos << e.what()).c_str());
             }
-        } catch (const yexception& e) {
-            UdfTerminate((TStringBuilder() << CastCtx_->PyCtx->Pos << e.what()).c_str());
         }
         return NUdf::TUnboxedValue();
     }
@@ -644,60 +679,60 @@ private:
     TPyObjectPtr PySeq_;
 };
 
-} // namspace
+} // namespace
 
 NUdf::TUnboxedValue FromPyDict(
-        const TPyCastContext::TPtr& castCtx,
-        const NUdf::TType* keyType,
-        const NUdf::TType* payType,
-        PyObject* dict)
+    const TPyCastContext::TPtr& castCtx,
+    const NUdf::TType* keyType,
+    const NUdf::TType* payType,
+    PyObject* dict)
 {
     return NUdf::TUnboxedValuePod(new TLazyDict(castCtx, keyType, payType, dict));
 }
 
 NUdf::TUnboxedValue FromPyMapping(
-        const TPyCastContext::TPtr& castCtx,
-        const NUdf::TType* keyType,
-        const NUdf::TType* payType,
-        PyObject* map)
+    const TPyCastContext::TPtr& castCtx,
+    const NUdf::TType* keyType,
+    const NUdf::TType* payType,
+    PyObject* map)
 {
     return NUdf::TUnboxedValuePod(new TLazyMapping(castCtx, keyType, payType, map));
 }
 
 NUdf::TUnboxedValue FromPySet(
-        const TPyCastContext::TPtr& castCtx,
-        const NUdf::TType* itemType,
-        PyObject* set)
+    const TPyCastContext::TPtr& castCtx,
+    const NUdf::TType* itemType,
+    PyObject* set)
 {
     return NUdf::TUnboxedValuePod(new TLazySet(castCtx, itemType, set));
 }
 
 NUdf::TUnboxedValue FromPySequence(
-        const TPyCastContext::TPtr& castCtx,
-        const NUdf::TType* keyType,
-        PyObject* set)
+    const TPyCastContext::TPtr& castCtx,
+    const NUdf::TType* keyType,
+    PyObject* set)
 {
     return NUdf::TUnboxedValuePod(new TLazySequenceAsSet(castCtx, keyType, set));
 }
 
 NUdf::TUnboxedValue FromPySequence(
-        const TPyCastContext::TPtr& castCtx,
-        const NUdf::TType* itemType,
-        const NUdf::TDataTypeId keyType,
-        PyObject* sequence)
+    const TPyCastContext::TPtr& castCtx,
+    const NUdf::TType* itemType,
+    const NUdf::TDataTypeId keyType,
+    PyObject* sequence)
 {
     if (TPyObjectPtr fast = PySequence_Fast(sequence, "Can't get fast sequence.")) {
-    const auto size = PySequence_Fast_GET_SIZE(fast.Get());
-    if (size >= 0) {
-        switch (keyType) {
+        const auto size = PySequence_Fast_GET_SIZE(fast.Get());
+        if (size >= 0) {
+            switch (keyType) {
 #define MAKE_PRIMITIVE_TYPE_SIZE(type) \
-            case NUdf::TDataType<type>::Id: \
-                return NUdf::TUnboxedValuePod(new TLazySequenceAsDict<type>(castCtx, itemType, std::move(fast), size));
-            INTEGRAL_VALUE_TYPES(MAKE_PRIMITIVE_TYPE_SIZE)
+    case NUdf::TDataType<type>::Id:    \
+        return NUdf::TUnboxedValuePod(new TLazySequenceAsDict<type>(castCtx, itemType, std::move(fast), size));
+                INTEGRAL_VALUE_TYPES(MAKE_PRIMITIVE_TYPE_SIZE)
 #undef MAKE_PRIMITIVE_TYPE_SIZE
+            }
+            Y_ABORT("Invalid key type.");
         }
-        Y_ABORT("Invalid key type.");
-    }
     }
     UdfTerminate((TStringBuilder() << castCtx->PyCtx->Pos << GetLastErrorAsString()).c_str());
 }