Merge pull request #3609 from s-mx/CLICKHOUSE-4090

Add primary_key, order_key, partition_key and sample_key to system.tables.

dbms/src/Core/Names.h

@@ -2,6 +2,7 @@
 
 #include <vector>
 #include <string>
+#include <set>
 #include <unordered_set>
 #include <unordered_map>
 
@@ -11,6 +12,7 @@ namespace DB
 
 using Names = std::vector<std::string>;
 using NameSet = std::unordered_set<std::string>;
+using NameOrderedSet = std::set<std::string>;
 using NameToNameMap = std::unordered_map<std::string, std::string>;
 using NameToNameSetMap = std::unordered_map<std::string, NameSet>;
 

dbms/src/Interpreters/InterpreterCreateQuery.cpp

@@ -350,7 +350,6 @@ ASTPtr InterpreterCreateQuery::formatColumns(const ColumnsDescription & columns)
     return columns_list;
 }
 
-
 ColumnsDescription InterpreterCreateQuery::getColumnsDescription(const ASTExpressionList & columns, const Context & context)
 {
     ColumnsDescription res;

dbms/src/Storages/IStorage.h

@@ -3,6 +3,7 @@
 #include <Core/Names.h>
 #include <Common/Exception.h>
 #include <Common/RWLockFIFO.h>
+#include <Core/Names.h>
 #include <Core/QueryProcessingStage.h>
 #include <Storages/ITableDeclaration.h>
 #include <Storages/SelectQueryInfo.h>
@@ -349,6 +350,24 @@ public:
     /// Returns primary expression for storage or nullptr if there is no.
     virtual ASTPtr getPrimaryExpression() const { return nullptr; }
 
+    /// Returns partition expression for storage or nullptr if there is no.
+    virtual ASTPtr getPartitionExpression() const { return nullptr; }
+
+    /// Returns secondary expression for storage or nullptr if there is no.
+    virtual ASTPtr getOrderExpression() const { return nullptr; }
+
+    /// Returns sampling key names for storage or empty vector if there is no.
+    virtual Names getSamplingExpressionNames() const { return {}; }
+
+    /// Returns primary key names for storage or empty vector if there is no.
+    virtual Names getPrimaryExpressionNames() const { return {}; }
+
+    /// Returns partition key names for storage or empty vector if there is no.
+    virtual Names getPartitionExpressionNames() const { return {}; }
+
+    /// Returns order key names for storage or empty vector if there is no.
+    virtual Names getOrderExpressionNames() const { return {}; }
+
     using ITableDeclaration::ITableDeclaration;
     using std::enable_shared_from_this<IStorage>::shared_from_this;
 

dbms/src/Storages/StorageMergeTree.cpp

@@ -938,4 +938,34 @@ ActionLock StorageMergeTree::getActionLock(StorageActionBlockType action_type)
     return {};
 }
 
+Names StorageMergeTree::getSamplingExpressionNames() const
+{
+    NameOrderedSet names;
+    const auto & expr = data.sampling_expression;
+    if (expr)
+        expr->collectIdentifierNames(names);
+
+    return Names(names.begin(), names.end());
+}
+
+Names StorageMergeTree::getPrimaryExpressionNames() const
+{
+    return data.getPrimarySortColumns();
+}
+
+Names StorageMergeTree::getPartitionExpressionNames() const
+{
+    NameOrderedSet names;
+    const auto & expr = data.partition_expr_ast;
+    if (expr)
+        expr->collectIdentifierNames(names);
+
+    return Names(names.cbegin(), names.cend());
+}
+
+Names StorageMergeTree::getOrderExpressionNames() const
+{
+    return data.getSortColumns();
+}
+
 }

dbms/src/Storages/StorageMergeTree.h

@@ -2,6 +2,7 @@
 
 #include <ext/shared_ptr_helper.h>
 
+#include <Core/Names.h>
 #include <Storages/IStorage.h>
 #include <Storages/MergeTree/MergeTreeData.h>
 #include <Storages/MergeTree/MergeTreeDataSelectExecutor.h>
@@ -98,6 +99,17 @@ public:
 
     ASTPtr getPrimaryExpression() const override { return data.primary_expr_ast; }
 
