Skip to content

Commit 0ff030e

Browse files
mrdrivingduckcodex
andcommitted
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>
1 parent 802c9ce commit 0ff030e

6 files changed

Lines changed: 361 additions & 28 deletions

File tree

src/include/paimon_insert.hpp

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,12 +40,13 @@ class PhysicalPaimonInsert : public PhysicalOperator {
4040

4141
PhysicalPaimonInsert(PhysicalPlan &physical_plan, LogicalOperator &op, SchemaCatalogEntry &schema,
4242
unique_ptr<BoundCreateTableInfo> info, string table_path, map<string, string> paimon_options,
43-
idx_t estimated_cardinality);
43+
vector<string> part_keys, idx_t estimated_cardinality);
4444

4545
SchemaCatalogEntry *schema;
4646
unique_ptr<BoundCreateTableInfo> info;
4747
string table_path;
4848
map<string, string> paimon_options;
49+
vector<string> part_keys;
4950

5051
public:
5152
bool IsSink() const override {

src/paimon_storage/paimon_catalog.cpp

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636
#include "duckdb/planner/parsed_data/bound_create_table_info.hpp"
3737

3838
#include "paimon/catalog/catalog.h"
39+
#include "paimon/schema/schema.h"
3940

4041
#include "paimon_catalog.hpp"
4142
#include "paimon_insert.hpp"
@@ -236,8 +237,10 @@ PhysicalOperator &PaimonCatalog::PlanCreateTableAs(ClientContext &context, Physi
236237
auto table_path = path + "/" + base.schema + paimon::Catalog::DB_SUFFIX + "/" + base.table;
237238
auto paimon_options = GetPaimonOptions(context, path, attached_options);
238239

240+
vector<string> part_keys;
241+
239242
auto &insert = planner.Make<PhysicalPaimonInsert>(op, op.schema, std::move(op.info), std::move(table_path),
240-
std::move(paimon_options), 0U);
243+
std::move(paimon_options), std::move(part_keys), 0U);
241244
insert.children.push_back(plan);
242245
return insert;
243246
}
@@ -252,12 +255,23 @@ PhysicalOperator &PaimonCatalog::PlanInsert(ClientContext &context, PhysicalPlan
252255
auto table_path = path + "/" + table.schema.name + paimon::Catalog::DB_SUFFIX + "/" + table.name;
253256
auto paimon_options = GetPaimonOptions(context, path, attached_options);
254257

258+
vector<string> part_keys;
259+
auto schema_result = paimon_catalog->LoadTableSchema(paimon::Identifier(table.schema.name, table.name));
260+
if (!schema_result.ok()) {
261+
throw IOException(schema_result.status().ToString());
262+
}
263+
auto data_schema = std::dynamic_pointer_cast<paimon::DataSchema>(schema_result.value());
264+
if (data_schema) {
265+
auto &schema_part_keys = data_schema->PartitionKeys();
266+
part_keys.assign(schema_part_keys.begin(), schema_part_keys.end());
267+
}
268+
255269
if (plan && !op.column_index_map.empty()) {
256270
plan = planner.ResolveDefaultsProjection(op, *plan);
257271
}
258272

259273
auto &insert = planner.Make<PhysicalPaimonInsert>(op, table.schema, nullptr, std::move(table_path),
260-
std::move(paimon_options), 0U);
274+
std::move(paimon_options), std::move(part_keys), 0U);
261275
if (plan) {
262276
insert.children.push_back(*plan);
263277
}

src/paimon_storage/paimon_insert.cpp

Lines changed: 118 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -26,12 +26,16 @@
2626
#include "duckdb/common/arrow/arrow_converter.hpp"
2727
#include "duckdb/main/client_context.hpp"
2828

29+
#include "paimon/catalog/catalog.h"
30+
#include "paimon/catalog/identifier.h"
2931
#include "paimon/commit_context.h"
3032
#include "paimon/file_store_commit.h"
3133
#include "paimon/file_store_write.h"
3234
#include "paimon/record_batch.h"
35+
#include "paimon/schema/schema.h"
3336
#include "paimon/write_context.h"
3437

38+
#include "paimon_catalog.hpp"
3539
#include "paimon_insert.hpp"
3640
#include "paimon_schema_entry.hpp"
3741

