Merge

dbms/include/DB/DataStreams/SummingSortedBlockInputStream.h

@@ -0,0 +1,99 @@
+#pragma once
+
+#include <Yandex/logger_useful.h>
+
+#include <DB/Core/Row.h>
+#include <DB/Core/ColumnNumbers.h>
+#include <DB/DataStreams/MergingSortedBlockInputStream.h>
+
+
+namespace DB
+{
+
+/** Соединяет несколько сортированных потоков в один.
+  * При этом, для каждой группы идущих подряд одинаковых значений первичного ключа (столбцов, по которым сортируются данные),
+  *  схлопывает их в одну строку, суммируя все числовые столбцы кроме первичного ключа.
+  */
+class SummingSortedBlockInputStream : public MergingSortedBlockInputStream
+{
+public:
+	SummingSortedBlockInputStream(BlockInputStreams inputs_, SortDescription & description_, size_t max_block_size_)
+		: MergingSortedBlockInputStream(inputs_, description_, max_block_size_),
+		log(&Logger::get("SummingSortedBlockInputStream"))
+	{
+	}
+
+	String getName() const { return "SummingSortedBlockInputStream"; }
+
+	String getID() const
+	{
+		std::stringstream res;
+		res << "SummingSorted(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_sum;
+
+	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);
+
+
+	/** Реализует операцию += */
+	class FieldVisitorSum : public StaticVisitor<>
+	{
+	private:
+		const Field & rhs;
+	public:
+		FieldVisitorSum(const Field & rhs_) : rhs(rhs_) {}
+
+		void operator() (UInt64 	& x) const { x += get<UInt64>(rhs); }
+		void operator() (Int64 		& x) const { x += get<Int64>(rhs); }
+		void operator() (Float64 	& x) const { x += get<Float64>(rhs); }
+
+		void operator() (Null 		& x) const { throw Exception("Cannot sum Nulls", ErrorCodes::LOGICAL_ERROR); }
+		void operator() (String 	& x) const { throw Exception("Cannot sum Strings", ErrorCodes::LOGICAL_ERROR); }
+		void operator() (Array 		& x) const { throw Exception("Cannot sum Arrays", ErrorCodes::LOGICAL_ERROR); }
+	};
+
+	/// Прибавить строчку под курсором к row.
+	template<class TSortCursor>
+	void addRow(Row & row, TSortCursor & cursor)
+	{
+		for (size_t i = 0, size = column_numbers_to_sum.size(); i < size; ++i)
+		{
+			size_t j = column_numbers_to_sum[i];
+			apply_visitor(FieldVisitorSum((*cursor->all_columns[j])[cursor->pos]), row[j]);
+		}
+	}
+};
+
+}

dbms/src/DataStreams/SummingSortedBlockInputStream.cpp

@@ -0,0 +1,118 @@
+#include <DB/DataStreams/SummingSortedBlockInputStream.h>
+
+
+namespace DB
+{
+
+
+void SummingSortedBlockInputStream::insertCurrentRow(ColumnPlainPtrs & merged_columns)
+{
+	for (size_t i = 0; i < num_columns; ++i)
+		merged_columns[i]->insert(current_row[i]);
+}
+	
+
+Block SummingSortedBlockInputStream::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 (!column.type->isNumeric() || column.type->getName() == "Date" || column.type->getName() == "DateTime")
+				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_sum.push_back(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 SummingSortedBlockInputStream::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();
+		queue.pop();
+
+		setPrimaryKey(next_key, current);
+
+		if (next_key != current_key)
+		{
+			/// Запишем данные для предыдущей группы.
+			++merged_rows;
+			insertCurrentRow(merged_columns);
+
+			current_key = next_key;
+			next_key.resize(description.size());
+
+			setRow(current_row, current);
+		}
+		else
+		{
+			addRow(current_row, current);
+		}
+
+		if (!current->isLast())
+		{
+			current->next();
+			queue.push(current);
+		}
+		else
+		{
+			/// Достаём из соответствующего источника следующий блок, если есть.
+			fetchNextBlock(current, queue);
+		}
+
+		if (merged_rows >= max_block_size)
+			return;;
+	}
+
+	/// Запишем данные для последней группы.
+	++merged_rows;
+	insertCurrentRow(merged_columns);
+
+	children.clear();
+}
+
+}