ClickHouse/dbms/src/Interpreters/Aggregator.cpp

#include <openssl/md5.h>

#include <DB/DataTypes/DataTypeAggregateFunction.h>
#include <DB/Columns/ColumnAggregateFunction.h>
#include <DB/Columns/ColumnString.h>
#include <DB/Columns/ColumnFixedString.h>
#include <DB/Columns/ColumnsNumber.h>

#include <DB/Interpreters/Aggregator.h>


namespace DB
{

class FieldVisitorHash : public boost::static_visitor<>
{
public:
	MD5_CTX state;
	
	FieldVisitorHash()
	{
		MD5_Init(&state);
	}

	void finalize(unsigned char * place)
	{
		MD5_Final(place, &state);
	}

	void operator() (const Null 	& x)
	{
		char type = FieldType::Null;
		MD5_Update(&state, reinterpret_cast<const char *>(&type), sizeof(type));
	}
	
	void operator() (const UInt64 	& x)
	{
		char type = FieldType::UInt64;
		MD5_Update(&state, reinterpret_cast<const char *>(&type), sizeof(type));
		MD5_Update(&state, reinterpret_cast<const char *>(&x), sizeof(x));
	}
	
 	void operator() (const Int64 	& x)
	{
		char type = FieldType::Int64;
		MD5_Update(&state, reinterpret_cast<const char *>(&type), sizeof(type));
		MD5_Update(&state, reinterpret_cast<const char *>(&x), sizeof(x));
	}

	void operator() (const Float64 	& x)
	{
		char type = FieldType::Float64;
		MD5_Update(&state, reinterpret_cast<const char *>(&type), sizeof(type));
		MD5_Update(&state, reinterpret_cast<const char *>(&x), sizeof(x));
	}

	void operator() (const String 	& x)
	{
		char type = FieldType::String;
		MD5_Update(&state, reinterpret_cast<const char *>(&type), sizeof(type));
		/// Используем ноль на конце.
		MD5_Update(&state, x.c_str(), x.size() + 1);
	}

	void operator() (const Array 	& x)
	{
		throw Exception("Cannot aggregate by array", ErrorCodes::ILLEGAL_KEY_OF_AGGREGATION);
	}

	void operator() (const SharedPtr<IAggregateFunction> & x)
	{
		throw Exception("Cannot aggregate by state of aggregate function", ErrorCodes::ILLEGAL_KEY_OF_AGGREGATION);
	}
};


/** Преобразование значения в 64 бита. Для чисел - однозначное, для строк - некриптографический хэш. */
class FieldVisitorToUInt64 : public boost::static_visitor<UInt64>
{
public:
	FieldVisitorToUInt64() {}
	
	UInt64 operator() (const Null 		& x) const { return 0; }
	UInt64 operator() (const UInt64 	& x) const { return x; }
	UInt64 operator() (const Int64 		& x) const { return x; }

	UInt64 operator() (const Float64 	& x) const
	{
		UInt64 res = 0;
		memcpy(reinterpret_cast<char *>(&res), reinterpret_cast<const char *>(&x), sizeof(x));
		return res;
	}

	UInt64 operator() (const String 	& x) const
	{
		return std::tr1::hash<String>()(x);
	}

	UInt64 operator() (const Array 	& x) const
	{
		throw Exception("Cannot aggregate by array", ErrorCodes::ILLEGAL_KEY_OF_AGGREGATION);
	}

