HappenLee commented on code in PR #68500:
URL: https://github.com/apache/doris/pull/68500#discussion_r4358063696
##########
be/test/exec/operator/analytic_sink_operator_test.cpp:
##########
@@ -310,6 +351,1160 @@ TEST_F(AnalyticSinkOperatorTest, AggFunction) {
std::cout << "######### AggFunction with sum test end #########" <<
std::endl;
}
+TEST_F(AnalyticSinkOperatorTest,
SpillEligibilityKeepsStreamingAggregateOnExistingPath) {
+ Initialize(1);
+ state->_query_options.__set_enable_spill(true);
+ create_operator(true, 1, "sum", {std::make_shared<DataTypeInt64>()},
+ std::make_shared<DataTypeInt64>());
+ TAnalyticWindow window;
+ window.type = TAnalyticWindowType::ROWS;
+ TAnalyticWindowBoundary window_end;
+ window_end.type = TAnalyticWindowBoundaryType::CURRENT_ROW;
+ window.__set_window_end(window_end);
+ create_window_type(false, true, window);
+
+ sink->_prepare_spill(state.get());
+
+ EXPECT_FALSE(sink->_enable_spill_analytic);
+ EXPECT_FALSE(sink->_spillable);
+ EXPECT_EQ(sink->_window_spill_unsupported_reason, "Streaming");
+}
+
+TEST_F(AnalyticSinkOperatorTest,
SpillEligibilityAllowsPartitionCardinalityFunction) {
+ Initialize(1);
+ state->_query_options.__set_enable_spill(true);
+ create_operator(true, 1, "ntile", {std::make_shared<DataTypeInt64>()},
+ std::make_shared<DataTypeInt64>());
+ TAnalyticWindow window;
+ window.type = TAnalyticWindowType::ROWS;
+ TAnalyticWindowBoundary window_end;
+ window_end.type = TAnalyticWindowBoundaryType::CURRENT_ROW;
+ window.__set_window_end(window_end);
+ create_window_type(false, true, window);
+
+ sink->_prepare_spill(state.get());
+
+ EXPECT_TRUE(sink->_enable_spill_analytic);
+ EXPECT_TRUE(sink->_spillable);
+ ASSERT_EQ(sink->_window_spill_strategies.size(), 1);
+ EXPECT_EQ(sink->_window_spill_strategies[0],
WindowSpillStrategy::PARTITION_CARDINALITY);
+}
+
+TEST_F(AnalyticSinkOperatorTest,
SpillEligibilityTreatsRangeToCurrentRowWithoutOrderAsFullFrame) {
+ Initialize(1);
+ state->_query_options.__set_enable_spill(true);
+ create_operator(true, 1, "sum", {std::make_shared<DataTypeInt64>()},
+ std::make_shared<DataTypeInt64>());
+ // `sum(v) OVER (PARTITION BY k)` reaches BE as RANGE UNBOUNDED PRECEDING
.. CURRENT ROW
+ // without ORDER BY expressions.
+ TAnalyticWindow window;
+ window.type = TAnalyticWindowType::RANGE;
+ TAnalyticWindowBoundary window_end;
+ window_end.type = TAnalyticWindowBoundaryType::CURRENT_ROW;
+ window.__set_window_end(window_end);
+ create_window_type(false, true, window);
+ sink->_has_range_window = true;
+
+ sink->_prepare_spill(state.get());
+
+ EXPECT_TRUE(sink->_enable_spill_analytic);
+ ASSERT_EQ(sink->_window_spill_strategies.size(), 1);
+ EXPECT_EQ(sink->_window_spill_strategies[0],
WindowSpillStrategy::PARTITION_REDUCE);
+}
+
+TEST_F(AnalyticSinkOperatorTest,
SpillEligibilityKeepsRangeToCurrentRowWithOrderStreaming) {
+ Initialize(1);
+ state->_query_options.__set_enable_spill(true);
+ create_operator(true, 1, "sum", {std::make_shared<DataTypeInt64>()},
+ std::make_shared<DataTypeInt64>());
+ TAnalyticWindow window;
+ window.type = TAnalyticWindowType::RANGE;
+ TAnalyticWindowBoundary window_end;
+ window_end.type = TAnalyticWindowBoundaryType::CURRENT_ROW;
+ window.__set_window_end(window_end);
+ create_window_type(false, true, window);
+ sink->_has_range_window = true;
+ sink->_order_by_eq_expr_ctxs =
+ MockSlotRef::create_mock_contexts(0,
std::make_shared<DataTypeInt64>());
+
+ sink->_prepare_spill(state.get());
+
+ EXPECT_FALSE(sink->_enable_spill_analytic);
+ EXPECT_EQ(sink->_window_spill_unsupported_reason, "Streaming");
+}
+
+TEST_F(AnalyticSinkOperatorTest, SpillPathFullPartitionSumInMemory) {
+ Initialize(10);
+ create_operator(false, 1, "sum", {std::make_shared<DataTypeInt64>()},
+ std::make_shared<DataTypeInt64>());
+ sink->_agg_expr_ctxs.resize(1);
+ sink->_agg_expr_ctxs[0] =
+ MockSlotRef::create_mock_contexts(0,
std::make_shared<DataTypeInt64>());
+ sink->_enable_spill_analytic = true;
+ sink->_window_spill_strategies = {WindowSpillStrategy::PARTITION_REDUCE};
+ sink->_window_spill_peer_functions = {WindowSpillPeerFunction::NONE};
+ create_local_state();
+
+ Block input = ColumnHelper::create_block<DataTypeInt64>(_data_vals);
+ auto status = sink->sink(state.get(), &input, true);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+ EXPECT_TRUE(input.empty());
+
+ Block output = ColumnHelper::create_block<DataTypeInt64>({});
+ bool eos = false;
+ status = source->get_block(state.get(), &output, &eos);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+ EXPECT_FALSE(eos);
+ EXPECT_TRUE(ColumnHelper::block_equal(
+ output,
+ ColumnHelper::create_block<DataTypeInt64>(_data_vals,
std::vector<int64_t>(10, 45))))
+ << output.dump_data();
+
+ output.clear();
+ status = source->get_block(state.get(), &output, &eos);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+ EXPECT_TRUE(eos);
+ EXPECT_TRUE(output.empty());
+}
+
+TEST_F(AnalyticSinkOperatorTest, SpillSchedulingBeforeFirstWrite) {
+ Initialize(10);
+ // Scheduler submission can precede local-state open(), where
_spill_enabled is set.
+ sink->_enable_spill_analytic = true;
+ EXPECT_TRUE(sink->is_blockable(state.get()));
+ prepare_spilled_full_partition_sum();
+
+ // Even an ordinary first input can cross the proactive limit or flush a
peer sidecar.
+ // Waiting for is_spilled() would dispatch that synchronous write to the
simple pool.
+ ASSERT_EQ(sink_local_state->_batch_store, nullptr);
+ EXPECT_TRUE(sink->is_blockable(state.get()));
+ state->_query_options.__set_enable_force_spill(true);
+ EXPECT_TRUE(sink->is_blockable(state.get()));
+
+ sink_local_state->_spill_enabled = false;
+ EXPECT_TRUE(sink->is_blockable(state.get()));
+ sink->_enable_spill_analytic = false;
+ EXPECT_FALSE(sink->is_blockable(state.get()));
+}
+
+TEST_F(AnalyticSinkOperatorTest, SpillSchedulingReplayUsesOwnSharedState) {
+ Initialize(10);
+ prepare_spilled_full_partition_sum();
+
+ // The fixture puts source and sink in one RuntimeState. Disable only its
sink's spill
+ // flag to model replay feeding an ordinary sink; the source's shared flag
stays enabled.
+ sink_local_state->_spill_enabled = false;
+ sink->_enable_spill_analytic = false;
+ EXPECT_FALSE(sink->is_blockable(state.get()));
+ {
+ LockGuard lock(source_local_state->_shared_state->buffer_mutex);
+ EXPECT_TRUE(source_local_state->_shared_state->spill_batches.empty());
+ }
+ EXPECT_TRUE(source->is_blockable(state.get()));
+
+ source_local_state->_shared_state->spill_enabled.store(false);
+ EXPECT_FALSE(source->is_blockable(state.get()));
+ // Unsupported/streaming windows must still preserve blocking expression
classification.
+ source->_blockable = true;
+ EXPECT_TRUE(source->is_blockable(state.get()));
+}
+
+TEST_F(AnalyticSinkOperatorTest,
SpillAdmissionDefersBatchPublishedAfterEmptySnapshot) {
+ Initialize(10);
+ prepare_spilled_full_partition_sum();
+
+ // Admission observes an empty queue, then the sink publishes before
get_block runs.
+ EXPECT_EQ(source->get_reserve_mem_size(state.get()),
+ state->minimum_operator_memory_required_bytes());
+ Block input = ColumnHelper::create_block<DataTypeInt64>(_data_vals);
+ auto status = sink->sink(state.get(), &input, true);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+
+ Block output;
+ bool eos = false;
+ status = source->get_block(state.get(), &output, &eos);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+ EXPECT_FALSE(eos);
+ EXPECT_TRUE(output.empty());
+ EXPECT_GT(source->get_reserve_mem_size(state.get()),
+ state->minimum_operator_memory_required_bytes());
+ status = source->get_block(state.get(), &output, &eos);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+ EXPECT_EQ(output.rows(), 10);
+}
+
+TEST_F(AnalyticSinkOperatorTest,
SpillAdmissionIncludesFirstInMemoryResultColumns) {
+ Initialize(65535);
+ prepare_spilled_full_partition_sum();
+ auto batch = std::make_shared<AnalyticSpillBatch>();
+ batch->rows = state->batch_size();
+
batch->blocks.push_back(ColumnHelper::create_block<DataTypeInt64>(_data_vals));
+ batch->result_types.assign(32, std::make_shared<DataTypeInt64>());
+ batch->partition_results.resize(32);
+ batch->change_to_nullable_flags.assign(32, false);
+ {
+ LockGuard lock(sink_local_state->_shared_state->buffer_mutex);
+ sink_local_state->_shared_state->spill_batches.push(batch);
+ }
+ EXPECT_EQ(source->get_reserve_mem_size(state.get()),
+ state->minimum_operator_memory_required_bytes() + 32 * 65535 *
sizeof(int64_t));
+}
+
+TEST_F(AnalyticSinkOperatorTest,
SpillAdmissionIncludesFirstAndNextInMemoryCuts) {
+ Initialize(1024);
+ prepare_spilled_full_partition_sum();
+ auto batch = std::make_shared<AnalyticSpillBatch>();
+ batch->rows = 4096;
+ batch->partition_ends = {4096};
+ const std::vector<std::string> values(2048, std::string(8 * 1024, 'x'));
+
batch->blocks.push_back(ColumnHelper::create_block<DataTypeString>(values));
+
batch->blocks.push_back(ColumnHelper::create_block<DataTypeString>(values));
+ {
+ LockGuard lock(sink_local_state->_shared_state->buffer_mutex);
+ sink_local_state->_shared_state->spill_batches.push(batch);
+ }
+ const auto first_cut = source->get_reserve_mem_size(state.get());
+ EXPECT_GT(first_cut, state->minimum_operator_memory_required_bytes() + 8 *
1024 * 1024);
+
+ Block output;
+ bool eos = false;
+ auto status = source->get_block(state.get(), &output, &eos);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+ ASSERT_EQ(output.rows(), 1024);
+ output.clear();
+ status = source->get_block(state.get(), &output, &eos);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+ ASSERT_EQ(output.rows(), 1024);
+ EXPECT_EQ(source_local_state->_replay_block_position, 2048);
+ EXPECT_EQ(source->get_reserve_mem_size(state.get()), first_cut);
+}
+
+TEST_F(AnalyticSinkOperatorTest,
SpillCardinalityKeepsConstantArgumentsCompact) {
+ Initialize(65535);
+ auto local_state = AnalyticSinkLocalState::create_unique(sink.get(),
state.get());
+ local_state->_agg_functions_size = 32;
+ local_state->_spill_strategies.assign(32,
WindowSpillStrategy::PARTITION_CARDINALITY);
+ local_state->_agg_expr_ctxs.resize(32);
+ for (size_t i = 0; i < 32; ++i) {
+ local_state->_agg_expr_ctxs[i].push_back(
+ MockLiteral::create_const<DataTypeInt64>(2147483648LL + i,
65535));
+ }
+ Block input = ColumnHelper::create_block<DataTypeInt64>(_data_vals);
+ std::vector<std::vector<ColumnPtr>> agg_columns;
+ std::vector<ColumnPtr> partition_columns;
+ std::vector<ColumnPtr> order_columns;
+ auto status = local_state->_materialize_spill_columns(&input,
&agg_columns, &partition_columns,
+ &order_columns);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+ size_t argument_bytes = 0;
+ for (const auto& columns : agg_columns) {
+ ASSERT_EQ(columns.size(), 1);
+ EXPECT_TRUE(is_column_const(*columns[0]));
+ EXPECT_EQ(columns[0]->get_int(65534), columns[0]->get_int(0));
+ argument_bytes += columns[0]->allocated_bytes();
+ }
+ EXPECT_LT(argument_bytes, 32 * 1024);
+}
+
+TEST_F(AnalyticSinkOperatorTest, SpillPathFullPartitionSumFromDisk) {
+ ASSERT_TRUE(spill_file_manager.init().ok());
+
+ Initialize(10);
+ state->_query_options.__set_enable_force_spill(true);
+ create_operator(false, 1, "sum", {std::make_shared<DataTypeInt64>()},
+ std::make_shared<DataTypeInt64>());
+ sink->_agg_expr_ctxs.resize(1);
+ sink->_agg_expr_ctxs[0] =
+ MockSlotRef::create_mock_contexts(0,
std::make_shared<DataTypeInt64>());
+ sink->_enable_spill_analytic = true;
+ sink->_window_spill_strategies = {WindowSpillStrategy::PARTITION_REDUCE};
+ sink->_window_spill_peer_functions = {WindowSpillPeerFunction::NONE};
+ create_local_state();
+
+ Block first = ColumnHelper::create_block<DataTypeInt64>({0, 1, 2, 3, 4});
+ auto status = sink->sink(state.get(), &first, false);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+ Block second = ColumnHelper::create_block<DataTypeInt64>({5, 6, 7, 8, 9});
+ status = sink->sink(state.get(), &second, true);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+
+ Block output = ColumnHelper::create_block<DataTypeInt64>({});
+ bool eos = false;
+ status = source->get_block(state.get(), &output, &eos);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+ EXPECT_FALSE(eos);
+ EXPECT_TRUE(ColumnHelper::block_equal(
+ output, ColumnHelper::create_block<DataTypeInt64>({0, 1, 2, 3, 4},
+
std::vector<int64_t>(5, 45))))
+ << output.dump_data();
+
+ output.clear();
+ status = source->get_block(state.get(), &output, &eos);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+ EXPECT_FALSE(eos);
+ EXPECT_TRUE(ColumnHelper::block_equal(
+ output, ColumnHelper::create_block<DataTypeInt64>({5, 6, 7, 8, 9},
+
std::vector<int64_t>(5, 45))))
+ << output.dump_data();
+
+ output.clear();
+ status = source->get_block(state.get(), &output, &eos);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+ EXPECT_TRUE(eos);
+}
+
+TEST_F(AnalyticSinkOperatorTest, SpillPathNtileUsesSealedPartitionCardinality)
{
+ const std::vector<int64_t> bucket_arguments(10, 3);
+ Initialize(bucket_arguments.size());
+ create_operator(true, 1, "ntile", {std::make_shared<DataTypeInt64>()},
+ std::make_shared<DataTypeInt64>());
+ sink->_agg_expr_ctxs.resize(1);
+ sink->_agg_expr_ctxs[0] =
+ MockSlotRef::create_mock_contexts(0,
std::make_shared<DataTypeInt64>());
+ sink->_enable_spill_analytic = true;
+ sink->_window_spill_strategies =
{WindowSpillStrategy::PARTITION_CARDINALITY};
+ sink->_window_spill_peer_functions = {WindowSpillPeerFunction::NONE};
+ create_local_state();
+
+ Block input = ColumnHelper::create_block<DataTypeInt64>(bucket_arguments);
+ auto status = sink->sink(state.get(), &input, true);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+
+ Block output = ColumnHelper::create_block<DataTypeInt64>({});
+ bool eos = false;
+ status = source->get_block(state.get(), &output, &eos);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+ const std::vector<int64_t> expected {1, 1, 1, 1, 2, 2, 2, 3, 3, 3};
+ EXPECT_TRUE(ColumnHelper::block_equal(
+ output,
ColumnHelper::create_block<DataTypeInt64>(bucket_arguments, expected)))
+ << output.dump_data();
+}
+
+TEST_F(AnalyticSinkOperatorTest, SpillPathSplitsPartitionAcrossInputBlocks) {
+ Initialize(3);
+ create_operator(false, 1, "sum", {std::make_shared<DataTypeInt64>()},
+ std::make_shared<DataTypeInt64>());
+ sink->_agg_expr_ctxs.resize(1);
+ sink->_agg_expr_ctxs[0] =
+ MockSlotRef::create_mock_contexts(1,
std::make_shared<DataTypeInt64>());
+ sink->_partition_by_eq_expr_ctxs =
+ MockSlotRef::create_mock_contexts(0,
std::make_shared<DataTypeInt64>());
+ sink->_enable_spill_analytic = true;
+ sink->_window_spill_strategies = {WindowSpillStrategy::PARTITION_REDUCE};
+ sink->_window_spill_peer_functions = {WindowSpillPeerFunction::NONE};
+ create_local_state();
+
+ Block first = ColumnHelper::create_block<DataTypeInt64>({1, 1}, {1, 2});
+ auto status = sink->sink(state.get(), &first, false);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+ Block second = ColumnHelper::create_block<DataTypeInt64>({2, 2, 2}, {10,
20, 30});
+ status = sink->sink(state.get(), &second, true);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+
+ auto check_partition = [&](const std::vector<int64_t>& keys, const
std::vector<int64_t>& values,
+ int64_t sum) {
+ Block output = ColumnHelper::create_block<DataTypeInt64>({});
+ bool eos = false;
+ auto get_status = source->get_block(state.get(), &output, &eos);
+ ASSERT_TRUE(get_status.ok()) << get_status.to_string();
+ EXPECT_FALSE(eos);
+ auto expected = ColumnHelper::create_block<DataTypeInt64>(keys,
values);
+ expected.insert(
+
{ColumnHelper::create_column<DataTypeInt64>(std::vector<int64_t>(keys.size(),
sum)),
+ std::make_shared<DataTypeInt64>(), "sum"});
+ EXPECT_TRUE(ColumnHelper::block_equal(output, expected)) <<
output.dump_data();
+ };
+ check_partition({1, 1}, {1, 2}, 3);
+ check_partition({2, 2, 2}, {10, 20, 30}, 60);
+}
+
+TEST_F(AnalyticSinkOperatorTest, SpillPathCumeDistUsesPeerGroupEnds) {
+ ASSERT_TRUE(spill_file_manager.init().ok());
+
+ Initialize(3);
+ state->_query_options.__set_enable_force_spill(true);
+ create_operator(true, 1, "cume_dist", {},
std::make_shared<DataTypeFloat64>());
+ sink->_agg_expr_ctxs.resize(1);
+ sink->_order_by_eq_expr_ctxs =
+ MockSlotRef::create_mock_contexts(0,
std::make_shared<DataTypeInt64>());
+ sink->_enable_spill_analytic = true;
+ sink->_window_spill_strategies = {WindowSpillStrategy::PEER_GROUP};
+ sink->_window_spill_peer_functions = {WindowSpillPeerFunction::CUME_DIST};
+ create_local_state();
+
+ Block first = ColumnHelper::create_block<DataTypeInt64>({1, 1, 2});
+ auto status = sink->sink(state.get(), &first, false);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+ Block second = ColumnHelper::create_block<DataTypeInt64>({2, 4});
+ status = sink->sink(state.get(), &second, true);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+
+ auto check_block = [&](const std::vector<int64_t>& order_values,
+ const std::vector<double>& expected) {
+ Block output = ColumnHelper::create_block<DataTypeInt64>({});
+ bool eos = false;
+ auto get_status = source->get_block(state.get(), &output, &eos);
+ ASSERT_TRUE(get_status.ok()) << get_status.to_string();
+ auto expected_block =
ColumnHelper::create_block<DataTypeInt64>(order_values);
+
expected_block.insert({ColumnHelper::create_column<DataTypeFloat64>(expected),
+ std::make_shared<DataTypeFloat64>(),
"cume_dist"});
+ EXPECT_TRUE(ColumnHelper::block_equal(output, expected_block)) <<
output.dump_data();
+ };
+ check_block({1, 1, 2}, {0.4, 0.4, 0.8});
+ check_block({2, 4}, {0.8, 1.0});
+
+ Block output = ColumnHelper::create_block<DataTypeInt64>({});
+ bool eos = false;
+ status = source->get_block(state.get(), &output, &eos);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+ EXPECT_TRUE(eos);
+}
+
+TEST_F(AnalyticSinkOperatorTest, SpillPathPercentRankUsesPeerGroupStarts) {
+ Initialize(5);
+ create_operator(true, 1, "percent_rank", {},
std::make_shared<DataTypeFloat64>());
+ sink->_agg_expr_ctxs.resize(1);
+ sink->_order_by_eq_expr_ctxs =
+ MockSlotRef::create_mock_contexts(0,
std::make_shared<DataTypeInt64>());
+ sink->_enable_spill_analytic = true;
+ sink->_window_spill_strategies = {WindowSpillStrategy::PEER_GROUP};
+ sink->_window_spill_peer_functions =
{WindowSpillPeerFunction::PERCENT_RANK};
+ create_local_state();
+
+ Block input = ColumnHelper::create_block<DataTypeInt64>({1, 1, 2, 2, 4});
+ auto status = sink->sink(state.get(), &input, true);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+
+ Block output = ColumnHelper::create_block<DataTypeInt64>({});
+ bool eos = false;
+ status = source->get_block(state.get(), &output, &eos);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+ auto expected = ColumnHelper::create_block<DataTypeInt64>({1, 1, 2, 2, 4});
+ expected.insert({ColumnHelper::create_column<DataTypeFloat64>({0.0, 0.0,
0.5, 0.5, 1.0}),
+ std::make_shared<DataTypeFloat64>(), "percent_rank"});
+ EXPECT_TRUE(ColumnHelper::block_equal(output, expected)) <<
output.dump_data();
+}
+
+namespace {
+
+void expect_next_spill_block(AnalyticSourceOperatorX* source, RuntimeState*
state,
+ const Block& expected) {
+ Block output;
+ bool eos = false;
+ auto status = source->get_block(state, &output, &eos);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+ EXPECT_FALSE(eos);
+ EXPECT_TRUE(ColumnHelper::block_equal(output, expected))
+ << "output: " << output.dump_data() << "expected: " <<
expected.dump_data();
+}
+
+void expect_spill_source_eos(AnalyticSourceOperatorX* source, RuntimeState*
state) {
+ Block output;
+ bool eos = false;
+ auto status = source->get_block(state, &output, &eos);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+ EXPECT_TRUE(eos);
+ EXPECT_TRUE(output.empty()) << output.dump_data();
+}
+
+Block int64_block_with_result(const std::vector<int64_t>& keys, const
std::vector<int64_t>& values,
+ const std::vector<int64_t>& results) {
+ auto block = ColumnHelper::create_block<DataTypeInt64>(keys, values);
+ block.insert({ColumnHelper::create_column<DataTypeInt64>(results),
+ std::make_shared<DataTypeInt64>(), "result"});
+ return block;
+}
+
+Block int64_block_with_double_result(const std::vector<int64_t>& keys,
+ const std::vector<int64_t>& values,
+ const std::vector<double>& results) {
+ auto block = ColumnHelper::create_block<DataTypeInt64>(keys, values);
+ block.insert({ColumnHelper::create_column<DataTypeFloat64>(results),
+ std::make_shared<DataTypeFloat64>(), "result"});
+ return block;
+}
+
+Block order_block_with_double_result(const std::vector<int64_t>& order_values,
+ const std::vector<double>& results) {
+ auto block = ColumnHelper::create_block<DataTypeInt64>(order_values);
+ block.insert({ColumnHelper::create_column<DataTypeFloat64>(results),
+ std::make_shared<DataTypeFloat64>(), "result"});
+ return block;
+}
+
+} // namespace
+
+TEST_F(AnalyticSinkOperatorTest,
SpillPathBatchesSmallPartitionsIntoOneOutputBlock) {
+ Initialize(8);
+ create_operator(false, 1, "sum", {std::make_shared<DataTypeInt64>()},
+ std::make_shared<DataTypeInt64>());
+ sink->_agg_expr_ctxs.resize(1);
+ sink->_agg_expr_ctxs[0] =
+ MockSlotRef::create_mock_contexts(1,
std::make_shared<DataTypeInt64>());
+ sink->_partition_by_eq_expr_ctxs =
+ MockSlotRef::create_mock_contexts(0,
std::make_shared<DataTypeInt64>());
+ sink->_enable_spill_analytic = true;
+ sink->_window_spill_strategies = {WindowSpillStrategy::PARTITION_REDUCE};
+ sink->_window_spill_peer_functions = {WindowSpillPeerFunction::NONE};
+ create_local_state();
+
+ Block first = ColumnHelper::create_block<DataTypeInt64>({1, 2, 2, 3, 4, 4,
4, 5},
+ {1, 2, 3, 4, 5, 6,
7, 8});
+ auto status = sink->sink(state.get(), &first, false);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+ // Four partitions finished inside one input Block and are published as a
single batch.
+ {
+ LockGuard lock(sink_local_state->_shared_state->buffer_mutex);
+ ASSERT_EQ(sink_local_state->_shared_state->spill_batches.size(), 1);
+
EXPECT_EQ(sink_local_state->_shared_state->spill_batches.front()->partition_ends,
+ (std::vector<int64_t> {1, 3, 4, 7}));
+ }
+ EXPECT_EQ(sink_local_state->_open_partition_rows, 1);
+
+ Block second = ColumnHelper::create_block<DataTypeInt64>({5, 6}, {10, 20});
+ status = sink->sink(state.get(), &second, true);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+
+ expect_next_spill_block(source.get(), state.get(),
+ int64_block_with_result({1, 2, 2, 3, 4, 4, 4}, {1,
2, 3, 4, 5, 6, 7},
+ {1, 5, 5, 4, 18, 18, 18}));
+ expect_next_spill_block(source.get(), state.get(),
int64_block_with_result({5}, {8}, {18}));
+ expect_next_spill_block(source.get(), state.get(),
int64_block_with_result({5}, {10}, {18}));
+ expect_next_spill_block(source.get(), state.get(),
int64_block_with_result({6}, {20}, {20}));
+ expect_spill_source_eos(source.get(), state.get());
+}
+
+TEST_F(AnalyticSinkOperatorTest,
SpillPathNullableSumKeepsAllNullPartitionNull) {
+ Initialize(4);
+ create_operator(false, 1, "sum", {std::make_shared<DataTypeInt64>()},
+ std::make_shared<DataTypeInt64>());
+ const auto nullable_int64 =
make_nullable(std::make_shared<DataTypeInt64>());
+ sink->_agg_functions[0] =
+ create_agg_fn(pool, "sum", {nullable_int64}, nullable_int64, true,
true);
+ sink->_agg_expr_ctxs.resize(1);
+ sink->_agg_expr_ctxs[0] = MockSlotRef::create_mock_contexts(1,
nullable_int64);
+ sink->_partition_by_eq_expr_ctxs =
+ MockSlotRef::create_mock_contexts(0,
std::make_shared<DataTypeInt64>());
+ sink->_enable_spill_analytic = true;
+ sink->_window_spill_strategies = {WindowSpillStrategy::PARTITION_REDUCE};
+ sink->_window_spill_peer_functions = {WindowSpillPeerFunction::NONE};
+ create_local_state();
+
+ Block input;
+ input.insert({ColumnHelper::create_column<DataTypeInt64>({1, 1, 2, 2, 3}),
+ std::make_shared<DataTypeInt64>(), "key"});
+ input.insert(
+ {ColumnHelper::create_nullable_column<DataTypeInt64>({0, 0, 1, 0,
7}, {1, 1, 0, 1, 0}),
+ nullable_int64, "value"});
+ auto status = sink->sink(state.get(), &input, true);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+
+ Block expected;
+ expected.insert({ColumnHelper::create_column<DataTypeInt64>({1, 1, 2, 2}),
+ std::make_shared<DataTypeInt64>(), "key"});
+ expected.insert(
+ {ColumnHelper::create_nullable_column<DataTypeInt64>({0, 0, 1, 0},
{1, 1, 0, 1}),
+ nullable_int64, "value"});
+ expected.insert(
+ {ColumnHelper::create_nullable_column<DataTypeInt64>({0, 0, 1, 1},
{1, 1, 0, 0}),
+ nullable_int64, "result"});
+ expect_next_spill_block(source.get(), state.get(), expected);
+
+ Block expected_tail;
+ expected_tail.insert({ColumnHelper::create_column<DataTypeInt64>({3}),
+ std::make_shared<DataTypeInt64>(), "key"});
+
expected_tail.insert({ColumnHelper::create_nullable_column<DataTypeInt64>({7},
{0}),
+ nullable_int64, "value"});
+
expected_tail.insert({ColumnHelper::create_nullable_column<DataTypeInt64>({7},
{0}),
+ nullable_int64, "result"});
+ expect_next_spill_block(source.get(), state.get(), expected_tail);
+ expect_spill_source_eos(source.get(), state.get());
+}
+
+TEST_F(AnalyticSinkOperatorTest,
SpillPathNtileUsesPartitionCardinalityInsideBatch) {
+ ASSERT_TRUE(spill_file_manager.init().ok());
+
+ Initialize(4);
+ state->_query_options.__set_enable_force_spill(true);
+ create_operator(true, 1, "ntile", {std::make_shared<DataTypeInt64>()},
+ std::make_shared<DataTypeInt64>());
+ sink->_agg_expr_ctxs.resize(1);
+ sink->_agg_expr_ctxs[0] =
+ MockSlotRef::create_mock_contexts(1,
std::make_shared<DataTypeInt64>());
+ sink->_partition_by_eq_expr_ctxs =
+ MockSlotRef::create_mock_contexts(0,
std::make_shared<DataTypeInt64>());
+ sink->_enable_spill_analytic = true;
+ sink->_window_spill_strategies =
{WindowSpillStrategy::PARTITION_CARDINALITY};
+ sink->_window_spill_peer_functions = {WindowSpillPeerFunction::NONE};
+ create_local_state();
+
+ Block first = ColumnHelper::create_block<DataTypeInt64>({1, 1, 1}, {2, 2,
2});
+ auto status = sink->sink(state.get(), &first, false);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+ EXPECT_TRUE(sink_local_state->_batch_store->is_spilled());
+ Block second = ColumnHelper::create_block<DataTypeInt64>({1, 2, 3, 3, 3},
{2, 2, 2, 2, 2});
+ status = sink->sink(state.get(), &second, true);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+
+ // The spilled partition 1 and the small partition 2 are replayed from one
spilled batch.
+ expect_next_spill_block(source.get(), state.get(),
+ int64_block_with_result({1, 1, 1}, {2, 2, 2}, {1,
1, 2}));
+ expect_next_spill_block(source.get(), state.get(),
+ int64_block_with_result({1, 2}, {2, 2}, {2, 1}));
+ expect_next_spill_block(source.get(), state.get(),
+ int64_block_with_result({3, 3, 3}, {2, 2, 2}, {1,
1, 2}));
+ expect_spill_source_eos(source.get(), state.get());
+}
+
+void AnalyticSinkOperatorTest::check_cume_dist_peer_groups_inside_batch(bool
force_spill) {
+ ASSERT_TRUE(spill_file_manager.init().ok());
+
+ Initialize(6);
+ state->_query_options.__set_enable_force_spill(force_spill);
+ create_operator(true, 1, "cume_dist", {},
std::make_shared<DataTypeFloat64>());
+ sink->_agg_expr_ctxs.resize(1);
+ sink->_partition_by_eq_expr_ctxs =
+ MockSlotRef::create_mock_contexts(0,
std::make_shared<DataTypeInt64>());
+ sink->_order_by_eq_expr_ctxs =
+ MockSlotRef::create_mock_contexts(1,
std::make_shared<DataTypeInt64>());
+ sink->_enable_spill_analytic = true;
+ sink->_window_spill_strategies = {WindowSpillStrategy::PEER_GROUP};
+ sink->_window_spill_peer_functions = {WindowSpillPeerFunction::CUME_DIST};
+ create_local_state();
+
+ Block input = ColumnHelper::create_block<DataTypeInt64>({1, 1, 1, 2, 2,
3}, {1, 1, 2, 2, 4, 4});
+ auto status = sink->sink(state.get(), &input, true);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+
+ // Equal order keys in adjacent partitions must not be merged into one
peer group.
+ expect_next_spill_block(source.get(), state.get(),
+ int64_block_with_double_result({1, 1, 1, 2, 2},
{1, 1, 2, 2, 4},
+ {2.0 / 3.0, 2.0 /
3.0, 1.0, 0.5, 1.0}));
+ expect_next_spill_block(source.get(), state.get(),
+ int64_block_with_double_result({3}, {4}, {1.0}));
+ expect_spill_source_eos(source.get(), state.get());
+}
+
+TEST_F(AnalyticSinkOperatorTest,
SpillPathPeerGroupsStayInsidePartitionOfInMemoryBatch) {
+ check_cume_dist_peer_groups_inside_batch(false);
+}
+
+TEST_F(AnalyticSinkOperatorTest,
SpillPathPeerGroupsStayInsidePartitionOfSpilledBatch) {
+ check_cume_dist_peer_groups_inside_batch(true);
+}
+
+TEST_F(AnalyticSinkOperatorTest, SpillPathPercentRankAcrossPartitionsInBatch) {
+ Initialize(6);
+ create_operator(true, 1, "percent_rank", {},
std::make_shared<DataTypeFloat64>());
+ sink->_agg_expr_ctxs.resize(1);
+ sink->_partition_by_eq_expr_ctxs =
+ MockSlotRef::create_mock_contexts(0,
std::make_shared<DataTypeInt64>());
+ sink->_order_by_eq_expr_ctxs =
+ MockSlotRef::create_mock_contexts(1,
std::make_shared<DataTypeInt64>());
+ sink->_enable_spill_analytic = true;
+ sink->_window_spill_strategies = {WindowSpillStrategy::PEER_GROUP};
+ sink->_window_spill_peer_functions =
{WindowSpillPeerFunction::PERCENT_RANK};
+ create_local_state();
+
+ Block first = ColumnHelper::create_block<DataTypeInt64>({1, 1, 1, 2}, {1,
1, 2, 3});
+ auto status = sink->sink(state.get(), &first, false);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+ Block second = ColumnHelper::create_block<DataTypeInt64>({2, 2, 3}, {3, 4,
5});
+ status = sink->sink(state.get(), &second, true);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+
+ expect_next_spill_block(source.get(), state.get(),
+ int64_block_with_double_result({1, 1, 1}, {1, 1,
2}, {0.0, 0.0, 1.0}));
+ // Partition 2 continues from the first input Block into the second one.
+ expect_next_spill_block(source.get(), state.get(),
+ int64_block_with_double_result({2}, {3}, {0.0}));
+ expect_next_spill_block(source.get(), state.get(),
+ int64_block_with_double_result({2, 2}, {3, 4},
{0.0, 1.0}));
+ expect_next_spill_block(source.get(), state.get(),
+ int64_block_with_double_result({3}, {5}, {0.0}));
+ expect_spill_source_eos(source.get(), state.get());
+}
+
+TEST_F(AnalyticSinkOperatorTest, SpillPathRevokeMemorySpillsOpenPartition) {
+ ASSERT_TRUE(spill_file_manager.init().ok());
+
+ Initialize(5);
+ create_operator(false, 1, "sum", {std::make_shared<DataTypeInt64>()},
+ std::make_shared<DataTypeInt64>());
+ sink->_agg_expr_ctxs.resize(1);
+ sink->_agg_expr_ctxs[0] =
+ MockSlotRef::create_mock_contexts(0,
std::make_shared<DataTypeInt64>());
+ sink->_enable_spill_analytic = true;
+ sink->_window_spill_strategies = {WindowSpillStrategy::PARTITION_REDUCE};
+ sink->_window_spill_peer_functions = {WindowSpillPeerFunction::NONE};
+ create_local_state();
+
+ Block first = ColumnHelper::create_block<DataTypeInt64>({1, 2});
+ auto status = sink->sink(state.get(), &first, false);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+ ASSERT_FALSE(sink_local_state->_batch_store->is_spilled());
+ EXPECT_GT(sink_local_state->_batch_store->revocable_mem_size(), 0);
+
+ status = sink->revoke_memory(state.get());
+ ASSERT_TRUE(status.ok()) << status.to_string();
+ EXPECT_TRUE(sink_local_state->_batch_store->is_spilled());
+ EXPECT_EQ(sink_local_state->_batch_store->revocable_mem_size(), 0);
+
+ Block second = ColumnHelper::create_block<DataTypeInt64>({3, 4, 5});
+ status = sink->sink(state.get(), &second, true);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+
+ auto expected_block = [](const std::vector<int64_t>& values) {
+ auto block = ColumnHelper::create_block<DataTypeInt64>(values);
+ block.insert({ColumnHelper::create_column<DataTypeInt64>(
+ std::vector<int64_t>(values.size(), 15)),
+ std::make_shared<DataTypeInt64>(), "sum"});
+ return block;
+ };
+ expect_next_spill_block(source.get(), state.get(), expected_block({1, 2}));
+ expect_next_spill_block(source.get(), state.get(), expected_block({3, 4,
5}));
+ expect_spill_source_eos(source.get(), state.get());
+}
+
+void AnalyticSinkOperatorTest::prepare_spilled_full_partition_sum() {
+ create_operator(false, 1, "sum", {std::make_shared<DataTypeInt64>()},
+ std::make_shared<DataTypeInt64>());
+ sink->_agg_expr_ctxs.resize(1);
+ sink->_agg_expr_ctxs[0] =
+ MockSlotRef::create_mock_contexts(0,
std::make_shared<DataTypeInt64>());
+ sink->_enable_spill_analytic = true;
+ sink->_window_spill_strategies = {WindowSpillStrategy::PARTITION_REDUCE};
+ sink->_window_spill_peer_functions = {WindowSpillPeerFunction::NONE};
+ create_local_state();
+}
+
+namespace {
+
+void sink_small_blocks(AnalyticSinkOperatorX* sink, RuntimeState* state,
int64_t first_value,
+ size_t block_count) {
+ for (size_t i = 0; i < block_count; ++i) {
+ const auto base = first_value + static_cast<int64_t>(i * 3);
+ Block block = ColumnHelper::create_block<DataTypeInt64>({base, base +
1, base + 2});
+ auto status = sink->sink(state, &block, false);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+ }
+}
+
+Block values_with_sum(int64_t first_value, size_t rows, int64_t sum) {
+ std::vector<int64_t> values(rows);
+ for (size_t i = 0; i < rows; ++i) {
+ values[i] = first_value + static_cast<int64_t>(i);
+ }
+ auto block = ColumnHelper::create_block<DataTypeInt64>(values);
+
block.insert({ColumnHelper::create_column<DataTypeInt64>(std::vector<int64_t>(rows,
sum)),
+ std::make_shared<DataTypeInt64>(), "sum"});
+ return block;
+}
+
+int64_t spill_write_block_count(AnalyticSinkLocalState* local_state) {
+ return
local_state->custom_profile()->get_counter("SpillWriteBlockCount")->value();
+}
+
+size_t replay_read_reserve(const AnalyticSpillBatch::RecordMemory& memory) {
+ return std::max(memory.serialized_bytes, memory.footer_bytes) +
memory.footer_bytes +
+ memory.serialized_bytes + 2 * memory.decoded_bytes;
+}
+
+} // namespace
+
+TEST_F(AnalyticSinkOperatorTest, SpillPathCoalescesSmallBlocksAndSlicesReplay)
{
+ ASSERT_TRUE(spill_file_manager.init().ok());
+
+ Initialize(8);
+ state->_query_options.__set_enable_force_spill(true);
+ prepare_spilled_full_partition_sum();
+
+ // The first Block triggers the spill and is written on its own. The next
nine small Blocks
+ // are coalesced and written as one Block when the batch is sealed.
+ sink_small_blocks(sink.get(), state.get(), 0, 10);
+ ASSERT_TRUE(sink_local_state->_batch_store->is_spilled());
+ EXPECT_EQ(spill_write_block_count(sink_local_state), 1);
+ EXPECT_GT(sink_local_state->_batch_store->revocable_mem_size(), 0);
+ Block eos_block = ColumnHelper::create_block<DataTypeInt64>({});
+ auto status = sink->sink(state.get(), &eos_block, true);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+ EXPECT_EQ(spill_write_block_count(sink_local_state), 2);
+
+ // Opening the batch reserves the data reader and its first record,
reading the second record
+ // one more Block, and slicing a retained Block nothing.
+ AnalyticSpillBatch::RecordMemory record_memory;
+ {
+ LockGuard lock(sink_local_state->_shared_state->buffer_mutex);
+ record_memory =
sink_local_state->_shared_state->spill_batches.front()->data_record_memory;
+ }
+ EXPECT_GT(record_memory.serialized_bytes, 0);
+ EXPECT_EQ(source_local_state->_spill_replay_reserve_bytes(state.get()),
+ replay_read_reserve(record_memory));
+ // The 27-row coalesced Block is replayed in batch-size slices.
+ expect_next_spill_block(source.get(), state.get(), values_with_sum(0, 3,
435));
+ EXPECT_EQ(source_local_state->_spill_replay_reserve_bytes(state.get()),
+ replay_read_reserve(record_memory));
+ expect_next_spill_block(source.get(), state.get(), values_with_sum(3, 8,
435));
+ EXPECT_EQ(source_local_state->_spill_replay_reserve_bytes(state.get()), 0);
+ source_local_state->estimate_memory_usage() = 100 * 1024 * 1024;
+ EXPECT_LT(source->get_reserve_mem_size(state.get()), 100 * 1024 * 1024);
+ expect_next_spill_block(source.get(), state.get(), values_with_sum(11, 8,
435));
+ expect_next_spill_block(source.get(), state.get(), values_with_sum(19, 8,
435));
+ expect_next_spill_block(source.get(), state.get(), values_with_sum(27, 3,
435));
+ expect_spill_source_eos(source.get(), state.get());
+}
+
+TEST_F(AnalyticSinkOperatorTest, SpillPathRevokeMemoryFlushesCoalescedBlocks) {
+ ASSERT_TRUE(spill_file_manager.init().ok());
+
+ Initialize(8);
+ state->_query_options.__set_enable_force_spill(true);
+ prepare_spilled_full_partition_sum();
+
+ sink_small_blocks(sink.get(), state.get(), 0, 5);
+ EXPECT_EQ(spill_write_block_count(sink_local_state), 1);
+ ASSERT_GT(sink_local_state->_batch_store->revocable_mem_size(), 0);
+
+ auto status = sink->revoke_memory(state.get());
+ ASSERT_TRUE(status.ok()) << status.to_string();
+ EXPECT_EQ(sink_local_state->_batch_store->revocable_mem_size(), 0);
+ EXPECT_EQ(spill_write_block_count(sink_local_state), 2);
+
+ sink_small_blocks(sink.get(), state.get(), 15, 5);
+ Block eos_block = ColumnHelper::create_block<DataTypeInt64>({});
+ status = sink->sink(state.get(), &eos_block, true);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+ EXPECT_EQ(spill_write_block_count(sink_local_state), 3);
+
+ expect_next_spill_block(source.get(), state.get(), values_with_sum(0, 3,
435));
+ expect_next_spill_block(source.get(), state.get(), values_with_sum(3, 8,
435));
+ expect_next_spill_block(source.get(), state.get(), values_with_sum(11, 4,
435));
+ expect_next_spill_block(source.get(), state.get(), values_with_sum(15, 8,
435));
+ expect_next_spill_block(source.get(), state.get(), values_with_sum(23, 7,
435));
+ expect_spill_source_eos(source.get(), state.get());
+}
+
+TEST_F(AnalyticSinkOperatorTest, SpillAdmissionIncludesPerPartitionResults) {
+ Initialize(60000);
+ prepare_spilled_full_partition_sum();
+
+ Block input = ColumnHelper::create_block<DataTypeInt64>(_data_vals);
+ const size_t input_bytes = input.allocated_bytes();
+ const size_t per_row = sizeof(size_t) + sizeof(int64_t) + sizeof(int64_t);
+ EXPECT_GE(sink->get_reserve_mem_size(state.get(), false, &input),
+ 2 * input_bytes + input.rows() * per_row);
+}
+
+TEST_F(AnalyticSinkOperatorTest,
SpillAdmissionReservesFirstWriteAndDropsEmptyEosPeak) {
+ ASSERT_TRUE(spill_file_manager.init().ok());
+ Initialize(4096);
+ state->_query_options.__set_enable_force_spill(true);
+ state->_query_options.__set_spill_buffer_size_bytes(1024 * 1024);
+ prepare_spilled_full_partition_sum();
+
+ std::vector<int64_t> values(100000, 1);
+ Block first = ColumnHelper::create_block<DataTypeInt64>(values);
+ EXPECT_GE(sink->get_reserve_mem_size(state.get(), false, &first), 3 *
first.allocated_bytes());
+ auto status = sink->sink(state.get(), &first, false);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+ ASSERT_TRUE(sink_local_state->_batch_store->is_spilled());
+
+ // The second Block remains staged below the default 32 MiB threshold but
is large enough
+ // for PipelineTask's 512 KiB revoke floor.
+ Block second = ColumnHelper::create_block<DataTypeInt64>(values);
+ status = sink->sink(state.get(), &second, false);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+ EXPECT_GE(sink->revocable_mem_size(state.get()),
SpillFile::MIN_SPILL_WRITE_BATCH_MEM);
+ EXPECT_LT(sink->revocable_mem_size(state.get()),
+ static_cast<size_t>(state->spill_min_revocable_mem()));
+ status = sink->revoke_memory(state.get());
+ ASSERT_TRUE(status.ok()) << status.to_string();
+ EXPECT_EQ(sink->revocable_mem_size(state.get()), 0);
+
+ Block eos_block = ColumnHelper::create_block<DataTypeInt64>({});
+ EXPECT_EQ(sink->get_reserve_mem_size(state.get(), true, &eos_block), 0);
+ status = sink->sink(state.get(), &eos_block, true);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+}
+
+TEST_F(AnalyticSinkOperatorTest, SpillAdmissionUsesActualOversizeRecord) {
+ ASSERT_TRUE(spill_file_manager.init().ok());
+ Initialize(256);
+ state->_query_options.__set_enable_force_spill(true);
+ state->_query_options.__set_spill_buffer_size_bytes(1024 * 1024);
+ prepare_spilled_full_partition_sum();
+
+ // A single wide row cannot be split at the effective 1 MiB spill-buffer
boundary.
+ std::string payload(2 * 1024 * 1024, '\0');
+ uint32_t random = 0x12345678;
+ for (char& byte : payload) {
+ random ^= random << 13;
+ random ^= random >> 17;
+ random ^= random << 5;
+ byte = static_cast<char>(random);
+ }
+ Block input = ColumnHelper::create_block<DataTypeInt64>({1});
+ input.insert({ColumnHelper::create_column<DataTypeString>({payload}),
+ std::make_shared<DataTypeString>(), "payload"});
+ auto status = sink->sink(state.get(), &input, true);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+ AnalyticSpillBatch::RecordMemory memory;
+ {
+ LockGuard lock(sink_local_state->_shared_state->buffer_mutex);
+ memory =
sink_local_state->_shared_state->spill_batches.front()->data_record_memory;
+ }
+ EXPECT_GT(memory.serialized_bytes, state->spill_buffer_size_bytes());
+ EXPECT_GT(
+ source->get_reserve_mem_size(state.get()),
+ state->minimum_operator_memory_required_bytes() + 2 *
state->spill_buffer_size_bytes());
+}
+
+TEST_F(AnalyticSinkOperatorTest,
SpillAdmissionLocalReaderIgnoresLargeBufferSetting) {
+ ASSERT_TRUE(spill_file_manager.init().ok());
+ Initialize(4096);
+ state->_query_options.__set_spill_buffer_size_bytes(256 * 1024 * 1024);
+ state->_query_options.__set_spill_analytic_sink_mem_limit_bytes(1024 *
1024);
+ prepare_spilled_full_partition_sum();
+
+ std::vector<int64_t> values(200000);
+ std::iota(values.begin(), values.end(), 0);
+ Block input = ColumnHelper::create_block<DataTypeInt64>(values);
+ auto status = sink->sink(state.get(), &input, true);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+ AnalyticSpillBatch::RecordMemory memory;
+ {
+ LockGuard lock(sink_local_state->_shared_state->buffer_mutex);
+ const auto& batch =
sink_local_state->_shared_state->spill_batches.front();
+ ASSERT_NE(batch->data_file, nullptr);
+ EXPECT_FALSE(batch->data_file->is_remote());
+ memory = batch->data_record_memory;
+ }
+ EXPECT_GT(memory.footer_bytes, 0);
+ EXPECT_LT(source->get_reserve_mem_size(state.get()), 16 * 1024 * 1024);
+}
+
+void AnalyticSinkOperatorTest::prepare_cume_dist_without_partition() {
+ create_operator(true, 1, "cume_dist", {},
std::make_shared<DataTypeFloat64>());
+ sink->_agg_expr_ctxs.resize(1);
+ sink->_order_by_eq_expr_ctxs =
+ MockSlotRef::create_mock_contexts(0,
std::make_shared<DataTypeInt64>());
+ sink->_enable_spill_analytic = true;
+ sink->_window_spill_strategies = {WindowSpillStrategy::PEER_GROUP};
+ sink->_window_spill_peer_functions = {WindowSpillPeerFunction::CUME_DIST};
+ create_local_state();
+}
+
+TEST_F(AnalyticSinkOperatorTest, SpillPathReplaysPeerGroupsFromSidecarFile) {
+ ASSERT_TRUE(spill_file_manager.init().ok());
+
+ Initialize(8);
+ prepare_cume_dist_without_partition();
+
+ // Peer group ends collected before the revoke go to the sidecar file, the
rest at seal.
+ Block first = ColumnHelper::create_block<DataTypeInt64>({1, 1, 2});
+ auto status = sink->sink(state.get(), &first, false);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+ status = sink->revoke_memory(state.get());
+ ASSERT_TRUE(status.ok()) << status.to_string();
+ ASSERT_TRUE(sink_local_state->_batch_store->has_spilled_peer_groups());
+ Block second = ColumnHelper::create_block<DataTypeInt64>({2, 4});
+ status = sink->sink(state.get(), &second, true);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+
+ // Opening the batch reserves the reader buffer and one deserialized
record for the data file
+ // and for the peer group sidecar file.
+ std::shared_ptr<AnalyticSpillBatch> batch;
+ {
+ LockGuard lock(sink_local_state->_shared_state->buffer_mutex);
+ batch = sink_local_state->_shared_state->spill_batches.front();
+ }
+ const auto data_open = replay_read_reserve(batch->data_record_memory);
+ const auto peer_open = replay_read_reserve(batch->peer_record_memory);
+ EXPECT_EQ(source_local_state->_spill_replay_reserve_bytes(state.get()),
data_open + peer_open);
+ EXPECT_GE(source->get_reserve_mem_size(state.get()),
+ data_open + peer_open +
state->minimum_operator_memory_required_bytes());
+ expect_next_spill_block(source.get(), state.get(),
+ order_block_with_double_result({1, 1, 2}, {0.4,
0.4, 0.8}));
+ ASSERT_NE(source_local_state->_peer_group_reader, nullptr);
+ // The second sidecar record ([4, 5]) was read while replaying row 2 and
covers the rest of
+ // the batch, so only the next data record is reserved.
+ EXPECT_EQ(source_local_state->_spill_replay_reserve_bytes(state.get()),
+ replay_read_reserve(batch->data_record_memory));
+ expect_next_spill_block(source.get(), state.get(),
+ order_block_with_double_result({2, 4}, {0.8,
1.0}));
+ // Nothing is queued behind the batch, so finishing it reserves nothing.
+ EXPECT_EQ(source_local_state->_spill_replay_reserve_bytes(state.get()), 0);
+ expect_spill_source_eos(source.get(), state.get());
+}
+
+TEST_F(AnalyticSinkOperatorTest, SpillPathFlushesPeerGroupsAtSpillBufferSize) {
+ ASSERT_TRUE(spill_file_manager.init().ok());
+
+ Initialize(4096);
+ state->_query_options.__set_spill_buffer_size_bytes(1024 * 1024);
+ prepare_cume_dist_without_partition();
+
+ // Every row starts a peer group, so the peer group ends exceed the 1MB
buffer
+ // (131072 entries) and are flushed to the sidecar file before the batch
is sealed.
+ constexpr int64_t rows = 200000;
+ std::vector<int64_t> order_values(rows);
+ std::iota(order_values.begin(), order_values.end(), 0);
+ Block input = ColumnHelper::create_block<DataTypeInt64>(order_values);
+ auto status = sink->sink(state.get(), &input, false);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+ EXPECT_TRUE(sink_local_state->_batch_store->has_spilled_peer_groups());
+ EXPECT_FALSE(sink_local_state->_batch_store->is_spilled());
+ Block eos_block = ColumnHelper::create_block<DataTypeInt64>({});
+ status = sink->sink(state.get(), &eos_block, true);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+
+ // Opening the batch reserves the sidecar reader and its first record; the
rows are in
+ // memory and reserve nothing. Afterwards only the slice that starts at
the first row not
+ // covered by the first sidecar record reads the second record.
+ AnalyticSpillBatch::RecordMemory peer_record_memory;
+ {
+ LockGuard lock(sink_local_state->_shared_state->buffer_mutex);
+ peer_record_memory =
+
sink_local_state->_shared_state->spill_batches.front()->peer_record_memory;
+ }
+ constexpr int64_t first_record_peer_groups = 1024 * 1024 / sizeof(int64_t);
+ int64_t output_rows = 0;
+ bool eos = false;
+ while (!eos) {
+ size_t expected_reserve_bytes = 0;
+ if (output_rows == 0) {
+ expected_reserve_bytes = replay_read_reserve(peer_record_memory);
+ } else if (output_rows == first_record_peer_groups) {
+ expected_reserve_bytes = replay_read_reserve(peer_record_memory);
+ }
+ EXPECT_EQ(source_local_state->_spill_replay_reserve_bytes(state.get()),
+ expected_reserve_bytes)
+ << "output_rows=" << output_rows;
+ Block output;
+ status = source->get_block(state.get(), &output, &eos);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+ if (eos) {
+ break;
+ }
+ ASSERT_NE(source_local_state->_peer_group_reader, nullptr);
+ ASSERT_LE(output.rows(), 4096);
+ const auto& order = assert_cast<const
ColumnInt64&>(*output.get_by_position(0).column);
+ const auto& cume_dist =
+ assert_cast<const
ColumnFloat64&>(*output.get_by_position(1).column);
+ for (size_t i = 0; i < output.rows(); ++i) {
+ ASSERT_EQ(order.get_data()[i], output_rows);
+ ASSERT_EQ(cume_dist.get_data()[i],
+ static_cast<double>(output_rows + 1) /
static_cast<double>(rows));
+ ++output_rows;
+ }
+ if (output_rows == 4096) {
+ // A data Block ending exactly at the last loaded peer end needs
no sidecar read,
+ // even though batch_size would extend into the following data
Block.
+ Block two_rows = ColumnHelper::create_block<DataTypeInt64>({1, 2});
+ source_local_state->_replay_block.swap(two_rows);
+ const auto saved_replay_pos =
source_local_state->_replay_block_position;
+ const auto saved_output_pos =
source_local_state->_batch_output_position;
+ source_local_state->_replay_block_position = 0;
+ source_local_state->_batch_output_position =
first_record_peer_groups - 2;
+
EXPECT_FALSE(source_local_state->_next_slice_reads_peer_group_record(state.get()));
+ source_local_state->_replay_block.swap(two_rows);
+ source_local_state->_replay_block_position = saved_replay_pos;
+ source_local_state->_batch_output_position = saved_output_pos;
+ }
+ }
+ EXPECT_EQ(output_rows, rows);
+}
+
+TEST_F(AnalyticSinkOperatorTest,
SpillPathKeepsSealedBatchBelowSinkLimitInMemory) {
+ ASSERT_TRUE(spill_file_manager.init().ok());
+
+ Initialize(4096);
+ // A sealed batch is only spilled when it reaches
spill_analytic_sink_mem_limit_bytes
+ // (64MB by default), not when it exceeds the much smaller
min_revocable_mem.
+ state->_query_options.__set_min_revocable_mem(1024 * 1024);
+ create_operator(false, 1, "sum", {std::make_shared<DataTypeInt64>()},
+ std::make_shared<DataTypeInt64>());
+ sink->_agg_expr_ctxs.resize(1);
+ sink->_agg_expr_ctxs[0] =
+ MockSlotRef::create_mock_contexts(1,
std::make_shared<DataTypeInt64>());
+ sink->_partition_by_eq_expr_ctxs =
+ MockSlotRef::create_mock_contexts(0,
std::make_shared<DataTypeInt64>());
+ sink->_enable_spill_analytic = true;
+ sink->_window_spill_strategies = {WindowSpillStrategy::PARTITION_REDUCE};
+ sink->_window_spill_peer_functions = {WindowSpillPeerFunction::NONE};
+ create_local_state();
+
+ constexpr size_t rows = 200000;
+ std::vector<int64_t> keys(rows, 1);
+ keys.back() = 2;
+ std::vector<int64_t> values(rows, 1);
+ Block input = ColumnHelper::create_block<DataTypeInt64>(keys, values);
+ auto status = sink->sink(state.get(), &input, false);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+
+ {
+ LockGuard lock(sink_local_state->_shared_state->buffer_mutex);
+ ASSERT_EQ(sink_local_state->_shared_state->spill_batches.size(), 1);
+ const auto& batch =
sink_local_state->_shared_state->spill_batches.front();
+ EXPECT_EQ(batch->rows, rows - 1);
+ EXPECT_EQ(batch->data_file, nullptr);
+ size_t batch_bytes = 0;
+ for (const auto& block : batch->blocks) {
+ batch_bytes += block.allocated_bytes();
+ }
+ EXPECT_GT(batch_bytes, state->spill_min_revocable_mem());
+ }
+ EXPECT_EQ(spill_write_block_count(sink_local_state), 0);
+ // Replaying the in-memory batch opens no spill file, but still allocates
one result column.
+ EXPECT_EQ(source_local_state->_spill_replay_reserve_bytes(state.get()), 0);
+ EXPECT_EQ(source->get_reserve_mem_size(state.get()),
Review Comment:
Rechecked this expectation against the unmodified PR head `204d20fb9fdd`.
The queued batch contains a 199,999-row Block with two Int64 columns, while
the next output has 4,096 rows. `_next_output_reserve_bytes()` therefore
includes both:
- First input slice: `4096 * 2 columns * 2 * sizeof(int64_t) = 131072` bytes.
- SUM result column: `4096 * sizeof(int64_t) = 32768` bytes.
The expected total should be `minimum_operator_memory_required_bytes() +
163840` bytes (160 KiB), rather than the current minimum plus 32768 bytes.
Please include the first-slice reservation and retain the separate assertion
that `_spill_replay_reserve_bytes()` is zero, since this batch opens no spill
file.
For validation context, I checked the [BE UT build for this exact
head](http://43.132.222.7:8111/viewLog.html?buildId=1078005&buildTypeId=Doris_DorisBeUt_BeUt).
It stops while compiling `function_string_test.cpp` (`Field` conversion and
`FunctionContext::FRAGMENT_LOCAL/THREAD_LOCAL` errors), before the UT
executable runs. Thus the 128 KiB mismatch above is established by tracing the
current code; it is not an observed assertion failure from that CI run. Please
rerun the focused test once that build blocker is resolved.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
