Fix optional columns (#7875)

zverevgeny · ivanmorozov333 · web-flow · commit d9d39c96cabc · 2024-08-16T13:23:43.000+03:00
Co-authored-by: ivanmorozov333 &lt;ivanmorozov@ydb.tech&gt;
diff --git a/ydb/core/formats/arrow/common/container.cpp b/ydb/core/formats/arrow/common/container.cpp
@@ -2,6 +2,7 @@
 #include <ydb/library/actors/core/log.h>
 #include <ydb/core/formats/arrow/arrow_helpers.h>
 #include <ydb/core/formats/arrow/simple_arrays_cache.h>
+#include <ydb/library/accessor/validator.h>
 
 namespace NKikimr::NArrow {
 
@@ -93,10 +94,10 @@ TGeneralContainer::TGeneralContainer(const std::shared_ptr<arrow::Schema>& schem
     Initialize();
 }
 
-TGeneralContainer::TGeneralContainer(const std::shared_ptr<arrow::Table>& table) {
+TGeneralContainer::TGeneralContainer(const std::shared_ptr<arrow::Table>& table)
+    : RecordsCount(TValidator::CheckNotNull(table)->num_rows())
+    , Schema(std::make_shared<NModifier::TSchema>(TValidator::CheckNotNull(table)->schema())) {
     AFL_VERIFY(table);
-    Schema = std::make_shared<NModifier::TSchema>(table->schema());
-    RecordsCount = table->num_rows();
     for (auto&& i : table->columns()) {
         if (i->num_chunks() == 1) {
             Columns.emplace_back(std::make_shared<NAccessor::TTrivialArray>(i->chunk(0)));
@@ -107,10 +108,10 @@ TGeneralContainer::TGeneralContainer(const std::shared_ptr<arrow::Table>& table)
     Initialize();
 }
 
-TGeneralContainer::TGeneralContainer(const std::shared_ptr<arrow::RecordBatch>& table) {
+TGeneralContainer::TGeneralContainer(const std::shared_ptr<arrow::RecordBatch>& table)
+    : RecordsCount(TValidator::CheckNotNull(table)->num_rows())
+    , Schema(std::make_shared<NModifier::TSchema>(TValidator::CheckNotNull(table)->schema())) {
     AFL_VERIFY(table);
-    Schema = std::make_shared<NModifier::TSchema>(table->schema());
-    RecordsCount = table->num_rows();
     for (auto&& i : table->columns()) {
         Columns.emplace_back(std::make_shared<NAccessor::TTrivialArray>(i));
     }
diff --git a/ydb/core/formats/arrow/common/container.h b/ydb/core/formats/arrow/common/container.h
@@ -24,14 +24,14 @@ class IFieldsConstructor {
 
 class TGeneralContainer {
 private:
-    YDB_READONLY_DEF(std::optional<ui64>, RecordsCount);
+    std::optional<ui64> RecordsCount;
     YDB_READONLY_DEF(std::shared_ptr<NModifier::TSchema>, Schema);
     std::vector<std::shared_ptr<NAccessor::IChunkedArray>> Columns;
     void Initialize();
 public:
     TGeneralContainer(const ui32 recordsCount);
 
-    ui32 GetRecordsCountVerified() const {
+    ui32 GetRecordsCount() const {
         AFL_VERIFY(RecordsCount);
         return *RecordsCount;
     }
diff --git a/ydb/core/formats/arrow/modifier/subset.h b/ydb/core/formats/arrow/modifier/subset.h
@@ -14,6 +14,12 @@ class TSchemaSubset {
     TSchemaSubset() = default;
     TSchemaSubset(const std::set<ui32>& fieldsIdx, const ui32 fieldsCount);
 
+    static TSchemaSubset AllFieldsAccepted() {
+        TSchemaSubset result;
+        result.Exclude = true;
+        return result;
+    }
+
     template <class T>
     std::vector<T> Apply(const std::vector<T>& fullSchema) const {
         if (FieldIdx.empty()) {
diff --git a/ydb/core/kqp/ut/olap/kqp_olap_ut.cpp b/ydb/core/kqp/ut/olap/kqp_olap_ut.cpp
@@ -2624,6 +2624,43 @@ Y_UNIT_TEST_SUITE(KqpOlap) {
 
     }
 
+    Y_UNIT_TEST(NormalizeAbsentColumn) {
+        auto settings = TKikimrSettings().SetWithSampleTables(false);
+        TKikimrRunner kikimr(settings);
+        TLocalHelper testHelper(kikimr);
+
+        auto csController = NYDBTest::TControllers::RegisterCSControllerGuard<NYDBTest::NColumnShard::TController>();
+        csController->SetPeriodicWakeupActivationPeriod(TDuration::Seconds(1));
+        csController->SetLagForCompactionBeforeTierings(TDuration::Seconds(1));
+        csController->SetOverrideReduceMemoryIntervalLimit(1LLU << 30);
+        csController->DisableBackground(NKikimr::NYDBTest::ICSController::EBackground::Indexation);
+
+        testHelper.CreateTestOlapTable();
+        auto tableClient = kikimr.GetTableClient();
+
+        Tests::NCommon::TLoggerInit(kikimr).SetComponents({ NKikimrServices::TX_COLUMNSHARD, NKikimrServices::TX_COLUMNSHARD_SCAN }, "CS").SetPriority(NActors::NLog::PRI_DEBUG).Initialize();
+
+        auto session = tableClient.CreateSession().GetValueSync().GetSession();
+
+        {
+            auto alterQuery = TStringBuilder() << "ALTER TABLESTORE `/Root/olapStore` ADD COLUMN new_column1 Uint64;";
+            auto alterResult = session.ExecuteSchemeQuery(alterQuery).GetValueSync();
+            UNIT_ASSERT_VALUES_EQUAL_C(alterResult.GetStatus(), EStatus::SUCCESS, alterResult.GetIssues().ToString());
+        }
+        WriteTestData(kikimr, "/Root/olapStore/olapTable", 1000000, 300000000, 1000);
+        WriteTestData(kikimr, "/Root/olapStore/olapTable", 1100000, 300100000, 1000);
+
+        {
+            auto alterQuery = TStringBuilder() << "ALTER TABLESTORE `/Root/olapStore` ADD COLUMN new_column2 Uint64;";
+            auto alterResult = session.ExecuteSchemeQuery(alterQuery).GetValueSync();
+            UNIT_ASSERT_VALUES_EQUAL_C(alterResult.GetStatus(), EStatus::SUCCESS, alterResult.GetIssues().ToString());
+        }
+
+        csController->EnableBackground(NKikimr::NYDBTest::ICSController::EBackground::Indexation);
+        csController->WaitIndexation(TDuration::Seconds(5));
+
+    }
+
     Y_UNIT_TEST(MultiInsertWithSinks) {
         NKikimrConfig::TAppConfig appConfig;
         appConfig.MutableTableServiceConfig()->SetEnableOlapSink(true);
diff --git a/ydb/core/kqp/ut/service/kqp_qs_queries_ut.cpp b/ydb/core/kqp/ut/service/kqp_qs_queries_ut.cpp
@@ -2,6 +2,8 @@
 #include <ydb/core/kqp/ut/common/kqp_ut_common.h>
 #include <ydb/core/kqp/ut/common/columnshard.h>
 #include <ydb/core/testlib/common_helper.h>
+#include <ydb/core/tx/columnshard/hooks/abstract/abstract.h>
+#include <ydb/core/tx/columnshard/hooks/testing/controller.h>
 #include <ydb/public/lib/ut_helpers/ut_helpers_query.h>
 #include <ydb/public/sdk/cpp/client/ydb_operation/operation.h>
 #include <ydb/public/sdk/cpp/client/ydb_proto/accessor.h>
@@ -3039,6 +3041,11 @@ Y_UNIT_TEST_SUITE(KqpQueryService) {
 
             auto session = Kikimr->GetTableClient().CreateSession().GetValueSync().GetSession();
 
+            auto csController = NYDBTest::TControllers::RegisterCSControllerGuard<NYDBTest::NColumnShard::TController>();
+            csController->SetPeriodicWakeupActivationPeriod(TDuration::Seconds(1));
+            csController->SetLagForCompactionBeforeTierings(TDuration::Seconds(1));
+            csController->DisableBackground(NKikimr::NYDBTest::ICSController::EBackground::Indexation);
+
             const TString query = Sprintf(R"(
                 CREATE TABLE `/Root/DataShard` (
                     Col1 Uint64 NOT NULL,
@@ -3054,6 +3061,8 @@ Y_UNIT_TEST_SUITE(KqpQueryService) {
             auto result = session.ExecuteSchemeQuery(query).GetValueSync();
             UNIT_ASSERT_C(result.GetStatus() == NYdb::EStatus::SUCCESS, result.GetIssues().ToString());
             DoExecute();
+            csController->EnableBackground(NKikimr::NYDBTest::ICSController::EBackground::Indexation);
+            csController->WaitIndexation(TDuration::Seconds(5));
         }
 
     };
diff --git a/ydb/core/kqp/ut/service/ya.make b/ydb/core/kqp/ut/service/ya.make
@@ -24,6 +24,7 @@ PEERDIR(
     library/cpp/threading/local_executor
     ydb/core/kqp
     ydb/core/kqp/ut/common
+    ydb/core/tx/columnshard/hooks/testing
     ydb/library/yql/sql/pg
     ydb/library/yql/parser/pg_wrapper
     ydb/public/lib/ut_helpers
diff --git a/ydb/core/protos/feature_flags.proto b/ydb/core/protos/feature_flags.proto
@@ -146,4 +146,6 @@ message TFeatureFlags {
     optional bool EnableSingleCompositeActionGroup = 131 [default = false];
     optional bool EnableResourcePoolsOnServerless = 132 [default = false];
     optional bool EnableChangefeedsOnIndexTables = 134 [default = false];
+    optional bool EnableResourcePoolsCounters = 135 [default = false];
+    optional bool EnableOptionalColumnsInColumnShard = 136 [default = false];
 }
diff --git a/ydb/core/tx/columnshard/engines/changes/compaction/abstract/merger.cpp b/ydb/core/tx/columnshard/engines/changes/compaction/abstract/merger.cpp
@@ -2,4 +2,14 @@
 
 namespace NKikimr::NOlap::NCompaction {
 
+void IColumnMerger::Start(const std::vector<std::shared_ptr<NArrow::NAccessor::IChunkedArray>>& input) {
+    AFL_VERIFY(!Started);
+    Started = true;
+    // for (auto&& i : input) {
+    //     AFL_VERIFY(i->GetDataType()->id() == Context.GetResultField()->type()->id())("input", i->GetDataType()->ToString())(
+    //                                              "result", Context.GetResultField()->ToString());
+    // }
+    return DoStart(input);
+}
+
 }
diff --git a/ydb/core/tx/columnshard/engines/changes/compaction/abstract/merger.h b/ydb/core/tx/columnshard/engines/changes/compaction/abstract/merger.h
@@ -21,11 +21,7 @@ class IColumnMerger {
 
     virtual ~IColumnMerger() = default;
 
-    void Start(const std::vector<std::shared_ptr<NArrow::NAccessor::IChunkedArray>>& input) {
-        AFL_VERIFY(!Started);
-        Started = true;
-        return DoStart(input);
-    }
+    void Start(const std::vector<std::shared_ptr<NArrow::NAccessor::IChunkedArray>>& input);
 
     std::vector<TColumnPortionResult> Execute(
         const NCompaction::TColumnMergeContext& context, const std::shared_ptr<arrow::RecordBatch>& remap) {
diff --git a/ydb/core/tx/columnshard/engines/changes/compaction/merger.cpp b/ydb/core/tx/columnshard/engines/changes/compaction/merger.cpp
@@ -30,6 +30,8 @@ std::vector<NKikimr::NOlap::TWritePortionInfoWithBlobsResult> TMerger::Execute(c
 
         ui32 idx = 0;
         for (auto&& batch : Batches) {
+            AFL_VERIFY(batch->GetColumnsCount() == resultFiltered->GetColumnsCount())("data", batch->GetColumnsCount())(
+                                                       "schema", resultFiltered->GetColumnsCount());
             {
                 NArrow::NConstruction::IArrayBuilder::TPtr column =
                     std::make_shared<NArrow::NConstruction::TSimpleArrayConstructor<NArrow::NConstruction::TIntConstFiller<arrow::UInt16Type>>>(
@@ -54,7 +56,6 @@ std::vector<NKikimr::NOlap::TWritePortionInfoWithBlobsResult> TMerger::Execute(c
         NActors::TLogContextGuard logGuard(
             NActors::TLogContextBuilder::Build()("field_name", resultFiltered->GetIndexInfo().GetColumnName(columnId)));
         auto columnInfo = stats->GetColumnInfo(columnId);
-        auto resultField = resultFiltered->GetIndexInfo().GetColumnFieldVerified(columnId);
         std::shared_ptr<IColumnMerger> merger = std::make_shared<TPlainMerger>();
         //        resultFiltered->BuildColumnMergerVerified(columnId);
 
diff --git a/ydb/core/tx/columnshard/engines/changes/indexation.cpp b/ydb/core/tx/columnshard/engines/changes/indexation.cpp
@@ -130,17 +130,7 @@ TConclusionStatus TInsertColumnEngineChanges::DoConstructBlobs(TConstructionCont
     Y_ABORT_UNLESS(!DataToIndex.empty());
     Y_ABORT_UNLESS(AppendedPortions.empty());
 
-    auto maxSnapshot = TSnapshot::Zero();
-    for (auto& inserted : DataToIndex) {
-        TSnapshot insertSnap = inserted.GetSnapshot();
-        Y_ABORT_UNLESS(insertSnap.Valid());
-        if (insertSnap > maxSnapshot) {
-            maxSnapshot = insertSnap;
-        }
-    }
-    Y_ABORT_UNLESS(maxSnapshot.Valid());
-
-    auto resultSchema = context.SchemaVersions.GetSchema(maxSnapshot);
+    auto resultSchema = context.SchemaVersions.GetLastSchema();
     Y_ABORT_UNLESS(resultSchema->GetIndexInfo().IsSorted());
 
     TPathesData pathBatches;
diff --git a/ydb/core/tx/columnshard/engines/scheme/versions/abstract_scheme.cpp b/ydb/core/tx/columnshard/engines/scheme/versions/abstract_scheme.cpp
@@ -29,23 +29,32 @@ std::set<ui32> ISnapshotSchema::GetPkColumnsIds() const {
 
 TConclusion<std::shared_ptr<NArrow::TGeneralContainer>> ISnapshotSchema::NormalizeBatch(
     const ISnapshotSchema& dataSchema, const std::shared_ptr<NArrow::TGeneralContainer>& batch, const std::set<ui32>& restoreColumnIds) const {
+    AFL_VERIFY(dataSchema.GetSnapshot() <= GetSnapshot());
     if (dataSchema.GetSnapshot() == GetSnapshot()) {
-        return batch;
+        if (batch->GetColumnsCount() == GetColumnsCount()) {
+            return batch;
+        }
     }
-    AFL_VERIFY(dataSchema.GetSnapshot() < GetSnapshot());
     const std::shared_ptr<arrow::Schema>& resultArrowSchema = GetSchema();
 
-    std::shared_ptr<NArrow::TGeneralContainer> result = std::make_shared<NArrow::TGeneralContainer>(batch->GetRecordsCountVerified());
+    std::shared_ptr<NArrow::TGeneralContainer> result = std::make_shared<NArrow::TGeneralContainer>(batch->GetRecordsCount());
     for (size_t i = 0; i < resultArrowSchema->fields().size(); ++i) {
         auto& resultField = resultArrowSchema->fields()[i];
         auto columnId = GetIndexInfo().GetColumnId(resultField->name());
         auto oldField = dataSchema.GetFieldByColumnIdOptional(columnId);
         if (oldField) {
-            auto conclusion = result->AddField(resultField, batch->GetAccessorByNameVerified(oldField->name()));
-            if (conclusion.IsFail()) {
-                return conclusion;
+            auto fAccessor = batch->GetAccessorByNameOptional(oldField->name());
+            if (fAccessor) {
+                auto conclusion = result->AddField(resultField, fAccessor);
+                if (conclusion.IsFail()) {
+                    return conclusion;
+                }
+                continue;
             }
-        } else if (restoreColumnIds.contains(columnId)) {
+        }
+        if (restoreColumnIds.contains(columnId)) {
+            AFL_VERIFY(!!GetExternalDefaultValueVerified(columnId) || GetIndexInfo().IsNullableVerified(columnId))("column_name",
+                                                                          GetIndexInfo().GetColumnName(columnId, false))("id", columnId);
             result->AddField(resultField,
                     NArrow::TThreadSimpleArraysCache::Get(resultField->type(), GetExternalDefaultValueVerified(columnId), batch->num_rows()))
                 .Validate();
@@ -107,8 +116,9 @@ TConclusion<std::shared_ptr<arrow::RecordBatch>> ISnapshotSchema::PrepareForModi
     Y_DEBUG_ABORT_UNLESS(NArrow::IsSortedAndUnique(batch, GetIndexInfo().GetPrimaryKey()));
 
     switch (mType) {
+        case NEvWrite::EModificationType::Replace:
         case NEvWrite::EModificationType::Upsert: {
-            AFL_VERIFY(batch->num_columns() <= dstSchema->num_fields());
+                AFL_VERIFY(batch->num_columns() <= dstSchema->num_fields());
             if (batch->num_columns() < dstSchema->num_fields()) {
                 for (auto&& f : dstSchema->fields()) {
                     if (GetIndexInfo().IsNullableVerified(f->name())) {
@@ -125,7 +135,6 @@ TConclusion<std::shared_ptr<arrow::RecordBatch>> ISnapshotSchema::PrepareForModi
             return batch;
         }
         case NEvWrite::EModificationType::Delete:
-        case NEvWrite::EModificationType::Replace:
         case NEvWrite::EModificationType::Insert:
         case NEvWrite::EModificationType::Update:
             return batch;
diff --git a/ydb/core/tx/columnshard/operations/slice_builder/builder.cpp b/ydb/core/tx/columnshard/operations/slice_builder/builder.cpp
@@ -55,7 +55,14 @@ TConclusionStatus TBuildSlicesTask::DoExecute(const std::shared_ptr<ITask>& /*ta
     if (OriginalBatch->num_columns() != indexSchema->num_fields()) {
         AFL_VERIFY(OriginalBatch->num_columns() < indexSchema->num_fields())("original", OriginalBatch->num_columns())(
                                                       "index", indexSchema->num_fields());
-
+        if (HasAppData() && !AppDataVerified().FeatureFlags.GetEnableOptionalColumnsInColumnShard()) {
+            subset = NArrow::TSchemaSubset::AllFieldsAccepted();
+            const std::vector<ui32>& columnIdsVector = ActualSchema->GetIndexInfo().GetColumnIds(false);
+            const std::set<ui32> columnIdsSet(columnIdsVector.begin(), columnIdsVector.end());
+            auto normalized =
+                ActualSchema->NormalizeBatch(*ActualSchema, std::make_shared<NArrow::TGeneralContainer>(OriginalBatch), columnIdsSet).DetachResult();
+            OriginalBatch = NArrow::ToBatch(normalized->BuildTableVerified(), true);
+        }
     }
     WriteData.MutableWriteMeta().SetWriteMiddle2StartInstant(TMonotonic::Now());
     auto batches = BuildSlices();

Original file line number	Diff line number	Diff line change
`@@ -146,4 +146,6 @@ message TFeatureFlags {`
`146`	`146`	`optional bool EnableSingleCompositeActionGroup = 131 [default = false];`
`147`	`147`	`optional bool EnableResourcePoolsOnServerless = 132 [default = false];`
`148`	`148`	`optional bool EnableChangefeedsOnIndexTables = 134 [default = false];`
	`149`	`+ optional bool EnableResourcePoolsCounters = 135 [default = false];`
	`150`	`+ optional bool EnableOptionalColumnsInColumnShard = 136 [default = false];`
`149`	`151`	`}`