remove the need to rewrite expressions

Author: Tishj

src/deletes/equality_delete.cpp

@@ -61,26 +61,28 @@ void IcebergMultiFileList::ScanEqualityDeleteFile(const IcebergManifestEntry &en
 		id_to_global_column[col.identifier.GetValue<int32_t>()] = i;
 	}
 
-	auto new_column_indexes = column_indexes;
-	for (auto field_id : entry.equality_ids) {
-		auto global_column_id = id_to_global_column[field_id];
-		ColumnIndex equality_index(global_column_id);
-		//! Check if the column needed by the equality delete is present
-		if (std::find(column_indexes.begin(), column_indexes.end(), equality_index) == column_indexes.end()) {
-			//! Column isn't being selected, add the column so it can be used for the equality delete
-			new_column_indexes.push_back(equality_index);
-		}
-	}
-
 	unordered_map<idx_t, idx_t> global_id_to_result_id;
-	for (idx_t i = 0; i < new_column_indexes.size(); i++) {
-		auto &column_index = new_column_indexes[i];
+	for (idx_t i = 0; i < column_indexes.size(); i++) {
+		auto &column_index = column_indexes[i];
 		if (column_index.IsVirtualColumn()) {
 			continue;
 		}
 		auto global_id = column_index.GetPrimaryIndex();
 		global_id_to_result_id[global_id] = i;
 	}
+	//! For the column(s) that are needed but aren't referenced, add them to the map
+	for (auto field_id : entry.equality_ids) {
+		auto global_column_id = id_to_global_column[field_id];
+		ColumnIndex equality_index(global_column_id);
+		//! Check if the column needed by the equality delete is present
+		if (std::find(column_indexes.begin(), column_indexes.end(), equality_index) != column_indexes.end()) {
+			continue;
+		}
+		auto new_result_id = column_indexes.size() + equality_id_to_result_id.size();
+		//! Create or get the result id mapping for this equality id
+		auto result_id = equality_id_to_result_id.emplace(field_id, new_result_id).first->second;
+		global_id_to_result_id[global_column_id] = result_id;
+	}
 
 	//! Take only the relevant columns from the result
 	InitializeFromOtherChunk(result, result_p, column_ids);
@@ -97,13 +99,13 @@ void IcebergMultiFileList::ScanEqualityDeleteFile(const IcebergManifestEntry &en
 
 		auto it = global_id_to_result_id.find(global_column_id);
 		D_ASSERT(it != global_id_to_result_id.end());
-		global_column_id = it->second;
+		auto result_column_id = it->second;
 
 		for (idx_t i = 0; i < count; i++) {
 			auto &row = rows[i];
 			auto constant = vec.GetValue(i);
 			unique_ptr<Expression> equality_filter;
-			auto bound_ref = make_uniq<BoundReferenceExpression>(col.type, global_column_id);
+			auto bound_ref = make_uniq<BoundReferenceExpression>(col.type, result_column_id);
 			if (!constant.IsNull()) {
 				//! Create a COMPARE_NOT_EQUAL expression
 				equality_filter =

src/iceberg_functions/iceberg_multi_file_reader.cpp

@@ -286,12 +286,18 @@ ReaderInitializeType IcebergMultiFileReader::InitializeReader(MultiFileReaderDat
 
 	//! Add the columns needed by the equality deletes if not present
 	auto new_global_column_ids = global_column_ids;
+	auto &equality_to_result_id = multi_file_list.equality_id_to_result_id;
+	new_global_column_ids.resize(global_column_ids.size() + equality_to_result_id.size());
 	for (auto field_id : equality_delete_ids) {
+		auto it = equality_to_result_id.find(field_id);
+		if (it == equality_to_result_id.end()) {
+			//! Already selected, no need to add
+			continue;
+		}
 		auto global_column_id = id_to_global_column[field_id];
 		ColumnIndex equality_index(global_column_id);
-		if (std::find(global_column_ids.begin(), global_column_ids.end(), equality_index) == global_column_ids.end()) {
-			new_global_column_ids.push_back(equality_index);
-		}
+		//! FIXME: is this correct?
+		new_global_column_ids[it->second] = equality_index;
 	}
 
 	return CreateMapping(context, reader_data, global_columns, new_global_column_ids, table_filters, gstate.file_list,
@@ -339,8 +345,7 @@ void IcebergMultiFileReader::FinalizeBind(MultiFileReaderData &reader_data, cons
 void IcebergMultiFileReader::ApplyEqualityDeletes(ClientContext &context, DataChunk &output_chunk,
                                                   const IcebergMultiFileList &multi_file_list,
                                                   const IcebergManifestEntry &data_file,
-                                                  const vector<MultiFileColumnDefinition> &local_columns,
-                                                  const unordered_map<idx_t, idx_t> &field_id_to_result_id) {
+                                                  const vector<MultiFileColumnDefinition> &local_columns) {
 	auto delete_rows = multi_file_list.GetEqualityDeletesForFile(data_file);
 
 	if (delete_rows.empty()) {
@@ -384,25 +389,7 @@ void IcebergMultiFileReader::ApplyEqualityDeletes(ClientContext &context, DataCh
 				}
 				continue;
 			}
-			if (field_id_to_result_id.empty()) {
-				equalities.push_back(expression->Copy());
-				continue;
-			}
-			idx_t index = field_id_to_result_id.at(field_id);
-			if (expression->type == ExpressionType::COMPARE_NOTEQUAL) {
-				auto &expr = expression->Cast<BoundComparisonExpression>();
-				auto bound_ref = make_uniq<BoundReferenceExpression>(expr.left->return_type, index);
-				unique_ptr<Expression> equality_filter = make_uniq<BoundComparisonExpression>(
-				    ExpressionType::COMPARE_NOTEQUAL, std::move(bound_ref), expr.right->Copy());
-				equalities.push_back(std::move(equality_filter));
-			} else if (expression->type == ExpressionType::OPERATOR_IS_NOT_NULL) {
-				auto &expr = expression->Cast<BoundOperatorExpression>();
-				auto bound_ref = make_uniq<BoundReferenceExpression>(expr.children[0]->return_type, index);
-				auto is_not_null =
-				    make_uniq<BoundOperatorExpression>(ExpressionType::OPERATOR_IS_NOT_NULL, LogicalType::BOOLEAN);
-				is_not_null->children.push_back(std::move(bound_ref));
-				equalities.push_back(std::move(is_not_null));
-			}
+			equalities.push_back(expression->Copy());
 		}
 
 		unique_ptr<Expression> filter;
@@ -440,43 +427,29 @@ void IcebergMultiFileReader::FinalizeChunk(ClientContext &context, const MultiFi
                                            DataChunk &input_chunk, DataChunk &output_chunk,
                                            ExpressionExecutor &executor,
                                            optional_ptr<MultiFileReaderGlobalState> global_state) {
+	D_ASSERT(global_state);
+	// Get the metadata for this file
+	const auto &multi_file_list = global_state->file_list->Cast<IcebergMultiFileList>();
 
 	//! Add the extra equality delete fields to output chunk.
-	int32_t diff = 0;
-	if (executor.expressions.size() != output_chunk.ColumnCount()) {
-		diff = executor.expressions.size() - output_chunk.ColumnCount();
-		for (int32_t i = diff; i > 0; i--) {
-			int32_t index = input_chunk.ColumnCount() - i;
-			output_chunk.data.emplace_back(input_chunk.data[index]);
+	idx_t diff = executor.expressions.size() - output_chunk.ColumnCount();
+	(void)diff;
+	D_ASSERT(diff == multi_file_list.equality_id_to_result_id.size());
+	if (diff > 0) {
+		int32_t start = input_chunk.ColumnCount() - diff;
+		for (int32_t i = 0; i < diff; i++) {
+			output_chunk.data.emplace_back(input_chunk.data[start + i]);
 		}
 	}
 
 	//! Base class finalization first
 	MultiFileReader::FinalizeChunk(context, bind_data, reader, reader_data, input_chunk, output_chunk, executor,
 	                               global_state);
 
-	//! Map from index into local_columns -> field_id
 	auto &local_columns = reader.columns;
-	unordered_map<idx_t, idx_t> column_index_to_field_id;
-	for (idx_t i = 0; i < local_columns.size(); i++) {
-		auto &col = local_columns[i];
-		column_index_to_field_id[i] = col.identifier.GetValue<int32_t>();
-	}
-
-	//! Map from field_id -> index in 'output_chunk'
-	unordered_map<idx_t, idx_t> field_id_to_result_id;
-	auto &column_indexes = reader.column_indexes;
-	auto result_id = executor.expressions.size() - column_indexes.size();
-	for (auto &column_index : column_indexes) {
-		field_id_to_result_id[column_index_to_field_id[column_index.GetPrimaryIndex()]] = result_id++;
-	}
-
-	D_ASSERT(global_state);
-	// Get the metadata for this file
-	const auto &multi_file_list = dynamic_cast<const IcebergMultiFileList &>(*global_state->file_list);
 	auto file_id = reader.file_list_idx.GetIndex();
 	auto &data_file = multi_file_list.data_files[file_id];
-	ApplyEqualityDeletes(context, output_chunk, multi_file_list, data_file, local_columns, field_id_to_result_id);
+	ApplyEqualityDeletes(context, output_chunk, multi_file_list, data_file, local_columns);
 
 	//! Remove the extra columns we added to perform the equality delete filtering
 	for (idx_t i = 0; i < diff; i++) {

src/include/iceberg_multi_file_list.hpp

@@ -119,6 +119,8 @@ struct IcebergMultiFileList : public MultiFileList {
 	//! All equality deletes with sequence numbers higher than that of the data_file apply to that data_file
 	mutable map<sequence_number_t, unique_ptr<IcebergEqualityDeleteData>> equality_delete_data;
 	mutable mutex delete_lock;
+	//! The columns needed by the equality deletes that aren't referenced by the scan
+	mutable unordered_map<int32_t, column_t> equality_id_to_result_id;
 
 	bool initialized = false;
 	const IcebergOptions &options;

src/include/iceberg_multi_file_reader.hpp

@@ -64,8 +64,7 @@ struct IcebergMultiFileReader : public MultiFileReader {
 	                   ExpressionExecutor &executor, optional_ptr<MultiFileReaderGlobalState> global_state) override;
 	void ApplyEqualityDeletes(ClientContext &context, DataChunk &output_chunk,
 	                          const IcebergMultiFileList &multi_file_list, const IcebergManifestEntry &data_file,
-	                          const vector<MultiFileColumnDefinition> &local_columns,
-	                          const unordered_map<idx_t, idx_t> &field_id_to_result_id);
+	                          const vector<MultiFileColumnDefinition> &local_columns);
 	bool ParseOption(const string &key, const Value &val, MultiFileOptions &options, ClientContext &context) override;
 
 public: