Merge bd7af75 into 8dde59c

ydb-platform · Jul 31, 2024 · 53a3c5e · 53a3c5e
2 parents 8dde59c + bd7af75
commit 53a3c5e
Show file tree

Hide file tree

Showing 5 changed files with 28 additions and 17 deletions.
diff --git a/ydb/core/tx/columnshard/test_helper/columnshard_ut_common.cpp b/ydb/core/tx/columnshard/test_helper/columnshard_ut_common.cpp
@@ -125,26 +125,25 @@ bool WriteDataImpl(TTestBasicRuntime& runtime, TActorId& sender, const ui64 shar
 }
 
 bool WriteData(TTestBasicRuntime& runtime, TActorId& sender, const ui64 shardId, const ui64 writeId, const ui64 tableId, const TString& data,
-    const std::vector<NArrow::NTest::TTestColumn>& ydbSchema, std::vector<ui64>* writeIds, const NEvWrite::EModificationType mType,
-    const std::set<std::string>& notNullColumns) {
+    const std::vector<NArrow::NTest::TTestColumn>& ydbSchema, std::vector<ui64>* writeIds, const NEvWrite::EModificationType mType) {
     NLongTxService::TLongTxId longTxId;
     UNIT_ASSERT(longTxId.ParseString("ydb://long-tx/01ezvvxjdk2hd4vdgjs68knvp8?node_id=1"));
     return WriteDataImpl(
-        runtime, sender, shardId, tableId, longTxId, writeId, data, NArrow::MakeArrowSchema(ydbSchema, notNullColumns), writeIds, mType);
+        runtime, sender, shardId, tableId, longTxId, writeId, data, NArrow::MakeArrowSchema(ydbSchema), writeIds, mType);
 }
 
 bool WriteData(TTestBasicRuntime& runtime, TActorId& sender, const ui64 writeId, const ui64 tableId, const TString& data,
     const std::vector<NArrow::NTest::TTestColumn>& ydbSchema, bool waitResult, std::vector<ui64>* writeIds,
-    const NEvWrite::EModificationType mType, const std::set<std::string>& notNullColumns) {
+    const NEvWrite::EModificationType mType) {
     NLongTxService::TLongTxId longTxId;
     UNIT_ASSERT(longTxId.ParseString("ydb://long-tx/01ezvvxjdk2hd4vdgjs68knvp8?node_id=1"));
     if (writeIds) {
         return WriteDataImpl(runtime, sender, TTestTxConfig::TxTablet0, tableId, longTxId, writeId, data,
-            NArrow::MakeArrowSchema(ydbSchema, notNullColumns), writeIds, mType);
+            NArrow::MakeArrowSchema(ydbSchema), writeIds, mType);
     }
     std::vector<ui64> ids;
     return WriteDataImpl(runtime, sender, TTestTxConfig::TxTablet0, tableId, longTxId, writeId, data,
-        NArrow::MakeArrowSchema(ydbSchema, notNullColumns), waitResult ? &ids : nullptr, mType);
+        NArrow::MakeArrowSchema(ydbSchema), waitResult ? &ids : nullptr, mType);
 }
 
 std::optional<ui64> WriteData(TTestBasicRuntime& runtime, TActorId& sender, const NLongTxService::TLongTxId& longTxId,

diff --git a/ydb/core/tx/columnshard/test_helper/columnshard_ut_common.h b/ydb/core/tx/columnshard/test_helper/columnshard_ut_common.h
@@ -407,11 +407,11 @@ void PlanWriteTx(TTestBasicRuntime& runtime, TActorId& sender, NOlap::TSnapshot
 
 bool WriteData(TTestBasicRuntime& runtime, TActorId& sender, const ui64 shardId, const ui64 writeId, const ui64 tableId, const TString& data,
     const std::vector<NArrow::NTest::TTestColumn>& ydbSchema, std::vector<ui64>* writeIds,
-    const NEvWrite::EModificationType mType = NEvWrite::EModificationType::Upsert, const std::set<std::string>& notNullColumns = {});
+    const NEvWrite::EModificationType mType = NEvWrite::EModificationType::Upsert);
 
 bool WriteData(TTestBasicRuntime& runtime, TActorId& sender, const ui64 writeId, const ui64 tableId, const TString& data,
     const std::vector<NArrow::NTest::TTestColumn>& ydbSchema, bool waitResult = true, std::vector<ui64>* writeIds = nullptr,
-    const NEvWrite::EModificationType mType = NEvWrite::EModificationType::Upsert, const std::set<std::string>& notNullColumns = {});
+    const NEvWrite::EModificationType mType = NEvWrite::EModificationType::Upsert);
 
 std::optional<ui64> WriteData(TTestBasicRuntime& runtime, TActorId& sender, const NLongTxService::TLongTxId& longTxId,
                               ui64 tableId, const ui64 writePartId, const TString& data,

diff --git a/ydb/core/tx/columnshard/test_helper/helper.cpp b/ydb/core/tx/columnshard/test_helper/helper.cpp
@@ -36,6 +36,16 @@ NKikimrSchemeOp::TOlapColumnDescription TTestColumn::CreateColumn(const ui32 id)
     return col;
 }
 