+    ASTPtr getPartitionExpression() const override { return data.partition_expr_ast; }
+
+    ASTPtr getOrderExpression() const override { return data.secondary_sort_expr_ast; }
+
+    Names getSamplingExpressionNames() const override;
+
+    Names getPrimaryExpressionNames() const override;
+
+    Names getPartitionExpressionNames() const override;
+
+    Names getOrderExpressionNames() const override;
 private:
     String path;
     String database_name;

dbms/src/Storages/StorageReplicatedMergeTree.cpp

@@ -4856,4 +4856,34 @@ bool StorageReplicatedMergeTree::dropPartsInPartition(
     return true;
 }
 
+Names StorageReplicatedMergeTree::getSamplingExpressionNames() const
+{
+    NameOrderedSet names;
+    const auto & expr = data.sampling_expression;
+    if (expr)
+        expr->collectIdentifierNames(names);
+
+    return Names(names.begin(), names.end());
+}
+
+Names StorageReplicatedMergeTree::getPrimaryExpressionNames() const
+{
+    return data.getPrimarySortColumns();
+}
+
+Names StorageReplicatedMergeTree::getOrderExpressionNames() const
+{
+    return data.getSortColumns();
+}
+
+Names StorageReplicatedMergeTree::getPartitionExpressionNames() const
+{
+    NameOrderedSet names;
+    const auto & expr = data.partition_expr_ast;
+    if (expr)
+        expr->collectIdentifierNames(names);
+
+    return Names(names.cbegin(), names.cend());
+}
+
 }

dbms/src/Storages/StorageReplicatedMergeTree.h

@@ -197,6 +197,14 @@ public:
 
     ASTPtr getPrimaryExpression() const override { return data.primary_expr_ast; }
 
+    Names getSamplingExpressionNames() const override;
+
+    Names getPrimaryExpressionNames() const override;
+
+    Names getOrderExpressionNames() const override;
+
+    Names getPartitionExpressionNames() const override;
+
 private:
     /// Delete old parts from disk and from ZooKeeper.
     void clearOldPartsAndRemoveFromZK();

dbms/src/Storages/System/StorageSystemColumns.cpp

@@ -36,6 +36,10 @@ StorageSystemColumns::StorageSystemColumns(const std::string & name_)
         { "data_compressed_bytes",      std::make_shared<DataTypeUInt64>() },
         { "data_uncompressed_bytes",    std::make_shared<DataTypeUInt64>() },
         { "marks_bytes",                std::make_shared<DataTypeUInt64>() },
+        { "is_in_primary_key", std::make_shared<DataTypeUInt8>() },
+        { "is_in_order_key", std::make_shared<DataTypeUInt8>() },
+        { "is_in_partition_key", std::make_shared<DataTypeUInt8>() },
+        { "is_in_sample_key", std::make_shared<DataTypeUInt8>() },
     }));
 }
 
@@ -80,6 +84,10 @@ protected:
 
             NamesAndTypesList columns;
             ColumnDefaults column_defaults;
+            Names partition_key_names;
+            Names order_key_names;
+            Names primary_key_names;
+            Names sampling_key_names;
             MergeTreeData::ColumnSizeByName column_sizes;
 
             {
@@ -106,6 +114,11 @@ protected:
                 columns = storage->getColumns().getAll();
                 column_defaults = storage->getColumns().defaults;
 
+                partition_key_names = storage->getPartitionExpressionNames();
+                order_key_names = storage->getOrderExpressionNames();
+                primary_key_names = storage->getPrimaryExpressionNames();
+                sampling_key_names = storage->getSamplingExpressionNames();
+
                 /** Info about sizes of columns for tables of MergeTree family.
                 * NOTE: It is possible to add getter for this info to IStorage interface.
                 */
@@ -173,6 +186,22 @@ protected:
                     }
                 }
 
+                {
+                    auto find_in_vector = [&key = column.name](const Names& names)
+                    {
+                        return std::find(names.cbegin(), names.cend(), key) != names.end();
+                    };
+
+                    if (columns_mask[src_index++])
+                        res_columns[res_index++]->insert(find_in_vector(primary_key_names));
+                    if (columns_mask[src_index++])
+                        res_columns[res_index++]->insert(find_in_vector(order_key_names));
+                    if (columns_mask[src_index++])
+                        res_columns[res_index++]->insert(find_in_vector(partition_key_names));
+                    if (columns_mask[src_index++])
+                        res_columns[res_index++]->insert(find_in_vector(sampling_key_names));
+                }
+
                 ++rows_count;
             }
         }

