Merge 62c1bea into 5cb33a5

ydb-platform · Sep 2, 2024 · 53ac89c · 53ac89c
2 parents 5cb33a5 + 62c1bea
commit 53ac89c
Show file tree

Hide file tree

Showing 8 changed files with 37 additions and 12 deletions.
diff --git a/ydb/core/tx/columnshard/engines/portions/portion_info.cpp b/ydb/core/tx/columnshard/engines/portions/portion_info.cpp
@@ -592,11 +592,10 @@ template <class TExternalBlobInfo>
 TPortionInfo::TPreparedBatchData PrepareForAssembleImpl(const TPortionInfo& portion, const ISnapshotSchema& dataSchema, const ISnapshotSchema& resultSchema,
     THashMap<TChunkAddress, TExternalBlobInfo>& blobsData) {
     std::vector<TPortionInfo::TColumnAssemblingInfo> columns;
-    auto arrowResultSchema = resultSchema.GetSchema();
-    columns.reserve(arrowResultSchema->num_fields());
+    columns.reserve(resultSchema.GetColumnIds().size());
     const ui32 rowsCount = portion.GetRecordsCount();
-    for (auto&& i : arrowResultSchema->fields()) {
-        columns.emplace_back(rowsCount, dataSchema.GetColumnLoaderOptional(i->name()), resultSchema.GetColumnLoaderVerified(i->name()));
+    for (auto&& i : resultSchema.GetColumnIds()) {
+        columns.emplace_back(rowsCount, dataSchema.GetColumnLoaderOptional(i), resultSchema.GetColumnLoaderVerified(i));
     }
     {
         int skipColumnId = -1;
@@ -628,7 +627,7 @@ TPortionInfo::TPreparedBatchData PrepareForAssembleImpl(const TPortionInfo& port
         preparedColumns.emplace_back(c.Compile());
     }
 
-    return TPortionInfo::TPreparedBatchData(std::move(preparedColumns), arrowResultSchema, rowsCount);
+    return TPortionInfo::TPreparedBatchData(std::move(preparedColumns), resultSchema.GetSchema(), rowsCount);
 }
 
 }

diff --git a/ydb/core/tx/columnshard/engines/scheme/abstract/index_info.h b/ydb/core/tx/columnshard/engines/scheme/abstract/index_info.h
@@ -109,9 +109,15 @@ class IIndexInfo {
 
     [[nodiscard]] static std::vector<ui32> AddSpecialFieldIds(const std::vector<ui32>& baseColumnIds) {
         std::vector<ui32> result = baseColumnIds;
-        for (auto&& i : GetSystemColumnIds()) {
-            result.emplace_back(i);
-        }
+        const auto& cIds = GetSystemColumnIds();
+        result.insert(result.end(), cIds.begin(), cIds.end());
+        return result;
+    }
+
+    [[nodiscard]] static std::set<ui32> AddSpecialFieldIds(const std::set<ui32>& baseColumnIds) {
+        std::set<ui32> result = baseColumnIds;
+        const auto& cIds = GetSystemColumnIds();
+        result.insert(cIds.begin(), cIds.end());
         return result;
     }
 

diff --git a/ydb/core/tx/columnshard/engines/scheme/index_info.cpp b/ydb/core/tx/columnshard/engines/scheme/index_info.cpp
@@ -334,6 +334,7 @@ void TIndexInfo::InitializeCaches(const std::shared_ptr<IStoragesManager>& opera
     }
     SchemaWithSpecials = IIndexInfo::AddSpecialFields(Schema);
     SchemaColumnIdsWithSpecials = IIndexInfo::AddSpecialFieldIds(SchemaColumnIds);
+    SchemaColumnIdsWithSpecialsSet = IIndexInfo::AddSpecialFieldIds(std::set<ui32>(SchemaColumnIds.begin(), SchemaColumnIds.end()));
 
     for (auto&& c : Columns) {
         AFL_VERIFY(ArrowColumnByColumnIdCache.emplace(c.first, GetColumnFieldVerified(c.first)).second);

diff --git a/ydb/core/tx/columnshard/engines/scheme/index_info.h b/ydb/core/tx/columnshard/engines/scheme/index_info.h
@@ -255,6 +255,9 @@ struct TIndexInfo : public NTable::TScheme::TTableSchema, public IIndexInfo {
     std::vector<TString> GetColumnNames(const std::vector<ui32>& ids) const;
     std::vector<std::string> GetColumnSTLNames(const std::vector<ui32>& ids) const;
     const std::vector<ui32>& GetColumnIds(const bool withSpecial = true) const;
+    const std::set<ui32>& GetColumnIdsSet() const {
+        return SchemaColumnIdsWithSpecialsSet;
+    }
     const std::vector<ui32>& GetPKColumnIds() const {
         AFL_VERIFY(PKColumnIds.size());
         return PKColumnIds;
@@ -324,6 +327,7 @@ struct TIndexInfo : public NTable::TScheme::TTableSchema, public IIndexInfo {
     TString Name;
     std::vector<ui32> SchemaColumnIds;
     std::vector<ui32> SchemaColumnIdsWithSpecials;
+    std::set<ui32> SchemaColumnIdsWithSpecialsSet;
     std::vector<ui32> PKColumnIds;
     std::shared_ptr<arrow::Schema> Schema;
     std::shared_ptr<arrow::Schema> SchemaWithSpecials;

diff --git a/ydb/core/tx/columnshard/engines/scheme/versions/abstract_scheme.h b/ydb/core/tx/columnshard/engines/scheme/versions/abstract_scheme.h
@@ -27,6 +27,7 @@ class ISnapshotSchema {
     std::shared_ptr<NArrow::NAccessor::TColumnLoader> GetColumnLoaderVerified(const std::string& columnName) const;
 
     bool IsSpecialColumnId(const ui32 columnId) const;
+    virtual const std::set<ui32>& GetColumnIds() const = 0;
 
     virtual NArrow::NAccessor::TColumnSaver GetColumnSaver(const ui32 columnId) const = 0;
     NArrow::NAccessor::TColumnSaver GetColumnSaver(const TString& columnName) const {

diff --git a/ydb/core/tx/columnshard/engines/scheme/versions/filtered_scheme.cpp b/ydb/core/tx/columnshard/engines/scheme/versions/filtered_scheme.cpp
@@ -5,8 +5,8 @@
 namespace NKikimr::NOlap {
 
 TFilteredSnapshotSchema::TFilteredSnapshotSchema(ISnapshotSchema::TPtr originalSnapshot, const std::vector<ui32>& columnIds)
-    : TFilteredSnapshotSchema(originalSnapshot, std::set(columnIds.begin(), columnIds.end()))
-{}
+    : TFilteredSnapshotSchema(originalSnapshot, std::set(columnIds.begin(), columnIds.end())) {
+}
 
 TFilteredSnapshotSchema::TFilteredSnapshotSchema(ISnapshotSchema::TPtr originalSnapshot, const std::set<ui32>& columnIds)
     : OriginalSnapshot(originalSnapshot)
@@ -48,7 +48,14 @@ std::shared_ptr<TColumnLoader> TFilteredSnapshotSchema::GetColumnLoaderOptional(
 }
 
 std::optional<ui32> TFilteredSnapshotSchema::GetColumnIdOptional(const std::string& columnName) const {
-    return OriginalSnapshot->GetColumnIdOptional(columnName);
+    auto result = OriginalSnapshot->GetColumnIdOptional(columnName);
+    if (!result) {
+        return result;
+    }
+    if (!ColumnIds.contains(*result)) {
+        return std::nullopt;
+    }
+    return result;
 }
 
 int TFilteredSnapshotSchema::GetFieldIndex(const ui32 columnId) const {

diff --git a/ydb/core/tx/columnshard/engines/scheme/versions/filtered_scheme.h b/ydb/core/tx/columnshard/engines/scheme/versions/filtered_scheme.h
@@ -9,14 +9,17 @@ namespace NKikimr::NOlap {
 class TFilteredSnapshotSchema: public ISnapshotSchema {
     ISnapshotSchema::TPtr OriginalSnapshot;
     std::shared_ptr<arrow::Schema> Schema;
-    YDB_READONLY_DEF(std::set<ui32>, ColumnIds);
+    std::set<ui32> ColumnIds;
 protected:
     virtual TString DoDebugString() const override;
 public:
     TFilteredSnapshotSchema(ISnapshotSchema::TPtr originalSnapshot, const std::vector<ui32>& columnIds);
     TFilteredSnapshotSchema(ISnapshotSchema::TPtr originalSnapshot, const std::set<ui32>& columnIds);
     TFilteredSnapshotSchema(ISnapshotSchema::TPtr originalSnapshot, const std::set<std::string>& columnNames);
 
+    virtual const std::set<ui32>& GetColumnIds() const override {
+        return ColumnIds;
+    }
     TColumnSaver GetColumnSaver(const ui32 columnId) const override;
     std::shared_ptr<TColumnLoader> GetColumnLoaderOptional(const ui32 columnId) const override;
     std::optional<ui32> GetColumnIdOptional(const std::string& columnName) const override;

diff --git a/ydb/core/tx/columnshard/engines/scheme/versions/snapshot_scheme.h b/ydb/core/tx/columnshard/engines/scheme/versions/snapshot_scheme.h
@@ -23,6 +23,10 @@ class TSnapshotSchema: public ISnapshotSchema {
 public:
     TSnapshotSchema(TIndexInfo&& indexInfo, const TSnapshot& snapshot);
 
+    virtual const std::set<ui32>& GetColumnIds() const override {
+        return IndexInfo.GetColumnIdsSet();
+    }
+
     TColumnSaver GetColumnSaver(const ui32 columnId) const override;
     std::shared_ptr<TColumnLoader> GetColumnLoaderOptional(const ui32 columnId) const override;
     std::optional<ui32> GetColumnIdOptional(const std::string& columnName) const override;