	UInt64 operator() (const SharedPtr<IAggregateFunction> & x) const
	{
		throw Exception("Cannot aggregate by state of aggregate function", ErrorCodes::ILLEGAL_KEY_OF_AGGREGATION);
	}
};


/** Результат хранится в оперативке и должен полностью помещаться в оперативку.
  */
void Aggregator::execute(BlockInputStreamPtr stream, AggregatedDataVariants & result)
{
	size_t keys_size = keys.empty() ? key_names.size() : keys.size();
	size_t aggregates_size = aggregates.size();
	Row key(keys_size);
	Columns key_columns(keys_size);

	typedef std::vector<Columns> AggregateColumns;
	AggregateColumns aggregate_columns(aggregates_size);

	typedef std::vector<Row> Rows;
	Rows aggregate_arguments(aggregates_size);

	/// Читаем все данные
	while (Block block = stream->read())
	{
		/// Преобразуем имена столбцов в номера, если номера не заданы
		if (keys.empty() && !key_names.empty())
			for (Names::const_iterator it = key_names.begin(); it != key_names.end(); ++it)
				keys.push_back(block.getPositionByName(*it));

		for (AggregateDescriptions::iterator it = aggregates.begin(); it != aggregates.end(); ++it)
			if (it->arguments.empty() && !it->argument_names.empty())
				for (Names::const_iterator jt = it->argument_names.begin(); jt != it->argument_names.end(); ++jt)
					it->arguments.push_back(block.getPositionByName(*jt));

		for (size_t i = 0; i < aggregates_size; ++i)
		{
			aggregate_arguments[i].resize(aggregates[i].arguments.size());
			aggregate_columns[i].resize(aggregates[i].arguments.size());
		}
		
		/// Запоминаем столбцы, с которыми будем работать
		for (size_t i = 0, size = keys_size; i < size; ++i)
			key_columns[i] = block.getByPosition(keys[i]).column;

		for (size_t i = 0; i < aggregates_size; ++i)
			for (size_t j = 0; j < aggregate_columns[i].size(); ++j)
				aggregate_columns[i][j] = block.getByPosition(aggregates[i].arguments[j]).column;

		/// Создадим пример блока, описывающего результат
		if (!sample)
		{
			for (size_t i = 0, size = keys_size; i < size; ++i)
				sample.insert(block.getByPosition(keys[i]).cloneEmpty());

			for (size_t i = 0; i < aggregates_size; ++i)
			{
				ColumnWithNameAndType col;
				col.name = aggregates[i].column_name;
				col.type = new DataTypeAggregateFunction;
				col.column = new ColumnAggregateFunction;

				sample.insert(col);
			}

			/// Вставим в блок результата все столбцы-константы из исходного блока, так как они могут ещё пригодиться.
			size_t columns = block.columns();
			for (size_t i = 0; i < columns; ++i)
				if (block.getByPosition(i).column->isConst())
					sample.insert(block.getByPosition(i).cloneEmpty());
		}

		size_t rows = block.rows();

		/// Каким способом выполнять агрегацию?
		bool has_strings = false;
		for (size_t j = 0; j < keys_size; ++j)
			if (dynamic_cast<const ColumnString *>(&*key_columns[j]) || dynamic_cast<const ColumnFixedString *>(&*key_columns[j]))
				has_strings = true;

		bool keys_fit_128_bits = true;
		size_t keys_bytes = 0;
		typedef std::vector<size_t> Sizes;
		Sizes key_sizes(keys_size);
		for (size_t j = 0; j < keys_size; ++j)
		{
			if (!key_columns[j]->isNumeric())
			{
				keys_fit_128_bits = false;
				break;
			}
			key_sizes[j] = key_columns[j]->sizeOfField();
			keys_bytes += key_sizes[j];
		}
		if (keys_bytes > 16)
			keys_fit_128_bits = false;

		if (keys_size == 0)
		{
			/// Если ключей нет
			AggregatedDataWithoutKey & res = result.without_key;
			if (res.empty())
			{
				res.resize(aggregates_size);
				for (size_t i = 0; i < aggregates_size; ++i)
					res[i] = aggregates[i].function->cloneEmpty();
			}
			
			for (size_t i = 0; i < rows; ++i)
			{
				/// Добавляем значения
				for (size_t j = 0; j < aggregates_size; ++j)
				{
					for (size_t k = 0, size = aggregate_arguments[j].size(); k < size; ++k)
						aggregate_arguments[j][k] = (*aggregate_columns[j][k])[i];

					res[j]->add(aggregate_arguments[j]);
				}
			}
		}
		else if (keys_size == 1 && key_columns[0]->isNumeric()
			&& !dynamic_cast<ColumnFloat32 *>(&*key_columns[0]) && !dynamic_cast<ColumnFloat64 *>(&*key_columns[0]))
		{
			/// Если есть один ключ, который помещается в 64 бита, и это не число с плавающей запятой
			AggregatedDataWithUInt64Key & res = result.key64;
			const FieldVisitorToUInt64 visitor;
			IColumn & column = *key_columns[0];

			/// Для всех строчек
			for (size_t i = 0; i < rows; ++i)
			{
				/// Строим ключ
				Field field = column[i];
				UInt64 key = boost::apply_visitor(visitor, field);

				AggregatedDataWithUInt64Key::iterator it;
				bool inserted;
				res.emplace(key, it, inserted);
				
				if (inserted)
				{
					new(&it->second) AggregateFunctions(aggregates_size);
					
					for (size_t j = 0; j < aggregates_size; ++j)
						it->second[j] = aggregates[j].function->cloneEmpty();
				}

				/// Добавляем значения
				for (size_t j = 0; j < aggregates_size; ++j)
				{
					for (size_t k = 0, size = aggregate_arguments[j].size(); k < size; ++k)
						aggregate_arguments[j][k] = (*aggregate_columns[j][k])[i];

					it->second[j]->add(aggregate_arguments[j]);
				}
			}
		}
		else if (keys_size == 1
			&& (dynamic_cast<ColumnString *>(&*key_columns[0]) || dynamic_cast<ColumnFixedString *>(&*key_columns[0])))
		{
			/// Если есть один строковый ключ, то используем хэш-таблицу с ним
			AggregatedDataWithStringKey & res = result.key_string;
			IColumn & column = *key_columns[0];

			/// Для всех строчек
			for (size_t i = 0; i < rows; ++i)
			{
				/// Строим ключ
				String key = boost::get<String>(column[i]);

				AggregatedDataWithStringKey::iterator it = res.find(key);
				if (it == res.end())
				{
					it = res.insert(std::make_pair(key, AggregateFunctions(aggregates_size))).first;

					for (size_t j = 0; j < aggregates_size; ++j)
						it->second[j] = aggregates[j].function->cloneEmpty();
				}

				/// Добавляем значения
				for (size_t j = 0; j < aggregates_size; ++j)
				{
					for (size_t k = 0, size = aggregate_arguments[j].size(); k < size; ++k)
						aggregate_arguments[j][k] = (*aggregate_columns[j][k])[i];

					it->second[j]->add(aggregate_arguments[j]);
				}
			}
		}
		else
		{
			/// Если много ключей - будем агрегировать по хэшу от них
			AggregatedDataHashed & res = result.hashed;
			const FieldVisitorToUInt64 to_uint64_visitor;

			/// Для всех строчек
			for (size_t i = 0; i < rows; ++i)
			{
				/// Строим ключ
				union
				{
					UInt128 key_hash;
					unsigned char bytes[16];
				} key_hash_union;

				/// Если все ключи числовые и помещаются в 128 бит
				if (keys_fit_128_bits)
				{
					memset(key_hash_union.bytes, 0, 16);
					size_t offset = 0;
					for (size_t j = 0; j < keys_size; ++j)
					{
						key[j] = (*key_columns[j])[i];
						UInt64 tmp = boost::apply_visitor(to_uint64_visitor, key[j]);
						/// Работает только на little endian
						memcpy(key_hash_union.bytes + offset, reinterpret_cast<const char *>(&tmp), key_sizes[j]);
						offset += key_sizes[j];
					}
				}
				else	/// Иначе используем md5.
				{
					FieldVisitorHash key_hash_visitor;

					for (size_t j = 0; j < keys_size; ++j)
					{
						key[j] = (*key_columns[j])[i];
						boost::apply_visitor(key_hash_visitor, key[j]);
					}

					key_hash_visitor.finalize(key_hash_union.bytes);
				}

				AggregatedDataHashed::iterator it;
				bool inserted;
				res.emplace(key_hash_union.key_hash, it, inserted);

				if (inserted)
				{
					new(&it->second) AggregatedDataHashed::mapped_type(key, AggregateFunctions(aggregates_size));

					for (size_t j = 0; j < aggregates_size; ++j)
						it->second.second[j] = aggregates[j].function->cloneEmpty();
				}

				/// Добавляем значения
				for (size_t j = 0; j < aggregates_size; ++j)
				{
					for (size_t k = 0, size = aggregate_arguments[j].size(); k < size; ++k)
						aggregate_arguments[j][k] = (*aggregate_columns[j][k])[i];

					it->second.second[j]->add(aggregate_arguments[j]);
				}
			}
		}
/*		else
		{
			/// Общий способ
			AggregatedData & res = result.generic;
			
			/// Для всех строчек
			for (size_t i = 0; i < rows; ++i)
			{
				/// Строим ключ
				for (size_t j = 0; j < keys_size; ++j)
					key[j] = (*key_columns[j])[i];

				AggregatedData::iterator it = res.find(key);
				if (it == res.end())
				{
					it = res.insert(std::make_pair(key, AggregateFunctions(aggregates_size))).first;

					for (size_t j = 0; j < aggregates_size; ++j)
						it->second[j] = aggregates[j].function->cloneEmpty();
				}

				/// Добавляем значения
				for (size_t j = 0; j < aggregates_size; ++j)
				{
					for (size_t k = 0, size = aggregate_arguments[j].size(); k < size; ++k)
						aggregate_arguments[j][k] = (*aggregate_columns[j][k])[i];

					it->second[j]->add(aggregate_arguments[j]);
				}
			}
		}*/
	}
}


}