fix query without equality delete fields

张军 · 张军 · commit cb13c01a84f0 · 2025-09-17T16:23:43.000+08:00
diff --git a/src/deletes/equality_delete.cpp b/src/deletes/equality_delete.cpp
@@ -61,9 +61,18 @@ void IcebergMultiFileList::ScanEqualityDeleteFile(const IcebergManifestEntry &en
 		id_to_global_column[col.identifier.GetValue<int32_t>()] = i;
 	}
 
+	std::vector<ColumnIndex> 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);
+		if (std::find(column_indexes.begin(), column_indexes.end(), equality_index) == column_indexes.end()) {
+			new_column_indexes.push_back(equality_index);
+		}
+	}
+
 	unordered_map<idx_t, idx_t> global_id_to_result_id;
-	for (idx_t i = 0; i < column_indexes.size(); i++) {
-		auto &column_index = column_indexes[i];
+	for (idx_t i = 0; i < new_column_indexes.size(); i++) {
+		auto &column_index = new_column_indexes[i];
 		if (column_index.IsVirtualColumn()) {
 			continue;
 		}
@@ -85,9 +94,6 @@ void IcebergMultiFileList::ScanEqualityDeleteFile(const IcebergManifestEntry &en
 		auto &vec = result.data[col_idx];
 
 		auto it = global_id_to_result_id.find(global_column_id);
-		if (it == global_id_to_result_id.end()) {
-			throw NotImplementedException("Equality deletes need the relevant columns to be selected");
-		}
 		global_column_id = it->second;
 
 		for (idx_t i = 0; i < count; i++) {
diff --git a/src/iceberg_functions/iceberg_multi_file_reader.cpp b/src/iceberg_functions/iceberg_multi_file_reader.cpp
@@ -250,6 +250,56 @@ static void ApplyPartitionConstants(const IcebergMultiFileList &multi_file_list,
 	}
 }
 
+ReaderInitializeType IcebergMultiFileReader::InitializeReader(MultiFileReaderData &reader_data,
+                                                              const MultiFileBindData &bind_data,
+                                                              const vector<MultiFileColumnDefinition> &global_columns,
+                                                              const vector<ColumnIndex> &global_column_ids,
+                                                              optional_ptr<TableFilterSet> table_filters,
+                                                              ClientContext &context, MultiFileGlobalState &gstate) {
+
+	FinalizeBind(reader_data, bind_data.file_options, bind_data.reader_bind, global_columns, global_column_ids, context,
+	             gstate.multi_file_reader_state.get());
+
+	unordered_map<int32_t, column_t> id_to_global_column;
+	for (column_t i = 0; i < global_columns.size(); i++) {
+		auto &col = global_columns[i];
+		D_ASSERT(!col.identifier.IsNull());
+		id_to_global_column[col.identifier.GetValue<int32_t>()] = i;
+	}
+
+	set<int32_t> equality_delete_ids;
+	const auto &multi_file_list = dynamic_cast<const IcebergMultiFileList &>(gstate.file_list);
+	auto &reader = *reader_data.reader;
+	auto file_id = reader.file_list_idx.GetIndex();
+	auto &data_file = multi_file_list.data_files[file_id];
+
+	auto delete_data_it = multi_file_list.equality_delete_data.upper_bound(data_file.sequence_number);
+	for (; delete_data_it != multi_file_list.equality_delete_data.end(); delete_data_it++) {
+		auto &files = delete_data_it->second->files;
+		for (auto &file : files) {
+			auto &rows = file.rows;
+			for (auto &row : rows) {
+				auto &filters = row.filters;
+				for (auto &filter : filters) {
+					equality_delete_ids.insert(filter.first);
+				}
+			}
+		}
+	}
+
+	vector<ColumnIndex> new_global_column_ids = global_column_ids;
+	for (auto field_id : equality_delete_ids) {
+		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);
+		}
+	}
+
+	return CreateMapping(context, reader_data, global_columns, new_global_column_ids, table_filters, gstate.file_list,
+	                     bind_data.reader_bind, bind_data.virtual_columns);
+}
+
 void IcebergMultiFileReader::FinalizeBind(MultiFileReaderData &reader_data, const MultiFileOptions &file_options,
                                           const MultiFileReaderBindData &options,
                                           const vector<MultiFileColumnDefinition> &global_columns,
@@ -291,7 +341,8 @@ 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 vector<MultiFileColumnDefinition> &local_columns,
+                                                  unordered_map<idx_t, idx_t> field_id_to_result_id) {
 	vector<reference<IcebergEqualityDeleteRow>> delete_rows;
 
 	auto &metadata = multi_file_list.GetMetadata();
@@ -360,7 +411,25 @@ void IcebergMultiFileReader::ApplyEqualityDeletes(ClientContext &context, DataCh
 					equalities.push_back(make_uniq<BoundConstantExpression>(Value::BOOLEAN(true)));
 				}
 			} else {
-				equalities.push_back(expression->Copy());
+				if (field_id_to_result_id.empty()) {
+					equalities.push_back(expression->Copy());
+				} else {
+					idx_t index = field_id_to_result_id[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));
+					}
+				}
 			}
 		}
 