dbms/src/Storages/System/StorageSystemTables.cpp

@@ -39,7 +39,11 @@ StorageSystemTables::StorageSystemTables(const std::string & name_)
         {"dependencies_database", std::make_shared<DataTypeArray>(std::make_shared<DataTypeString>())},
         {"dependencies_table", std::make_shared<DataTypeArray>(std::make_shared<DataTypeString>())},
         {"create_table_query", std::make_shared<DataTypeString>()},
-        {"engine_full", std::make_shared<DataTypeString>()}
+        {"engine_full", std::make_shared<DataTypeString>()},
+        {"primary_key", std::make_shared<DataTypeString>()},
+        {"order_key", std::make_shared<DataTypeString>()},
+        {"partition_key", std::make_shared<DataTypeString>()},
+        {"sample_key", std::make_shared<DataTypeString>()},
     }));
 }
 
@@ -65,7 +69,7 @@ public:
         size_t max_block_size,
         ColumnPtr databases,
         const Context & context)
-        : columns_mask(columns_mask), header(header), max_block_size(max_block_size), databases(std::move(databases)), context(context) {}
+        : columns_mask(std::move(columns_mask)), header(std::move(header)), max_block_size(max_block_size), databases(std::move(databases)), context(context) {}
 
     String getName() const override { return "Tables"; }
     Block getHeader() const override { return header; }
@@ -144,6 +148,18 @@ protected:
 
                         if (columns_mask[src_index++])
                             res_columns[res_index++]->insert(table.second->getName());
+
+                        if (columns_mask[src_index++])
+                            res_columns[res_index++]->insertDefault();
+
+                        if (columns_mask[src_index++])
+                            res_columns[res_index++]->insertDefault();
+
+                        if (columns_mask[src_index++])
+                            res_columns[res_index++]->insertDefault();
+
+                        if (columns_mask[src_index++])
+                            res_columns[res_index++]->insertDefault();
                     }
                 }
 
@@ -173,7 +189,7 @@ protected:
                     res_columns[res_index++]->insert(tables_it->table()->getName());
 
                 if (columns_mask[src_index++])
-                    res_columns[res_index++]->insert(0u);
+                    res_columns[res_index++]->insert(0u);  // is_temporary
 
                 if (columns_mask[src_index++])
                     res_columns[res_index++]->insert(tables_it->table()->getDataPath());
@@ -234,6 +250,42 @@ protected:
                         res_columns[res_index++]->insert(engine_full);
                     }
                 }
+                else
+                    src_index += 2;
+
+                const auto table_it = context.getTable(database_name, table_name);
+                ASTPtr expression_ptr;
+                if (columns_mask[src_index++])
+                {
+                    if ((expression_ptr = table_it->getPrimaryExpression()))
+                        res_columns[res_index++]->insert(queryToString(expression_ptr));
+                    else
+                        res_columns[res_index++]->insertDefault();
+                }
+
+                if (columns_mask[src_index++])
+                {
+                    if ((expression_ptr = table_it->getOrderExpression()))
+                        res_columns[res_index++]->insert(queryToString(expression_ptr));
+                    else
+                        res_columns[res_index++]->insertDefault();
+                }
+
+                if (columns_mask[src_index++])
+                {
+                    if ((expression_ptr = table_it->getPartitionExpression()))
+                        res_columns[res_index++]->insert(queryToString(expression_ptr));
+                    else
+                        res_columns[res_index++]->insertDefault();
+                }
+
+                if (columns_mask[src_index++])
+                {
+                    if ((expression_ptr = table_it->getSamplingExpression()))
+                        res_columns[res_index++]->insert(queryToString(expression_ptr));
+                    else
+                        res_columns[res_index++]->insertDefault();
+                }
             }
         }
 

dbms/tests/queries/0_stateless/00753_system_columns_and_system_tables.reference

