Fix optimize_read_in_order with IN and a subquery.

dbms/src/Interpreters/InterpreterSelectQuery.cpp

@@ -2083,19 +2083,8 @@ void InterpreterSelectQuery::executeOrder(Pipeline & pipeline, SortingInfoPtr so
             });
         }
 
-        if (pipeline.hasMoreThanOneStream())
-        {
-            pipeline.transform([&](auto & stream)
-            {
-                stream = std::make_shared<AsynchronousBlockInputStream>(stream);
-            });
-
-            UInt64 limit_for_merging = (need_finish_sorting ? 0 : limit);
-            pipeline.firstStream() = std::make_shared<MergingSortedBlockInputStream>(
-                pipeline.streams, sorting_info->prefix_order_descr,
-                settings.max_block_size, limit_for_merging);
-            pipeline.streams.resize(1);
-        }
+        UInt64 limit_for_merging = (need_finish_sorting ? 0 : limit);
+        executeMergeSorted(pipeline, sorting_info->prefix_order_descr, limit_for_merging);
 
         if (need_finish_sorting)
         {
@@ -2217,12 +2206,20 @@ void InterpreterSelectQuery::executeMergeSorted(Pipeline & pipeline)
     SortDescription order_descr = getSortDescription(query, context);
     UInt64 limit = getLimitForSorting(query, context);
 
-    const Settings & settings = context.getSettingsRef();
-
     /// If there are several streams, then we merge them into one
     if (pipeline.hasMoreThanOneStream())
     {
         unifyStreams(pipeline, pipeline.firstStream()->getHeader());
+        executeMergeSorted(pipeline, order_descr, limit);
+    }
+}
+
+
+void InterpreterSelectQuery::executeMergeSorted(Pipeline & pipeline, const SortDescription & sort_description, UInt64 limit)
+{
+    if (pipeline.hasMoreThanOneStream())
+    {
+        const Settings & settings = context.getSettingsRef();
 
         /** MergingSortedBlockInputStream reads the sources sequentially.
           * To make the data on the remote servers prepared in parallel, we wrap it in AsynchronousBlockInputStream.
@@ -2232,8 +2229,8 @@ void InterpreterSelectQuery::executeMergeSorted(Pipeline & pipeline)
             stream = std::make_shared<AsynchronousBlockInputStream>(stream);
         });
 
-        /// Merge the sorted sources into one sorted source.
-        pipeline.firstStream() = std::make_shared<MergingSortedBlockInputStream>(pipeline.streams, order_descr, settings.max_block_size, limit);
+        pipeline.firstStream() = std::make_shared<MergingSortedBlockInputStream>(
+            pipeline.streams, sort_description, settings.max_block_size, limit);
         pipeline.streams.resize(1);
     }
 }
@@ -2244,15 +2241,20 @@ void InterpreterSelectQuery::executeMergeSorted(QueryPipeline & pipeline)
     SortDescription order_descr = getSortDescription(query, context);
     UInt64 limit = getLimitForSorting(query, context);
 
-    const Settings & settings = context.getSettingsRef();
+    executeMergeSorted(pipeline, order_descr, limit);
+}
 
+void InterpreterSelectQuery::executeMergeSorted(QueryPipeline & pipeline, const SortDescription & sort_description, UInt64 limit)
+{
     /// If there are several streams, then we merge them into one
     if (pipeline.getNumStreams() > 1)
     {
+        const Settings & settings = context.getSettingsRef();
+
         auto transform = std::make_shared<MergingSortedTransform>(
             pipeline.getHeader(),
             pipeline.getNumStreams(),
-            order_descr,
+            sort_description,
             settings.max_block_size, limit);
 
         pipeline.addPipe({ std::move(transform) });
@@ -2615,13 +2617,29 @@ void InterpreterSelectQuery::executeExtremes(QueryPipeline & pipeline)
 
 void InterpreterSelectQuery::executeSubqueriesInSetsAndJoins(Pipeline & pipeline, SubqueriesForSets & subqueries_for_sets)
 {
-    executeUnion(pipeline, {});
+    /// Merge streams to one. Use MergeSorting if data was read in sorted order, Union otherwise.
+    if (query_info.sorting_info)
+    {
+        if (pipeline.stream_with_non_joined_data)
+            throw Exception("Using read in order optimization, but has stream with non-joined data in pipeline", ErrorCodes::LOGICAL_ERROR);
+        executeMergeSorted(pipeline, query_info.sorting_info->prefix_order_descr, 0);
+    }
+    else
+        executeUnion(pipeline, {});
+
     pipeline.firstStream() = std::make_shared<CreatingSetsBlockInputStream>(
         pipeline.firstStream(), subqueries_for_sets, context);
 }
 
 void InterpreterSelectQuery::executeSubqueriesInSetsAndJoins(QueryPipeline & pipeline, SubqueriesForSets & subqueries_for_sets)
 {
+    if (query_info.sorting_info)
+    {
+        if (pipeline.hasDelayedStream())
+            throw Exception("Using read in order optimization, but has delayed stream in pipeline", ErrorCodes::LOGICAL_ERROR);
+        executeMergeSorted(pipeline, query_info.sorting_info->prefix_order_descr, 0);
+    }
+
     const Settings & settings = context.getSettingsRef();
 
     auto creating_sets = std::make_shared<CreatingSetsTransform>(

dbms/src/Interpreters/InterpreterSelectQuery.h

@@ -214,6 +214,7 @@ class InterpreterSelectQuery : public IInterpreter
     void executeDistinct(Pipeline & pipeline, bool before_order, Names columns);
     void executeExtremes(Pipeline & pipeline);
     void executeSubqueriesInSetsAndJoins(Pipeline & pipeline, std::unordered_map<String, SubqueryForSet> & subqueries_for_sets);
+    void executeMergeSorted(Pipeline & pipeline, const SortDescription & sort_description, UInt64 limit);
 
     void executeWhere(QueryPipeline & pipeline, const ExpressionActionsPtr & expression, bool remove_fiter);
     void executeAggregation(QueryPipeline & pipeline, const ExpressionActionsPtr & expression, bool overflow_row, bool final);
@@ -231,6 +232,7 @@ class InterpreterSelectQuery : public IInterpreter
     void executeDistinct(QueryPipeline & pipeline, bool before_order, Names columns);
     void executeExtremes(QueryPipeline & pipeline);
     void executeSubqueriesInSetsAndJoins(QueryPipeline & pipeline, std::unordered_map<String, SubqueryForSet> & subqueries_for_sets);
+    void executeMergeSorted(QueryPipeline & pipeline, const SortDescription & sort_description, UInt64 limit);
 
     /// Add ConvertingBlockInputStream to specified header.
     void unifyStreams(Pipeline & pipeline, Block header);

dbms/tests/queries/0_stateless/01018_optimize_read_in_order_with_in_subquery.reference

@@ -0,0 +1,4 @@
+1
+0
+1
+0

dbms/tests/queries/0_stateless/01018_optimize_read_in_order_with_in_subquery.sql

@@ -0,0 +1,15 @@
+SET max_threads = 2;
+SET optimize_read_in_order = 1;
+
+DROP TABLE IF EXISTS TESTTABLE4;
+CREATE TABLE TESTTABLE4 (_id UInt64, pt String, l String ) 
+ENGINE = MergeTree() PARTITION BY (pt) ORDER BY (_id);
+INSERT INTO TESTTABLE4 VALUES (0,'1','1'), (1,'0','1');
+
+SELECT _id FROM TESTTABLE4 PREWHERE l IN (select '1') ORDER BY _id DESC LIMIT 10;
+
+SET experimental_use_processors=1;
+
+SELECT _id FROM TESTTABLE4 PREWHERE l IN (select '1') ORDER BY _id DESC LIMIT 10;
+
+DROP TABLE TESTTABLE4;