allow int8 int16 uint16 types to be used in db columns

6 files changed, 211 insertions, 8 deletions

diff --git a/ydb/core/kqp/ut/kqp_scheme_ut.cpp b/ydb/core/kqp/ut/kqp_scheme_ut.cpp
index e72501799f8..7e1d58f3693 100644
--- a/ydb/core/kqp/ut/kqp_scheme_ut.cpp
+++ b/ydb/core/kqp/ut/kqp_scheme_ut.cpp
@@ -1,4 +1,5 @@
 #include <ydb/core/kqp/ut/common/kqp_ut_common.h>
+#include <ydb/core/formats/arrow_helpers.h>
 #include <ydb/public/sdk/cpp/client/ydb_proto/accessor.h>
 #include <ydb/public/sdk/cpp/client/ydb_scheme/scheme.h>
 #include <ydb/public/sdk/cpp/client/ydb_topic/topic.h>
@@ -3315,6 +3316,169 @@ Y_UNIT_TEST_SUITE(KqpScheme) {
         UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
     }
 #endif
+
+    Y_UNIT_TEST(Int8Int16) {
+        TKikimrRunner kikimr;
+
+        auto db = kikimr.GetTableClient();
+        auto session = db.CreateSession().GetValueSync().GetSession();
+
+        TString tableName("/Root/Types");
+        auto createTable = Sprintf(R"(
+            CREATE TABLE `%s` (
+                CUint8 Uint8,
+                CInt8 Int8,
+                CUint16 Uint16,
+                CInt16 Int16,
+                PRIMARY KEY (CUint8, CInt8, CUint16, CInt16))
+        )", tableName.c_str());
+
+        auto result = session.ExecuteSchemeQuery(createTable).ExtractValueSync();
+        UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+
+        NYdb::TValueBuilder rows;
+        rows.BeginList();
+        for (size_t i = 0; i < 10; ++i) {
+            rows.AddListItem()
+                .BeginStruct()
+                .AddMember("CUint8").Uint8(i)
+                .AddMember("CInt8").Int8(i)
+                .AddMember("CUint16").Uint16(i)
+                .AddMember("CInt16").Int16(i)
+                .EndStruct();
+        }
+        rows.EndList();
+
+        result = db.BulkUpsert(tableName, rows.Build()).GetValueSync();
+        UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+
+        auto readSettings = TReadTableSettings()
+            .AppendColumns("CUint8")
+            .AppendColumns("CInt8")
+            .AppendColumns("CUint16")
+            .AppendColumns("CInt16");
+
+        auto it = session.ReadTable(tableName, readSettings).GetValueSync();
+        UNIT_ASSERT_C(it.IsSuccess(), result.GetIssues().ToString());
+
+        bool eos = false;
+        while (!eos) {
+            auto part = it.ReadNext().ExtractValueSync();
+            if (!part.IsSuccess()) {
+                eos = true;
+                UNIT_ASSERT_C(part.EOS(), result.GetIssues().ToString());
+                continue;
+            }
+            auto resultSet = part.ExtractPart();
+            TResultSetParser parser(resultSet);
+            for (size_t i = 0; parser.TryNextRow(); ++i) {
+                {
+                    auto& c = parser.ColumnParser("CUint8");
+                    UNIT_ASSERT_VALUES_EQUAL(i, *c.GetOptionalUint8().Get());
+                }
+                {
+                    auto& c = parser.ColumnParser("CInt8");
+                    UNIT_ASSERT_VALUES_EQUAL(i, *c.GetOptionalInt8().Get());
+                }
+                {
+                    auto& c = parser.ColumnParser("CUint16");
+                    UNIT_ASSERT_VALUES_EQUAL(i, *c.GetOptionalUint16().Get());
+                }
+                {
+                    auto& c = parser.ColumnParser("CInt16");
+                    UNIT_ASSERT_VALUES_EQUAL(i, *c.GetOptionalInt16().Get());
+                }
+            }
+        }
+
+        session.Close().GetValueSync();
+    }
+
+    Y_UNIT_TEST(Int8Int16Olap) {
+        TKikimrRunner kikimr;
+
+        auto db = kikimr.GetTableClient();
+        auto session = db.CreateSession().GetValueSync().GetSession();
+
+        TString tableName("/Root/Types");
+        auto createTable = Sprintf(R"(
+            CREATE TABLE `%s` (
+                CUint8 Uint8 NOT NULL,
+                CInt8 Int8 NOT NULL,
+                CUint16 Uint16 NOT NULL,
+                CInt16 Int16 NOT NULL,
+                PRIMARY KEY (CUint8, CInt8, CUint16, CInt16))
+            PARTITION BY HASH(CUint8, CInt8, CUint16, CInt16)
+            WITH (
+                STORE = COLUMN,
+                AUTO_PARTITIONING_MIN_PARTITIONS_COUNT = 1
+            )
+        )", tableName.c_str());
+
+        auto result = session.ExecuteSchemeQuery(createTable).ExtractValueSync();
+        UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+
+        auto schema = std::make_shared<arrow::Schema>(
+            std::vector<std::shared_ptr<arrow::Field>>{
+                arrow::field("CUint8", arrow::uint8()),
+                arrow::field("CInt8", arrow::int8()),
+                arrow::field("CUint16", arrow::uint16()),
+                arrow::field("CInt16", arrow::int16())
+            });
+
+        size_t rowsCount = 10;
+        auto builders = NArrow::MakeBuilders(schema, rowsCount);
+        for (size_t i = 0; i < rowsCount; ++i) {
+            Y_VERIFY(NArrow::Append<arrow::UInt8Type>(*builders[0], i));
+            Y_VERIFY(NArrow::Append<arrow::Int8Type>(*builders[1], i));
+            Y_VERIFY(NArrow::Append<arrow::UInt16Type>(*builders[2], i));
+            Y_VERIFY(NArrow::Append<arrow::Int16Type>(*builders[3], i));
+        }
+        auto batch = arrow::RecordBatch::Make(schema, rowsCount, NArrow::Finish(std::move(builders)));
+
+        TString strSchema = NArrow::SerializeSchema(*schema);
+        TString strBatch = NArrow::SerializeBatchNoCompression(batch);
+
+        result = db.BulkUpsert(tableName, NYdb::NTable::EDataFormat::ApacheArrow,
+            strBatch, strSchema).GetValueSync();
+        UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+
+        auto scan = Sprintf("SELECT * FROM `%s` ORDER BY CUint8", tableName.c_str());
+        auto it = db.StreamExecuteScanQuery(scan).GetValueSync();
+        UNIT_ASSERT_C(it.IsSuccess(), result.GetIssues().ToString());
+
+        bool eos = false;
+        while (!eos) {
+            auto part = it.ReadNext().ExtractValueSync();
+            if (!part.IsSuccess()) {
+                eos = true;
+                UNIT_ASSERT_C(part.EOS(), result.GetIssues().ToString());
+                continue;
+            }
+            auto resultSet = part.ExtractResultSet();
+            TResultSetParser parser(resultSet);
+            for (size_t i = 0; parser.TryNextRow(); ++i) {
+                {
+                    auto& c = parser.ColumnParser("CUint8");
+                    UNIT_ASSERT_VALUES_EQUAL(i, *c.GetOptionalUint8().Get());
+                }
+                {
+                    auto& c = parser.ColumnParser("CInt8");
+                    UNIT_ASSERT_VALUES_EQUAL(i, *c.GetOptionalInt8().Get());
+                }
+                {
+                    auto& c = parser.ColumnParser("CUint16");
+                    UNIT_ASSERT_VALUES_EQUAL(i, *c.GetOptionalUint16().Get());
+                }
+                {
+                    auto& c = parser.ColumnParser("CInt16");
+                    UNIT_ASSERT_VALUES_EQUAL(i, *c.GetOptionalInt16().Get());
+                }
+            }
+        }
+
+        session.Close().GetValueSync();
+    }
 }
 
 } // namespace NKqp
