reuse arrow to ydb conversion types

3 files changed, 37 insertions, 26 deletions

diff --git a/ydb/services/ext_index/service/add_index.cpp b/ydb/services/ext_index/service/add_index.cpp
index d8c2c1dd519..c15eb50a738 100644
--- a/ydb/services/ext_index/service/add_index.cpp
+++ b/ydb/services/ext_index/service/add_index.cpp
@@ -55,33 +55,13 @@ void TIndexUpsertActor::Bootstrap() {
             return;
         }
         pkColumns.emplace_back(Data->GetColumnByName(i));
-        switch (f->type()->id()) {
-            case arrow::Type::INT64:
-                types->emplace_back(i, NMetadata::NInternal::TYDBType::Primitive(Ydb::Type::INT64));
-                break;
-            case arrow::Type::INT32:
-                types->emplace_back(i, NMetadata::NInternal::TYDBType::Primitive(Ydb::Type::INT32));
-                break;
-            case arrow::Type::STRING:
-                types->emplace_back(i, NMetadata::NInternal::TYDBType::Primitive(Ydb::Type::UTF8));
-                break;
-            case arrow::Type::BINARY:
-                types->emplace_back(i, NMetadata::NInternal::TYDBType::Primitive(Ydb::Type::STRING));
-                break;
-            case arrow::Type::UINT64:
-                types->emplace_back(i, NMetadata::NInternal::TYDBType::Primitive(Ydb::Type::UINT64));
-                break;
-            case arrow::Type::UINT32:
-                types->emplace_back(i, NMetadata::NInternal::TYDBType::Primitive(Ydb::Type::UINT32));
-                break;
-            case arrow::Type::TIMESTAMP:
-                types->emplace_back(i, NMetadata::NInternal::TYDBType::Primitive(Ydb::Type::TIMESTAMP));
-                break;
-            default:
-                ExternalController->OnIndexUpsertionFailed("incorrect type for pk column");
-                PassAway();
-                return;
+        auto ydbType = NMetadata::NInternal::TYDBType::ConvertArrowToYDB(f->type()->id());
+        if (!ydbType) {
+            ExternalController->OnIndexUpsertionFailed("incorrect arrow type for ydb field");
+            PassAway();
+            return;
         }
+        types->emplace_back(i, NMetadata::NInternal::TYDBType::Primitive(*ydbType));
     }
 
     const std::vector<ui64> hashes = IndexInfo.GetExtractor()->ExtractIndex(Data);
diff --git a/ydb/services/metadata/manager/ydb_value_operator.cpp b/ydb/services/metadata/manager/ydb_value_operator.cpp
index e448526475b..b66100971a4 100644
--- a/ydb/services/metadata/manager/ydb_value_operator.cpp
+++ b/ydb/services/metadata/manager/ydb_value_operator.cpp
@@ -173,6 +173,35 @@ Ydb::Type TYDBType::Primitive(const Ydb::Type::PrimitiveTypeId type) {
     return result;
 }
 
+std::optional<Ydb::Type::PrimitiveTypeId> TYDBType::ConvertArrowToYDB(const arrow::Type::type type) {
+    switch (type) {
+        case arrow::Type::INT64:
+            return Ydb::Type::INT64;
+        case arrow::Type::INT32:
+            return Ydb::Type::INT32;
+        case arrow::Type::INT16:
+            return Ydb::Type::INT16;
+        case arrow::Type::INT8:
+            return Ydb::Type::INT8;
+        case arrow::Type::STRING:
+            return Ydb::Type::UTF8;
+        case arrow::Type::BINARY:
+            return Ydb::Type::STRING;
+        case arrow::Type::UINT64:
+            return Ydb::Type::UINT64;
+        case arrow::Type::UINT32:
+            return Ydb::Type::UINT32;
+        case arrow::Type::UINT16:
+            return Ydb::Type::UINT16;
+        case arrow::Type::UINT8:
+            return Ydb::Type::UINT8;
+        case arrow::Type::TIMESTAMP:
+            return Ydb::Type::TIMESTAMP;
+        default:
+            return {};
+    }
+}
+
 std::optional<Ydb::Type::PrimitiveTypeId> TYDBType::ConvertYQLToYDB(const NScheme::TTypeId type) {
     switch (type) {
         case NScheme::NTypeIds::Int8:
diff --git a/ydb/services/metadata/manager/ydb_value_operator.h b/ydb/services/metadata/manager/ydb_value_operator.h
index c5de2d402d0..28983a16f4e 100644
--- a/ydb/services/metadata/manager/ydb_value_operator.h
+++ b/ydb/services/metadata/manager/ydb_value_operator.h
@@ -4,6 +4,7 @@
 #include <ydb/public/lib/scheme_types/scheme_type_id.h>
 #include <ydb/core/scheme_types/scheme_type_info.h>
 #include <util/generic/vector.h>
+#include <contrib/libs/apache/arrow/cpp/src/arrow/type_fwd.h>
 
 namespace NKikimr::NMetadata::NInternal {
 
@@ -21,6 +22,7 @@ public:
     static Ydb::Type Primitive(const Ydb::Type::PrimitiveTypeId type);
     static std::optional<NScheme::TTypeId> ConvertYDBToYQL(const Ydb::Type::PrimitiveTypeId type);
     static std::optional<Ydb::Type::PrimitiveTypeId> ConvertYQLToYDB(const NScheme::TTypeId type);
+    static std::optional<Ydb::Type::PrimitiveTypeId> ConvertArrowToYDB(const arrow::Type::type type);
 
     static std::optional<TVector<std::pair<TString, NScheme::TTypeInfo>>> ConvertYDBToYQL(const std::vector<std::pair<TString, Ydb::Type>>& input);
 };