Trivial count optimization is disabled

src/Storages/StorageLoop.cpp

@@ -0,0 +1,63 @@
+#include "StorageLoop.h"
+#include <Interpreters/DatabaseCatalog.h>
+#include <Databases/IDatabase.h>
+#include <Interpreters/Context.h>
+#include <Common/Exception.h>
+#include <Storages/StorageFactory.h>
+#include <Processors/QueryPlan/QueryPlan.h>
+#include <Processors/QueryPlan/BuildQueryPipelineSettings.h>
+#include <Processors/QueryPlan/Optimizations/QueryPlanOptimizationSettings.h>
+#include <Storages/checkAndGetLiteralArgument.h>
+#include <Parsers/ASTSelectQuery.h>
+#include <Storages/SelectQueryInfo.h>
+
+
+namespace DB
+{
+namespace ErrorCodes
+{
+extern const int UNKNOWN_TABLE;
+}
+StorageLoop::StorageLoop(
+    const StorageID & table_id_,
+    StoragePtr inner_storage_)
+    : IStorage(table_id_)
+    , inner_storage(std::move(inner_storage_))
+{
+    StorageInMemoryMetadata storage_metadata = inner_storage->getInMemoryMetadata();
+    setInMemoryMetadata(storage_metadata);
+}
+
+
+void StorageLoop::read(
+    QueryPlan & query_plan,
+    const Names & column_names,
+    const StorageSnapshotPtr & storage_snapshot,
+    SelectQueryInfo & query_info,
+    ContextPtr context,
+    QueryProcessingStage::Enum processed_stage,
+    size_t max_block_size,
+    size_t num_streams)
+{
+
+    query_info.optimize_trivial_count = false;
+
+    inner_storage->read(query_plan,
+                        column_names,
+                        storage_snapshot,
+                        query_info,
+                        context,
+                        processed_stage,
+                        max_block_size,
+                        num_streams);
+}
+
+void registerStorageLoop(StorageFactory & factory)
+{
+    factory.registerStorage("Loop", [](const StorageFactory::Arguments & args)
+                            {
+                                StoragePtr inner_storage;
+                                return std::make_shared<StorageLoop>(args.table_id, inner_storage);
+                            });
+}
+}

src/Storages/StorageLoop.h

@@ -0,0 +1,33 @@
+#pragma once
+#include "config.h"
+#include <Storages/IStorage.h>
+
+
+namespace DB
+{
+
+class StorageLoop final : public IStorage
+{
+public:
+    StorageLoop(
+        const StorageID & table_id,
+        StoragePtr inner_storage_);
+
+    std::string getName() const override { return "Loop"; }
+
+    void read(
+        QueryPlan & query_plan,
+        const Names & column_names,
+        const StorageSnapshotPtr & storage_snapshot,
+        SelectQueryInfo & query_info,
+        ContextPtr context,
+        QueryProcessingStage::Enum processed_stage,
+        size_t max_block_size,
+        size_t num_streams) override;
+
+    bool supportsTrivialCountOptimization(const StorageSnapshotPtr &, ContextPtr) const override { return false; }
+
+private:
+    StoragePtr inner_storage;
+};
+}

src/Storages/registerStorages.cpp

@@ -25,6 +25,7 @@ void registerStorageLiveView(StorageFactory & factory);
 void registerStorageGenerateRandom(StorageFactory & factory);
 void registerStorageExecutable(StorageFactory & factory);
 void registerStorageWindowView(StorageFactory & factory);
+void registerStorageLoop(StorageFactory & factory);
 #if USE_RAPIDJSON || USE_SIMDJSON
 void registerStorageFuzzJSON(StorageFactory & factory);
 #endif
@@ -126,6 +127,7 @@ void registerStorages()
     registerStorageGenerateRandom(factory);
     registerStorageExecutable(factory);
     registerStorageWindowView(factory);
+    registerStorageLoop(factory);
 #if USE_RAPIDJSON || USE_SIMDJSON
     registerStorageFuzzJSON(factory);
 #endif

src/TableFunctions/TableFunctionLoop.cpp

@@ -1,20 +1,17 @@
 #include "config.h"
 #include <TableFunctions/ITableFunction.h>
 #include <TableFunctions/TableFunctionFactory.h>
-#include <Storages/StorageMaterializedView.h>
 #include <Interpreters/Context.h>
 #include <Interpreters/DatabaseCatalog.h>
