dbms: limits for transfering external tables, only_external mode in set [METR-11023]

dbms/include/DB/Interpreters/Limits.h

@@ -69,6 +69,11 @@ struct Limits
 	M(SettingUInt64, max_bytes_in_set, 0) \
 	M(SettingOverflowMode<false>, set_overflow_mode, OverflowMode::THROW) \
 	\
+	/** Ограничения для максимального размера передаваемой внешней таблицы, получающейся при выполнении секции Global IN. */ \
+	M(SettingUInt64, max_rows_to_transfer, 0) \
+	M(SettingUInt64, max_bytes_to_transfer, 0) \
+	M(SettingOverflowMode<false>, transfer_overflow_mode, OverflowMode::THROW) \
+	\
 	/** Ограничения для максимального размера запоминаемого состояния при выполнении DISTINCT. */ \
 	M(SettingUInt64, max_rows_in_distinct, 0) \
 	M(SettingUInt64, max_bytes_in_distinct, 0) \

dbms/include/DB/Interpreters/Set.h

@@ -33,7 +33,12 @@ class Set
 {
 public:
 	Set(const Limits & limits)
-		: log(&Logger::get("Set")),
+		: max_bytes_to_transfer(limits.max_bytes_to_transfer),
+		max_rows_to_transfer(limits.max_rows_to_transfer),
+		transfer_overflow_mode(limits.transfer_overflow_mode),
+		bytes_in_external_table(0),
+		rows_in_external_table(0),
+		log(&Logger::get("Set")),
 		max_rows(limits.max_rows_in_set),
 		max_bytes(limits.max_bytes_in_set),
 		overflow_mode(limits.set_overflow_mode)
@@ -54,6 +59,7 @@ public:
 	void setSource(BlockInputStreamPtr stream) { source = stream; }
 
 	void setExternalOutput(StoragePtr storage) { external_table = storage; }
+	void setOnlyExternal(bool flag) { only_external = flag; }
 
 	BlockInputStreamPtr getSource() { return source; }
 
@@ -98,7 +104,14 @@ private:
 
 	BlockInputStreamPtr source;
 
+	/// Информация о внешней таблице, заполняемой этим классом
 	StoragePtr external_table;
+	size_t max_bytes_to_transfer;
+	size_t max_rows_to_transfer;
+	OverflowMode transfer_overflow_mode;
+	size_t bytes_in_external_table;
+	size_t rows_in_external_table;
+	bool only_external;
 
 	/// Специализация для случая, когда есть один числовой ключ.
 	std::unique_ptr<SetUInt64> key64;
@@ -165,6 +178,8 @@ private:
 	
 	/// Проверить не превышены ли допустимые размеры множества ключей
 	bool checkSetSizeLimits() const;
+	/// Проверить не превышены ли допустимые размеры внешней таблицы для передачи данных
+	bool checkExternalSizeLimits() const;
 
 	/// Считает суммарное число ключей во всех Set'ах
 	size_t getTotalRowCount() const;

dbms/src/Interpreters/ExpressionAnalyzer.cpp

@@ -661,6 +661,7 @@ void ExpressionAnalyzer::addExternalStorage(ASTFunction * node)
 		ast_set->set = new Set(settings.limits);
 		ast_set->set->setSource(external_data[external_table_name]);
 		ast_set->set->setExternalOutput(external_tables[external_table_name]);
+		ast_set->set->setOnlyExternal();
 		sets_with_subqueries[ast_set->getColumnName()] = ast_set->set;
 	}
 	else

dbms/src/Interpreters/Set.cpp

@@ -58,6 +58,16 @@ bool Set::checkSetSizeLimits() const
 }
 	
 
+bool Set::checkExternalSizeLimits() const
+{
+	if (max_rows_to_transfer && rows_in_external_table > max_rows_to_transfer)
+		return false;
+	if (max_bytes_to_transfer && bytes_in_external_table > max_bytes_to_transfer)
+		return false;
+	return true;
+}
+
+
 Set::Type Set::chooseMethod(const ConstColumnPlainPtrs & key_columns, bool & keys_fit_128_bits, Sizes & key_sizes)
 {
 	size_t keys_size = key_columns.size();
@@ -99,7 +109,27 @@ bool Set::insertFromBlock(Block & block, bool create_ordered_set)
 	{
 		BlockOutputStreamPtr output = external_table->write(ASTPtr());
 		output->write(block);
+		bytes_in_external_table += block.bytes();
+		rows_in_external_table += block.rows();
+
+		if (!checkExternalSizeLimits())
+		{
+			if (transfer_overflow_mode == OverflowMode::THROW)
+				throw Exception("IN-external_table size exceeded."
+					" Rows: " + toString(rows_in_external_table) +
+					", limit: " + toString(max_rows_to_transfer) +
+					". Bytes: " + toString(bytes_in_external_table) +
+					", limit: " + toString(max_bytes_to_transfer) + ".",
+					ErrorCodes::SET_SIZE_LIMIT_EXCEEDED);
+
+			if (transfer_overflow_mode == OverflowMode::BREAK)
+				return false;
+
+			throw Exception("Logical error: unknown overflow mode", ErrorCodes::LOGICAL_ERROR);
 		}
+	}
+
+	if (only_external) return true;
 
 	size_t keys_size = block.columns();
 	Row key(keys_size);