Initial support for TAutoMap in block udfs

2 files changed, 18 insertions, 1 deletions

diff --git a/ydb/library/yql/public/udf/udf_type_builder.h b/ydb/library/yql/public/udf/udf_type_builder.h
index ef27298ef4..3a593045b3 100644
--- a/ydb/library/yql/public/udf/udf_type_builder.h
+++ b/ydb/library/yql/public/udf/udf_type_builder.h
@@ -855,6 +855,13 @@ struct TTypeBuilderHelper<TOptional<T>> {
     }
 };
 
+template <typename T>
+struct TTypeBuilderHelper<TAutoMap<T>> {
+    static TType* Build(const IFunctionTypeInfoBuilder& builder) {
+        return TTypeBuilderHelper<T>::Build(builder);
+    }
+};
+
 template <typename TKey, typename TValue>
 struct TTypeBuilderHelper<TDict<TKey, TValue>> {
     static TType* Build(const IFunctionTypeInfoBuilder& builder) {
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
index ad39f67fd1..abfe02c1c5 100644
--- 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
@@ -59,11 +59,21 @@ SIMPLE_UDF(TNormalize, TOptional<char*>(TOptional<char*>)) {
                 : TUnboxedValue();
 }
 
-SIMPLE_UDF(TGetScheme, char*(TAutoMap<char*>)) {
+BEGIN_SIMPLE_STRICT_ARROW_UDF(TGetScheme, char*(TAutoMap<char*>)) {
     const std::string_view url(args[0].AsStringRef());
     const std::string_view prefix(GetSchemePrefix(url));
     return valueBuilder->SubString(args[0], std::distance(url.begin(), prefix.begin()), prefix.size());
 }
+struct TGetSchemeKernelExec : public TUnaryKernelExec<TGetSchemeKernelExec> {
+    template <typename TSink>
+    static void Process(TBlockItem arg, const TSink& sink) {
+        const std::string_view url(arg.AsStringRef());
+        const std::string_view prefix(GetSchemePrefix(url));
+        const std::string_view scheme = url.substr(std::distance(url.begin(), prefix.begin()), prefix.size());
+        sink(TBlockItem(scheme));
+    }
+};
+END_SIMPLE_ARROW_UDF(TGetScheme, TGetSchemeKernelExec::Do);
 
 ARROW_UDF_SINGLE_STRING_FUNCTION_FOR_URL(TGetHost, GetOnlyHost)