+#include <Databases/IDatabase.h>
 #include <Parsers/ASTFunction.h>
+#include <Parsers/ASTIdentifier.h>
 #include <Common/Exception.h>
-#include <Common/parseAddress.h>
 #include <Interpreters/evaluateConstantExpression.h>
 #include <Storages/checkAndGetLiteralArgument.h>
-#include <Common/quoteString.h>
-#include <Parsers/ASTLiteral.h>
+#include <Storages/StorageLoop.h>
 #include "registerTableFunctions.h"
 
-#include <Common/parseRemoteDescription.h>
-
 namespace DB
 {
 namespace ErrorCodes
@@ -55,30 +52,45 @@ void TableFunctionLoop::parseArguments(const ASTPtr & ast_function, ContextPtr c
     if (args.empty())
         throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "No arguments provided for table function 'loop'");
 
+    if (args.size() == 1)
+    {
+        if (const auto * id = args[0]->as<ASTIdentifier>())
+        {
+            String id_name = id->name();
+
+            size_t dot_pos = id_name.find('.');
+            if (dot_pos != String::npos)
+            {
+                database_name_ = id_name.substr(0, dot_pos);
+                table_name_ = id_name.substr(dot_pos + 1);
+            }
+            else
+            {
+                table_name_ = id_name;
+            }
+        }
+        else if (const auto * func = args[0]->as<ASTFunction>())
+        {
+            inner_table_function_ast = args[0];
+        }
+        else
+        {
+            throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "Expected identifier or function for argument 1 of function 'loop', got {}", args[0]->getID());
+        }
+    }
     // loop(database, table)
-    if (args.size() == 2)
+    else if (args.size() == 2)
     {
         args[0] = evaluateConstantExpressionForDatabaseName(args[0], context);
         args[1] = evaluateConstantExpressionOrIdentifierAsLiteral(args[1], context);
 
         database_name_ = checkAndGetLiteralArgument<String>(args[0], "database");
         table_name_ = checkAndGetLiteralArgument<String>(args[1], "table");
-        /*if (const auto * lit = args[0]->as<ASTLiteral>())
-            database_name_ = lit->value.safeGet<String>();
-        else
-            throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "Expected literal for argument 1 of function 'loop', got {}", args[0]->getID());
-
-        if (const auto * lit = args[1]->as<ASTLiteral>())
-            table_name_ = lit->value.safeGet<String>();
-        else
-            throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "Expected literal for argument 2 of function 'loop', got {}", args[1]->getID());*/
     }
-    // loop(other_table_function(...))
-    else if (args.size() == 1)
-        inner_table_function_ast = args[0];
-
     else
+    {
         throw Exception(ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH, "Table function 'loop' must have 1 or 2 arguments.");
+    }
 }
 
 ColumnsDescription TableFunctionLoop::getActualTableStructure(ContextPtr context, bool is_insert_query) const
@@ -97,13 +109,18 @@ StoragePtr TableFunctionLoop::executeImpl(
     bool is_insert_query) const
 {
     StoragePtr storage;
-    if (!database_name_.empty() && !table_name_.empty())
+    if (!table_name_.empty())
     {
-        auto database = DatabaseCatalog::instance().getDatabase(database_name_);
+        String database_name = database_name_;
+        if (database_name.empty())
+            database_name = context->getCurrentDatabase();
+
+        auto database = DatabaseCatalog::instance().getDatabase(database_name);
         storage = database->tryGetTable(table_name_ ,context);
         if (!storage)
-            throw Exception(ErrorCodes::UNKNOWN_TABLE, "Table '{}' not found in database '{}'", table_name_, database_name_);
+            throw Exception(ErrorCodes::UNKNOWN_TABLE, "Table '{}' not found in database '{}'", table_name_, database_name);
     }
+
     else
     {
         auto inner_table_function = TableFunctionFactory::instance().get(inner_table_function_ast, context);
@@ -114,8 +131,13 @@ StoragePtr TableFunctionLoop::executeImpl(
             std::move(cached_columns),
             is_insert_query);
     }
-
-    return storage;
+    auto res = std::make_shared<StorageLoop>(
+        StorageID(getDatabaseName(), table_name),
+        storage
+        );
+    res->startup();
+    return res;
+    // return storage;
 }
 
 void registerTableFunctionLoop(TableFunctionFactory & factory)