Merge

dbms/include/DB/AggregateFunctions/AggregateFunctionMerge.h

@@ -81,10 +81,7 @@ public:
 
 	void add(AggregateDataPtr place, const IColumn ** columns, size_t row_num) const
 	{
-		Field field;
-		columns[0]->get(row_num, field);
-		ReadBufferFromString read_buffer(field.safeGet<String &>());
-		nested_func->deserializeMerge(place, read_buffer);
+		merge(place, columns[0]->getDataAt(row_num).data);
 	}
 
 	void merge(AggregateDataPtr place, ConstAggregateDataPtr rhs) const

dbms/include/DB/Columns/ColumnAggregateFunction.h

@@ -103,18 +103,26 @@ public:
 		return StringRef(reinterpret_cast<const char *>(&data[n]), sizeof(data[n]));
 	}
 
+	/// Объединить состояние в последней строке с заданным
+	void insertMerge(const Field & x)
+	{
+		ReadBufferFromString read_buffer(x.safeGet<const String &>());
+		func->deserializeMerge(data.back(), read_buffer);
+	}
+
 	void insert(const Field & x)
 	{
-		insertDefault();
-		ReadBufferFromString read_buffer(x.safeGet<const String &>());
-		func->deserializeMerge(data[data.size()-1], read_buffer);
+		data.push_back(AggregateDataPtr());
+		func->create(data.back());
+		insertMerge(x);
 	}
 
 	void insertData(const char * pos, size_t length)
 	{
-		insertDefault();
+		data.push_back(AggregateDataPtr());
+		func->create(data.back());
 		ReadBuffer read_buffer(const_cast<char *>(pos), length);
-		func->deserializeMerge(data[data.size()-1], read_buffer);
+		func->deserializeMerge(data.back(), read_buffer);
 	}
 	
 	ColumnPtr cut(size_t start, size_t length) const

dbms/include/DB/DataStreams/AggregatingSortedBlockInputStream.h

@@ -0,0 +1,87 @@
+#pragma once
+
+#include <Yandex/logger_useful.h>
+
+#include <DB/Core/Row.h>
+#include <DB/Core/ColumnNumbers.h>
+#include <DB/DataStreams/MergingSortedBlockInputStream.h>
+#include <DB/AggregateFunctions/IAggregateFunction.h>
+#include <DB/Columns/ColumnAggregateFunction.h>
+
+
+namespace DB
+{
+
+/** Соединяет несколько сортированных потоков в один.
+  * При этом, для каждой группы идущих подряд одинаковых значений первичного ключа (столбцов, по которым сортируются данные),
+  * сливает их в одну строку. При слиянии, производится доагрегация данных - слияние состояний агрегатных функций,
+  * соответствующих одному значению первичного ключа. Для столбцов, не входящих в первичный ключ, и не имеющих тип AggregateFunction,
+  * при слиянии, выбирается первое попавшееся значение.
+  */
+class AggregatingSortedBlockInputStream : public MergingSortedBlockInputStream
+{
+public:
+	AggregatingSortedBlockInputStream(BlockInputStreams inputs_, const SortDescription & description_, size_t max_block_size_)
+		: MergingSortedBlockInputStream(inputs_, description_, max_block_size_),
+		log(&Logger::get("SummingSortedBlockInputStream"))
+	{
+	}
+
+	String getName() const { return "AggregatingSortedBlockInputStream"; }
+
+	String getID() const
+	{
+		std::stringstream res;
+		res << "AggregatingSorted(inputs";
+
+		for (size_t i = 0; i < children.size(); ++i)
+			res << ", " << children[i]->getID();
+
+		res << ", description";
+
+		for (size_t i = 0; i < description.size(); ++i)
+			res << ", " << description[i].getID();
+
+		res << ")";
+		return res.str();
+	}
+
+protected:
+	/// Может возвращаться на 1 больше записей, чем max_block_size.
+	Block readImpl();
+
+private:
+	Logger * log;
+
+	/// Столбцы с какими номерами надо аггрегировать.
+	ColumnNumbers column_numbers_to_aggregate;
+	std::vector<ColumnAggregateFunction *> column_to_aggregate;
+
+	Row current_key;		/// Текущий первичный ключ.
+	Row next_key;			/// Первичный ключ следующей строки.
+
+	Row current_row;
+
+	/** Делаем поддержку двух разных курсоров - с Collation и без.
+	 *  Шаблоны используем вместо полиморфных SortCursor'ов и вызовов виртуальных функций.
+	 */
+	template<class TSortCursor>
+	void merge(Block & merged_block, ColumnPlainPtrs & merged_columns, std::priority_queue<TSortCursor> & queue);
+
+	/// Вставить в результат первую строку для текущей группы.
+	void insertCurrentRow(ColumnPlainPtrs & merged_columns);
+
+	/** Извлечь все состояния аггрегатных функции и объединить с текущей группой.
+	  */
+	template<class TSortCursor>
+	void addRow(TSortCursor & cursor)
+	{
+		for (size_t i = 0, size = column_numbers_to_aggregate.size(); i < size; ++i)
+		{
+			size_t j = column_numbers_to_aggregate[i];
+			column_to_aggregate[i]->insertMerge((*cursor->all_columns[j])[cursor->pos]);
+		}
+	}
+};
+
+}

dbms/include/DB/DataStreams/MergingSortedBlockInputStream.h

@@ -85,7 +85,7 @@ protected:
 	QueueWithCollation queue_with_collation;
 
 
-	/// Эти методы используются в Collapsing/Summing SortedBlockInputStream-ах.
+	/// Эти методы используются в Collapsing/Summing/Aggregating SortedBlockInputStream-ах.
 
 	/// Сохранить строчку, на которую указывает cursor, в row.
 	template<class TSortCursor>

