pg_duckdb/include/pgduckdb/scan/postgres_scan.hpp at f075a41518eea71f00b07386a7a3b14d6dbf91cc · duckdb/pg_duckdb · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
#pragma once

#include "duckdb.hpp"

#include "pgduckdb/pg/declarations.hpp"
#include "pgduckdb/utility/allocator.hpp"

#include "pgduckdb/scan/postgres_table_reader.hpp"

#include "pgduckdb/utility/cpp_only_file.hpp" // Must be last include.

namespace pgduckdb {

// Global State

struct PostgresScanGlobalState : public duckdb::GlobalTableFunctionState {
	explicit PostgresScanGlobalState(Snapshot, Relation rel, const duckdb::TableFunctionInitInput &input);
	~PostgresScanGlobalState();
	idx_t
	MaxThreads() const override {
		return 1;
	}
	void ConstructTableScanQuery(const duckdb::TableFunctionInitInput &input);

private:
	int ExtractQueryFilters(duckdb::TableFilter *filter, const char *column_name, duckdb::string &filters,
	                        bool is_optional_filter_parent);

public:
	Snapshot snapshot;
	Relation rel;
	TupleDesc table_tuple_desc;
	bool count_tuples_only;
	duckdb::vector<AttrNumber> output_columns;
	std::atomic<std::uint32_t> total_row_count;
	std::ostringstream scan_query;
	duckdb::shared_ptr<PostgresTableReader> table_reader_global_state;
};

// Local State

struct PostgresScanLocalState : public duckdb::LocalTableFunctionState {
	PostgresScanLocalState(PostgresScanGlobalState *global_state);
	~PostgresScanLocalState() override;

	PostgresScanGlobalState *global_state;

	size_t output_vector_size;
	bool exhausted_scan;
};

// PostgresScanFunctionData

struct PostgresScanFunctionData : public duckdb::TableFunctionData {
	PostgresScanFunctionData(Relation rel, uint64_t cardinality, Snapshot snapshot);
	~PostgresScanFunctionData() override;
	duckdb::vector<duckdb::string> complex_filters;
	Relation rel;
	uint64_t cardinality;
	Snapshot snapshot;
};

// PostgresScanTableFunction

struct PostgresScanTableFunction : public duckdb::TableFunction {
	PostgresScanTableFunction();

	static duckdb::unique_ptr<duckdb::GlobalTableFunctionState>
	PostgresScanInitGlobal(duckdb::ClientContext &context, duckdb::TableFunctionInitInput &input);

	static duckdb::unique_ptr<duckdb::LocalTableFunctionState>
	PostgresScanInitLocal(duckdb::ExecutionContext &context, duckdb::TableFunctionInitInput &input,
	                      duckdb::GlobalTableFunctionState *gstate);

	static void PostgresScanFunction(duckdb::ClientContext &context, duckdb::TableFunctionInput &data,
	                                 duckdb::DataChunk &output);

	static duckdb::unique_ptr<duckdb::NodeStatistics> PostgresScanCardinality(duckdb::ClientContext &context,
	                                                                          const duckdb::FunctionData *data);
	static duckdb::InsertionOrderPreservingMap<duckdb::string> ToString(duckdb::TableFunctionToStringInput &input);
};

} // namespace pgduckdb