feat(scan): support partition filter pushdown

Push partition filters into Paimon planning. Scans can prune
irrelevant partitions before row evaluation, keeping pruning close
to partition-key queries.

Support partitioned writes by grouping rows by partition value and
attaching metadata for each group. Keep SQL NULL row values while
partition metadata uses Paimon default markers.

Co-authored-by: GPT-5.5 <codex@users.noreply.github.com>

Author: mrdrivingduck

src/include/paimon_insert.hpp

@@ -40,12 +40,13 @@ class PhysicalPaimonInsert : public PhysicalOperator {
 
 	PhysicalPaimonInsert(PhysicalPlan &physical_plan, LogicalOperator &op, SchemaCatalogEntry &schema,
 	                     unique_ptr<BoundCreateTableInfo> info, string table_path, map<string, string> paimon_options,
-	                     idx_t estimated_cardinality);
+	                     vector<string> part_keys, idx_t estimated_cardinality);
 
 	SchemaCatalogEntry *schema;
 	unique_ptr<BoundCreateTableInfo> info;
 	string table_path;
 	map<string, string> paimon_options;
+	vector<string> part_keys;
 
 public:
 	bool IsSink() const override {

src/paimon_storage/paimon_catalog.cpp

@@ -30,12 +30,14 @@
 #include "duckdb/parser/parsed_data/create_schema_info.hpp"
 #include "duckdb/parser/parsed_data/drop_info.hpp"
 #include "duckdb/parser/parsed_data/attach_info.hpp"
+#include "duckdb/parser/expression/columnref_expression.hpp"
 #include "duckdb/parser/parsed_data/create_table_info.hpp"
 #include "duckdb/planner/operator/logical_create_table.hpp"
 #include "duckdb/planner/operator/logical_insert.hpp"
 #include "duckdb/planner/parsed_data/bound_create_table_info.hpp"
 
 #include "paimon/catalog/catalog.h"
+#include "paimon/schema/schema.h"
 
 #include "paimon_catalog.hpp"
 #include "paimon_insert.hpp"
@@ -236,8 +238,16 @@ PhysicalOperator &PaimonCatalog::PlanCreateTableAs(ClientContext &context, Physi
 	auto table_path = path + "/" + base.schema + paimon::Catalog::DB_SUFFIX + "/" + base.table;
 	auto paimon_options = GetPaimonOptions(context, path, attached_options);
 
+	vector<string> part_keys;
+	for (auto &part_expr : base.partition_keys) {
+		if (part_expr->GetExpressionType() != ExpressionType::COLUMN_REF) {
+			throw InvalidInputException("Paimon partition key must be a column reference");
+		}
+		part_keys.push_back(part_expr->Cast<ColumnRefExpression>().GetColumnName());
+	}
+
 	auto &insert = planner.Make<PhysicalPaimonInsert>(op, op.schema, std::move(op.info), std::move(table_path),
-	                                                  std::move(paimon_options), 0U);
+	                                                  std::move(paimon_options), std::move(part_keys), 0U);
 	insert.children.push_back(plan);
 	return insert;
 }
@@ -252,12 +262,23 @@ PhysicalOperator &PaimonCatalog::PlanInsert(ClientContext &context, PhysicalPlan
 	auto table_path = path + "/" + table.schema.name + paimon::Catalog::DB_SUFFIX + "/" + table.name;
 	auto paimon_options = GetPaimonOptions(context, path, attached_options);
 
+	vector<string> part_keys;
+	auto schema_result = paimon_catalog->LoadTableSchema(paimon::Identifier(table.schema.name, table.name));
+	if (!schema_result.ok()) {
+		throw IOException(schema_result.status().ToString());
+	}
+	auto data_schema = std::dynamic_pointer_cast<paimon::DataSchema>(schema_result.value());
+	if (data_schema) {
+		auto &schema_part_keys = data_schema->PartitionKeys();
+		part_keys.assign(schema_part_keys.begin(), schema_part_keys.end());
+	}
+
 	if (plan && !op.column_index_map.empty()) {
 		plan = planner.ResolveDefaultsProjection(op, *plan);
 	}
 
 	auto &insert = planner.Make<PhysicalPaimonInsert>(op, table.schema, nullptr, std::move(table_path),
-	                                                  std::move(paimon_options), 0U);
+	                                                  std::move(paimon_options), std::move(part_keys), 0U);
 	if (plan) {
 		insert.children.push_back(*plan);
 	}

