24-3: Partition at decimal keys (ydb-platform#10696)

azevaykin · Oct 23, 2024 · f4dcfe0 · f4dcfe0
1 parent 7ff1c86
commit f4dcfe0
Show file tree

Hide file tree

Showing 3 changed files with 54 additions and 2 deletions.
diff --git a/ydb/core/engine/mkql_proto.cpp b/ydb/core/engine/mkql_proto.cpp
@@ -223,6 +223,19 @@ bool CellsFromTuple(const NKikimrMiniKQL::TType* tupleType,
             }
             break;
         }
+        case NScheme::NTypeIds::Decimal:
+        {
+            if (v.HasLow128() && v.HasHi128()) {
+                NYql::NDecimal::TInt128 int128 = NYql::NDecimal::FromProto(v);
+                auto &data = memoryOwner.emplace_back();
+                data.resize(sizeof(NYql::NDecimal::TInt128));
+                std::memcpy(data.Detach(), &int128, sizeof(NYql::NDecimal::TInt128));
+                c = TCell(data);                
+            } else {
+                CHECK_OR_RETURN_ERROR(false, Sprintf("Cannot parse value of type Decimal in tuple at position %" PRIu32, i));
+            }
+            break;
+        }        
         default:
             CHECK_OR_RETURN_ERROR(false, Sprintf("Unsupported typeId %" PRIu16 " at index %" PRIu32, typeId, i));
             break;
@@ -328,6 +341,13 @@ bool CellToValue(NScheme::TTypeInfo type, const TCell& c, NKikimrMiniKQL::TValue
         val.MutableOptional()->SetText(c.Data(), c.Size());
         break;
 
+    case NScheme::NTypeIds::Decimal: {
+        const auto loHi = c.AsValue<std::pair<ui64, i64>>();
+        val.MutableOptional()->SetLow128(loHi.first);
+        val.MutableOptional()->SetHi128(loHi.second);
+        break;
+    }
+
     case NScheme::NTypeIds::Pg: {
         auto convert = NPg::PgNativeTextFromNativeBinary(c.AsBuf(), type.GetTypeDesc());
         if (convert.Error) {

diff --git a/ydb/core/grpc_services/rpc_object_storage.cpp b/ydb/core/grpc_services/rpc_object_storage.cpp
@@ -144,6 +144,19 @@ bool CellFromTuple(NScheme::TTypeInfo type,
         }
         break;
     }
+    case NScheme::NTypeIds::Decimal:
+    {
+        if (tupleValue.Haslow_128()) {
+            NYql::NDecimal::TInt128 int128 = NYql::NDecimal::FromHalfs(tupleValue.Getlow_128(), tupleValue.Gethigh_128());
+            auto &data = memoryOwner.emplace_back();
+            data.resize(sizeof(NYql::NDecimal::TInt128));
+            std::memcpy(data.Detach(), &int128, sizeof(NYql::NDecimal::TInt128));
+            c = TCell(data);                
+        } else {
+            CHECK_OR_RETURN_ERROR(false, Sprintf("Cannot parse value of type Decimal in tuple at position %" PRIu32, position));
+        }
+        break;
+    }    
     default:
         CHECK_OR_RETURN_ERROR(false, Sprintf("Unsupported typeId %" PRIu16 " at index %" PRIu32, typeId, position));
         break;

diff --git a/ydb/core/kqp/ut/scan/kqp_scan_ut.cpp b/ydb/core/kqp/ut/scan/kqp_scan_ut.cpp
@@ -237,15 +237,34 @@ Y_UNIT_TEST_SUITE(KqpScan) {
                 .BeginTuple().AddElement().BeginOptional().Decimal(TDecimalValue("1.5", 22, 9)).EndOptional().EndTuple()
                 .Build());
 
-        auto ret = session.CreateTable("/Root/DecimalTest",
+       auto ret = session.CreateTable("/Root/DecimalTest",
                 TTableBuilder()
                     .AddNullableColumn("Key", TDecimalType(22, 9))
                     .AddNullableColumn("Value", TDecimalType(22, 9))
                     .SetPrimaryKeyColumn("Key")
-                    // .SetPartitionAtKeys(partitions)  // Error at split boundary 0: Unsupported typeId 4865 at index 0
+                    .SetPartitionAtKeys(partitions)
                     .Build()).GetValueSync();
         UNIT_ASSERT_C(ret.IsSuccess(), ret.GetIssues().ToString());
 
+        {
+            auto describeResult = session.DescribeTable("/Root/DecimalTest" , TDescribeTableSettings().WithKeyShardBoundary(true)).ExtractValueSync();
+            UNIT_ASSERT_VALUES_EQUAL(describeResult.GetStatus(), NYdb::EStatus::SUCCESS);
+            const NYdb::NTable::TTableDescription& tableDescription = describeResult.GetTableDescription();
+            const TVector<NYdb::NTable::TKeyRange>& keyRanges = tableDescription.GetKeyRanges();
+            const TVector<NYdb::NTable::TTableColumn>& columns = tableDescription.GetTableColumns();
+            UNIT_ASSERT_VALUES_EQUAL(columns.size(), 2);
+            UNIT_ASSERT_STRINGS_EQUAL(columns[0].Type.ToString(), "Decimal(22,9)?");
+            UNIT_ASSERT_STRINGS_EQUAL(columns[1].Type.ToString(), "Decimal(22,9)?");
+            auto extractValue = [](const TValue& val) {
+                auto parser = TValueParser(val);
+                parser.OpenTuple();
+                UNIT_ASSERT(parser.TryNextElement());
+                return parser.GetOptionalDecimal()->ToString();
+            };
+            UNIT_ASSERT_VALUES_EQUAL(keyRanges.size(), 2);
+            UNIT_ASSERT_STRINGS_EQUAL(extractValue(keyRanges[0].To()->GetValue()), "1.5");
+        }
+
         auto params = TParamsBuilder().AddParam("$in").BeginList()
                 .AddListItem().BeginStruct()
                     .AddMember("Key").Decimal(TDecimalValue("1.0"))