@@ -42,10 +46,11 @@ namespace duckdb {
4246

4347
PhysicalPaimonInsert::PhysicalPaimonInsert(PhysicalPlan &physical_plan, LogicalOperator &op, SchemaCatalogEntry &schema,
4448
unique_ptr<BoundCreateTableInfo> info, string table_path,
45-
map<string, string> paimon_options, idx_t estimated_cardinality)
49+
map<string, string> paimon_options, vector<string> part_keys,
50+
idx_t estimated_cardinality)
4651
: PhysicalOperator(physical_plan, PhysicalOperatorType::EXTENSION, {LogicalType::BIGINT}, estimated_cardinality),
4752
schema(&schema), info(std::move(info)), table_path(std::move(table_path)),
48-
paimon_options(std::move(paimon_options)) {
53+
paimon_options(std::move(paimon_options)), part_keys(std::move(part_keys)) {
4954
}
5055

5156
// ---------------------------------------------------------------------------
@@ -58,6 +63,10 @@ struct PaimonInsertGlobalState : public GlobalSinkState {
5863
std::vector<std::shared_ptr<paimon::CommitMessage>> all_commit_messages;
5964
idx_t insert_count = 0;
6065
bool finished = false;
66+
67+
vector<string> part_key_names;
68+
vector<idx_t> part_col_idxs;
69+
string null_part_name = "__DEFAULT_PARTITION__";
6170
};
6271

6372
struct PaimonInsertLocalState : public LocalSinkState {
@@ -78,6 +87,45 @@ unique_ptr<GlobalSinkState> PhysicalPaimonInsert::GetGlobalSinkState(ClientConte
7887
paimon_schema.CreateTable(txn, *info);
7988
}
8089

90+
if (!part_keys.empty()) {
91+
auto &paimon_catalog = schema->catalog.Cast<PaimonCatalog>();
92+
auto &catalog = paimon_catalog.GetPaimonCatalog();
93+
auto schema_name = info ? info->Base().schema : schema->name;
94+
auto table_name = info ? info->Base().table : string();
95+
96+
if (table_name.empty()) {
97+
auto last_slash = table_path.rfind('/');
98+
table_name = (last_slash != string::npos) ? table_path.substr(last_slash + 1) : table_path;
99+
}
100+
101+
auto schema_result = catalog.LoadTableSchema(paimon::Identifier(schema_name, table_name));
102+
if (!schema_result.ok()) {
103+
throw IOException(schema_result.status().ToString());
104+
}
105+
106+
auto table_schema = schema_result.value();
107+
auto data_schema = std::dynamic_pointer_cast<paimon::DataSchema>(table_schema);
108+
if (!data_schema) {
109+
throw IOException("Failed to resolve Paimon data schema for partitioned insert");
110+
}
111+
112+
auto &table_options = data_schema->Options();
113+
auto default_name = table_options.find(paimon::Options::PARTITION_DEFAULT_NAME);
114+
if (default_name != table_options.end()) {
115+
state->null_part_name = default_name->second;
116+
}
117+
118+
auto field_names = table_schema->FieldNames();
119+
state->part_key_names = part_keys;
120+
for (auto &part_key : part_keys) {
121+
auto it = std::find(field_names.begin(), field_names.end(), part_key);
122+
if (it == field_names.end()) {
123+
throw IOException("Partition key \"%s\" not found in Paimon table schema", part_key);
124+
}
125+
state->part_col_idxs.push_back(std::distance(field_names.begin(), it));
126+
}
127+
}
128+
81129
return std::move(state);
82130
}
83131

@@ -102,22 +150,16 @@ unique_ptr<LocalSinkState> PhysicalPaimonInsert::GetLocalSinkState(ExecutionCont
102150
return std::move(lstate);
103151
}
104152

105-
SinkResultType PhysicalPaimonInsert::Sink(ExecutionContext &context, DataChunk &chunk, OperatorSinkInput &input) const {
106-
auto &lstate = input.local_state.Cast<PaimonInsertLocalState>();
107-
108-
ArrowArray out_array {};
109-
auto client_props = context.client.GetClientProperties();
110-
ArrowConverter::ToArrowArray(chunk, &out_array, client_props, {});
111-
struct ArrowArrayGuard {
112-
ArrowArray &array;
113-
~ArrowArrayGuard() {
114-
if (array.release) {
115-
array.release(&array);
116-
}
117-
}
118-
} arrow_guard {out_array};
153+
static void WriteChunkWithPartition(PaimonInsertLocalState &lstate, DataChunk &chunk, ClientContext &client,
154+
const std::map<string, string> &partition) {
155+
ArrowArrayWrapper arrow_wrapper;
156+
auto client_props = client.GetClientProperties();
157+
ArrowConverter::ToArrowArray(chunk, &arrow_wrapper.arrow_array, client_props, {});
119158

120-
paimon::RecordBatchBuilder batch_builder(&out_array);
159+
paimon::RecordBatchBuilder batch_builder(&arrow_wrapper.arrow_array);
160+
if (!partition.empty()) {
161+
batch_builder.SetPartition(partition).SetBucket(0);
162+
}
121163
auto batch_result = batch_builder.Finish();
122164
if (!batch_result.ok()) {
123165
throw IOException(batch_result.status().ToString());
@@ -127,8 +169,66 @@ SinkResultType PhysicalPaimonInsert::Sink(ExecutionContext &context, DataChunk &
127169
if (!status.ok()) {
128170
throw IOException(status.ToString());
129171
}
172+
}
173+
174+
SinkResultType PhysicalPaimonInsert::Sink(ExecutionContext &context, DataChunk &chunk, OperatorSinkInput &input) const {
175+
auto &gstate = input.global_state.Cast<PaimonInsertGlobalState>();
176+
auto &lstate = input.local_state.Cast<PaimonInsertLocalState>();
177+
178+
if (gstate.part_col_idxs.empty()) {
179+
WriteChunkWithPartition(lstate, chunk, context.client, {});
180+
lstate.local_count += chunk.size();
181+
return SinkResultType::NEED_MORE_INPUT;
182+
}
183+
184+
auto num_rows = chunk.size();
185+
auto &part_idxs = gstate.part_col_idxs;
186+
auto &part_names = gstate.part_key_names;
187+
D_ASSERT(part_idxs.size() == part_names.size());
188+
189+
vector<UnifiedVectorFormat> part_formats(part_idxs.size());
190+
for (idx_t i = 0; i < part_idxs.size(); i++) {
191+
chunk.data[part_idxs[i]].ToUnifiedFormat(num_rows, part_formats[i]);
192+
}
193+
194+
std::map<std::vector<string>, vector<idx_t>> partition_groups;
195+
for (idx_t row = 0; row < num_rows; row++) {
196+
std::vector<string> key;
197+
key.reserve(part_idxs.size());
198+
for (idx_t i = 0; i < part_idxs.size(); i++) {
199+
auto idx = part_formats[i].sel->get_index(row);
200+
if (!part_formats[i].validity.RowIsValid(idx)) {
201+
key.push_back(gstate.null_part_name);
202+
} else {
203+
key.push_back(chunk.GetValue(part_idxs[i], row).ToString());
204+
}
205+
}
206+
auto entry = partition_groups.emplace(std::move(key), vector<idx_t>());
207+
entry.first->second.push_back(row);
208+
}
209+
210+
for (auto &entry : partition_groups) {
211+
auto &key_values = entry.first;
212+
auto &row_indices = entry.second;
213+
214+
std::map<string, string> partition;
215+
for (idx_t i = 0; i < part_names.size(); i++) {
216+
partition[part_names[i]] = key_values[i];
217+
}
218+
219+
SelectionVector sel(row_indices.size());
220+
for (idx_t i = 0; i < row_indices.size(); i++) {
221+
sel.set_index(i, row_indices[i]);
222+
}
223+
224+
DataChunk sub_chunk;
225+
sub_chunk.Initialize(Allocator::DefaultAllocator(), chunk.GetTypes());
226+
sub_chunk.Slice(chunk, sel, row_indices.size());
227+
228+
WriteChunkWithPartition(lstate, sub_chunk, context.client, partition);
229+
}
130230

131-
lstate.local_count += chunk.size();
231+
lstate.local_count += num_rows;
132232
return SinkResultType::NEED_MORE_INPUT;
133233
}
134234

0 commit comments

Comments
 (0)