src/paimon_storage/paimon_insert.cpp

@@ -24,13 +24,18 @@
 
 #include "duckdb/catalog/catalog.hpp"
 #include "duckdb/common/arrow/arrow_converter.hpp"
+#include "duckdb/common/arrow/arrow_wrapper.hpp"
 
+#include "paimon/catalog/catalog.h"
+#include "paimon/catalog/identifier.h"
 #include "paimon/commit_context.h"
 #include "paimon/file_store_commit.h"
 #include "paimon/file_store_write.h"
 #include "paimon/record_batch.h"
+#include "paimon/schema/schema.h"
 #include "paimon/write_context.h"
 
+#include "paimon_catalog.hpp"
 #include "paimon_insert.hpp"
 #include "paimon_schema_entry.hpp"
 
@@ -41,10 +46,11 @@ namespace duckdb {
 
 PhysicalPaimonInsert::PhysicalPaimonInsert(PhysicalPlan &physical_plan, LogicalOperator &op, SchemaCatalogEntry &schema,
                                            unique_ptr<BoundCreateTableInfo> info, string table_path,
-                                           map<string, string> paimon_options, idx_t estimated_cardinality)
+                                           map<string, string> paimon_options, vector<string> part_keys,
+                                           idx_t estimated_cardinality)
     : PhysicalOperator(physical_plan, PhysicalOperatorType::EXTENSION, {LogicalType::BIGINT}, estimated_cardinality),
       schema(&schema), info(std::move(info)), table_path(std::move(table_path)),
