github-actions[bot] commented on code in PR #68500: URL: https://github.com/apache/doris/pull/68500#discussion_r4348511961
########## be/src/exec/operator/analytic_spill.h: ########## @@ -0,0 +1,129 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#pragma once + +#include <cstddef> +#include <cstdint> +#include <memory> +#include <vector> + +#include "common/status.h" +#include "core/block/block.h" +#include "core/custom_allocator.h" +#include "exprs/aggregate/aggregate_function.h" + +namespace doris { + +class RuntimeProfile; +class RuntimeState; +class SpillFile; +class SpillFileWriter; + +using SpillFileSPtr = std::shared_ptr<SpillFile>; +using SpillFileWriterSPtr = std::shared_ptr<SpillFileWriter>; + +/// An immutable run of consecutive analytic partitions handed from the sink to the source. Rows +/// keep the input order. partition_ends[p] is the exclusive end row of partition p inside the batch +/// and peer_group_ends are batch-relative. Small partitions are batched together so that the source +/// replays ordinary sized Blocks even when every partition contains only a few rows. +struct AnalyticSpillBatch { + struct RecordMemory { + size_t serialized_bytes = 0; + size_t decoded_bytes = 0; + size_t footer_bytes = 0; + }; + + int64_t rows = 0; + std::vector<Block> blocks; + SpillFileSPtr data_file; + RecordMemory data_record_memory; + + DorisVector<int64_t> peer_group_ends; + SpillFileSPtr peer_group_file; + RecordMemory peer_record_memory; + + std::vector<int64_t> partition_ends; + std::vector<WindowSpillStrategy> strategies; + std::vector<WindowSpillPeerFunction> peer_functions; + // Row p holds the result of partition p for PARTITION_REDUCE functions. + std::vector<ColumnPtr> partition_results; + // Element p holds the parameter of partition p for PARTITION_CARDINALITY functions. + std::vector<std::vector<int64_t>> function_parameters; Review Comment: [P2] Track the retained NTILE parameter buffers Each finished partition appends an int64 per NTILE function to `_batch_function_parameters`, then moves those standard vectors here. With 65,535 distinct partition keys and 32 distinct NTILE literals, one batch retains about 16 MiB that never reaches MemTracker: these vectors use the standard allocator, neither side accounts their capacity, and the temporary admission reservation is released after the sink call. Consequently query/workload-group accounting misses that memory throughout replay. Use allocator-aware inner vectors in both the producer and descriptor, or explicitly transfer balanced capacity accounting with ownership. This remains an issue after keeping constant arguments compact. ########## be/src/exec/operator/analytic_sink_operator.cpp: ########## @@ -905,9 +1389,102 @@ void AnalyticSinkLocalState::_remove_unused_rows() { DCHECK_GE(_partition_by_pose.end, 0); } +bool AnalyticSinkOperatorX::is_blockable(RuntimeState* state) const { + // The scheduler can classify the first task before local-state open initializes its flag. + return _enable_spill_analytic || DataSinkOperatorX<AnalyticSinkLocalState>::is_blockable(state); +} + size_t AnalyticSinkOperatorX::get_reserve_mem_size(RuntimeState* state, bool eos) { + return get_reserve_mem_size(state, eos, nullptr); +} + +size_t AnalyticSinkOperatorX::get_reserve_mem_size(RuntimeState* state, bool eos, + const Block* block) { auto& local_state = get_local_state(state); - return local_state._reserve_mem_size; + if (!local_state._spill_enabled) { + return local_state._reserve_mem_size; + } + const size_t incoming_bytes = block ? block->allocated_bytes() : 0; + const size_t incoming_rows = block ? block->rows() : 0; + const size_t staged_bytes = + local_state._batch_store ? local_state._batch_store->pending_write_bytes() : 0; + const size_t pending_data_bytes = + local_state._batch_store ? local_state._batch_store->pending_data_bytes() : 0; + const size_t pending_peer_bytes = + local_state._batch_store ? local_state._batch_store->pending_peer_bytes() : 0; + const size_t spill_limit = static_cast<size_t>(state->spill_analytic_sink_mem_limit_bytes()); + const bool data_write = + (local_state._batch_store && local_state._batch_store->is_spilled() && + (incoming_rows > 0 || pending_data_bytes > 0)) || + ((incoming_rows > 0 || eos) && staged_bytes + incoming_bytes >= spill_limit); + const bool first_forced_write = incoming_rows > 0 && state->enable_force_spill(); + // An input Block may close the preceding peer group at its first row; empty EOS closes + // the final group and can flush a sidecar without writing a data Block. + const size_t possible_peer_ends = + incoming_rows + (local_state._open_partition_rows > 0 ? 1 : 0); + const bool peer_write = local_state._has_peer_group_functions && + ((pending_peer_bytes > 0 && local_state._batch_store && + local_state._batch_store->has_spilled_peer_groups()) || + pending_peer_bytes + possible_peer_ends * sizeof(int64_t) >= + static_cast<size_t>(state->spill_buffer_size_bytes())); + if (eos && incoming_rows == 0 && !data_write && !peer_write) { + // An empty EOS merely publishes an in-memory batch; the previous write peak is stale. + return 0; + } + // In the worst case every input row ends a partition. Its boundary and each reduction + // result remain in the batch while the input Block is collected, even without a write. + size_t partition_result_bytes = incoming_rows * (sizeof(size_t) + sizeof(int64_t)); + for (size_t i = 0; i < local_state._spill_strategies.size(); ++i) { + switch (local_state._spill_strategies[i]) { + case WindowSpillStrategy::PARTITION_REDUCE: { + const auto& type = local_state._agg_functions[i]->data_type(); + partition_result_bytes += type->have_maximum_size_of_value() + ? incoming_rows * type->get_size_of_value_in_memory() + : incoming_bytes; + if (type->is_nullable()) { + partition_result_bytes += incoming_rows; + } + break; + } + case WindowSpillStrategy::PARTITION_CARDINALITY: + partition_result_bytes += incoming_rows * sizeof(int64_t); + break; + case WindowSpillStrategy::PEER_GROUP: + partition_result_bytes += incoming_rows * sizeof(int64_t); + break; + case WindowSpillStrategy::UNSUPPORTED: + DORIS_CHECK(false); + } + } + if (data_write || first_forced_write || peer_write) { + // The writer can hold a PBlock and serialized/compressed bytes at the same time. + // A wide input row is not splittable, so do not cap this at spill_buffer_size_bytes. + // Small records only need a fixed metadata allowance, not a full spill buffer. + return 3 * std::max({incoming_bytes, staged_bytes, size_t {64 * 1024}}) + Review Comment: [P2] Reserve spill scratch for the combined coalesced record After a data writer exists, `_write_data_block()` merges a sub-buffer input into the staged Block before checking whether to flush. Two roughly 4 MiB allocations can therefore produce one roughly 8 MiB record with an 8 MiB spill buffer, while this expression admits only roughly 12 MiB. For poorly compressible payloads, `Block::serialize()` simultaneously holds the combined serialization and compression buffers, exceeding that process-memory allowance even before accounting for merged-column growth. This concerns process headroom, including ordinary string allocations; it does not require exceeding the query MemTracker. Bound the pending record from staged plus incoming data and include merge growth when reserving a later coalesced write. ########## 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: [P2] Include first-slice copies in the reservation expectation This test publishes a 199,999-row in-memory Block with two Int64 columns and uses `batch_size=4096`. `_next_output_reserve_bytes()` now includes its first cut: `4096 * 2 columns * 2 * sizeof(int64_t) = 131072` bytes, in addition to the 32768 bytes for the SUM result. This expectation still includes only the result, so `SpillPathKeepsSealedBatchBelowSinkLimitInMemory` deterministically fails by 128 KiB. Update the expected reservation to include the first input slice while retaining the separate assertion that no spill-reader memory is reserved. -- 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]
