Add decimal tests for columnshard (#8959)

6 files changed, 228 insertions, 12 deletions

diff --git a/.github/config/muted_ya.txt b/.github/config/muted_ya.txt
index 2e772972ae1..c22e363f669 100644
--- a/.github/config/muted_ya.txt
+++ b/.github/config/muted_ya.txt
@@ -340,6 +340,8 @@ ydb/tests/postgres_integrations/go-libpq  docker_wrapper_test.py.test_pg_generat
 ydb/tests/tools/pq_read/test test_timeout.py.TestTimeout.test_timeout
 ydb/core/kqp/ut/query KqpStats.SysViewClientLost
 ydb/core/kqp/ut/olap KqpOlap.ManyColumnShards
+ydb/core/kqp/ut/olap KqpDecimalColumnShard.TestFilterCompare
+ydb/core/kqp/ut/olap KqpDecimalColumnShard.TestAggregation
 ydb/public/sdk/cpp/client/ydb_topic/ut TxUsage.WriteToTopic_Invalid_Session
 ydb/library/actors/interconnect/ut_huge_cluster HugeCluster.AllToAll
 ydb/tests/functional/hive test_drain.py.TestHive.test_drain_on_stop
diff --git a/ydb/core/grpc_services/rpc_load_rows.cpp b/ydb/core/grpc_services/rpc_load_rows.cpp
index e780de7eea0..4b5cab1c5b1 100644
--- a/ydb/core/grpc_services/rpc_load_rows.cpp
+++ b/ydb/core/grpc_services/rpc_load_rows.cpp
@@ -75,9 +75,12 @@ bool ConvertArrowToYdbPrimitive(const arrow::DataType& type, Ydb::Type& toType)
         case arrow::Type::DURATION:
             toType.set_type_id(Ydb::Type::INTERVAL);
             return true;
-        case arrow::Type::DECIMAL:
-            // TODO
-            return false;
+        case arrow::Type::DECIMAL: {
+            Ydb::DecimalType* decimalType = toType.mutable_decimal_type();
+            decimalType->set_precision(22);
+            decimalType->set_scale(9);
+            return true;
+        }
         case arrow::Type::NA:
         case arrow::Type::HALF_FLOAT:
         case arrow::Type::FIXED_SIZE_BINARY:
diff --git a/ydb/core/kqp/ut/common/columnshard.cpp b/ydb/core/kqp/ut/common/columnshard.cpp
index 955e260b598..6318a8f1e7e 100644
--- a/ydb/core/kqp/ut/common/columnshard.cpp
+++ b/ydb/core/kqp/ut/common/columnshard.cpp
@@ -246,6 +246,8 @@ namespace NKqp {
             return arrow::field(name, arrow::int64(), nullable);
         case NScheme::NTypeIds::JsonDocument:
             return arrow::field(name, arrow::binary(), nullable);
+        case NScheme::NTypeIds::Decimal:
+            return arrow::field(name, arrow::decimal(22, 9));
         case NScheme::NTypeIds::Pg:
             switch (NPg::PgTypeIdFromTypeDesc(typeDesc)) {
                 case INT2OID:
diff --git a/ydb/core/kqp/ut/olap/decimal_ut.cpp b/ydb/core/kqp/ut/olap/decimal_ut.cpp
new file mode 100644
index 00000000000..151226bd01d
--- /dev/null
+++ b/ydb/core/kqp/ut/olap/decimal_ut.cpp
@@ -0,0 +1,198 @@
+#include <ydb/core/formats/arrow/arrow_helpers.h>
+#include <ydb/core/kqp/ut/common/columnshard.h>
+#include <ydb/core/kqp/ut/common/kqp_ut_common.h>
+#include <ydb/core/testlib/common_helper.h>
+#include <ydb/core/testlib/cs_helper.h>
+#include <ydb/core/tx/columnshard/hooks/testing/controller.h>
+#include <ydb/core/tx/tx_proxy/proxy.h>
+
+#include <ydb/library/binary_json/write.h>
+#include <ydb/library/uuid/uuid.h>
+#include <ydb/public/sdk/cpp/client/draft/ydb_replication.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>
+
+#include <library/cpp/threading/local_executor/local_executor.h>
+#include <util/generic/serialized_enum.h>
+#include <util/string/printf.h>
+
+namespace NKikimr {
+namespace NKqp {
+
+using namespace NYdb;
+using namespace NYdb::NTable;
+
+Y_UNIT_TEST_SUITE(KqpDecimalColumnShard) {
+    class TDecimalTestCase {
+    public:
+        TDecimalTestCase()
+            : TestHelper(TKikimrSettings().SetWithSampleTables(false)) {
+        }
+
+        TTestHelper::TUpdatesBuilder Inserter() {
+            return TTestHelper::TUpdatesBuilder(TestTable.GetArrowSchema(Schema));
+        }
+
+        void Upsert(TTestHelper::TUpdatesBuilder& inserter) {
+            TestHelper.BulkUpsert(TestTable, inserter);
+        }
+
+        void CheckQuery(const TString& query, const TString& expected) {
+            TestHelper.ReadData(query, expected);
+        }
+
+        void PrepareTable1() {
+            Schema = {
+                TTestHelper::TColumnSchema().SetName("id").SetType(NScheme::NTypeIds::Int32).SetNullable(false),
+                TTestHelper::TColumnSchema().SetName("int").SetType(NScheme::NTypeIds::Int64),
+                TTestHelper::TColumnSchema().SetName("dec").SetType(NScheme::NTypeIds::Decimal),
+            };
+            TestTable.SetName("/Root/Table1").SetPrimaryKey({ "id" }).SetSharding({ "id" }).SetSchema(Schema);
+            TestHelper.CreateTable(TestTable);
+
+            {
+                TTestHelper::TUpdatesBuilder inserter = Inserter();
+                inserter.AddRow().Add(1).Add(4).Add(TDecimalValue("3.14"));
+                inserter.AddRow().Add(2).Add(3).Add(TDecimalValue("8.16"));
+                Upsert(inserter);
+            }
+            {
+                TTestHelper::TUpdatesBuilder inserter = Inserter();
+                inserter.AddRow().Add(4).Add(1).Add(TDecimalValue("12.46"));
+                inserter.AddRow().Add(3).Add(2).Add(TDecimalValue("8.492"));
+
+                Upsert(inserter);
+            }
+        }
+
+        void PrepareTable2() {
+            Schema = {
+                TTestHelper::TColumnSchema().SetName("id").SetType(NScheme::NTypeIds::Int32).SetNullable(false),
+                TTestHelper::TColumnSchema().SetName("table1_id").SetType(NScheme::NTypeIds::Int64),
+                TTestHelper::TColumnSchema().SetName("dec").SetType(NScheme::NTypeIds::Decimal),
+            };
+            TestTable.SetName("/Root/Table2").SetPrimaryKey({ "id" }).SetSharding({ "id" }).SetSchema(Schema);
+            TestHelper.CreateTable(TestTable);
+
+            {
+                TTestHelper::TUpdatesBuilder inserter = Inserter();
+                inserter.AddRow().Add(1).Add(1).Add(TDecimalValue("12.46"));
+                inserter.AddRow().Add(2).Add(1).Add(TDecimalValue("8.16"));
+                inserter.AddRow().Add(3).Add(2).Add(TDecimalValue("12.46"));
+                inserter.AddRow().Add(4).Add(2).Add(TDecimalValue("8.16"));
+                Upsert(inserter);
+            }
+        }
+
+    private:
+        TTestHelper TestHelper;
+
+        TVector<TTestHelper::TColumnSchema> Schema;
+        TTestHelper::TColumnTable TestTable;
+    };
+
+    Y_UNIT_TEST(TestSimpleQueries) {
+        TDecimalTestCase tester;
+        tester.PrepareTable1();
+
+        tester.CheckQuery("SELECT * FROM `/Root/Table1` WHERE id=1", "[[[\"3.14\"];1;[4]]]");
+        tester.CheckQuery(
+            "SELECT * FROM `/Root/Table1` order by id", "[[[\"3.14\"];1;[4]];[[\"8.16\"];2;[3]];[[\"8.492\"];3;[2]];[[\"12.46\"];4;[1]]]");
+    }
+
+    Y_UNIT_TEST(TestFilterEqual) {
+        TDecimalTestCase tester;
+        tester.PrepareTable1();
+
+        tester.CheckQuery("SELECT * FROM `/Root/Table1` WHERE dec == cast(\"3.14\" as decimal(22,9))", "[[[\"3.14\"];1;[4]]]");
+
+        tester.CheckQuery("SELECT * FROM `/Root/Table1` WHERE dec != cast(\"3.14\" as decimal(22,9)) order by id",
+            "[[[\"8.16\"];2;[3]];[[\"8.492\"];3;[2]];[[\"12.46\"];4;[1]]]");
+    }
+
+    Y_UNIT_TEST(TestFilterNulls) {
+        TDecimalTestCase tester;
+        tester.PrepareTable1();
+
+        TTestHelper::TUpdatesBuilder inserter = tester.Inserter();
+        inserter.AddRow().Add(5).Add(5).AddNull();
+        inserter.AddRow().Add(6).Add(6).AddNull();
+        tester.Upsert(inserter);
+
+        tester.CheckQuery("SELECT * FROM `/Root/Table1` WHERE dec is NULL order by id", "[[#;5;[5]];[#;6;[6]]]");
+
+        tester.CheckQuery("SELECT * FROM `/Root/Table1` WHERE dec is not NULL order by id",
+            "[[[\"3.14\"];1;[4]];[[\"8.16\"];2;[3]];[[\"8.492\"];3;[2]];[[\"12.46\"];4;[1]]]");
+    }
+
+    Y_UNIT_TEST(TestFilterCompare) {
+        TDecimalTestCase tester;
+        tester.PrepareTable1();
+
+        tester.CheckQuery("SELECT * FROM `/Root/Table1` WHERE dec < cast(\"12.46\" as decimal(22,9)) order by id",
+            "[[[\"3.14\"];1;[4]];[[\"8.16\"];2;[3]];[[\"8.492\"];3;[2]]]");
+
+        tester.CheckQuery(
+            "SELECT * FROM `/Root/Table1` WHERE dec > cast(\"8.16\" as decimal(22,9)) order by id", "[[[\"8.492\"];3;[2]];[[\"12.46\"];4;[1]]]");
+
+        tester.CheckQuery("SELECT * FROM `/Root/Table1` WHERE dec <= cast(\"12.46\" as decimal(22,9)) order by id",
+            "[[[\"3.14\"];1;[4]];[[\"8.16\"];2;[3]];[[\"8.492\"];3;[2]];[[\"12.46\"];4;[1]]]");
+
+        tester.CheckQuery("SELECT * FROM `/Root/Table1` WHERE dec >= cast(\"8.492\" as decimal(22,9)) order by id",
+            "[[[\"8.16\"];2;[3]];[[\"8.492\"];3;[2]];[[\"12.46\"];4;[1]]]");
+    }
+
+    Y_UNIT_TEST(TestOrderByDecimal) {
+        TDecimalTestCase tester;
+        tester.PrepareTable1();
+
+        tester.CheckQuery(
+            "SELECT * FROM `/Root/Table1` order by dec", "[[[\"3.14\"];1;[4]];[[\"8.16\"];2;[3]];[[\"8.492\"];3;[2]];[[\"12.46\"];4;[1]]]");
+    }
+
+    Y_UNIT_TEST(TestGroupByDecimal) {
+        TDecimalTestCase tester;
+        tester.PrepareTable1();
+
+        TTestHelper::TUpdatesBuilder inserter = tester.Inserter();
+        inserter.AddRow().Add(5).Add(12).Add(TDecimalValue("8.492"));
+        inserter.AddRow().Add(6).Add(30).Add(TDecimalValue("12.46"));
+        tester.Upsert(inserter);
+
+        tester.CheckQuery("SELECT dec, count(*) FROM `/Root/Table1` group by dec order by dec",
+            "[[[\"3.14\"];1u];[[\"8.16\"];1u];[[\"8.492\"];2u];[[\"12.46\"];2u]]");
+    }
+
+    Y_UNIT_TEST(TestAggregation) {
+        TDecimalTestCase tester;
+        tester.PrepareTable1();
+        tester.CheckQuery("SELECT min(dec) FROM `/Root/Table1`", "[[[\"3.14\"]]]");
+        tester.CheckQuery("SELECT max(dec) FROM `/Root/Table1`", "[[[\"12.46\"]]]");
+        tester.CheckQuery("SELECT sum(dec) FROM `/Root/Table1`", "[[[\"32.252\"]]]");
+    }
+
+    Y_UNIT_TEST(TestJoinById) {
+        TDecimalTestCase tester;
+        tester.PrepareTable1();
+        tester.PrepareTable2();
+
+        tester.CheckQuery(
+            "SELECT t1.id, t1.dec, t2.dec FROM `/Root/Table1` as t1 join `/Root/Table2` as t2 on t1.id = t2.table1_id order by t1.id, t1.dec, "
+            "t2.dec",
+            R"([[1;["3.14"];["8.16"]];[1;["3.14"];["12.46"]];[2;["8.16"];["8.16"]];[2;["8.16"];["12.46"]]])");
+    }
+
+    Y_UNIT_TEST(TestJoinByDecimal) {
+        TDecimalTestCase tester;
+        tester.PrepareTable1();
+        tester.PrepareTable2();
+
+        tester.CheckQuery(
+            "SELECT t1.id, t2.id, t1.dec FROM `/Root/Table1` as t1 join `/Root/Table2` as t2 on t1.dec = t2.dec order by t1.id, t2.id, t1.dec",
+            R"([[2;2;["8.16"]];[2;4;["8.16"]];[4;1;["12.46"]];[4;3;["12.46"]]])");
+    }
+}
+
+}   // namespace NKqp
+}   // namespace NKikimr
diff --git a/ydb/core/kqp/ut/olap/ya.make b/ydb/core/kqp/ut/olap/ya.make
index eae368ee159..e324116597e 100644
--- a/ydb/core/kqp/ut/olap/ya.make
+++ b/ydb/core/kqp/ut/olap/ya.make
@@ -25,6 +25,7 @@ SRCS(
     write_ut.cpp
     sparsed_ut.cpp
     tiering_ut.cpp
+    decimal_ut.cpp
 )
 
 PEERDIR(
diff --git a/ydb/core/tx/columnshard/test_helper/columnshard_ut_common.h b/ydb/core/tx/columnshard/test_helper/columnshard_ut_common.h
index 86060d21e7b..7594be5da95 100644
--- a/ydb/core/tx/columnshard/test_helper/columnshard_ut_common.h
+++ b/ydb/core/tx/columnshard/test_helper/columnshard_ut_common.h
@@ -1,20 +1,21 @@
 #pragma once
 
-#include <ydb/core/tx/columnshard/blob_cache.h>
-#include <ydb/core/tx/columnshard/common/snapshot.h>
-
 #include <ydb/core/formats/arrow/arrow_batch_builder.h>
-#include <ydb/core/tx/columnshard/test_helper/helper.h>
+#include <ydb/core/protos/tx_columnshard.pb.h>
 #include <ydb/core/scheme/scheme_tabledefs.h>
 #include <ydb/core/scheme/scheme_types_proto.h>
 #include <ydb/core/testlib/tablet_helpers.h>
 #include <ydb/core/testlib/test_client.h>
-#include <ydb/core/protos/tx_columnshard.pb.h>
+#include <ydb/core/tx/columnshard/blob_cache.h>
+#include <ydb/core/tx/columnshard/common/snapshot.h>
+#include <ydb/core/tx/columnshard/test_helper/helper.h>
+#include <ydb/core/tx/data_events/common/modification_type.h>
+#include <ydb/core/tx/long_tx_service/public/types.h>
+
+#include <ydb/public/sdk/cpp/client/ydb_value/value.h>
 #include <ydb/services/metadata/abstract/fetcher.h>
 
 #include <library/cpp/testing/unittest/registar.h>
-#include <ydb/core/tx/long_tx_service/public/types.h>
-#include <ydb/core/tx/data_events/common/modification_type.h>
 
 namespace NKikimr::NOlap {
 struct TIndexInfo;
@@ -476,11 +477,13 @@ namespace NKikimr::NColumnShard {
                 auto& builder = Owner.Builders[Index];
                 auto type = builder->type();
 
-                NArrow::SwitchType(type->id(), [&](const auto& t) {
+                Y_ABORT_UNLESS(NArrow::SwitchType(type->id(), [&](const auto& t) {
                     using TWrap = std::decay_t<decltype(t)>;
                     using T = typename TWrap::T;
                     using TBuilder = typename arrow::TypeTraits<typename TWrap::T>::BuilderType;
 
+                    AFL_NOTICE(NKikimrServices::TX_COLUMNSHARD)("T", typeid(T).name());
+
                     auto& typedBuilder = static_cast<TBuilder&>(*builder);
                     if constexpr (std::is_arithmetic<TData>::value) {
                         if constexpr (arrow::has_c_type<T>::value) {
@@ -495,9 +498,16 @@ namespace NKikimr::NColumnShard {
                             return true;
                         }
                     }
+
+                    if constexpr (std::is_same<TData, NYdb::TDecimalValue>::value) {
+                        if constexpr (arrow::is_decimal128_type<T>::value) {
+                            Y_ABORT_UNLESS(typedBuilder.Append(arrow::Decimal128(data.Hi_, data.Low_)).ok());
+                            return true;
+                        }
+                    }
                     Y_ABORT("Unknown type combination");
                     return false;
-                });
+                }));
                 return TRowBuilder(Index + 1, Owner);
             }