-      paimon_options(std::move(paimon_options)) {
+      paimon_options(std::move(paimon_options)), part_keys(std::move(part_keys)) {
 }
 
 // ---------------------------------------------------------------------------
@@ -57,6 +63,10 @@ struct PaimonInsertGlobalState : public GlobalSinkState {
 	std::vector<std::shared_ptr<paimon::CommitMessage>> all_commit_messages;
 	idx_t insert_count = 0;
 	bool finished = false;
+
+	vector<string> part_key_names;
+	vector<idx_t> part_col_idxs;
+	string null_part_name = "__DEFAULT_PARTITION__";
 };
 
 struct PaimonInsertLocalState : public LocalSinkState {
@@ -77,6 +87,45 @@ unique_ptr<GlobalSinkState> PhysicalPaimonInsert::GetGlobalSinkState(ClientConte
 		paimon_schema.CreateTable(txn, *info);
 	}
 
+	if (!part_keys.empty()) {
+		auto &paimon_catalog = schema->catalog.Cast<PaimonCatalog>();
+		auto &catalog = paimon_catalog.GetPaimonCatalog();
+		auto schema_name = info ? info->Base().schema : schema->name;
+		auto table_name = info ? info->Base().table : string();
+
+		if (table_name.empty()) {
+			auto last_slash = table_path.rfind('/');
+			table_name = (last_slash != string::npos) ? table_path.substr(last_slash + 1) : table_path;
+		}
+
+		auto schema_result = catalog.LoadTableSchema(paimon::Identifier(schema_name, table_name));
+		if (!schema_result.ok()) {
+			throw IOException(schema_result.status().ToString());
+		}
+
+		auto table_schema = schema_result.value();
+		auto data_schema = std::dynamic_pointer_cast<paimon::DataSchema>(table_schema);
+		if (!data_schema) {
+			throw IOException("Failed to resolve Paimon data schema for partitioned insert");
+		}
+
+		auto &table_options = data_schema->Options();
+		auto default_name = table_options.find(paimon::Options::PARTITION_DEFAULT_NAME);
+		if (default_name != table_options.end()) {
+			state->null_part_name = default_name->second;
+		}
+
+		auto field_names = table_schema->FieldNames();
+		state->part_key_names = part_keys;
+		for (auto &part_key : part_keys) {
+			auto it = std::find(field_names.begin(), field_names.end(), part_key);
+			if (it == field_names.end()) {
+				throw IOException("Partition key \"%s\" not found in Paimon table schema", part_key);
+			}
+			state->part_col_idxs.push_back(std::distance(field_names.begin(), it));
+		}
+	}
+
 	return std::move(state);
 }
 
@@ -101,22 +150,16 @@ unique_ptr<LocalSinkState> PhysicalPaimonInsert::GetLocalSinkState(ExecutionCont
 	return std::move(lstate);
 }
 
-SinkResultType PhysicalPaimonInsert::Sink(ExecutionContext &context, DataChunk &chunk, OperatorSinkInput &input) const {
-	auto &lstate = input.local_state.Cast<PaimonInsertLocalState>();
-
-	ArrowArray out_array {};
-	auto client_props = context.client.GetClientProperties();
-	ArrowConverter::ToArrowArray(chunk, &out_array, client_props, {});
-	struct ArrowArrayGuard {
-		ArrowArray &array;
-		~ArrowArrayGuard() {
-			if (array.release) {
-				array.release(&array);
-			}
-		}
-	} arrow_guard {out_array};
+static void WriteChunkWithPartition(PaimonInsertLocalState &lstate, DataChunk &chunk, ClientContext &client,
+                                    const std::map<string, string> &partition) {
+	ArrowArrayWrapper arrow_wrapper;
+	auto client_props = client.GetClientProperties();
+	ArrowConverter::ToArrowArray(chunk, &arrow_wrapper.arrow_array, client_props, {});
 
-	paimon::RecordBatchBuilder batch_builder(&out_array);
+	paimon::RecordBatchBuilder batch_builder(&arrow_wrapper.arrow_array);
+	if (!partition.empty()) {
+		batch_builder.SetPartition(partition).SetBucket(0);
+	}
 	auto batch_result = batch_builder.Finish();
 	if (!batch_result.ok()) {
 		throw IOException(batch_result.status().ToString());
@@ -126,8 +169,66 @@ SinkResultType PhysicalPaimonInsert::Sink(ExecutionContext &context, DataChunk &
 	if (!status.ok()) {
 		throw IOException(status.ToString());
 	}
+}
+
+SinkResultType PhysicalPaimonInsert::Sink(ExecutionContext &context, DataChunk &chunk, OperatorSinkInput &input) const {
+	auto &gstate = input.global_state.Cast<PaimonInsertGlobalState>();
+	auto &lstate = input.local_state.Cast<PaimonInsertLocalState>();
+
+	if (gstate.part_col_idxs.empty()) {
+		WriteChunkWithPartition(lstate, chunk, context.client, {});
+		lstate.local_count += chunk.size();
+		return SinkResultType::NEED_MORE_INPUT;
+	}
+
+	auto num_rows = chunk.size();
+	auto &part_idxs = gstate.part_col_idxs;
+	auto &part_names = gstate.part_key_names;
+	D_ASSERT(part_idxs.size() == part_names.size());
+
+	vector<UnifiedVectorFormat> part_formats(part_idxs.size());
+	for (idx_t i = 0; i < part_idxs.size(); i++) {
+		chunk.data[part_idxs[i]].ToUnifiedFormat(num_rows, part_formats[i]);
+	}
+
+	std::map<std::vector<string>, vector<idx_t>> partition_groups;
+	for (idx_t row = 0; row < num_rows; row++) {
+		std::vector<string> key;
+		key.reserve(part_idxs.size());
+		for (idx_t i = 0; i < part_idxs.size(); i++) {
+			auto idx = part_formats[i].sel->get_index(row);
+			if (!part_formats[i].validity.RowIsValid(idx)) {
+				key.push_back(gstate.null_part_name);
+			} else {
+				key.push_back(chunk.GetValue(part_idxs[i], row).ToString());
+			}
+		}
+		auto entry = partition_groups.emplace(std::move(key), vector<idx_t>());
+		entry.first->second.push_back(row);
+	}
+
+	for (auto &entry : partition_groups) {
+		auto &key_values = entry.first;
+		auto &row_indices = entry.second;
+
+		std::map<string, string> partition;
+		for (idx_t i = 0; i < part_names.size(); i++) {
+			partition[part_names[i]] = key_values[i];
+		}
+
+		SelectionVector sel(row_indices.size());
+		for (idx_t i = 0; i < row_indices.size(); i++) {
+			sel.set_index(i, row_indices[i]);
+		}
+
+		DataChunk sub_chunk;
+		sub_chunk.Initialize(Allocator::DefaultAllocator(), chunk.GetTypes());
+		sub_chunk.Slice(chunk, sel, row_indices.size());
+
+		WriteChunkWithPartition(lstate, sub_chunk, context.client, partition);
+	}
 
-	lstate.local_count += chunk.size();
+	lstate.local_count += num_rows;
 	return SinkResultType::NEED_MORE_INPUT;
 }