Merge pull request #54610 from CurtizJ/fix-after-refactoring

Fix reading of virtual columns in reverse order

src/Storages/MergeTree/MergeTreeSelectAlgorithms.cpp

@@ -1,5 +1,4 @@
 #include <Storages/MergeTree/MergeTreeSelectAlgorithms.h>
-#include <Storages/MergeTree/IMergeTreeReadPool.h>
 
 namespace DB
 {
@@ -9,57 +8,29 @@ namespace ErrorCodes
     extern const int LOGICAL_ERROR;
 }
 
-MergeTreeThreadSelectAlgorithm::TaskResult MergeTreeThreadSelectAlgorithm::getNewTask(IMergeTreeReadPool & pool, MergeTreeReadTask * previous_task)
-{
-    TaskResult res;
-    res.first = pool.getTask(thread_idx, previous_task);
-    res.second = !!res.first;
-    return res;
-}
-
-MergeTreeReadTask::BlockAndProgress MergeTreeThreadSelectAlgorithm::readFromTask(MergeTreeReadTask * task, const MergeTreeReadTask::BlockSizeParams & params)
-{
-    if (!task)
-        return {};
-
-    return task->read(params);
-}
-
-IMergeTreeSelectAlgorithm::TaskResult MergeTreeInOrderSelectAlgorithm::getNewTask(IMergeTreeReadPool & pool, MergeTreeReadTask * previous_task)
+MergeTreeReadTaskPtr MergeTreeInOrderSelectAlgorithm::getNewTask(IMergeTreeReadPool & pool, MergeTreeReadTask * previous_task)
 {
     if (!pool.preservesOrderOfRanges())
         throw Exception(ErrorCodes::LOGICAL_ERROR,
             "MergeTreeInOrderSelectAlgorithm requires read pool that preserves order of ranges, got: {}", pool.getName());
 
-    TaskResult res;
-    res.first = pool.getTask(part_idx, previous_task);
-    res.second = !!res.first;
-    return res;
+    return pool.getTask(part_idx, previous_task);
 }
 
-MergeTreeReadTask::BlockAndProgress MergeTreeInOrderSelectAlgorithm::readFromTask(MergeTreeReadTask * task, const BlockSizeParams & params)
-{
-    if (!task)
-        return {};
-
-    return task->read(params);
-}
-
-IMergeTreeSelectAlgorithm::TaskResult MergeTreeInReverseOrderSelectAlgorithm::getNewTask(IMergeTreeReadPool & pool, MergeTreeReadTask * previous_task)
+MergeTreeReadTaskPtr MergeTreeInReverseOrderSelectAlgorithm::getNewTask(IMergeTreeReadPool & pool, MergeTreeReadTask * previous_task)
 {
     if (!pool.preservesOrderOfRanges())
         throw Exception(ErrorCodes::LOGICAL_ERROR,
             "MergeTreeInReverseOrderSelectAlgorithm requires read pool that preserves order of ranges, got: {}", pool.getName());
 
-    TaskResult res;
-    res.first = pool.getTask(part_idx, previous_task);
-    /// We may have some chunks to return in buffer.
-    /// Set continue_reading to true but actually don't create a new task.
-    res.second = !!res.first || !chunks.empty();
-    return res;
+    if (!chunks.empty())
+        throw Exception(ErrorCodes::LOGICAL_ERROR,
+            "Cannot get new task for reading in reverse order because there are {} buffered chunks", chunks.size());
+
+    return pool.getTask(part_idx, previous_task);
 }
 
-MergeTreeReadTask::BlockAndProgress MergeTreeInReverseOrderSelectAlgorithm::readFromTask(MergeTreeReadTask * task, const BlockSizeParams & params)
+MergeTreeReadTask::BlockAndProgress MergeTreeInReverseOrderSelectAlgorithm::readFromTask(MergeTreeReadTask & task, const BlockSizeParams & params)
 {
     MergeTreeReadTask::BlockAndProgress res;
 
@@ -70,11 +41,8 @@ MergeTreeReadTask::BlockAndProgress MergeTreeInReverseOrderSelectAlgorithm::read
         return res;
     }
 
-    if (!task)
-        return {};
-
-    while (!task->isFinished())
-        chunks.push_back(task->read(params));
+    while (!task.isFinished())
+        chunks.push_back(task.read(params));
 
     if (chunks.empty())
         return {};

src/Storages/MergeTree/MergeTreeSelectAlgorithms.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #include <Storages/MergeTree/MergeTreeReadTask.h>
+#include <Storages/MergeTree/IMergeTreeReadPool.h>
 #include <boost/core/noncopyable.hpp>
 
 namespace DB
@@ -11,15 +12,16 @@ class IMergeTreeReadPool;
 class IMergeTreeSelectAlgorithm : private boost::noncopyable
 {
 public:
-    /// The pair of {task, continue_reading}.
-    using TaskResult = std::pair<MergeTreeReadTaskPtr, bool>;
     using BlockSizeParams = MergeTreeReadTask::BlockSizeParams;
+    using BlockAndProgress = MergeTreeReadTask::BlockAndProgress;
 
     virtual ~IMergeTreeSelectAlgorithm() = default;
 
     virtual String getName() const = 0;
-    virtual TaskResult getNewTask(IMergeTreeReadPool & pool, MergeTreeReadTask * previous_task) = 0;
-    virtual MergeTreeReadTask::BlockAndProgress readFromTask(MergeTreeReadTask * task, const BlockSizeParams & params) = 0;
+    virtual bool needNewTask(const MergeTreeReadTask & task) const = 0;
+
+    virtual MergeTreeReadTaskPtr getNewTask(IMergeTreeReadPool & pool, MergeTreeReadTask * previous_task) = 0;
+    virtual BlockAndProgress readFromTask(MergeTreeReadTask & task, const BlockSizeParams & params) = 0;
 };
 
 using MergeTreeSelectAlgorithmPtr = std::unique_ptr<IMergeTreeSelectAlgorithm>;