+std::set<std::string> TTestColumn::GetNullableSet(const std::vector<TTestColumn>& columns) {
+    std::set<std::string> result;
+    for (auto&& i : columns) {
+        if (!i.GetNullable()) {
+            result.emplace(i.GetName());
+        }
+    }
+    return result;
+}
+
 std::vector<std::pair<TString, NKikimr::NScheme::TTypeInfo>> TTestColumn::ConvertToPairs(const std::vector<TTestColumn>& columns) {
     std::vector<std::pair<TString, NScheme::TTypeInfo>> result;
     for (auto&& i : columns) {
@@ -70,15 +80,15 @@ std::vector<NKikimr::NArrow::NTest::TTestColumn> TTestColumn::CropSchema(const s
 
 namespace NKikimr::NArrow {
 
-std::vector<std::shared_ptr<arrow::Field>> MakeArrowFields(
-    const std::vector<NTest::TTestColumn>& columns, const std::set<std::string>& notNullColumns /*= {}*/) {
+std::vector<std::shared_ptr<arrow::Field>> MakeArrowFields(const std::vector<NTest::TTestColumn>& columns) {
+    std::set<std::string> notNullColumns = NTest::TTestColumn::GetNullableSet(columns);
     auto result = MakeArrowFields(NTest::TTestColumn::ConvertToPairs(columns), notNullColumns);
     UNIT_ASSERT_C(result.ok(), result.status().ToString());
     return result.ValueUnsafe();
 }
 
-std::shared_ptr<arrow::Schema> MakeArrowSchema(
-    const std::vector<NTest::TTestColumn>& columns, const std::set<std::string>& notNullColumns /*= {}*/) {
+std::shared_ptr<arrow::Schema> MakeArrowSchema(const std::vector<NTest::TTestColumn>& columns) {
+    std::set<std::string> notNullColumns = NTest::TTestColumn::GetNullableSet(columns);
     auto result = MakeArrowSchema(NTest::TTestColumn::ConvertToPairs(columns), notNullColumns);
     UNIT_ASSERT_C(result.ok(), result.status().ToString());
     return result.ValueUnsafe();

diff --git a/ydb/core/tx/columnshard/test_helper/helper.h b/ydb/core/tx/columnshard/test_helper/helper.h
@@ -52,6 +52,7 @@ class TTestColumn {
     YDB_ACCESSOR_DEF(NScheme::TTypeInfo, Type);
     YDB_ACCESSOR_DEF(TString, StorageId);
     YDB_ACCESSOR_DEF(TString, AccessorClassName);
+    YDB_ACCESSOR(bool, Nullable, true);
 
 public:
     explicit TTestColumn(const TString& name, const NScheme::TTypeInfo& type)
@@ -65,13 +66,14 @@ class TTestColumn {
     static THashMap<TString, NScheme::TTypeInfo> ConvertToHash(const std::vector<TTestColumn>& columns);
     static std::vector<TTestColumn> BuildFromPairs(const std::vector<std::pair<TString, NScheme::TTypeInfo>>& columns);
     static std::vector<TTestColumn> CropSchema(const std::vector<TTestColumn>& input, const ui32 size);
+    static std::set<std::string> GetNullableSet(const std::vector<TTestColumn>& columns);
 };
 
 }
 
 namespace NKikimr::NArrow {
 
-std::vector<std::shared_ptr<arrow::Field>> MakeArrowFields(const std::vector<NTest::TTestColumn>& columns, const std::set<std::string>& notNullColumns = {});
-std::shared_ptr<arrow::Schema> MakeArrowSchema(const std::vector<NTest::TTestColumn>& columns, const std::set<std::string>& notNullColumns = {});
+std::vector<std::shared_ptr<arrow::Field>> MakeArrowFields(const std::vector<NTest::TTestColumn>& columns);
+std::shared_ptr<arrow::Schema> MakeArrowSchema(const std::vector<NTest::TTestColumn>& columns);
 
 }
diff --git a/ydb/core/tx/schemeshard/ut_olap/ut_olap.cpp b/ydb/core/tx/schemeshard/ut_olap/ut_olap.cpp
@@ -40,7 +40,7 @@ static const TString defaultTableSchema = R"(
 )";
 
 static const TVector<NArrow::NTest::TTestColumn> defaultYdbSchema = {
-    NArrow::NTest::TTestColumn("timestamp", TTypeInfo(NTypeIds::Timestamp) ),
+    NArrow::NTest::TTestColumn("timestamp", TTypeInfo(NTypeIds::Timestamp).SetNullable(false) ),
     NArrow::NTest::TTestColumn("data", TTypeInfo(NTypeIds::Utf8) )
 };
 
@@ -713,7 +713,7 @@ Y_UNIT_TEST_SUITE(TOlap) {
             TSet<ui64> txIds;
             for (ui32 i = 0; i < 10; ++i) {
                 std::vector<ui64> writeIds;
-                NTxUT::WriteData(runtime, sender, shardId, ++writeId, pathId, data, defaultYdbSchema, &writeIds, NEvWrite::EModificationType::Upsert, { "timestamp" });
+                NTxUT::WriteData(runtime, sender, shardId, ++writeId, pathId, data, defaultYdbSchema, &writeIds, NEvWrite::EModificationType::Upsert);
                 NTxUT::ProposeCommit(runtime, sender, shardId, ++txId, writeIds);
                 txIds.insert(txId);
             }
@@ -725,7 +725,7 @@ Y_UNIT_TEST_SUITE(TOlap) {
 
             // trigger periodic stats at shard (after timeout)
             std::vector<ui64> writeIds;
-            NTxUT::WriteData(runtime, sender, shardId, ++writeId, pathId, data, defaultYdbSchema, &writeIds, NEvWrite::EModificationType::Upsert, { "timestamp" });
+            NTxUT::WriteData(runtime, sender, shardId, ++writeId, pathId, data, defaultYdbSchema, &writeIds, NEvWrite::EModificationType::Upsert);
             NTxUT::ProposeCommit(runtime, sender, shardId, ++txId, writeIds);
             NTxUT::PlanCommit(runtime, sender, shardId, ++planStep, {txId});
         }