diff --git a/ydb/core/scheme/scheme_tablecell_ut.cpp b/ydb/core/scheme/scheme_tablecell_ut.cpp
index 0b01da964b3..4ff0b48c95f 100644
--- a/ydb/core/scheme/scheme_tablecell_ut.cpp
+++ b/ydb/core/scheme/scheme_tablecell_ut.cpp
@@ -294,6 +294,22 @@ Y_UNIT_TEST_SUITE(Scheme) {
         for (NScheme::TTypeId typeId : yqlIds) {
             NScheme::TTypeInfo typeInfo(typeId);
             switch (typeId) {
+            case NScheme::NTypeIds::Int8:
+                GetValueHash(typeInfo, TCell(charArr, sizeof(i8)));
+                CompareTypedCells(TCell(charArr, sizeof(i8)), TCell(charArr, sizeof(i8)), typeInfo);
+                break;
+            case NScheme::NTypeIds::Uint8:
+                GetValueHash(typeInfo, TCell(charArr, sizeof(ui8)));
+                CompareTypedCells(TCell(charArr, sizeof(ui8)), TCell(charArr, sizeof(ui8)), typeInfo);
+                break;
+            case NScheme::NTypeIds::Int16:
+                GetValueHash(typeInfo, TCell(charArr, sizeof(i16)));
+                CompareTypedCells(TCell(charArr, sizeof(i16)), TCell(charArr, sizeof(i16)), typeInfo);
+                break;
+            case NScheme::NTypeIds::Uint16:
+                GetValueHash(typeInfo, TCell(charArr, sizeof(ui16)));
+                CompareTypedCells(TCell(charArr, sizeof(ui16)), TCell(charArr, sizeof(ui16)), typeInfo);
+                break;
             case NScheme::NTypeIds::Int32:
                 GetValueHash(typeInfo, TCell(charArr, sizeof(i32)));
                 CompareTypedCells(TCell(charArr, sizeof(i32)), TCell(charArr, sizeof(i32)), typeInfo);
@@ -310,10 +326,6 @@ Y_UNIT_TEST_SUITE(Scheme) {
                 GetValueHash(typeInfo, TCell(charArr, sizeof(ui64)));
                 CompareTypedCells(TCell(charArr, sizeof(ui64)), TCell(charArr, sizeof(ui64)), typeInfo);
                 break;
-            case NScheme::NTypeIds::Byte:
-                GetValueHash(typeInfo, TCell(charArr, sizeof(ui8)));
-                CompareTypedCells(TCell(charArr, sizeof(ui8)), TCell(charArr, sizeof(ui8)), typeInfo);
-                break;
             case NScheme::NTypeIds::Bool:
                 GetValueHash(typeInfo, TCell(charArr, sizeof(ui8)));
                 CompareTypedCells(TCell(charArr, sizeof(ui8)), TCell(charArr, sizeof(ui8)), typeInfo);
diff --git a/ydb/core/scheme_types/scheme_type_registry.cpp b/ydb/core/scheme_types/scheme_type_registry.cpp
index d737254361e..0e41da264f1 100644
--- a/ydb/core/scheme_types/scheme_type_registry.cpp
+++ b/ydb/core/scheme_types/scheme_type_registry.cpp
@@ -14,13 +14,16 @@ namespace NScheme {
 TTypeRegistry::TTypeRegistry()
 {
     // move to 'init defaults?'
+    RegisterType<TInt8>();
+    RegisterType<TUint8>();
+    RegisterType<TInt16>();
+    RegisterType<TUint16>();
     RegisterType<TInt32>();
     RegisterType<TUint32>();
     RegisterType<TInt64>();
     RegisterType<TUint64>();
     RegisterType<TDouble>();
     RegisterType<TFloat>();
-    RegisterType<TUint8>();
     RegisterType<TBool>();
     RegisterType<TPairUi64Ui64>();
     RegisterType<TString>();
diff --git a/ydb/core/scheme_types/scheme_types_defs.cpp b/ydb/core/scheme_types/scheme_types_defs.cpp
index eeaba2e9af0..363daeea073 100644
--- a/ydb/core/scheme_types/scheme_types_defs.cpp
+++ b/ydb/core/scheme_types/scheme_types_defs.cpp
@@ -7,11 +7,14 @@ namespace NKikimr {
 namespace NScheme {
 
 namespace NNames {
+    DECLARE_TYPED_TYPE_NAME(Int8);
+    DECLARE_TYPED_TYPE_NAME(Uint8);
+    DECLARE_TYPED_TYPE_NAME(Int16);
+    DECLARE_TYPED_TYPE_NAME(Uint16);
     DECLARE_TYPED_TYPE_NAME(Int32);
     DECLARE_TYPED_TYPE_NAME(Uint32);
     DECLARE_TYPED_TYPE_NAME(Int64);
     DECLARE_TYPED_TYPE_NAME(Uint64);
-    DECLARE_TYPED_TYPE_NAME(Uint8);
     DECLARE_TYPED_TYPE_NAME(Bool);
     DECLARE_TYPED_TYPE_NAME(Double);
     DECLARE_TYPED_TYPE_NAME(Float);
diff --git a/ydb/core/scheme_types/scheme_types_defs.h b/ydb/core/scheme_types/scheme_types_defs.h
index 590f218056e..398bcffc906 100644
--- a/ydb/core/scheme_types/scheme_types_defs.h
+++ b/ydb/core/scheme_types/scheme_types_defs.h
@@ -84,13 +84,28 @@ class TUint64 : public IIntegerTypeWithKeyString<ui64, NTypeIds::Uint64, NNames:
 // upyachka to get around undefined tryfromstring for chars
 
 namespace NNames {
+    extern const char Int8[5];
     extern const char Uint8[6];
+    extern const char Int16[6];
+    extern const char Uint16[7];
 }
 
+class TInt8 : public TTypedType<i8, TInt8, NTypeIds::Int8, NNames::Int8> {
+public:
+};
+
 class TUint8 : public TTypedType<ui8, TUint8, NTypeIds::Uint8, NNames::Uint8> {
 public:
 };
 
+class TInt16 : public TTypedType<i16, TInt16, NTypeIds::Int16, NNames::Int16> {
+public:
+};
+
+class TUint16 : public TTypedType<ui16, TUint16, NTypeIds::Uint16, NNames::Uint16> {
+public:
+};
+
 namespace NNames {
     extern const char Bool[5];
 }
@@ -233,11 +248,14 @@ class TInterval : public IIntegerTypeWithKeyString<i64, NTypeIds::Interval, NNam
 #define KIKIMR_FOREACH_MINIKQL_TYPE_I(name, size, macro, ...) macro(name, T##name, __VA_ARGS__)
 
 #define KIKIMR_FOREACH_MINIKQL_TYPE(xx, ...) \
+    xx(Int8, TInt8, __VA_ARGS__) \
+    xx(Uint8, TUint8, __VA_ARGS__) \
+    xx(Int16, TInt16, __VA_ARGS__) \
+    xx(Uint16, TUint16, __VA_ARGS__) \
     xx(Int32, TInt32, __VA_ARGS__) \
     xx(Uint32, TUint32, __VA_ARGS__) \
     xx(Int64, TInt64, __VA_ARGS__) \
     xx(Uint64, TUint64, __VA_ARGS__) \
-    xx(Uint8, TUint8, __VA_ARGS__) \
     xx(Bool, TBool, __VA_ARGS__) \
     xx(Double, TDouble, __VA_ARGS__) \
     xx(Float, TFloat, __VA_ARGS__) \
diff --git a/ydb/public/lib/scheme_types/scheme_type_id.h b/ydb/public/lib/scheme_types/scheme_type_id.h
index d18e4475f01..8c2aa7d2d50 100644
--- a/ydb/public/lib/scheme_types/scheme_type_id.h
+++ b/ydb/public/lib/scheme_types/scheme_type_id.h
@@ -56,11 +56,14 @@ static constexpr TTypeId Decimal = NYql::NProto::Decimal;
 static constexpr TTypeId Pg = 0x3000;
 
 static constexpr TTypeId YqlIds[] = {
+    Int8,
+    Uint8,
+    Int16,
+    Uint16,
     Int32,
     Uint32,
     Int64,
     Uint64,
-    Byte,
     Bool,
     Double,
     Float,