Fix spill crash (#10004) (#10014)

close #9999

Fix sort spill crash in some corner cases.

Signed-off-by: xufei <[email protected]>

Co-authored-by: xufei <[email protected]>
Co-authored-by: xufei <[email protected]>
Co-authored-by: JaySon <[email protected]>

Author: 3 people

dbms/src/DataStreams/FilterTransformAction.cpp

@@ -102,16 +102,18 @@ bool FilterTransformAction::transform(Block & block, FilterPtr & res_filter, boo
         *  and now - are calculated. That is, not all cases are covered by the code above.
         * This happens if the function returns a constant for a non-constant argument.
         * For example, `ignore` function.
+        * use a local variable to avoid the case that function return a constant for a non-constant argument in one block
+        * but return a non-constant for the same argument in another block.
         */
-    constant_filter_description = ConstantFilterDescription(*column_of_filter);
+    auto current_constant_filter_description = ConstantFilterDescription(*column_of_filter);
 
-    if (constant_filter_description.always_false)
+    if (current_constant_filter_description.always_false)
     {
         block.clear();
         return true;
     }
 
-    if (constant_filter_description.always_true)
+    if (current_constant_filter_description.always_true)
     {
         if (return_filter)
             res_filter = nullptr;
@@ -140,7 +142,7 @@ bool FilterTransformAction::transform(Block & block, FilterPtr & res_filter, boo
     if (filtered_rows == rows)
     {
         /// Replace the column with the filter by a constant.
-        auto filter_column = block.safeGetByPosition(filter_column_position);
+        auto & filter_column = block.safeGetByPosition(filter_column_position);
         filter_column.column = filter_column.type->createColumnConst(filtered_rows, static_cast<UInt64>(1));
         /// No need to touch the rest of the columns.
         return true;

dbms/src/DataStreams/MergeSortingBlockInputStream.cpp

@@ -88,6 +88,11 @@ Block MergeSortingBlockInputStream::readImpl()
                 return block;
 
             SortHelper::removeConstantsFromBlock(block);
+            RUNTIME_CHECK_MSG(
+                block.columns() == header_without_constants.columns(),
+                "Unexpected number of constant columns in block in MergeSortingBlockInputStream, n_block={}, n_head={}",
+                block.columns(),
+                header_without_constants.columns());
 
             blocks.push_back(block);
             sum_bytes_in_blocks += block.estimateBytesForSpill();

dbms/src/Flash/tests/gtest_spill_sort.cpp

@@ -12,10 +12,12 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <Functions/FunctionHelpers.h>
 #include <Interpreters/Context.h>
 #include <TestUtils/ColumnGenerator.h>
 #include <TestUtils/ExecutorTestUtils.h>
 #include <TestUtils/mockExecutor.h>
+#include <common/types.h>
 
 namespace DB
 {
@@ -166,5 +168,58 @@ try
 }
 CATCH
 
+TEST_F(SpillSortTestRunner, SpillAfterFilter)
+try
+{
+    DB::MockColumnInfoVec column_infos{
+        {"a", TiDB::TP::TypeTiny},
+        {"b", TiDB::TP::TypeTiny},
+        {"c", TiDB::TP::TypeTiny},
+        {"d", TiDB::TP::TypeTiny},
+        {"e", TiDB::TP::TypeTiny}};
+    ColumnsWithTypeAndName column_data;
+    size_t table_rows = 102400;
+    UInt64 max_block_size = 64;
+    size_t total_data_size = 0;
+    size_t limit_size = table_rows / 10 * 9;
+    for (const auto & column_info : mockColumnInfosToTiDBColumnInfos(column_infos))
+    {
+        ColumnGeneratorOpts opts{
+            table_rows,
+            getDataTypeByColumnInfoForComputingLayer(column_info)->getName(),
+            RANDOM,
+            column_info.name};
+        column_data.push_back(ColumnGenerator::instance().generate(opts));
+        total_data_size += column_data.back().column->byteSize();
+    }
+    context.addMockTable("spill_sort_test", "simple_table", column_infos, column_data, 8);
+
+    MockOrderByItemVec order_by_items{
+        std::make_pair("a", true),
+        std::make_pair("b", true),
+        std::make_pair("c", true),
+        std::make_pair("d", true),
+        std::make_pair("e", true)};
+
+    auto request = context.scan("spill_sort_test", "simple_table")
+                       .filter(gt(col("d"), lit(toField(static_cast<Int8>(-128)))))
+                       .topN(order_by_items, limit_size)
+                       .project({gt(col("d"), lit(toField(static_cast<Int8>(-128))))})
+                       .build(context);
+    context.context->setSetting("max_block_size", Field(static_cast<UInt64>(max_block_size)));
+
+    /// disable spill
+    context.context->setSetting("max_bytes_before_external_sort", Field(static_cast<UInt64>(0)));
+    auto ref_columns = executeStreams(request, 1);
+
+    // The implementation of topN in the pipeline model is LocalSort, and the result of using multiple threads is unstable. Therefore, a single thread is used here instead.
+    enablePipeline(true);
+    context.context->setSetting("max_bytes_before_external_sort", Field(static_cast<UInt64>(total_data_size / 10)));
+    ASSERT_COLUMNS_EQ_R(ref_columns, executeStreams(request, 1));
+    context.context->setSetting("max_cached_data_bytes_in_spiller", Field(static_cast<UInt64>(total_data_size / 100)));
+    ASSERT_COLUMNS_EQ_R(ref_columns, executeStreams(request, 1));
+}
+CATCH
+
 } // namespace tests
 } // namespace DB

dbms/src/Operators/MergeSortTransformOp.cpp

@@ -149,6 +149,11 @@ OperatorStatus MergeSortTransformOp::transformImpl(Block & block)
 
         // store the sorted block in `sorted_blocks`.
         SortHelper::removeConstantsFromBlock(block);
+        RUNTIME_CHECK_MSG(
+            block.columns() == header_without_constants.columns(),
+            "Unexpected number of constant columns in block in MergeSortTransformOp, n_block={}, n_head={}",
+            block.columns(),
+            header_without_constants.columns());
         sum_bytes_in_blocks += block.estimateBytesForSpill();
         sorted_blocks.emplace_back(std::move(block));