Merge pull request #54477 from azat/view-sampling

Support SAMPLE BY for VIEW

src/Storages/StorageView.cpp

@@ -101,6 +101,19 @@ ContextPtr getViewContext(ContextPtr context)
     return view_context;
 }
 
+ASTTableExpression * getFirstTableExpression(ASTSelectQuery & select_query)
+{
+    if (!select_query.tables() || select_query.tables()->children.empty())
+        throw Exception(ErrorCodes::LOGICAL_ERROR, "Logical error: no table expression in view select AST");
+
+    auto * select_element = select_query.tables()->children[0]->as<ASTTablesInSelectQueryElement>();
+
+    if (!select_element->table_expression)
+        throw Exception(ErrorCodes::LOGICAL_ERROR, "Logical error: incorrect table expression");
+
+    return select_element->table_expression->as<ASTTableExpression>();
+}
+
 }
 
 StorageView::StorageView(
@@ -148,6 +161,21 @@ void StorageView::read(
         current_inner_query = query_info.view_query->clone();
     }
 
+    const auto & select_query = query_info.query->as<ASTSelectQuery &>();
+    if (auto sample_size = select_query.sampleSize(), sample_offset = select_query.sampleOffset(); sample_size || sample_offset)
+    {
+        for (auto & inner_select_query : current_inner_query->as<ASTSelectWithUnionQuery &>().list_of_selects->children)
+        {
+            if (auto * select = inner_select_query->as<ASTSelectQuery>(); select)
+            {
+                ASTTableExpression * table_expression = getFirstTableExpression(*select);
+
+                table_expression->sample_offset = sample_offset;
+                table_expression->sample_size = sample_size;
+            }
+        }
+    }
+
     auto options = SelectQueryOptions(QueryProcessingStage::Complete, 0, false, query_info.settings_limit_offset_done);
 
     if (context->getSettingsRef().allow_experimental_analyzer)
@@ -196,19 +224,6 @@ void StorageView::read(
     query_plan.addStep(std::move(converting));
 }
 
-static ASTTableExpression * getFirstTableExpression(ASTSelectQuery & select_query)
-{
-    if (!select_query.tables() || select_query.tables()->children.empty())
-        throw Exception(ErrorCodes::LOGICAL_ERROR, "Logical error: no table expression in view select AST");
-
-    auto * select_element = select_query.tables()->children[0]->as<ASTTablesInSelectQueryElement>();
-
-    if (!select_element->table_expression)
-        throw Exception(ErrorCodes::LOGICAL_ERROR, "Logical error: incorrect table expression");
-
-    return select_element->table_expression->as<ASTTableExpression>();
-}
-
 void StorageView::replaceQueryParametersIfParametrizedView(ASTPtr & outer_query, const NameToNameMap & parameter_values)
 {
     ReplaceQueryParameterVisitor visitor(parameter_values);

tests/queries/0_stateless/02881_view_sampling.reference

@@ -0,0 +1,7 @@
+-- { echo }
+select * from data_mt sample 0.1 order by key;
+9
+select * from view_mt sample 0.1 order by key;
+9
+select * from data_mem sample 0.1 order by key; -- { serverError SAMPLING_NOT_SUPPORTED }
+select * from view_mem sample 0.1 order by key; -- { serverError SAMPLING_NOT_SUPPORTED }

tests/queries/0_stateless/02881_view_sampling.sql

@@ -0,0 +1,18 @@
+drop table if exists data_mt;
+drop table if exists view_mt;
+drop table if exists data_mem;
+drop table if exists view_mem;
+
+create table data_mt (key Int) engine=MergeTree() order by (key, sipHash64(key)) sample by sipHash64(key);
+insert into data_mt select * from numbers(10);
+create view view_mt as select * from data_mt;
+
+create table data_mem (key Int) engine=Memory();
+insert into data_mem select * from numbers(10);
+create view view_mem as select * from data_mem;
+
+-- { echo }
+select * from data_mt sample 0.1 order by key;
+select * from view_mt sample 0.1 order by key;
+select * from data_mem sample 0.1 order by key; -- { serverError SAMPLING_NOT_SUPPORTED }
+select * from view_mem sample 0.1 order by key; -- { serverError SAMPLING_NOT_SUPPORTED }

tests/queries/1_stateful/00056_view.reference

@@ -19,13 +19,13 @@
 59183	85379
 33010362	77807
 800784	77492
-1704509	523264
-732797	475698
-598875	337212
-792887	252197
-3807842	196036
-25703952	147211
-716829	90109
-59183	85379
-33010362	77807
-800784	77492
+732797	46144
+1704509	40129
+598875	30482
+792887	22585
+3807842	18891
+11312316	13181
+25703952	11177
+59183	9935
+4379238	8593
+716829	8411