@@ -28,9 +30,12 @@ class MergeTreeThreadSelectAlgorithm : public IMergeTreeSelectAlgorithm
 {
 public:
     explicit MergeTreeThreadSelectAlgorithm(size_t thread_idx_) : thread_idx(thread_idx_) {}
+
     String getName() const override { return "Thread"; }
-    TaskResult getNewTask(IMergeTreeReadPool & pool, MergeTreeReadTask * previous_task) override;
-    MergeTreeReadTask::BlockAndProgress readFromTask(MergeTreeReadTask * task, const BlockSizeParams & params) override;
+    bool needNewTask(const MergeTreeReadTask & task) const override { return task.isFinished(); }
+
+    MergeTreeReadTaskPtr getNewTask(IMergeTreeReadPool & pool, MergeTreeReadTask * previous_task) override { return pool.getTask(thread_idx, previous_task); }
+    BlockAndProgress readFromTask(MergeTreeReadTask & task, const BlockSizeParams & params) override { return task.read(params); }
 
 private:
     const size_t thread_idx;
@@ -40,9 +45,12 @@ class MergeTreeInOrderSelectAlgorithm : public IMergeTreeSelectAlgorithm
 {
 public:
     explicit MergeTreeInOrderSelectAlgorithm(size_t part_idx_) : part_idx(part_idx_) {}
+
     String getName() const override { return "InOrder"; }
-    TaskResult getNewTask(IMergeTreeReadPool & pool, MergeTreeReadTask * previous_task) override;
-    MergeTreeReadTask::BlockAndProgress readFromTask(MergeTreeReadTask * task, const BlockSizeParams & params) override;
+    bool needNewTask(const MergeTreeReadTask & task) const override { return task.isFinished(); }
+
+    MergeTreeReadTaskPtr getNewTask(IMergeTreeReadPool & pool, MergeTreeReadTask * previous_task) override;
+    MergeTreeReadTask::BlockAndProgress readFromTask(MergeTreeReadTask & task, const BlockSizeParams & params) override { return task.read(params); }
 
 private:
     const size_t part_idx;
@@ -52,13 +60,16 @@ class MergeTreeInReverseOrderSelectAlgorithm : public IMergeTreeSelectAlgorithm
 {
 public:
     explicit MergeTreeInReverseOrderSelectAlgorithm(size_t part_idx_) : part_idx(part_idx_) {}
+
     String getName() const override { return "InReverseOrder"; }
-    TaskResult getNewTask(IMergeTreeReadPool & pool, MergeTreeReadTask * previous_task) override;
-    MergeTreeReadTask::BlockAndProgress readFromTask(MergeTreeReadTask * task, const BlockSizeParams & params) override;
+    bool needNewTask(const MergeTreeReadTask & task) const override { return chunks.empty() && task.isFinished(); }
+
+    MergeTreeReadTaskPtr getNewTask(IMergeTreeReadPool & pool, MergeTreeReadTask * previous_task) override;
+    BlockAndProgress readFromTask(MergeTreeReadTask & task, const BlockSizeParams & params) override;
 
 private:
     const size_t part_idx;
-    std::vector<MergeTreeReadTask::BlockAndProgress> chunks;
+    std::vector<BlockAndProgress> chunks;
 };
 
 }

src/Storages/MergeTree/MergeTreeSelectProcessor.cpp

@@ -139,11 +139,10 @@ ChunkAndProgress MergeTreeSelectProcessor::read()
     {
         try
         {
-            bool continue_reading = true;
-            if (!task || task->isFinished())
-                std::tie(task, continue_reading) = algorithm->getNewTask(*pool, task.get());
+            if (!task || algorithm->needNewTask(*task))
+                task = algorithm->getNewTask(*pool, task.get());
 
-            if (!continue_reading)
+            if (!task)
                 break;
         }
         catch (const Exception & e)
@@ -153,10 +152,10 @@ ChunkAndProgress MergeTreeSelectProcessor::read()
             throw;
         }
 
-        if (task && !task->getMainRangeReader().isInitialized())
+        if (!task->getMainRangeReader().isInitialized())
             initializeRangeReaders();
 
-        auto res = algorithm->readFromTask(task.get(), block_size_params);
+        auto res = algorithm->readFromTask(*task, block_size_params);
 
         if (res.row_count)
         {

tests/queries/0_stateless/02883_read_in_reverse_order_virtual_column.reference

@@ -0,0 +1 @@
+198401_1_1_0

tests/queries/0_stateless/02883_read_in_reverse_order_virtual_column.sql

@@ -0,0 +1,10 @@
+DROP TABLE IF EXISTS t_reverse_order_virt_col;
+
+CREATE TABLE t_reverse_order_virt_col (`order_0` Decimal(76, 53), `p_time` Date)
+ENGINE = MergeTree PARTITION BY toYYYYMM(p_time)
+ORDER BY order_0;
+
+INSERT INTO t_reverse_order_virt_col SELECT number, '1984-01-01' FROM numbers(1000000);
+SELECT DISTINCT _part FROM (SELECT _part FROM t_reverse_order_virt_col ORDER BY order_0 DESC);
+
+DROP TABLE IF EXISTS t_reverse_order_virt_col;