github-actions[bot] commented on code in PR #68500:
URL: https://github.com/apache/doris/pull/68500#discussion_r4373838899
##########
be/test/exec/operator/analytic_sink_operator_test.cpp:
##########
@@ -310,6 +383,1500 @@ 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,
SpillAdmissionCountsNullableFixedWidthCutStorage) {
+ Initialize(65535);
+ prepare_spilled_full_partition_sum();
+ constexpr size_t block_rows = 65536;
+ constexpr size_t columns = 32;
+ const auto nullable_type =
make_nullable(std::make_shared<DataTypeInt64>());
+ const std::vector<int64_t> values(block_rows, 7);
+
+ for (bool mixed_nulls : {false, true}) {
+ std::vector<UInt8> null_map(block_rows, 1);
+ if (mixed_nulls) {
+ for (size_t row = 0; row < block_rows; row += 2) {
+ null_map[row] = 0;
+ }
+ }
+ const auto column =
ColumnHelper::create_nullable_column<DataTypeInt64>(values, null_map);
+ Block block;
+ for (size_t i = 0; i < columns; ++i) {
+ block.insert({column, nullable_type, "nullable_input"});
+ }
+ auto batch = std::make_shared<AnalyticSpillBatch>();
+ batch->rows = block_rows;
+ batch->blocks.push_back(std::move(block));
+ source_local_state->_current_batch = batch;
+
+ const size_t rows = state->batch_size();
+ const size_t budget =
source_local_state->_next_output_reserve_bytes(state.get());
+ size_t cut_bytes = 0;
+ for (const auto& input : batch->blocks.front()) {
+ cut_bytes += input.column->cut(0, rows)->allocated_bytes();
+ }
+ EXPECT_EQ(budget, 2 * rows * columns *
nullable_type->get_size_of_value_in_memory());
+ EXPECT_GE(budget, cut_bytes);
+ }
+
+ // A constant NULL remains one stored value after cut, not 65,535 nested
Int64 values.
+ auto null_column = nullable_type->create_column();
+ null_column->insert_default();
+ ColumnPtr const_column = ColumnConst::create(std::move(null_column),
block_rows);
+ Block const_block;
+ const_block.insert({const_column, nullable_type, "constant_null"});
+ auto batch = std::make_shared<AnalyticSpillBatch>();
+ batch->rows = block_rows;
+ batch->blocks.push_back(std::move(const_block));
+ source_local_state->_current_batch = batch;
+ const size_t budget =
source_local_state->_next_output_reserve_bytes(state.get());
+ EXPECT_EQ(budget, 2 * static_cast<size_t>(state->batch_size()));
+ EXPECT_GE(budget, const_column->cut(0,
state->batch_size())->allocated_bytes());
+}
+
+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();
+}
+
+TEST_F(AnalyticSinkOperatorTest, SpillBoundarySearchSkipsLargeStringGroups) {
+ Initialize(1);
+ create_operator(false, 0, "", {}, nullptr);
+ create_local_state();
+ constexpr size_t rows = 65536;
+ constexpr size_t group_boundary = 60000;
+ const std::string prefix(256, 'x');
+ std::vector<std::string> values(rows, prefix + "b");
+ std::fill_n(values.begin(), group_boundary, prefix + "a");
+ auto counting = AnalyticCompareCountingColumn::create(
+ ColumnHelper::create_column<DataTypeString>(values));
+ auto* counter = counting.get();
+ const std::vector<ColumnPtr> columns {std::move(counting)};
+ PartitionStatistics statistics;
+ size_t boundary = 0;
+ ASSERT_TRUE(sink_local_state
+ ->_find_spill_boundary(state.get(), columns, 0, rows,
statistics, &boundary)
+ .ok());
+ EXPECT_EQ(boundary, group_boundary);
+ EXPECT_LE(counter->comparisons, 20);
+
+ counter->comparisons = 0;
+ ASSERT_TRUE(sink_local_state
+ ->_find_spill_boundary(state.get(), columns,
group_boundary, rows,
+ statistics, &boundary)
+ .ok());
+ EXPECT_EQ(boundary, rows);
+ EXPECT_EQ(counter->comparisons, 1);
+}
+
+TEST_F(AnalyticSinkOperatorTest, SpillBoundarySearchAdaptsToSingletonGroups) {
+ Initialize(1);
+ create_operator(false, 0, "", {}, nullptr);
+ create_local_state();
+ constexpr size_t rows = 4096;
+ std::vector<int64_t> values(rows);
+ std::iota(values.begin(), values.end(), 0);
+ auto counting = AnalyticCompareCountingColumn::create(
+ ColumnHelper::create_column<DataTypeInt64>(values));
+ auto* counter = counting.get();
+ const std::vector<ColumnPtr> columns {std::move(counting)};
+ PartitionStatistics statistics;
+ for (size_t row = 0; row < rows; ++row) {
+ size_t boundary = rows;
+ ASSERT_TRUE(sink_local_state
+ ->_find_spill_boundary(state.get(), columns, row,
rows, statistics,
+ &boundary)
+ .ok());
+ ASSERT_EQ(boundary, row + 1);
+ statistics.update(boundary - row);
+ EXPECT_EQ(statistics.is_high_cardinality(), row >= 16);
+ }
+ // The initial binary searches are bounded; the remaining singleton groups
cost one each.
+ EXPECT_LT(counter->comparisons, rows * 2);
+
+ counter->comparisons = 0;
+ std::vector<size_t> partition_starts;
+ ASSERT_TRUE(
+ sink_local_state
+ ->_find_spill_partition_starts(state.get(), columns, rows,
&partition_starts)
+ .ok());
+ ASSERT_EQ(partition_starts.size(), rows - 1);
+
EXPECT_TRUE(sink_local_state->_partition_column_statistics.is_high_cardinality());
+ EXPECT_LT(counter->comparisons, rows * 2);
+}
+
+TEST_F(AnalyticSinkOperatorTest,
SpillBoundarySearchHandlesCompositeNullableDescendingKeys) {
+ Initialize(1);
+ create_operator(false, 0, "", {}, nullptr);
+ create_local_state();
+ const std::string prefix(256, 'x');
+ const std::vector<ColumnPtr> columns {
+ ColumnHelper::create_nullable_column<DataTypeString>(
+ {"ignored1", "ignored2", "ignored3", prefix + "z", prefix
+ "z", prefix + "z",
+ prefix + "a", prefix + "a"},
+ {1, 1, 1, 0, 0, 0, 0, 0}),
+ ColumnHelper::create_column<DataTypeInt64>({1, 1, 2, 1, 1, 2, 1,
1})};
+ for (bool high_cardinality : {false, true}) {
+ PartitionStatistics statistics;
+ if (high_cardinality) {
+ for (int i = 0; i < 17; ++i) {
+ statistics.update(1);
+ }
+ }
+ // Exercise nonzero starts, singleton ranges, and boundaries exactly
at the range end.
+ for (size_t start = 0; start < 8; ++start) {
+ for (size_t end = start + 1; end <= 8; ++end) {
+ size_t expected = start + 1;
+ while (expected < end &&
+ sink_local_state->_keys_equal(columns, start, columns,
expected)) {
+ ++expected;
+ }
+ size_t boundary = 0;
+ ASSERT_TRUE(sink_local_state
+ ->_find_spill_boundary(state.get(),
columns, start, end,
+ statistics,
&boundary)
+ .ok());
+ EXPECT_EQ(boundary, expected)
+ << "range: " << start << ", " << end << ", linear: "
<< high_cardinality;
+ }
+ }
+ }
+}
+
+TEST_F(AnalyticSinkOperatorTest, SpillBoundarySearchCancelsInsideLinearRun) {
+ Initialize(1);
+ create_operator(false, 0, "", {}, nullptr);
+ create_local_state();
+ const std::vector<ColumnPtr> columns {
+
ColumnHelper::create_column<DataTypeInt64>(std::vector<int64_t>(20000, 1))};
+ PartitionStatistics statistics;
+ for (int i = 0; i < 17; ++i) {
+ statistics.update(1);
+ }
+ CancelAfterChecksRuntimeState cancelled_state;
+ cancelled_state.reset_cancel_after(0);
+ size_t boundary = 0;
+ expect_cancelled(sink_local_state->_find_spill_boundary(
+ &cancelled_state, columns, 0, columns[0]->size(), statistics,
&boundary));
+}
+
+TEST_F(AnalyticSinkOperatorTest,
SpillPathPeerStatisticsCountCompleteGroupsAcrossBlocks) {
+ Initialize(8);
+ prepare_cume_dist_without_partition();
+ for (int i = 0; i < 20; ++i) {
+ Block input = ColumnHelper::create_block<DataTypeInt64>({7});
+ ASSERT_TRUE(sink->sink(state.get(), &input, false).ok());
+ EXPECT_EQ(sink_local_state->_order_by_column_statistics._count, 0);
+ }
+ Block input = ColumnHelper::create_block<DataTypeInt64>({7, 8});
+ ASSERT_TRUE(sink->sink(state.get(), &input, true).ok());
+ EXPECT_EQ(sink_local_state->_order_by_column_statistics._count, 1);
+ EXPECT_EQ(sink_local_state->_order_by_column_statistics._average_size, 21);
+ LockGuard lock(sink_local_state->_shared_state->buffer_mutex);
+ const auto& batch = sink_local_state->_shared_state->spill_batches.front();
+ ASSERT_EQ(batch->peer_group_ends.size(), 2);
+ EXPECT_EQ(batch->peer_group_ends[0], 21);
+ EXPECT_EQ(batch->peer_group_ends[1], 22);
+}
+
+TEST_F(AnalyticSinkOperatorTest, SpillPathPartitionStatisticsAdaptWithinBlock)
{
+ Initialize(64);
+ 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();
+ for (int i = 0; i < 20; ++i) {
+ Block input = ColumnHelper::create_block<DataTypeInt64>({0}, {1});
+ ASSERT_TRUE(sink->sink(state.get(), &input, false).ok());
+ EXPECT_EQ(sink_local_state->_partition_column_statistics._count, 0);
+ }
+ std::vector<int64_t> keys(41);
+ std::iota(keys.begin(), keys.end(), 0);
+ Block input = ColumnHelper::create_block<DataTypeInt64>(keys,
std::vector<int64_t>(41, 1));
+ ASSERT_TRUE(sink->sink(state.get(), &input, true).ok());
+ EXPECT_EQ(sink_local_state->_partition_column_statistics._count, 40);
+ EXPECT_EQ(sink_local_state->_partition_column_statistics._cumulative_size,
60);
+
EXPECT_TRUE(sink_local_state->_partition_column_statistics.is_high_cardinality());
+ LockGuard lock(sink_local_state->_shared_state->buffer_mutex);
+ const auto& batch = sink_local_state->_shared_state->spill_batches.front();
+ ASSERT_EQ(batch->partition_ends.size(), 40);
+ for (size_t i = 0; i < batch->partition_ends.size(); ++i) {
+ EXPECT_EQ(batch->partition_ends[i], 21 + i);
+ EXPECT_EQ(batch->partition_results[0]->get_int(i), i == 0 ? 21 : 1);
+ }
+}
+
+TEST_F(AnalyticSinkOperatorTest,
SpillPathPeerStatisticsResetAtPartitionBoundary) {
+ Initialize(32);
+ 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();
+
+ std::vector<int64_t> keys(21, 1);
+ keys.back() = 2;
+ std::vector<int64_t> orders(21);
+ std::iota(orders.begin(), orders.end(), 0);
+ orders.back() = 19;
+ Block first = ColumnHelper::create_block<DataTypeInt64>(keys, orders);
+ ASSERT_TRUE(sink->sink(state.get(), &first, false).ok());
+ EXPECT_EQ(sink_local_state->_order_by_column_statistics._count, 0);
+ Block second = ColumnHelper::create_block<DataTypeInt64>({2, 2, 2}, {19,
19, 20});
+ ASSERT_TRUE(sink->sink(state.get(), &second, true).ok());
+ EXPECT_EQ(sink_local_state->_order_by_column_statistics._count, 1);
+ EXPECT_EQ(sink_local_state->_order_by_column_statistics._average_size, 3);
+ LockGuard lock(sink_local_state->_shared_state->buffer_mutex);
+ auto& batches = sink_local_state->_shared_state->spill_batches;
+ ASSERT_EQ(batches.size(), 2);
+ EXPECT_EQ(batches.front()->peer_group_ends.size(), 20);
+ EXPECT_EQ(batches.front()->peer_group_ends.back(), 20);
+ const auto& second_batch = batches.back();
+ ASSERT_EQ(second_batch->peer_group_ends.size(), 2);
+ EXPECT_EQ(second_batch->peer_group_ends[0], 3);
+ EXPECT_EQ(second_batch->peer_group_ends[1], 4);
+}
+
+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;
+}
+
+void expect_large_cume_dist_block(const Block& output, size_t* seen) {
+ constexpr size_t total_rows = 8002;
+ ASSERT_LE(output.rows(), 1024);
+ const auto& order = assert_cast<const
ColumnInt64&>(*output.get_by_position(0).column);
+ const auto& result = assert_cast<const
ColumnFloat64&>(*output.get_by_position(1).column);
+ for (size_t row = 0; row < output.rows(); ++row, ++(*seen)) {
+ int64_t expected_order = *seen < 5000 ? 1 : 2;
+ double expected_result = *seen < 5000 ? 5000.0 / total_rows : 8001.0 /
total_rows;
+ if (*seen == 8001) {
+ expected_order = 3;
+ expected_result = 1.0;
+ }
+ EXPECT_EQ(order.get_data()[row], expected_order);
+ EXPECT_EQ(result.get_data()[row], expected_result);
+ }
+}
+
+void expect_large_cume_dist_peer_output(AnalyticSourceOperatorX* source,
RuntimeState* state) {
+ size_t seen = 0;
+ bool eos = false;
+ while (!eos) {
+ Block output;
+ const auto status = source->get_block(state, &output, &eos);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+ if (eos) {
+ break;
+ }
+ expect_large_cume_dist_block(output, &seen);
+ }
+ EXPECT_EQ(seen, 8002);
+}
+
+} // 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);
Review Comment:
[P2] Keep the staged test block below the spill-buffer threshold
`ColumnHelper::create_column` inserts these 100,000 Int64 values into a
doubling `PaddedPODArray`, so the second Block allocates 1,048,576 bytes,
exactly the configured 1 MiB spill buffer. The first call has already created
the writer, so `_write_data_block()` takes its `allocated_bytes() >=
spill_buffer_bytes` direct-write branch for this second Block. No data or peer
buffer remains, `revocable_mem_size()` returns zero, and this assertion fails
instead of exercising the intended 512 KiB revocation case. Use a 2 MiB spill
buffer here, or a second Block whose allocation is at least 512 KiB and
strictly below 1 MiB. This is separate from the already-reported first-slice
reservation expectation at line 1832; the failure follows statically from the
allocation and write branches.
--
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]