@@ -399,17 +468,46 @@ void IcebergMultiFileReader::FinalizeChunk(ClientContext &context, const MultiFi
                                            DataChunk &input_chunk, DataChunk &output_chunk,
                                            ExpressionExecutor &executor,
                                            optional_ptr<MultiFileReaderGlobalState> global_state) {
+
+	// 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]);
+		}
+	}
+
 	// Base class finalization first
 	MultiFileReader::FinalizeChunk(context, bind_data, reader, reader_data, input_chunk, output_chunk, executor,
 	                               global_state);
 
+	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>();
+	}
+	unordered_map<idx_t, idx_t> field_id_to_result_id;
+	vector<ColumnIndex> column_indexes = reader.column_indexes;
+	int32_t result_id = executor.expressions.size() - 1;
+	for (int32_t i = column_indexes.size() - 1; i >= 0; i--) {
+		ColumnIndex column_index = column_indexes[i];
+		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];
-	auto &local_columns = reader.columns;
-	ApplyEqualityDeletes(context, output_chunk, multi_file_list, data_file, local_columns);
+	ApplyEqualityDeletes(context, output_chunk, multi_file_list, data_file, local_columns, field_id_to_result_id);
+
+	// delete the equality delete fields for result
+	for (idx_t i = 0; i < diff; i++) {
+		output_chunk.data.pop_back();
+	}
 }
 
 bool IcebergMultiFileReader::ParseOption(const string &key, const Value &val, MultiFileOptions &options,
diff --git a/src/include/iceberg_multi_file_reader.hpp b/src/include/iceberg_multi_file_reader.hpp
@@ -50,6 +50,11 @@ struct IcebergMultiFileReader : public MultiFileReader {
 	                      const MultiFileReaderBindData &bind_data, const MultiFileList &file_list,
 	                      const vector<MultiFileColumnDefinition> &global_columns,
 	                      const vector<ColumnIndex> &global_column_ids) override;
+	ReaderInitializeType InitializeReader(MultiFileReaderData &reader_data, const MultiFileBindData &bind_data,
+	                                      const vector<MultiFileColumnDefinition> &global_columns,
+	                                      const vector<ColumnIndex> &global_column_ids,
+	                                      optional_ptr<TableFilterSet> table_filters, ClientContext &context,
+	                                      MultiFileGlobalState &gstate) override;
 	void FinalizeBind(MultiFileReaderData &reader_data, const MultiFileOptions &file_options,
 	                  const MultiFileReaderBindData &options, const vector<MultiFileColumnDefinition> &global_columns,
 	                  const vector<ColumnIndex> &global_column_ids, ClientContext &context,
@@ -59,7 +64,8 @@ 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 vector<MultiFileColumnDefinition> &local_columns,
+	                          unordered_map<idx_t, idx_t> field_id_to_result_id);
 	bool ParseOption(const string &key, const Value &val, MultiFileOptions &options, ClientContext &context) override;
 
 public:
diff --git a/test/sql/local/equality_deletes.test b/test/sql/local/equality_deletes.test
@@ -44,7 +44,8 @@ SELECT filename[17:32], name, id FROM ICEBERG_SCAN('data/persistent/equality_del
 equality_deletes	b	2
 equality_deletes	b	1
 
-statement error
+query II
 SELECT filename[17:32], name FROM ICEBERG_SCAN('data/persistent/equality_deletes/warehouse/mydb/mytable');
 ----
-Not implemented Error: Equality deletes need the relevant columns to be selected
+equality_deletes	b
+equality_deletes	b
