Merge 2f5a4ad into cc81c7e

fedor-miron · web-flow · commit c78925fda163 · 2024-07-08T14:24:30.000Z
diff --git a/ydb/library/yql/minikql/comp_nodes/ut/mkql_blocks_ut.cpp b/ydb/library/yql/minikql/comp_nodes/ut/mkql_blocks_ut.cpp
@@ -16,7 +16,8 @@ BEGIN_SIMPLE_ARROW_UDF(TInc, i32(i32)) {
 
 struct TIncKernelExec : public NYql::NUdf::TUnaryKernelExec<TIncKernelExec> {
     template <typename TSink>
-    static void Process(NYql::NUdf::TBlockItem arg, const TSink& sink) {
+    static void Process(NYql::NUdf::TBlockItem arg, const NYql::NUdf::IValueBuilder& valueBuilder, const TSink& sink) {
+        Y_UNUSED(valueBuilder);
         sink(NYql::NUdf::TBlockItem(arg.As<i32>() + 1));
     }
 };
diff --git a/ydb/library/yql/public/udf/arrow/block_builder.h b/ydb/library/yql/public/udf/arrow/block_builder.h
@@ -55,6 +55,7 @@ class IScalarBuilder {
 public:
     virtual ~IScalarBuilder() = default;
     virtual arrow::Datum Build(TBlockItem value) const = 0;
+    virtual arrow::Datum Build(NUdf::TUnboxedValuePod value) const = 0;
 };
 
 inline std::shared_ptr<arrow::DataType> GetArrowType(const ITypeInfoHelper& typeInfoHelper, const TType* type) {
diff --git a/ydb/library/yql/public/udf/arrow/block_item.h b/ydb/library/yql/public/udf/arrow/block_item.h
@@ -59,6 +59,19 @@ class TBlockItem {
         Raw.Halfs[1] = high;
     }
 
+    inline static TBlockItem Embedded(const TStringRef& value) {
+        UDF_VERIFY(value.Size() <= sizeof(TRawEmbeddedValue::Buffer));
+
+        TBlockItem v;
+        v.Raw.Embedded.Size = value.Size();
+        v.Raw.Embedded.Meta = static_cast<ui8>(EMarkers::Embedded);
+        if (v.Raw.Embedded.Size) {
+            std::memcpy(v.Raw.Embedded.Buffer, value.Data(), v.Raw.Embedded.Size);
+        }
+
+        return v;
+    }
+
     inline ui64 Low() const {
         return Raw.Halfs[0];
     }
diff --git a/ydb/library/yql/public/udf/arrow/block_reader.h b/ydb/library/yql/public/udf/arrow/block_reader.h
@@ -40,6 +40,8 @@ class TFixedSizeBlockReaderBase : public IBlockReader {
             if (IsNull(data, index)) {
                 return {};
             }
+        } else {
+            Y_DEBUG_ABORT_UNLESS(!data.MayHaveNulls());
         }
         return static_cast<TDerived*>(this)->MakeBlockItem(data.GetValues<T>(1)[index]);
     }
@@ -129,6 +131,8 @@ class TStringBlockReader final : public IBlockReader {
             if (IsNull(data, index)) {
                 return {};
             }
+        } else {
+            Y_DEBUG_ABORT_UNLESS(!data.MayHaveNulls());
         }
 
         const TOffset* offsets = data.GetValues<TOffset>(1);
@@ -210,6 +214,8 @@ class TTupleBlockReaderBase : public IBlockReader {
         if constexpr (Nullable) {
             if (IsNull(data, index)) {
                 return {};
+            } else {
+                Y_DEBUG_ABORT_UNLESS(!data.MayHaveNulls());
             }
         }
         return static_cast<TDerived*>(this)->GetChildrenItems(data, index);
diff --git a/ydb/library/yql/public/udf/arrow/udf_arrow_helpers.h b/ydb/library/yql/public/udf/arrow/udf_arrow_helpers.h
@@ -27,12 +27,12 @@ using TExec = arrow::Status(*)(arrow::compute::KernelContext*, const arrow::comp
 
 class TUdfKernelState : public arrow::compute::KernelState {
 public:
-    TUdfKernelState(const TVector<const TType*>& argTypes, const TType* outputType, bool onlyScalars, const ITypeInfoHelper* typeInfoHelper, const IPgBuilder& pgBuilder)
+    TUdfKernelState(const TVector<const TType*>& argTypes, const TType* outputType, bool onlyScalars, const ITypeInfoHelper* typeInfoHelper, const IValueBuilder* valueBuilder)
         : ArgTypes_(argTypes)
         , OutputType_(outputType)
         , OnlyScalars_(onlyScalars)
         , TypeInfoHelper_(typeInfoHelper)
-        , PgBuilder_(pgBuilder)
+        , ValueBuilder_(valueBuilder)
     {
         Readers_.resize(ArgTypes_.size());
     }
@@ -48,7 +48,7 @@ class TUdfKernelState : public arrow::compute::KernelState {
     IArrayBuilder& GetArrayBuilder() {
         Y_ENSURE(!OnlyScalars_);
         if (!ArrayBuilder_) {
-            ArrayBuilder_ = MakeArrayBuilder(*TypeInfoHelper_, OutputType_, *GetYqlMemoryPool(), TypeInfoHelper_->GetMaxBlockLength(OutputType_), &PgBuilder_);
+            ArrayBuilder_ = MakeArrayBuilder(*TypeInfoHelper_, OutputType_, *GetYqlMemoryPool(), TypeInfoHelper_->GetMaxBlockLength(OutputType_), &ValueBuilder_->GetPgBuilder());
         }
 
         return *ArrayBuilder_;
@@ -62,13 +62,18 @@ class TUdfKernelState : public arrow::compute::KernelState {
 
         return *ScalarBuilder_;
     }
+    
+    const IValueBuilder& GetValueBuilder() {
+        Y_ENSURE(ValueBuilder_);
+        return *ValueBuilder_;
+    }
 
 private:
     const TVector<const TType*> ArgTypes_;
     const TType* OutputType_;
     const bool OnlyScalars_;
     const ITypeInfoHelper* TypeInfoHelper_;
-    const IPgBuilder& PgBuilder_;  
+    const IValueBuilder* ValueBuilder_;
     TVector<std::unique_ptr<IBlockReader>> Readers_;
     std::unique_ptr<IArrayBuilder> ArrayBuilder_;
     std::unique_ptr<IScalarBuilder> ScalarBuilder_;
@@ -157,7 +162,7 @@ class TSimpleArrowUdfImpl : public TBoxedValue {
                 }
             }
 
-            TUdfKernelState kernelState(ArgTypes_, OutputType_, OnlyScalars_, TypeInfoHelper_.Get(), valueBuilder->GetPgBuilder());
+            TUdfKernelState kernelState(ArgTypes_, OutputType_, OnlyScalars_, TypeInfoHelper_.Get(), valueBuilder);
             arrow::compute::ExecContext execContext(GetYqlMemoryPool());
             arrow::compute::KernelContext kernelContext(&execContext);
             kernelContext.SetState(&kernelState);
@@ -351,7 +356,8 @@ TReader* CastToBlockReaderImpl(IBlockReader& reader) {
 
 template <typename TDerived, typename TReader = IBlockReader, typename TArrayBuilderImpl = IArrayBuilder, typename TScalarBuilderImpl = IScalarBuilder>
 struct TUnaryKernelExec {
-    static arrow::Status Do(arrow::compute::KernelContext* ctx, const arrow::compute::ExecBatch& batch, arrow::Datum* res) {
+
+static arrow::Status Do(arrow::compute::KernelContext* ctx, const arrow::compute::ExecBatch& batch, arrow::Datum* res) {
         auto& state = dynamic_cast<TUdfKernelState&>(*ctx->state());
         auto& reader = state.GetReader(0);
         auto* readerImpl = CastToBlockReaderImpl<TReader>(reader);
@@ -362,9 +368,10 @@ struct TUnaryKernelExec {
             auto* builderImpl = CastToScalarBuilderImpl<TScalarBuilderImpl>(builder);
 
             auto item = readerImpl->GetScalarItem(*arg.scalar());
-            TDerived::Process(item, [&](TBlockItem out) {
+            auto sink = [&](auto out) {
                 *res = builderImpl->Build(out);
-            });
+            };
+            TDerived::Process(item, state.GetValueBuilder(), sink);
         }
         else {
             auto& array = *arg.array();
@@ -377,9 +384,10 @@ struct TUnaryKernelExec {
             for (int64_t i = 0; i < array.length;) {
                 for (size_t j = 0; j < maxBlockLength && i < array.length; ++j, ++i) {
                     auto item = readerImpl->GetItem(array, i);
-                    TDerived::Process(item, [&](TBlockItem out) {
+                    auto sink = [&](auto out) {
                         builderImpl->Add(out);
-                    });
+                    };
+                    TDerived::Process(item, state.GetValueBuilder(), sink);
                 }
                 auto outputDatum = builderImpl->Build(false);
                 ForEachArrayData(outputDatum, [&](const auto& arr) { outputArrays.push_back(arr); });
@@ -411,9 +419,11 @@ struct TBinaryKernelExec {
 
             auto item1 = reader1Impl->GetScalarItem(*arg1.scalar());
             auto item2 = reader2Impl->GetScalarItem(*arg2.scalar());
-            TDerived::Process(item1, item2, [&](TBlockItem out) {
+
+            auto sink = [&](TBlockItem out) {
                 *res = builderImpl->Build(out);
-            });
+            };
+            TDerived::Process(item1, item2, state.GetValueBuilder(), sink);
         }
         else if (arg1.is_scalar() && arg2.is_array()) {
             auto item1 = reader1Impl->GetScalarItem(*arg1.scalar());
@@ -427,9 +437,11 @@ struct TBinaryKernelExec {
             for (int64_t i = 0; i < array2.length;) {
                 for (size_t j = 0; j < maxBlockLength && i < array2.length; ++j, ++i) {
                     auto item2 = reader2Impl->GetItem(array2, i);
-                    TDerived::Process(item1, item2, [&](TBlockItem out) {
+
+                    auto sink = [&](TBlockItem out) {
                         builderImpl->Add(out);
-                    });
+                    };
+                    TDerived::Process(item1, item2, state.GetValueBuilder(), sink);
                 }
                 auto outputDatum = builder.Build(false);
                 ForEachArrayData(outputDatum, [&](const auto& arr) { outputArrays.push_back(arr); });
@@ -448,9 +460,11 @@ struct TBinaryKernelExec {
             for (int64_t i = 0; i < array1.length;) {
                 for (size_t j = 0; j < maxBlockLength && i < array1.length; ++j, ++i) {
                     auto item1 = reader1Impl->GetItem(array1, i);
-                    TDerived::Process(item1, item2, [&](TBlockItem out) {
+
+                    auto sink = [&](TBlockItem out) {
                         builderImpl->Add(out);
-                    });
+                    };
+                    TDerived::Process(item1, item2, state.GetValueBuilder(), sink);
                 }
                 auto outputDatum = builder.Build(false);
                 ForEachArrayData(outputDatum, [&](const auto& arr) { outputArrays.push_back(arr); });
@@ -470,11 +484,13 @@ struct TBinaryKernelExec {
             Y_ENSURE(array1.length == array2.length);
             for (int64_t i = 0; i < array1.length;) {
                 for (size_t j = 0; j < maxBlockLength && i < array1.length; ++j, ++i) {
-                    auto item1 = reader1.GetItem(array1, i);
-                    auto item2 = reader2.GetItem(array2, i);
-                    TDerived::Process(item1, item2, [&](TBlockItem out) {
+                    auto item1 = reader1Impl->GetItem(array1, i);
+                    auto item2 = reader2Impl->GetItem(array2, i);
+
+                    auto sink = [&](TBlockItem out) {
                         builderImpl->Add(out);
-                    });
+                    };
+                    TDerived::Process(item1, item2, state.GetValueBuilder(), sink);
                 }
                 auto outputDatum = builder.Build(false);
                 ForEachArrayData(outputDatum, [&](const auto& arr) { outputArrays.push_back(arr); });
@@ -529,9 +545,11 @@ struct TGenericKernelExec {
                 auto& reader = state.GetReader(k);
                 args[k] = reader.GetScalarItem(*batch[k].scalar());
             }
-            TDerived::Process(items, [&](TBlockItem out) {
+            
+            auto sink = [&](TBlockItem out) {
                 *res = builderImpl->Build(out);
-            });
+            };
+            TDerived::Process(items, state.GetValueBuilder(), sink);
         } else {
             auto& builder = state.GetArrayBuilder();
             auto* builderImpl = CastToArrayBuilderImpl<TArrayBuilderImpl>(builder);
@@ -559,9 +577,11 @@ struct TGenericKernelExec {
 
                         args[k] = reader.GetItem(*batch[k].array(), i);
                     }
-                    TDerived::Process(items, [&](TBlockItem out) {
+
+                    auto sink = [&](TBlockItem out) {
                         builderImpl->Add(out);
-                    });
+                    };
+                    TDerived::Process(items, state.GetValueBuilder(), sink);
                 }
                 auto outputDatum = builderImpl->Build(false);
                 ForEachArrayData(outputDatum, [&](const auto& arr) { outputArrays.push_back(arr); });
diff --git a/ydb/library/yql/udfs/common/string/string_udf.cpp b/ydb/library/yql/udfs/common/string/string_udf.cpp
diff --git a/ydb/library/yql/udfs/common/url_base/lib/url_base_udf.h b/ydb/library/yql/udfs/common/url_base/lib/url_base_udf.h

Original file line number	Diff line number	Diff line change
`@@ -16,7 +16,8 @@ BEGIN_SIMPLE_ARROW_UDF(TInc, i32(i32)) {`
`16`	`16`
`17`	`17`	`struct TIncKernelExec : public NYql::NUdf::TUnaryKernelExec<TIncKernelExec> {`
`18`	`18`	`template <typename TSink>`
`19`		`- static void Process(NYql::NUdf::TBlockItem arg, const TSink& sink) {`
	`19`	`+ static void Process(NYql::NUdf::TBlockItem arg, const NYql::NUdf::IValueBuilder& valueBuilder, const TSink& sink) {`
	`20`	`+ Y_UNUSED(valueBuilder);`
`20`	`21`	`sink(NYql::NUdf::TBlockItem(arg.As<i32>() + 1));`
`21`	`22`	`}`
`22`	`23`	`};`
Original file line number	Diff line number	Diff line change
`@@ -40,6 +40,8 @@ class TFixedSizeBlockReaderBase : public IBlockReader {`
`40`	`40`	`if (IsNull(data, index)) {`
`41`	`41`	`return {};`
`42`	`42`	`}`
	`43`	`+ } else {`
	`44`	`+ Y_DEBUG_ABORT_UNLESS(!data.MayHaveNulls());`
`43`	`45`	`}`
`44`	`46`	`return static_cast<TDerived*>(this)->MakeBlockItem(data.GetValues<T>(1)[index]);`
`45`	`47`	`}`
`@@ -129,6 +131,8 @@ class TStringBlockReader final : public IBlockReader {`
`129`	`131`	`if (IsNull(data, index)) {`
`130`	`132`	`return {};`
`131`	`133`	`}`
	`134`	`+ } else {`
	`135`	`+ Y_DEBUG_ABORT_UNLESS(!data.MayHaveNulls());`
`132`	`136`	`}`
`133`	`137`
`134`	`138`	`const TOffset* offsets = data.GetValues<TOffset>(1);`
`@@ -210,6 +214,8 @@ class TTupleBlockReaderBase : public IBlockReader {`
`210`	`214`	`if constexpr (Nullable) {`
`211`	`215`	`if (IsNull(data, index)) {`
`212`	`216`	`return {};`
	`217`	`+ } else {`
	`218`	`+ Y_DEBUG_ABORT_UNLESS(!data.MayHaveNulls());`
`213`	`219`	`}`
`214`	`220`	`}`
`215`	`221`	`return static_cast<TDerived*>(this)->GetChildrenItems(data, index);`