handle the case first prefix fixed

src/Processors/QueryPlan/Optimizations/optimizeReadInOrder.cpp

@@ -370,6 +370,7 @@ InputOrderInfoPtr buildInputOrderInfo(
     int read_direction = 0;
     size_t next_description_column = 0;
     size_t next_sort_key = 0;
+    bool first_prefix_fixed = false;
 
     while (next_description_column < description.size() && next_sort_key < sorting_key.column_names.size())
     {
@@ -447,6 +448,9 @@ InputOrderInfoPtr buildInputOrderInfo(
             }
             else if (fixed_key_columns.contains(sort_column_node))
             {
+                if (next_sort_key == 0)
+                    first_prefix_fixed = true;
+
                 //std::cerr << "+++++++++ Found fixed key by match" << std::endl;
                 ++next_sort_key;
             }
@@ -481,7 +485,7 @@ InputOrderInfoPtr buildInputOrderInfo(
     if (read_direction == 0 || order_key_prefix_descr.empty())
         return nullptr;
 
-    return std::make_shared<InputOrderInfo>(order_key_prefix_descr, next_sort_key, read_direction, limit);
+    return std::make_shared<InputOrderInfo>(order_key_prefix_descr, next_sort_key, read_direction, limit, first_prefix_fixed);
 }
 
 /// We really need three different sort descriptions here.
@@ -685,7 +689,7 @@ AggregationInputOrder buildInputOrderInfo(
     for (const auto & key : not_matched_group_by_keys)
         group_by_sort_description.emplace_back(SortColumnDescription(std::string(key)));
 
-    auto input_order = std::make_shared<InputOrderInfo>(order_key_prefix_descr, next_sort_key, /*read_direction*/ 1, /* limit */ 0);
+    auto input_order = std::make_shared<InputOrderInfo>(order_key_prefix_descr, next_sort_key, /*read_direction*/ 1, /* limit */ 0, false);
     return { std::move(input_order), std::move(sort_description_for_merging), std::move(group_by_sort_description) };
 }
 
@@ -823,7 +827,7 @@ InputOrderInfoPtr buildInputOrderInfo(SortingStep & sorting, QueryPlan::Node & n
 
             bool use_buffering = (order_info->limit == 0) && sorting.getSettings().read_in_order_use_buffering;
             /// Avoid conflict with buffering.
-            if (!use_buffering)
+            if (!use_buffering && !order_info->first_prefix_fixed)
                 reading->enableVirtualRow();
         }
 

src/Processors/QueryPlan/ReadFromMergeTree.cpp

@@ -1808,7 +1808,7 @@ bool ReadFromMergeTree::requestReadingInOrder(size_t prefix_size, int direction,
     if (direction != 1 && query_info.isFinal())
         return false;
 
-    query_info.input_order_info = std::make_shared<InputOrderInfo>(SortDescription{}, prefix_size, direction, read_limit);
+    query_info.input_order_info = std::make_shared<InputOrderInfo>(SortDescription{}, prefix_size, direction, read_limit, false);
     reader_settings.read_in_order = true;
 
     /// In case or read-in-order, don't create too many reading streams.

src/Storages/ReadInOrderOptimizer.cpp

@@ -249,7 +249,7 @@ InputOrderInfoPtr ReadInOrderOptimizer::getInputOrderImpl(
     if (sort_description_for_merging.empty())
         return {};
 
-    return std::make_shared<InputOrderInfo>(std::move(sort_description_for_merging), key_pos, read_direction, limit);
+    return std::make_shared<InputOrderInfo>(std::move(sort_description_for_merging), key_pos, read_direction, limit, false);
 }
 
 InputOrderInfoPtr ReadInOrderOptimizer::getInputOrder(

src/Storages/SelectQueryInfo.h

@@ -119,13 +119,22 @@ struct InputOrderInfo
     const int direction;
     const UInt64 limit;
 
+    /** For virtual row optimization only
+     * for example, when pk is (a,b), a = 1, order by b, virtual row should be
+     * disabled in the following case:
+     * 1st part (0, 100), (1, 2), (1, 3), (1, 4)
+     * 2nd part (0, 100), (1, 2), (1, 3), (1, 4).
+     */
+    bool first_prefix_fixed;
+
     InputOrderInfo(
         const SortDescription & sort_description_for_merging_,
         size_t used_prefix_of_sorting_key_size_,
-        int direction_, UInt64 limit_)
+        int direction_, UInt64 limit_, bool first_prefix_fixed_)
         : sort_description_for_merging(sort_description_for_merging_)
         , used_prefix_of_sorting_key_size(used_prefix_of_sorting_key_size_)
         , direction(direction_), limit(limit_)
+        , first_prefix_fixed(first_prefix_fixed_)
     {
     }
 

tests/queries/0_stateless/03031_read_in_order_optimization_with_virtual_row.reference

@@ -28,6 +28,12 @@
 1	3
 1	4
 1	4
+1	2
+1	2
+1	3
+1	3
+1	4
+1	4
 ========
 1	3
 1	2

tests/queries/0_stateless/03031_read_in_order_optimization_with_virtual_row.sql

@@ -147,6 +147,15 @@ read_in_order_use_buffering = false,
 optimize_read_in_order = 1,
 read_in_order_two_level_merge_threshold = 0;  --force preliminary merge
 
+SELECT a, b
+FROM fixed_prefix
+WHERE a = 1
+ORDER BY b
+SETTINGS max_threads = 1,
+read_in_order_use_buffering = false,
+optimize_read_in_order = 1,
+read_in_order_two_level_merge_threshold = 5;  --avoid preliminary merge
+
 DROP TABLE fixed_prefix;
 
 SELECT '========';