@@ -0,0 +1,25 @@
+┌─name────────────────┬─primary_key─┬─order_key─┬─partition_key─┬─sample_key─┐
+│ check_system_tables │ name1       │           │ name2         │ name1      │
+└─────────────────────┴─────────────┴───────────┴───────────────┴────────────┘
+┌─name──┬─is_in_primary_key─┬─is_in_order_key─┬─is_in_partition_key─┬─is_in_sample_key─┐
+│ name1 │                 1 │               1 │                   0 │                1 │
+│ name2 │                 0 │               0 │                   1 │                0 │
+│ name3 │                 0 │               0 │                   0 │                0 │
+└───────┴───────────────────┴─────────────────┴─────────────────────┴──────────────────┘
+┌─name────────────────┬─primary_key─┬─order_key─┬─partition_key─┬─sample_key─┐
+│ check_system_tables │ date        │ version   │ date          │            │
+└─────────────────────┴─────────────┴───────────┴───────────────┴────────────┘
+┌─name────┬─is_in_primary_key─┬─is_in_order_key─┬─is_in_partition_key─┬─is_in_sample_key─┐
+│ date    │                 1 │               1 │                   1 │                0 │
+│ value   │                 0 │               0 │                   0 │                0 │
+│ version │                 0 │               1 │                   0 │                0 │
+│ sign    │                 0 │               0 │                   0 │                0 │
+└─────────┴───────────────────┴─────────────────┴─────────────────────┴──────────────────┘
+┌─name────────────────┬─primary_key───────────────────────┬─order_key─┬─partition_key───┬─sample_key────────┐
+│ check_system_tables │ Counter, Event, intHash32(UserId) │           │ toYYYYMM(Event) │ intHash32(UserId) │
+└─────────────────────┴───────────────────────────────────┴───────────┴─────────────────┴───────────────────┘
+┌─name────┬─is_in_primary_key─┬─is_in_order_key─┬─is_in_partition_key─┬─is_in_sample_key─┐
+│ Event   │                 1 │               1 │                   1 │                0 │
+│ UserId  │                 0 │               0 │                   0 │                1 │
+│ Counter │                 1 │               1 │                   0 │                0 │
+└─────────┴───────────────────┴─────────────────┴─────────────────────┴──────────────────┘

dbms/tests/queries/0_stateless/00753_system_columns_and_system_tables.sql

@@ -0,0 +1,68 @@
+CREATE DATABASE IF NOT EXISTS test;
+DROP TABLE IF EXISTS test.check_system_tables;
+
+-- Check MergeTree declaration in new format
+CREATE TABLE test.check_system_tables
+  (
+    name1 UInt8,
+    name2 UInt8,
+    name3 UInt8
+  ) ENGINE = MergeTree()
+    ORDER BY name1
+    PARTITION BY name2
+    SAMPLE BY name1;
+
+SELECT name, primary_key, order_key, partition_key, sample_key
+FROM system.tables
+WHERE name = 'check_system_tables'
+FORMAT PrettyCompactNoEscapes;
+
+SELECT name, is_in_primary_key, is_in_order_key, is_in_partition_key, is_in_sample_key
+FROM system.columns
+WHERE table = 'check_system_tables'
+FORMAT PrettyCompactNoEscapes;
+
+DROP TABLE IF EXISTS test.check_system_tables;
+
+-- Check VersionedCollapsingMergeTree
+CREATE TABLE test.check_system_tables
+  (
+    date Date,
+    value String,
+    version UInt64,
+    sign Int8
+  ) ENGINE = VersionedCollapsingMergeTree(sign, version)
+    PARTITION BY date
+    ORDER BY date;
+
+SELECT name, primary_key, order_key, partition_key, sample_key
+FROM system.tables
+WHERE name = 'check_system_tables'
+FORMAT PrettyCompactNoEscapes;
+
+SELECT name, is_in_primary_key, is_in_order_key, is_in_partition_key, is_in_sample_key
+FROM system.columns
+WHERE table = 'check_system_tables'
+FORMAT PrettyCompactNoEscapes;
+
+DROP TABLE IF EXISTS test.check_system_tables;
+
+-- Check MergeTree declaration in old format
+CREATE TABLE test.check_system_tables
+  (
+    Event Date,
+    UserId UInt32,
+    Counter UInt32
+  ) ENGINE = MergeTree(Event, intHash32(UserId), (Counter, Event, intHash32(UserId)), 8192);
+
+SELECT name, primary_key, order_key, partition_key, sample_key
+FROM system.tables
+WHERE name = 'check_system_tables'
+FORMAT PrettyCompactNoEscapes;
+
+SELECT name, is_in_primary_key, is_in_order_key, is_in_partition_key, is_in_sample_key
+FROM system.columns
+WHERE table = 'check_system_tables'
+FORMAT PrettyCompactNoEscapes;
+
+DROP TABLE IF EXISTS test.check_system_tables;