|
1 | 1 | #pragma once
|
2 | 2 |
|
3 | 3 | #include "kqp_ut_common.h"
|
| 4 | + |
| 5 | +#include <ydb/core/tx/columnshard/test_helper/columnshard_ut_common.h> |
| 6 | + |
4 | 7 | #include <ydb/library/accessor/accessor.h>
|
5 |
| -#include <ydb/library/formats/arrow/simple_builder/filler.h> |
6 | 8 | #include <ydb/library/formats/arrow/simple_builder/array.h>
|
7 | 9 | #include <ydb/library/formats/arrow/simple_builder/batch.h>
|
| 10 | +#include <ydb/library/formats/arrow/simple_builder/filler.h> |
8 | 11 | #include <ydb/public/lib/scheme_types/scheme_type_id.h>
|
9 | 12 | #include <ydb/public/sdk/cpp/client/ydb_table/table.h>
|
10 | 13 | #include <ydb/public/sdk/cpp/client/ydb_types/status_codes.h>
|
11 |
| -#include <ydb/core/tx/columnshard/test_helper/columnshard_ut_common.h> |
12 | 14 |
|
13 | 15 | #include <contrib/libs/apache/arrow/cpp/src/arrow/type.h>
|
14 | 16 |
|
15 | 17 | namespace NKikimr {
|
16 | 18 | namespace NKqp {
|
17 |
| - class TTestHelper { |
| 19 | +class TTestHelper { |
| 20 | +public: |
| 21 | + class TCompression { |
| 22 | + YDB_ACCESSOR(TString, SerializerName, "ARROW_SERIALIZER"); |
| 23 | + YDB_ACCESSOR(arrow::Compression::type, Type, arrow::Compression::type::UNCOMPRESSED); |
| 24 | + YDB_ACCESSOR(i32, CompressionLevel, Max<i32>()); |
| 25 | + }; |
| 26 | + |
| 27 | + class TColumnSchema { |
| 28 | + YDB_ACCESSOR_DEF(TString, Name); |
| 29 | + YDB_ACCESSOR_DEF(NScheme::TTypeInfo, TypeInfo); |
| 30 | + YDB_FLAG_ACCESSOR(Nullable, true); |
| 31 | + |
18 | 32 | public:
|
19 |
| - class TColumnSchema { |
20 |
| - YDB_ACCESSOR_DEF(TString, Name); |
21 |
| - YDB_ACCESSOR_DEF(NScheme::TTypeInfo, TypeInfo); |
22 |
| - YDB_FLAG_ACCESSOR(Nullable, true); |
23 |
| - public: |
24 |
| - TString BuildQuery() const; |
25 |
| - |
26 |
| - TColumnSchema& SetType(NScheme::TTypeId typeId); |
27 |
| - }; |
28 |
| - |
29 |
| - using TUpdatesBuilder = NColumnShard::TTableUpdatesBuilder; |
30 |
| - |
31 |
| - class TColumnTableBase { |
32 |
| - YDB_ACCESSOR_DEF(TString, Name); |
33 |
| - YDB_ACCESSOR_DEF(TVector<TColumnSchema>, Schema); |
34 |
| - YDB_ACCESSOR_DEF(TVector<TString>, PrimaryKey); |
35 |
| - YDB_ACCESSOR_DEF(TVector<TString>, Sharding); |
36 |
| - YDB_ACCESSOR(ui32, MinPartitionsCount, 1); |
37 |
| - |
38 |
| - std::optional<std::pair<TString, TString>> TTLConf; |
39 |
| - public: |
40 |
| - TString BuildQuery() const; |
41 |
| - std::shared_ptr<arrow::Schema> GetArrowSchema(const TVector<TColumnSchema>& columns); |
42 |
| - |
43 |
| - TColumnTableBase& SetTTL(const TString& columnName, const TString& ttlConf) { |
44 |
| - TTLConf = std::make_pair(columnName, ttlConf); |
45 |
| - return *this; |
46 |
| - } |
47 |
| - |
48 |
| - private: |
49 |
| - virtual TString GetObjectType() const = 0; |
50 |
| - TString BuildColumnsStr(const TVector<TColumnSchema>& clumns) const; |
51 |
| - std::shared_ptr<arrow::Field> BuildField(const TString name, const NScheme::TTypeInfo& typeInfo, bool nullable) const; |
52 |
| - }; |
53 |
| - |
54 |
| - class TColumnTable : public TColumnTableBase { |
55 |
| - private: |
56 |
| - TString GetObjectType() const override; |
57 |
| - }; |
58 |
| - |
59 |
| - class TColumnTableStore : public TColumnTableBase { |
60 |
| - private: |
61 |
| - TString GetObjectType() const override; |
62 |
| - }; |
| 33 | + TString BuildQuery() const; |
63 | 34 |
|
64 |
| - private: |
65 |
| - std::unique_ptr<TKikimrRunner> Kikimr; |
66 |
| - std::unique_ptr<NYdb::NTable::TTableClient> TableClient; |
67 |
| - std::unique_ptr<NYdb::NTable::TSession> Session; |
| 35 | + TColumnSchema& SetType(NScheme::TTypeId typeId); |
| 36 | + }; |
| 37 | + |
| 38 | + using TUpdatesBuilder = NColumnShard::TTableUpdatesBuilder; |
| 39 | + |
| 40 | + class TColumnTableBase { |
| 41 | + YDB_ACCESSOR_DEF(TString, Name); |
| 42 | + YDB_ACCESSOR_DEF(TVector<TColumnSchema>, Schema); |
| 43 | + YDB_ACCESSOR_DEF(TVector<TString>, PrimaryKey); |
| 44 | + YDB_ACCESSOR_DEF(TVector<TString>, Sharding); |
| 45 | + YDB_ACCESSOR(ui32, MinPartitionsCount, 1); |
| 46 | + |
| 47 | + std::optional<std::pair<TString, TString>> TTLConf; |
68 | 48 |
|
69 | 49 | public:
|
70 |
| - TTestHelper(const TKikimrSettings& settings); |
71 |
| - TKikimrRunner& GetKikimr(); |
72 |
| - TTestActorRuntime& GetRuntime(); |
73 |
| - NYdb::NTable::TSession& GetSession(); |
74 |
| - void CreateTable(const TColumnTableBase& table, const NYdb::EStatus expectedStatus = NYdb::EStatus::SUCCESS); |
75 |
| - void DropTable(const TString& tableName); |
76 |
| - void CreateTier(const TString& tierName); |
77 |
| - TString CreateTieringRule(const TString& tierName, const TString& columnName); |
78 |
| - void SetTiering(const TString& tableName, const TString& ruleName); |
79 |
| - void ResetTiering(const TString& tableName); |
80 |
| - void BulkUpsert(const TColumnTable& table, TTestHelper::TUpdatesBuilder& updates, const Ydb::StatusIds_StatusCode& opStatus = Ydb::StatusIds::SUCCESS); |
81 |
| - void BulkUpsert(const TColumnTable& table, std::shared_ptr<arrow::RecordBatch> batch, const Ydb::StatusIds_StatusCode& opStatus = Ydb::StatusIds::SUCCESS); |
82 |
| - void ReadData(const TString& query, const TString& expected, const NYdb::EStatus opStatus = NYdb::EStatus::SUCCESS); |
83 |
| - void RebootTablets(const TString& tableName); |
84 |
| - void WaitTabletDeletionInHive(ui64 tabletId, TDuration duration); |
| 50 | + TString BuildQuery() const; |
| 51 | + TString BuildAlterCompressionQuery(const TString& columnName, const TCompression& compression) const; |
| 52 | + std::shared_ptr<arrow::Schema> GetArrowSchema(const TVector<TColumnSchema>& columns); |
| 53 | + |
| 54 | + TColumnTableBase& SetTTL(const TString& columnName, const TString& ttlConf) { |
| 55 | + TTLConf = std::make_pair(columnName, ttlConf); |
| 56 | + return *this; |
| 57 | + } |
| 58 | + |
| 59 | + private: |
| 60 | + virtual TString GetObjectType() const = 0; |
| 61 | + TString BuildColumnsStr(const TVector<TColumnSchema>& clumns) const; |
| 62 | + std::shared_ptr<arrow::Field> BuildField(const TString name, const NScheme::TTypeInfo& typeInfo, bool nullable) const; |
| 63 | + }; |
| 64 | + |
| 65 | + class TColumnTable: public TColumnTableBase { |
| 66 | + private: |
| 67 | + TString GetObjectType() const override; |
| 68 | + }; |
| 69 | + |
| 70 | + class TColumnTableStore: public TColumnTableBase { |
| 71 | + private: |
| 72 | + TString GetObjectType() const override; |
85 | 73 | };
|
86 | 74 |
|
| 75 | +private: |
| 76 | + std::unique_ptr<TKikimrRunner> Kikimr; |
| 77 | + std::unique_ptr<NYdb::NTable::TTableClient> TableClient; |
| 78 | + std::unique_ptr<NYdb::NTable::TSession> Session; |
| 79 | + |
| 80 | +public: |
| 81 | + TTestHelper(const TKikimrSettings& settings); |
| 82 | + TKikimrRunner& GetKikimr(); |
| 83 | + TTestActorRuntime& GetRuntime(); |
| 84 | + NYdb::NTable::TSession& GetSession(); |
| 85 | + void CreateTable(const TColumnTableBase& table, const NYdb::EStatus expectedStatus = NYdb::EStatus::SUCCESS); |
| 86 | + void DropTable(const TString& tableName); |
| 87 | + void CreateTier(const TString& tierName); |
| 88 | + TString CreateTieringRule(const TString& tierName, const TString& columnName); |
| 89 | + void SetTiering(const TString& tableName, const TString& ruleName); |
| 90 | + void ResetTiering(const TString& tableName); |
| 91 | + void BulkUpsert( |
| 92 | + const TColumnTable& table, TTestHelper::TUpdatesBuilder& updates, const Ydb::StatusIds_StatusCode& opStatus = Ydb::StatusIds::SUCCESS); |
| 93 | + void BulkUpsert(const TColumnTable& table, std::shared_ptr<arrow::RecordBatch> batch, |
| 94 | + const Ydb::StatusIds_StatusCode& opStatus = Ydb::StatusIds::SUCCESS); |
| 95 | + void ReadData(const TString& query, const TString& expected, const NYdb::EStatus opStatus = NYdb::EStatus::SUCCESS); |
| 96 | + void RebootTablets(const TString& tableName); |
| 97 | + void WaitTabletDeletionInHive(ui64 tabletId, TDuration duration); |
| 98 | + void SetCompression(const TColumnTableBase& columnTable, const TString& columnName, const TCompression& compression, |
| 99 | + const NYdb::EStatus expectedStatus = NYdb::EStatus::SUCCESS); |
| 100 | +}; |
87 | 101 | }
|
88 | 102 | }
|
0 commit comments