dbms/include/DB/Storages/MergeTree/MergeTreeData.h

@@ -350,6 +350,7 @@ public:
 		Ordinary,
 		Collapsing,
 		Summing,
+		Aggregating,
 	};
 
 	/** Подцепить таблицу с соответствующим именем, по соответствующему пути (с / на конце),

dbms/src/DataStreams/AggregatingSortedBlockInputStream.cpp

@@ -0,0 +1,117 @@
+#include <DB/DataStreams/AggregatingSortedBlockInputStream.h>
+
+
+namespace DB
+{
+
+
+void AggregatingSortedBlockInputStream::insertCurrentRow(ColumnPlainPtrs & merged_columns)
+{
+	for (size_t i = 0; i < num_columns; ++i)
+		merged_columns[i]->insert(current_row[i]);
+}
+
+
+Block AggregatingSortedBlockInputStream::readImpl()
+{
+	if (!children.size())
+		return Block();
+
+	if (children.size() == 1)
+		return children[0]->read();
+
+	Block merged_block;
+	ColumnPlainPtrs merged_columns;
+
+	init(merged_block, merged_columns);
+	if (merged_columns.empty())
+		return Block();
+
+	/// Дополнительная инициализация.
+	if (current_row.empty())
+	{
+		current_row.resize(num_columns);
+		current_key.resize(description.size());
+		next_key.resize(description.size());
+
+		/// Заполним номера столбцов, которые нужно доагрегировать.
+		for (size_t i = 0; i < num_columns; ++i)
+		{
+			ColumnWithNameAndType & column = merged_block.getByPosition(i);
+
+			/// Оставляем только состояния аггрегатных функций.
+			if (strncmp(column.type->getName().data(), "AggregateFunction", strlen("AggregateFunction")) != 0)
+				continue;
+
+			/// Входят ли в PK?
+			SortDescription::const_iterator it = description.begin();
+			for (; it != description.end(); ++it)
+				if (it->column_name == column.name || (it->column_name.empty() && it->column_number == i))
+					break;
+
+			if (it != description.end())
+				continue;
+
+			column_numbers_to_aggregate.push_back(i);
+			column_to_aggregate.push_back(dynamic_cast<ColumnAggregateFunction *>(merged_columns[i]));
+		}
+	}
+
+	if (has_collation)
+		merge(merged_block, merged_columns, queue_with_collation);
+	else
+		merge(merged_block, merged_columns, queue);
+
+	return merged_block;
+}
+
+
+template<class TSortCursor>
+void AggregatingSortedBlockInputStream::merge(Block & merged_block, ColumnPlainPtrs & merged_columns, std::priority_queue<TSortCursor> & queue)
+{
+	size_t merged_rows = 0;
+
+	/// Вынимаем строки в нужном порядке и кладём в merged_block, пока строк не больше max_block_size
+	while (!queue.empty())
+	{
+		TSortCursor current = queue.top();
+
+		setPrimaryKey(next_key, current);
+
+		/// если накопилось достаточно строк и последняя посчитана полностью
+		if (next_key != current_key && merged_rows >= max_block_size)
+			return;
+
+		queue.pop();
+
+		if (next_key != current_key)
+		{
+			current_key = std::move(next_key);
+			next_key.resize(description.size());
+
+			++merged_rows;
+			/// Запишем данные для очередной группы.
+			setRow(current_row, current);
+			insertCurrentRow(merged_columns);
+		}
+		else
+		{
+			addRow(current);
+		}
+
+		if (!current->isLast())
+		{
+			current->next();
+			queue.push(current);
+		}
+		else
+		{
+			/// Достаём из соответствующего источника следующий блок, если есть.
+			fetchNextBlock(current, queue);
+		}
+	}
+
+	children.clear();
+}
+
+}

dbms/src/Interpreters/Aggregator.cpp

@@ -530,8 +530,10 @@ Block Aggregator::convertToBlock(AggregatedDataVariants & data_variants, bool fi
 	try
 	{
 		for (size_t i = 0; i < aggregates_size; ++i)
-		{
 			is_final[i] = final && aggregate_functions[i]->canBeFinal();
+
+		for (size_t i = 0; i < aggregates_size; ++i)
+		{
 			if (!is_final[i])
 			{
 				/// Столбец ColumnAggregateFunction захватывает разделяемое владение ареной с состояниями агрегатных функций.

dbms/src/Storages/MergeTree/MergeTreeDataMerger.cpp

@@ -5,6 +5,7 @@
 #include <DB/DataStreams/MergingSortedBlockInputStream.h>
 #include <DB/DataStreams/CollapsingSortedBlockInputStream.h>
 #include <DB/DataStreams/SummingSortedBlockInputStream.h>
+#include <DB/DataStreams/AggregatingSortedBlockInputStream.h>
 
 
 namespace DB
@@ -291,6 +292,10 @@ MergeTreeData::DataPartPtr MergeTreeDataMerger::mergeParts(const MergeTreeData::
 			merged_stream = new SummingSortedBlockInputStream(src_streams, data.getSortDescription(), DEFAULT_MERGE_BLOCK_SIZE);
 			break;
 
+		case MergeTreeData::Aggregating:
+			merged_stream = new AggregatingSortedBlockInputStream(src_streams, data.getSortDescription(), DEFAULT_MERGE_BLOCK_SIZE);
+			break;
+
 		default:
 			throw Exception("Unknown mode of operation for MergeTreeData: " + toString(data.mode), ErrorCodes::LOGICAL_ERROR);
 	}