From 35ff478abd779a1f978dfde012d0237e7dce3656 Mon Sep 17 00:00:00 2001 From: mkaruza Date: Thu, 28 Nov 2024 13:21:43 +0100 Subject: [PATCH 01/53] Use PostgreSQL execution nodes to scan tables # Conflicts: # include/pgduckdb/pgduckdb_guc.h # Conflicts: # src/scan/heap_reader.cpp # src/scan/postgres_seq_scan.cpp --- include/pgduckdb/catalog/pgduckdb_table.hpp | 28 +-- include/pgduckdb/pg/declarations.hpp | 13 ++ include/pgduckdb/pgduckdb_guc.h | 2 +- include/pgduckdb/pgduckdb_types.hpp | 3 +- include/pgduckdb/scan/heap_reader.hpp | 60 ----- include/pgduckdb/scan/postgres_scan.hpp | 97 +++++--- include/pgduckdb/scan/postgres_seq_scan.hpp | 78 ------- .../pgduckdb/scan/postgres_table_reader.hpp | 34 +++ include/pgduckdb/scan/postgres_view.hpp | 13 ++ src/catalog/pgduckdb_table.cpp | 23 +- src/catalog/pgduckdb_transaction.cpp | 4 +- src/pgduckdb.cpp | 8 +- src/pgduckdb_duckdb.cpp | 10 +- src/pgduckdb_hooks.cpp | 8 +- src/pgduckdb_types.cpp | 152 ++----------- src/scan/heap_reader.cpp | 180 --------------- src/scan/postgres_scan.cpp | 208 ++++++++++++------ src/scan/postgres_seq_scan.cpp | 115 ---------- src/scan/postgres_table_reader.cpp | 173 +++++++++++++++ src/scan/postgres_view.cpp | 90 ++++++++ 20 files changed, 571 insertions(+), 728 deletions(-) delete mode 100644 include/pgduckdb/scan/heap_reader.hpp delete mode 100644 include/pgduckdb/scan/postgres_seq_scan.hpp create mode 100644 include/pgduckdb/scan/postgres_table_reader.hpp create mode 100644 include/pgduckdb/scan/postgres_view.hpp delete mode 100644 src/scan/heap_reader.cpp delete mode 100644 src/scan/postgres_seq_scan.cpp create mode 100644 src/scan/postgres_table_reader.cpp create mode 100644 src/scan/postgres_view.cpp diff --git a/include/pgduckdb/catalog/pgduckdb_table.hpp b/include/pgduckdb/catalog/pgduckdb_table.hpp index b5d1b186..8157d6e8 100644 --- a/include/pgduckdb/catalog/pgduckdb_table.hpp +++ b/include/pgduckdb/catalog/pgduckdb_table.hpp @@ -11,35 +11,27 @@ namespace pgduckdb { class PostgresTable : public duckdb::TableCatalogEntry { public: + PostgresTable(duckdb::Catalog &catalog, duckdb::SchemaCatalogEntry &schema, duckdb::CreateTableInfo &info, + Relation rel, Cardinality cardinality, Snapshot snapshot); + virtual ~PostgresTable(); +public: + duckdb::unique_ptr GetStatistics(duckdb::ClientContext &context, + duckdb::column_t column_id) override; + duckdb::TableFunction GetScanFunction(duckdb::ClientContext &context, + duckdb::unique_ptr &bind_data) override; + duckdb::TableStorageInfo GetStorageInfo(duckdb::ClientContext &context) override; + public: static Relation OpenRelation(Oid relid); static void SetTableInfo(duckdb::CreateTableInfo &info, Relation rel); static Cardinality GetTableCardinality(Relation rel); -protected: - PostgresTable(duckdb::Catalog &catalog, duckdb::SchemaCatalogEntry &schema, duckdb::CreateTableInfo &info, - Relation rel, Cardinality cardinality, Snapshot snapshot); - protected: Relation rel; Cardinality cardinality; Snapshot snapshot; }; -class PostgresHeapTable : public PostgresTable { -public: - PostgresHeapTable(duckdb::Catalog &catalog, duckdb::SchemaCatalogEntry &schema, duckdb::CreateTableInfo &info, - Relation rel, Cardinality cardinality, Snapshot snapshot); - -public: - // -- Table API -- - duckdb::unique_ptr GetStatistics(duckdb::ClientContext &context, - duckdb::column_t column_id) override; - duckdb::TableFunction GetScanFunction(duckdb::ClientContext &context, - duckdb::unique_ptr &bind_data) override; - duckdb::TableStorageInfo GetStorageInfo(duckdb::ClientContext &context) override; -}; - } // namespace pgduckdb diff --git a/include/pgduckdb/pg/declarations.hpp b/include/pgduckdb/pg/declarations.hpp index fb7a4043..46d1a6ee 100644 --- a/include/pgduckdb/pg/declarations.hpp +++ b/include/pgduckdb/pg/declarations.hpp @@ -66,4 +66,17 @@ struct TableAmRoutine; typedef uint32_t CommandId; typedef uint32_t SubTransactionId; + +struct QueryDesc; + +struct ParallelExecutorInfo; + +struct MinimalTupleData; +typedef MinimalTupleData *MinimalTuple; + +struct TupleQueueReader; + +struct PlanState; + +struct Plan; } diff --git a/include/pgduckdb/pgduckdb_guc.h b/include/pgduckdb/pgduckdb_guc.h index b1e0664d..35a8369a 100644 --- a/include/pgduckdb/pgduckdb_guc.h +++ b/include/pgduckdb/pgduckdb_guc.h @@ -8,7 +8,7 @@ extern bool duckdb_enable_external_access; extern bool duckdb_allow_unsigned_extensions; extern bool duckdb_autoinstall_known_extensions; extern bool duckdb_autoload_known_extensions; -extern int duckdb_max_threads_per_postgres_scan; +extern int max_workers_per_postgres_scan; extern char *duckdb_motherduck_postgres_database; extern int duckdb_motherduck_enabled; extern char *duckdb_motherduck_token; diff --git a/include/pgduckdb/pgduckdb_types.hpp b/include/pgduckdb/pgduckdb_types.hpp index 97028a46..a47cd3c6 100644 --- a/include/pgduckdb/pgduckdb_types.hpp +++ b/include/pgduckdb/pgduckdb_types.hpp @@ -21,7 +21,6 @@ int32_t GetPostgresDuckDBTypemod(const duckdb::LogicalType &type); duckdb::Value ConvertPostgresParameterToDuckValue(Datum value, Oid postgres_type); void ConvertPostgresToDuckValue(Oid attr_type, Datum value, duckdb::Vector &result, uint64_t offset); bool ConvertDuckToPostgresValue(TupleTableSlot *slot, duckdb::Value &value, uint64_t col); -void InsertTupleIntoChunk(duckdb::DataChunk &output, duckdb::shared_ptr scan_global_state, - duckdb::shared_ptr scan_local_state, HeapTupleData *tuple); +void InsertTupleIntoChunk(duckdb::DataChunk &output, PostgresScanLocalState &scan_local_state, TupleTableSlot *slot); } // namespace pgduckdb diff --git a/include/pgduckdb/scan/heap_reader.hpp b/include/pgduckdb/scan/heap_reader.hpp deleted file mode 100644 index a077dd13..00000000 --- a/include/pgduckdb/scan/heap_reader.hpp +++ /dev/null @@ -1,60 +0,0 @@ -#pragma once - -#include "duckdb.hpp" - -#include "pgduckdb/scan/postgres_scan.hpp" -#include "pgduckdb/pg/declarations.hpp" - -#include "pgduckdb/utility/cpp_only_file.hpp" // Must be last include. - -namespace pgduckdb { - -// HeapReaderGlobalState - -class HeapReaderGlobalState { -public: - HeapReaderGlobalState(Relation rel); - BlockNumber AssignNextBlockNumber(std::mutex &lock); - -private: - BlockNumber m_nblocks; - BlockNumber m_last_assigned_block_number; -}; - -// HeapReader - -class HeapReader { -public: - HeapReader(Relation rel, duckdb::shared_ptr heap_reader_global_state, - duckdb::shared_ptr global_state, - duckdb::shared_ptr local_state); - ~HeapReader(); - HeapReader(const HeapReader &other) = delete; - HeapReader &operator=(const HeapReader &other) = delete; - HeapReader &operator=(HeapReader &&other) = delete; - HeapReader(HeapReader &&other) = delete; - bool ReadPageTuples(duckdb::DataChunk &output); - BlockNumber - GetCurrentBlockNumber() { - return m_block_number; - } - -private: - Page PreparePageRead(); - - duckdb::shared_ptr m_global_state; - duckdb::shared_ptr m_heap_reader_global_state; - duckdb::shared_ptr m_local_state; - Relation m_rel; - bool m_inited; - bool m_read_next_page; - bool m_page_tuples_all_visible; - BlockNumber m_block_number; - Buffer m_buffer; - OffsetNumber m_current_tuple_index; - int m_page_tuples_left; - duckdb::unique_ptr m_tuple; - BufferAccessStrategy m_buffer_access_strategy; -}; - -} // namespace pgduckdb diff --git a/include/pgduckdb/scan/postgres_scan.hpp b/include/pgduckdb/scan/postgres_scan.hpp index 85fdf86e..a16f5ab8 100644 --- a/include/pgduckdb/scan/postgres_scan.hpp +++ b/include/pgduckdb/scan/postgres_scan.hpp @@ -5,52 +5,79 @@ #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 { -class PostgresScanGlobalState { -public: - PostgresScanGlobalState() : m_snapshot(nullptr), m_count_tuples_only(false), m_total_row_count(0) { +// Global State + +struct PostgresScanGlobalState : public duckdb::GlobalTableFunctionState { + explicit PostgresScanGlobalState(Snapshot, Relation rel, Cardinality cardinality, + duckdb::TableFunctionInitInput &input); + ~PostgresScanGlobalState(); + idx_t + MaxThreads() const override { + return 1; } + void ConstructTableScanQuery(duckdb::TableFunctionInitInput &input); - void InitGlobalState(duckdb::TableFunctionInitInput &input); - - void InitRelationMissingAttrs(TupleDesc tuple_desc); - - Snapshot m_snapshot; - TupleDesc m_tuple_desc; - std::mutex m_lock; // Lock for one replacement scan - bool m_count_tuples_only; - /* Postgres column id to duckdb scanned index. The scanned index is DuckDB - * its scan order of the columns. */ - std::vector> m_columns_to_scan; - /* These are indexed by the DuckDB scan index */ - std::vector m_column_filters; - /* Duckdb output vector idx with information about postgres column id */ - duckdb::vector> m_output_columns; - std::atomic m_total_row_count; - duckdb::map m_relation_missing_attrs; +public: + Snapshot snapshot; + Relation rel; + TupleDesc table_tuple_desc; + std::mutex lock; // Lock for one replacement scan + bool count_tuples_only; + duckdb::vector> output_columns; + std::atomic total_row_count; + std::ostringstream scan_query; + duckdb::shared_ptr table_reader_global_state; }; -class PostgresScanLocalState { +// Local State + +struct PostgresScanLocalState : public duckdb::LocalTableFunctionState { public: - PostgresScanLocalState(const PostgresScanGlobalState *psgs) : m_output_vector_size(0), m_exhausted_scan(false) { - if (!psgs->m_count_tuples_only) { - const auto s = psgs->m_columns_to_scan.size(); - values.resize(s); - nulls.resize(s); - } - } + PostgresScanLocalState(PostgresScanGlobalState *global_state); + ~PostgresScanLocalState() override; + +public: + PostgresScanGlobalState *global_state; + int output_vector_size; + bool exhausted_scan; +}; - uint32_t m_output_vector_size; - bool m_exhausted_scan; - std::vector> values; - std::vector> nulls; +// PostgresScanFunctionData + +struct PostgresScanFunctionData : public duckdb::TableFunctionData { +public: + PostgresScanFunctionData(Relation rel, uint64_t cardinality, Snapshot snapshot); + ~PostgresScanFunctionData() override; + +public: + duckdb::vector complex_filters; + Relation rel; + uint64_t cardinality; + Snapshot snapshot; }; -duckdb::unique_ptr PostgresReplacementScan(duckdb::ClientContext &context, - duckdb::ReplacementScanInput &input, - duckdb::optional_ptr data); +// PostgresScanTableFunction + +struct PostgresScanTableFunction : public duckdb::TableFunction { +public: + PostgresScanTableFunction(); + +public: + static duckdb::unique_ptr + PostgresScanInitGlobal(duckdb::ClientContext &context, duckdb::TableFunctionInitInput &input); + static duckdb::unique_ptr + 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 PostgresScanCardinality(duckdb::ClientContext &context, + const duckdb::FunctionData *data); +}; } // namespace pgduckdb diff --git a/include/pgduckdb/scan/postgres_seq_scan.hpp b/include/pgduckdb/scan/postgres_seq_scan.hpp deleted file mode 100644 index 287ee0f6..00000000 --- a/include/pgduckdb/scan/postgres_seq_scan.hpp +++ /dev/null @@ -1,78 +0,0 @@ -#pragma once - -#include "duckdb.hpp" - -#include "pgduckdb/pgduckdb_guc.h" -#include "pgduckdb/pg/declarations.hpp" - -#include "pgduckdb/utility/cpp_only_file.hpp" // Must be last include. - -namespace pgduckdb { - -class HeapReaderGlobalState; -class HeapReader; -class PostgresScanGlobalState; -class PostgresScanLocalState; - -// Global State - -struct PostgresSeqScanGlobalState : public duckdb::GlobalTableFunctionState { - explicit PostgresSeqScanGlobalState(Relation rel, duckdb::TableFunctionInitInput &input); - ~PostgresSeqScanGlobalState(); - idx_t - MaxThreads() const override { - return duckdb_max_threads_per_postgres_scan; - } - -public: - duckdb::shared_ptr m_global_state; - duckdb::shared_ptr m_heap_reader_global_state; - Relation m_rel; -}; - -// Local State - -struct PostgresSeqScanLocalState : public duckdb::LocalTableFunctionState { -public: - PostgresSeqScanLocalState(Relation rel, duckdb::shared_ptr heap_reader_global_state, - duckdb::shared_ptr global_state); - ~PostgresSeqScanLocalState() override; - -public: - duckdb::shared_ptr m_local_state; - duckdb::unique_ptr m_heap_table_reader; -}; - -// PostgresSeqScanFunctionData - -struct PostgresSeqScanFunctionData : public duckdb::TableFunctionData { -public: - PostgresSeqScanFunctionData(Relation rel, uint64_t cardinality, Snapshot snapshot); - ~PostgresSeqScanFunctionData() override; - -public: - Relation m_rel; - uint64_t m_cardinality; - Snapshot m_snapshot; -}; - -// PostgresSeqScanFunction - -struct PostgresSeqScanFunction : public duckdb::TableFunction { -public: - PostgresSeqScanFunction(); - -public: - static duckdb::unique_ptr - PostgresSeqScanInitGlobal(duckdb::ClientContext &context, duckdb::TableFunctionInitInput &input); - static duckdb::unique_ptr - PostgresSeqScanInitLocal(duckdb::ExecutionContext &context, duckdb::TableFunctionInitInput &input, - duckdb::GlobalTableFunctionState *gstate); - static void PostgresSeqScanFunc(duckdb::ClientContext &context, duckdb::TableFunctionInput &data, - duckdb::DataChunk &output); - - static duckdb::unique_ptr PostgresSeqScanCardinality(duckdb::ClientContext &context, - const duckdb::FunctionData *data); -}; - -} // namespace pgduckdb diff --git a/include/pgduckdb/scan/postgres_table_reader.hpp b/include/pgduckdb/scan/postgres_table_reader.hpp new file mode 100644 index 00000000..344f626f --- /dev/null +++ b/include/pgduckdb/scan/postgres_table_reader.hpp @@ -0,0 +1,34 @@ +#pragma once + +#include "duckdb.hpp" + +#include "pgduckdb/pg/declarations.hpp" + +#include "pgduckdb/utility/cpp_only_file.hpp" // Must be last include. + +namespace pgduckdb { + +// PostgresTableReader + +class PostgresTableReader { +public: + PostgresTableReader(Cardinality cardinality, const char *table_scan_query); + ~PostgresTableReader(); + TupleTableSlot *GetNextTuple(); + +private: + MinimalTuple GetNextWorkerTuple(); + int ParalleWorkerNumber(Cardinality cardinality); + +private: + QueryDesc *table_scan_query_desc; + PlanState *table_scan_planstate; + ParallelExecutorInfo *pei; + void **parallel_worker_readers; + TupleTableSlot *slot; + int nworkers_launched; + int nreaders; + int next_parallel_reader; +}; + +} // namespace pgduckdb diff --git a/include/pgduckdb/scan/postgres_view.hpp b/include/pgduckdb/scan/postgres_view.hpp new file mode 100644 index 00000000..78cfec56 --- /dev/null +++ b/include/pgduckdb/scan/postgres_view.hpp @@ -0,0 +1,13 @@ +#pragma once + +#include "duckdb.hpp" + +#include "pgduckdb/utility/cpp_only_file.hpp" // Must be last include. + +namespace pgduckdb { + +duckdb::unique_ptr PostgresViewScan(duckdb::ClientContext &context, + duckdb::ReplacementScanInput &input, + duckdb::optional_ptr data); + +} diff --git a/src/catalog/pgduckdb_table.cpp b/src/catalog/pgduckdb_table.cpp index 6bc3dce8..180fa51d 100644 --- a/src/catalog/pgduckdb_table.cpp +++ b/src/catalog/pgduckdb_table.cpp @@ -1,11 +1,11 @@ #include "pgduckdb/catalog/pgduckdb_table.hpp" +#include "pgduckdb/scan/postgres_scan.hpp" #include "pgduckdb/catalog/pgduckdb_schema.hpp" #include "pgduckdb/logger.hpp" #include "pgduckdb/pg/relations.hpp" #include "pgduckdb/pgduckdb_process_lock.hpp" #include "pgduckdb/pgduckdb_types.hpp" // ConvertPostgresToDuckColumnType -#include "pgduckdb/scan/postgres_seq_scan.hpp" #include "duckdb/parser/parsed_data/create_table_info.hpp" @@ -55,29 +55,20 @@ PostgresTable::GetTableCardinality(Relation rel) { return cardinality; } -//===--------------------------------------------------------------------===// -// PostgresHeapTable -//===--------------------------------------------------------------------===// - -PostgresHeapTable::PostgresHeapTable(duckdb::Catalog &_catalog, duckdb::SchemaCatalogEntry &_schema, - duckdb::CreateTableInfo &_info, Relation _rel, Cardinality _cardinality, - Snapshot _snapshot) - : PostgresTable(_catalog, _schema, _info, _rel, _cardinality, _snapshot) { -} - duckdb::unique_ptr -PostgresHeapTable::GetStatistics(duckdb::ClientContext &, duckdb::column_t) { +PostgresTable::GetStatistics(duckdb::ClientContext &context, duckdb::column_t column_id) { throw duckdb::NotImplementedException("GetStatistics not supported yet"); } duckdb::TableFunction -PostgresHeapTable::GetScanFunction(duckdb::ClientContext &, duckdb::unique_ptr &bind_data) { - bind_data = duckdb::make_uniq(rel, cardinality, snapshot); - return PostgresSeqScanFunction(); +PostgresTable::GetScanFunction(duckdb::ClientContext &context, + duckdb::unique_ptr &bind_data) { + bind_data = duckdb::make_uniq(rel, cardinality, snapshot); + return PostgresScanTableFunction(); } duckdb::TableStorageInfo -PostgresHeapTable::GetStorageInfo(duckdb::ClientContext &) { +PostgresTable::GetStorageInfo(duckdb::ClientContext &context) { throw duckdb::NotImplementedException("GetStorageInfo not supported yet"); } diff --git a/src/catalog/pgduckdb_transaction.cpp b/src/catalog/pgduckdb_transaction.cpp index a40eb3a6..bb0f6e83 100644 --- a/src/catalog/pgduckdb_transaction.cpp +++ b/src/catalog/pgduckdb_transaction.cpp @@ -55,8 +55,8 @@ SchemaItems::GetTable(const duckdb::string &entry_name) { PostgresTable::SetTableInfo(info, rel); auto cardinality = PostgresTable::GetTableCardinality(rel); - tables.emplace(entry_name, duckdb::make_uniq(schema->catalog, *schema, info, rel, cardinality, - schema->snapshot)); + tables.emplace(entry_name, duckdb::make_uniq(schema->catalog, *schema, info, rel, cardinality, + schema->snapshot)); return tables[entry_name].get(); } diff --git a/src/pgduckdb.cpp b/src/pgduckdb.cpp index f914ac2a..a9aa97ae 100644 --- a/src/pgduckdb.cpp +++ b/src/pgduckdb.cpp @@ -14,7 +14,7 @@ extern "C" { static void DuckdbInitGUC(void); bool duckdb_force_execution = false; -int duckdb_max_threads_per_postgres_scan = 1; +int max_workers_per_postgres_scan = 2; int duckdb_motherduck_enabled = MotherDuckEnabled::MOTHERDUCK_AUTO; char *duckdb_motherduck_token = strdup(""); char *duckdb_motherduck_postgres_database = strdup("postgres"); @@ -159,9 +159,9 @@ DuckdbInitGUC(void) { "Maximum number of DuckDB threads per Postgres backend, alias for duckdb.threads", &duckdb_maximum_threads, -1, 1024, PGC_SUSET); - DefineCustomVariable("duckdb.max_threads_per_postgres_scan", - "Maximum number of DuckDB threads used for a single Postgres scan", - &duckdb_max_threads_per_postgres_scan, 1, 64); + DefineCustomVariable("duckdb.max_workers_per_postgres_scan", + "Maximum number of PostgreSQL threads used for a single Postgres scan", + &max_workers_per_postgres_scan, 2, 8); DefineCustomVariable("duckdb.postgres_role", "Which postgres role should be allowed to use DuckDB execution, use the secrets and create " diff --git a/src/pgduckdb_duckdb.cpp b/src/pgduckdb_duckdb.cpp index 70331dfb..c2a67c0a 100644 --- a/src/pgduckdb_duckdb.cpp +++ b/src/pgduckdb_duckdb.cpp @@ -8,7 +8,7 @@ #include "pgduckdb/catalog/pgduckdb_storage.hpp" #include "pgduckdb/scan/postgres_scan.hpp" -#include "pgduckdb/scan/postgres_seq_scan.hpp" +#include "pgduckdb/scan/postgres_view.hpp" #include "pgduckdb/pg/transactions.hpp" #include "pgduckdb/pgduckdb_utils.hpp" @@ -136,7 +136,7 @@ DuckDBManager::Initialize() { config.SetOptionByName("custom_user_agent", "pg_duckdb"); config.SetOptionByName("extension_directory", CreateOrGetDirectoryPath("duckdb_extensions")); // Transforms VIEWs into their view definition - config.replacement_scans.emplace_back(pgduckdb::PostgresReplacementScan); + config.replacement_scans.emplace_back(pgduckdb::PostgresViewScan); SET_DUCKDB_OPTION(allow_unsigned_extensions); SET_DUCKDB_OPTION(enable_external_access); SET_DUCKDB_OPTION(autoinstall_known_extensions); @@ -208,14 +208,10 @@ DuckDBManager::Initialize() { void DuckDBManager::LoadFunctions(duckdb::ClientContext &context) { - pgduckdb::PostgresSeqScanFunction seq_scan_fun; - duckdb::CreateTableFunctionInfo seq_scan_info(seq_scan_fun); - - auto &catalog = duckdb::Catalog::GetSystemCatalog(context); + //auto &catalog = duckdb::Catalog::GetSystemCatalog(context); context.transaction.BeginTransaction(); auto &instance = *database->instance; duckdb::ExtensionUtil::RegisterType(instance, "UnsupportedPostgresType", duckdb::LogicalTypeId::VARCHAR); - catalog.CreateTableFunction(context, &seq_scan_info); context.transaction.Commit(); } diff --git a/src/pgduckdb_hooks.cpp b/src/pgduckdb_hooks.cpp index 6936edbe..c3144866 100644 --- a/src/pgduckdb_hooks.cpp +++ b/src/pgduckdb_hooks.cpp @@ -181,10 +181,10 @@ IsAllowedStatement(Query *query, bool throw_error = false) { /* * When accessing the partitioned table, we temporarily let PG handle it instead of DuckDB. */ - if (ContainsPartitionedTable(query->rtable)) { - elog(elevel, "DuckDB does not support querying PG partitioned table"); - return false; - } + // if (ContainsPartitionedTable(query->rtable)) { + // elog(elevel, "DuckDB does not support querying PG partitioned table"); + // return false; + // } /* Anything else is hopefully fine... */ return true; diff --git a/src/pgduckdb_types.cpp b/src/pgduckdb_types.cpp index c8ee60ac..3fe1dc29 100644 --- a/src/pgduckdb_types.cpp +++ b/src/pgduckdb_types.cpp @@ -1348,157 +1348,33 @@ ConvertPostgresToDuckValue(Oid attr_type, Datum value, duckdb::Vector &result, i } } -typedef struct HeapTupleReadState { - bool m_slow = 0; - int m_last_tuple_att = 0; - uint32 m_page_tuple_offset = 0; -} HeapTupleReadState; - -static Datum -HeapTupleFetchNextColumnDatum(TupleDesc tupleDesc, HeapTuple tuple, HeapTupleReadState &heap_tuple_read_state, - AttrNumber target_attr_num, bool *is_null, const duckdb::map &missing_attrs) { - HeapTupleHeader tup = tuple->t_data; - bool hasnulls = HeapTupleHasNulls(tuple); - int natts = HeapTupleHeaderGetNatts(tup); - bits8 *null_bitmap = tup->t_bits; - - if (natts < target_attr_num) { - if (auto missing_attr = missing_attrs.find(target_attr_num - 1); missing_attr != missing_attrs.end()) { - *is_null = false; - return missing_attr->second; - } else { - *is_null = true; - return PointerGetDatum(NULL); - } - } - - /* Which tuple are we currently reading */ - AttrNumber current_attr_num = heap_tuple_read_state.m_last_tuple_att + 1; - /* Either restore from previous fetch, or use the defaults of 0 and false */ - uint32 current_tuple_offset = heap_tuple_read_state.m_page_tuple_offset; - bool slow = heap_tuple_read_state.m_slow; - - /* Points to the start of the tuple data section, i.e. right after the - * tuple header */ - char *tuple_data = (char *)tup + tup->t_hoff; - - Datum value = (Datum)0; - for (; current_attr_num <= target_attr_num; current_attr_num++) { - Form_pg_attribute thisatt = TupleDescAttr(tupleDesc, current_attr_num - 1); - - if (hasnulls && att_isnull(current_attr_num - 1, null_bitmap)) { - value = (Datum)0; - *is_null = true; - /* - * Can't use attcacheoff anymore. The hardcoded attribute offset - * assumes all attribute before it are present in the tuple. If - * they are NULL, they are not present. - */ - slow = true; - continue; - } - - *is_null = false; - - if (!slow && thisatt->attcacheoff >= 0) { - current_tuple_offset = thisatt->attcacheoff; - } else if (thisatt->attlen == -1) { - if (!slow && current_tuple_offset == att_align_nominal(current_tuple_offset, thisatt->attalign)) { - thisatt->attcacheoff = current_tuple_offset; - } else { - current_tuple_offset = - att_align_pointer(current_tuple_offset, thisatt->attalign, -1, tuple_data + current_tuple_offset); - slow = true; - } - } else { - current_tuple_offset = att_align_nominal(current_tuple_offset, thisatt->attalign); - if (!slow) { - thisatt->attcacheoff = current_tuple_offset; - } - } - - value = fetchatt(thisatt, tuple_data + current_tuple_offset); - - current_tuple_offset = - att_addlength_pointer(current_tuple_offset, thisatt->attlen, tuple_data + current_tuple_offset); - - if (thisatt->attlen <= 0) { - slow = true; - } - } - - heap_tuple_read_state.m_last_tuple_att = target_attr_num; - heap_tuple_read_state.m_page_tuple_offset = current_tuple_offset; - heap_tuple_read_state.m_slow = slow; - - return value; -} - void -InsertTupleIntoChunk(duckdb::DataChunk &output, duckdb::shared_ptr scan_global_state, - duckdb::shared_ptr scan_local_state, HeapTupleData *tuple) { - HeapTupleReadState heap_tuple_read_state = {}; - - if (scan_global_state->m_count_tuples_only) { - scan_local_state->m_output_vector_size++; - return; - } +InsertTupleIntoChunk(duckdb::DataChunk &output, PostgresScanLocalState &scan_local_state, TupleTableSlot *slot) { - auto &values = scan_local_state->values; - auto &nulls = scan_local_state->nulls; - - /* First we are fetching all required columns ordered by column id - * and than we need to write this tuple into output vector. Output column id list - * could be out of order so we need to match column values from ordered list. - */ - - /* Read heap tuple with all required columns. */ - for (auto const &[attr_num, duckdb_scanned_index] : scan_global_state->m_columns_to_scan) { - bool is_null = false; - values[duckdb_scanned_index] = - HeapTupleFetchNextColumnDatum(scan_global_state->m_tuple_desc, tuple, heap_tuple_read_state, attr_num, - &is_null, scan_global_state->m_relation_missing_attrs); - nulls[duckdb_scanned_index] = is_null; - auto filter = scan_global_state->m_column_filters[duckdb_scanned_index]; - if (!filter) { - continue; - } + auto scan_global_state = scan_local_state.global_state; - const auto valid_tuple = ApplyValueFilter(*filter, values[duckdb_scanned_index], is_null, - scan_global_state->m_tuple_desc->attrs[attr_num - 1].atttypid); - if (!valid_tuple) { - return; - } + if (scan_global_state->count_tuples_only) { + scan_global_state->total_row_count += slot->tts_values[0]; + scan_local_state.output_vector_size += slot->tts_values[0]; + return; } - /* Write tuple columns in output vector. */ int duckdb_output_index = 0; - for (auto const &[duckdb_scanned_index, attr_num] : scan_global_state->m_output_columns) { + for (auto const &[duckdb_scanned_index, attr_num] : scan_global_state->output_columns) { auto &result = output.data[duckdb_output_index]; - if (nulls[duckdb_scanned_index]) { + if (slot->tts_isnull[duckdb_scanned_index]) { auto &array_mask = duckdb::FlatVector::Validity(result); - array_mask.SetInvalid(scan_local_state->m_output_vector_size); + array_mask.SetInvalid(scan_local_state.output_vector_size); } else { - auto attr = scan_global_state->m_tuple_desc->attrs[attr_num - 1]; - if (attr.attlen == -1) { - bool should_free = false; - values[duckdb_scanned_index] = - DetoastPostgresDatum(reinterpret_cast(values[duckdb_scanned_index]), &should_free); - ConvertPostgresToDuckValue(attr.atttypid, values[duckdb_scanned_index], result, - scan_local_state->m_output_vector_size); - if (should_free) { - duckdb_free(reinterpret_cast(values[duckdb_scanned_index])); - } - } else { - ConvertPostgresToDuckValue(attr.atttypid, values[duckdb_scanned_index], result, - scan_local_state->m_output_vector_size); - } + auto attr = slot->tts_tupleDescriptor->attrs[attr_num - 1]; + ConvertPostgresToDuckValue(attr.atttypid, slot->tts_values[duckdb_scanned_index], result, + scan_local_state.output_vector_size); } duckdb_output_index++; } - scan_local_state->m_output_vector_size++; - scan_global_state->m_total_row_count++; + scan_local_state.output_vector_size++; + scan_global_state->total_row_count++; } NumericVar diff --git a/src/scan/heap_reader.cpp b/src/scan/heap_reader.cpp deleted file mode 100644 index 95344441..00000000 --- a/src/scan/heap_reader.cpp +++ /dev/null @@ -1,180 +0,0 @@ -#include "duckdb.hpp" - -#include "pgduckdb/scan/heap_reader.hpp" -#include "pgduckdb/pgduckdb_types.hpp" -#include "pgduckdb/pgduckdb_utils.hpp" - -extern "C" { -#include "postgres.h" -#include "pgstat.h" -#include "access/heapam.h" -#include "storage/bufmgr.h" -#include "storage/bufpage.h" -#include "utils/rel.h" -} - -#include "pgduckdb/pgduckdb_process_lock.hpp" - -#include - -namespace pgduckdb { - -// -// HeapReaderGlobalState -// - -HeapReaderGlobalState::HeapReaderGlobalState(Relation rel) - : m_nblocks(RelationGetNumberOfBlocks(rel)), m_last_assigned_block_number(InvalidBlockNumber) { -} - -BlockNumber -HeapReaderGlobalState::AssignNextBlockNumber(std::mutex &lock) { - lock.lock(); - BlockNumber block_number = InvalidBlockNumber; - if (m_nblocks > 0 && m_last_assigned_block_number == InvalidBlockNumber) { - block_number = m_last_assigned_block_number = 0; - } else if (m_nblocks > 0 && m_last_assigned_block_number < m_nblocks - 1) { - block_number = ++m_last_assigned_block_number; - } - lock.unlock(); - return block_number; -} - -// -// HeapReader -// - -HeapReader::HeapReader(Relation rel, duckdb::shared_ptr heap_reader_global_state, - duckdb::shared_ptr global_state, - duckdb::shared_ptr local_state) - : m_global_state(global_state), m_heap_reader_global_state(heap_reader_global_state), m_local_state(local_state), - m_rel(rel), m_inited(false), m_read_next_page(true), m_block_number(InvalidBlockNumber), m_buffer(InvalidBuffer), - m_current_tuple_index(InvalidOffsetNumber), m_page_tuples_left(0) { - m_tuple = duckdb::make_uniq(); - m_tuple->t_data = NULL; - m_tuple->t_tableOid = RelationGetRelid(m_rel); - ItemPointerSetInvalid(&m_tuple->t_self); - DuckdbProcessLock::GetLock().lock(); - m_buffer_access_strategy = GetAccessStrategy(BAS_BULKREAD); - DuckdbProcessLock::GetLock().unlock(); -} - -HeapReader::~HeapReader() { - DuckdbProcessLock::GetLock().lock(); - /* If execution is interrupted and buffer is still opened close it now */ - if (m_buffer != InvalidBuffer) { - UnlockReleaseBuffer(m_buffer); - } - FreeAccessStrategy(m_buffer_access_strategy); - DuckdbProcessLock::GetLock().unlock(); -} - -Page -HeapReader::PreparePageRead() { - Page page = BufferGetPage(m_buffer); -#if PG_VERSION_NUM < 170000 - TestForOldSnapshot(m_global_state->m_snapshot, m_rel, page); -#endif - m_page_tuples_all_visible = PageIsAllVisible(page) && !m_global_state->m_snapshot->takenDuringRecovery; - m_page_tuples_left = PageGetMaxOffsetNumber(page) - FirstOffsetNumber + 1; - m_current_tuple_index = FirstOffsetNumber; - return page; -} - -bool -HeapReader::ReadPageTuples(duckdb::DataChunk &output) { - BlockNumber block = InvalidBlockNumber; - Page page = nullptr; - - if (!m_inited) { - block = m_block_number = m_heap_reader_global_state->AssignNextBlockNumber(m_global_state->m_lock); - if (m_block_number == InvalidBlockNumber) { - return false; - } - m_inited = true; - m_read_next_page = true; - } else { - block = m_block_number; - if (!m_read_next_page) { - page = BufferGetPage(m_buffer); - } - } - - while (block != InvalidBlockNumber) { - if (m_read_next_page) { - CHECK_FOR_INTERRUPTS(); - std::lock_guard lock(DuckdbProcessLock::GetLock()); - block = m_block_number; - - m_buffer = PostgresFunctionGuard(ReadBufferExtended, m_rel, MAIN_FORKNUM, block, RBM_NORMAL, - m_buffer_access_strategy); - - PostgresFunctionGuard(LockBuffer, m_buffer, BUFFER_LOCK_SHARE); - - page = PreparePageRead(); - m_read_next_page = false; - } - - for (; m_page_tuples_left > 0 && m_local_state->m_output_vector_size < STANDARD_VECTOR_SIZE; - m_page_tuples_left--, m_current_tuple_index++) { - bool visible = true; - ItemId lpp = PageGetItemId(page, m_current_tuple_index); - - if (!ItemIdIsNormal(lpp)) - continue; - - m_tuple->t_data = (HeapTupleHeader)PageGetItem(page, lpp); - m_tuple->t_len = ItemIdGetLength(lpp); - ItemPointerSet(&(m_tuple->t_self), block, m_current_tuple_index); - - if (!m_page_tuples_all_visible) { - std::lock_guard lock(DuckdbProcessLock::GetLock()); - visible = HeapTupleSatisfiesVisibility(m_tuple.get(), m_global_state->m_snapshot, m_buffer); - /* skip tuples not visible to this snapshot */ - if (!visible) - continue; - } - - pgstat_count_heap_getnext(m_rel); - InsertTupleIntoChunk(output, m_global_state, m_local_state, m_tuple.get()); - } - - /* No more items on current page */ - if (!m_page_tuples_left) { - DuckdbProcessLock::GetLock().lock(); - UnlockReleaseBuffer(m_buffer); - DuckdbProcessLock::GetLock().unlock(); - m_buffer = InvalidBuffer; - m_read_next_page = true; - /* Handle cancel request */ - if (QueryCancelPending) { - block = m_block_number = InvalidBlockNumber; - } else { - block = m_block_number = m_heap_reader_global_state->AssignNextBlockNumber(m_global_state->m_lock); - } - } - - /* We have collected STANDARD_VECTOR_SIZE */ - if (m_local_state->m_output_vector_size == STANDARD_VECTOR_SIZE) { - output.SetCardinality(m_local_state->m_output_vector_size); - output.Verify(); - m_local_state->m_output_vector_size = 0; - return true; - } - } - - /* Next assigned block number is InvalidBlockNumber so we check did we write any tuples in output vector */ - if (m_local_state->m_output_vector_size) { - output.SetCardinality(m_local_state->m_output_vector_size); - output.Verify(); - m_local_state->m_output_vector_size = 0; - } - - m_buffer = InvalidBuffer; - m_block_number = InvalidBlockNumber; - m_tuple->t_data = NULL; - m_read_next_page = false; - - return false; -} -} // namespace pgduckdb diff --git a/src/scan/postgres_scan.cpp b/src/scan/postgres_scan.cpp index ffa9603e..a1947050 100644 --- a/src/scan/postgres_scan.cpp +++ b/src/scan/postgres_scan.cpp @@ -13,9 +13,15 @@ #include "duckdb/common/enums/expression_type.hpp" #include "pgduckdb/scan/postgres_scan.hpp" +#include "pgduckdb/scan/postgres_table_reader.hpp" #include "pgduckdb/pgduckdb_types.hpp" #include "pgduckdb/pgduckdb_utils.hpp" +#include "pgduckdb/pgduckdb_process_lock.hpp" +#include "pgduckdb/pgduckdb_types.hpp" +#include "pgduckdb/pgduckdb_utils.hpp" +#include "pgduckdb/logger.hpp" + extern "C" { #include "postgres.h" #include "access/htup_details.h" @@ -25,8 +31,12 @@ extern "C" { #include "optimizer/planner.h" #include "utils/builtins.h" #include "utils/regproc.h" +#include "utils/rel.h" #include "utils/snapmgr.h" #include "utils/syscache.h" +#include "tcop/tcopprot.h" +#include "executor/execdesc.h" +#include "executor/executor.h" } #include "pgduckdb/pgduckdb_process_lock.hpp" @@ -34,13 +44,13 @@ extern "C" { namespace pgduckdb { void -PostgresScanGlobalState::InitGlobalState(duckdb::TableFunctionInitInput &input) { +PostgresScanGlobalState::ConstructTableScanQuery(duckdb::TableFunctionInitInput &input) { /* SELECT COUNT(*) FROM */ if (input.column_ids.size() == 1 && input.column_ids[0] == UINT64_MAX) { - m_count_tuples_only = true; + scan_query << "SELECT COUNT(*) FROM " << rel->rd_rel->relname.data; + count_tuples_only = true; return; } - /* * We need to read columns from the Postgres tuple in column order, but for * outputting them we care about the DuckDB order. A map automatically @@ -55,10 +65,12 @@ PostgresScanGlobalState::InitGlobalState(duckdb::TableFunctionInitInput &input) } auto table_filters = input.filters.get(); - m_column_filters.resize(input.column_ids.size(), 0); + + std::vector> columns_to_scan; + std::vector column_filters(input.column_ids.size(), 0); for (auto const &[att_num, duckdb_scanned_index] : pg_column_order) { - m_columns_to_scan.emplace_back(att_num, duckdb_scanned_index); + columns_to_scan.emplace_back(att_num, duckdb_scanned_index); if (!table_filters) { continue; @@ -66,7 +78,7 @@ PostgresScanGlobalState::InitGlobalState(duckdb::TableFunctionInitInput &input) auto column_filter_it = table_filters->filters.find(duckdb_scanned_index); if (column_filter_it != table_filters->filters.end()) { - m_column_filters[duckdb_scanned_index] = column_filter_it->second.get(); + column_filters[duckdb_scanned_index] = column_filter_it->second.get(); } } @@ -78,99 +90,159 @@ PostgresScanGlobalState::InitGlobalState(duckdb::TableFunctionInitInput &input) */ if (input.CanRemoveFilterColumns()) { for (const auto &projection_id : input.projection_ids) { - m_output_columns.emplace_back(projection_id, input.column_ids[projection_id] + 1); + output_columns.emplace_back(projection_id, input.column_ids[projection_id] + 1); } } else { duckdb::idx_t output_index = 0; for (const auto &column_id : input.column_ids) { - m_output_columns.emplace_back(output_index++, column_id + 1); + output_columns.emplace_back(output_index++, column_id + 1); } } -} -void -PostgresScanGlobalState::InitRelationMissingAttrs(TupleDesc tuple_desc) { - std::lock_guard lock(DuckdbProcessLock::GetLock()); - for (int attnum = 0; attnum < tuple_desc->natts; attnum++) { - bool is_null = false; - Datum attr = PostgresFunctionGuard(getmissingattr, tuple_desc, attnum + 1, &is_null); - /* Add missing attr datum if not null*/ - if (!is_null) { - m_relation_missing_attrs[attnum] = attr; + scan_query << "SELECT "; + + bool first = true; + for (auto const &[duckdb_scanned_index, attr_num] : output_columns) { + if (!first) { + scan_query << ", "; } + first = false; + auto attr = table_tuple_desc->attrs[attr_num - 1]; + scan_query << attr.attname.data; } -} -static Oid -FindMatchingRelation(const duckdb::string &schema, const duckdb::string &table) { - List *name_list = NIL; - if (!schema.empty()) { - name_list = lappend(name_list, makeString(pstrdup(schema.c_str()))); + scan_query << " FROM " << rel->rd_rel->relname.data; + + first = true; + + for (auto const &[attr_num, duckdb_scanned_index] : columns_to_scan) { + auto filter = column_filters[duckdb_scanned_index]; + if (filter) { + if (!first) { + scan_query << " AND "; + } else { + scan_query << " WHERE "; + } + first = false; + scan_query << " ("; + auto attr = table_tuple_desc->attrs[attr_num - 1]; + scan_query << filter->ToString(attr.attname.data); + scan_query << ") "; + } } - name_list = lappend(name_list, makeString(pstrdup(table.c_str()))); + scan_query << ";"; - RangeVar *table_range_var = makeRangeVarFromNameList(name_list); - return RangeVarGetRelid(table_range_var, AccessShareLock, true); + elog(DEBUG1, "scan_query: %s", scan_query.str().c_str()); } -const char * -pgduckdb_pg_get_viewdef(Oid view) { - auto oid = ObjectIdGetDatum(view); - Datum viewdef = DirectFunctionCall1(pg_get_viewdef, oid); - return text_to_cstring(DatumGetTextP(viewdef)); +// +// PostgresScanGlobalState +// + +PostgresScanGlobalState::PostgresScanGlobalState(Snapshot snapshot, Relation rel, Cardinality cardinality, duckdb::TableFunctionInitInput &input) + : snapshot(snapshot), rel(rel), table_tuple_desc(RelationGetDescr(rel)), count_tuples_only(false), + total_row_count(0) { + ConstructTableScanQuery(input); + table_reader_global_state = duckdb::make_shared_ptr(cardinality, scan_query.str().c_str()); + pd_log(DEBUG2, "(DuckDB/PostgresSeqScanGlobalState) Running %" PRIu64 " threads -- ", (uint64_t)MaxThreads()); } -duckdb::unique_ptr -ReplaceView(Oid view) { - const auto view_definition = PostgresFunctionGuard(pgduckdb_pg_get_viewdef, view); +PostgresScanGlobalState::~PostgresScanGlobalState() { +} - if (!view_definition) { - throw duckdb::InvalidInputException("Could not retrieve view definition for Relation with relid: %u", view); - } +// +// PostgresScanLocalState +// - duckdb::Parser parser; - parser.ParseQuery(view_definition); - auto &statements = parser.statements; - if (statements.size() != 1) { - throw duckdb::InvalidInputException("View definition contained more than 1 statement!"); - } +PostgresScanLocalState::PostgresScanLocalState(PostgresScanGlobalState *global_state) + : global_state(global_state), exhausted_scan(false) { +} - if (statements[0]->type != duckdb::StatementType::SELECT_STATEMENT) { - throw duckdb::InvalidInputException("View definition (%s) did not contain a SELECT statement!", - view_definition); - } +PostgresScanLocalState::~PostgresScanLocalState() { +} + +// +// PostgresSeqScanFunctionData +// - auto select = duckdb::unique_ptr_cast(std::move(statements[0])); - return duckdb::make_uniq(std::move(select)); +PostgresScanFunctionData::PostgresScanFunctionData(Relation rel, uint64_t cardinality, Snapshot snapshot) + : rel(rel), cardinality(cardinality), snapshot(snapshot) { } -duckdb::unique_ptr -PostgresReplacementScan(duckdb::ClientContext &, duckdb::ReplacementScanInput &input, - duckdb::optional_ptr) { +PostgresScanFunctionData::~PostgresScanFunctionData() { +} - auto &schema_name = input.schema_name; - auto &table_name = input.table_name; +// +// PostgresScanFunction +// + +PostgresScanTableFunction::PostgresScanTableFunction() + : TableFunction("postgres_scan", {}, PostgresScanFunction, nullptr, PostgresScanInitGlobal, PostgresScanInitLocal) { + named_parameters["cardinality"] = duckdb::LogicalType::UBIGINT; + named_parameters["relid"] = duckdb::LogicalType::UINTEGER; + named_parameters["snapshot"] = duckdb::LogicalType::POINTER; + projection_pushdown = true; + filter_pushdown = true; + filter_prune = true; + cardinality = PostgresScanCardinality; +} - auto relid = PostgresFunctionGuard(FindMatchingRelation, schema_name, table_name); - if (relid == InvalidOid) { - return nullptr; +duckdb::unique_ptr +PostgresScanTableFunction::PostgresScanInitGlobal(duckdb::ClientContext &context, + duckdb::TableFunctionInitInput &input) { + auto &bind_data = input.bind_data->CastNoConst(); + auto global_state = duckdb::make_uniq(bind_data.snapshot, bind_data.rel, bind_data.cardinality, input); + return std::move(global_state); +} + +duckdb::unique_ptr +PostgresScanTableFunction::PostgresScanInitLocal(duckdb::ExecutionContext &context, + duckdb::TableFunctionInitInput &input, + duckdb::GlobalTableFunctionState *gstate) { + auto global_state = reinterpret_cast(gstate); + return duckdb::make_uniq(global_state); +} + +void +PostgresScanTableFunction::PostgresScanFunction(duckdb::ClientContext &context, duckdb::TableFunctionInput &data, + duckdb::DataChunk &output) { + auto &local_state = data.local_state->Cast(); + + local_state.output_vector_size = 0; + + /* We have exhausted seq scan of heap table so we can return */ + if (local_state.exhausted_scan) { + output.SetCardinality(0); + return; + } + + int i = 0; + for (; i < STANDARD_VECTOR_SIZE; i++) { + TupleTableSlot *slot = local_state.global_state->table_reader_global_state->GetNextTuple(); + if (TupIsNull(slot)) { + local_state.exhausted_scan = true; + break; + } + slot_getallattrs(slot); + InsertTupleIntoChunk(output, local_state, slot); } - auto tuple = PostgresFunctionGuard(SearchSysCache1, RELOID, ObjectIdGetDatum(relid)); - if (!HeapTupleIsValid(tuple)) { - elog(WARNING, "(PGDuckDB/PostgresReplacementScan) Cache lookup failed for relation %u", relid); - return nullptr; + if (i != STANDARD_VECTOR_SIZE) { + local_state.exhausted_scan = true; } - auto relForm = (Form_pg_class)GETSTRUCT(tuple); - if (relForm->relkind != RELKIND_VIEW) { - PostgresFunctionGuard(ReleaseSysCache, tuple); - return nullptr; + if (local_state.global_state->count_tuples_only) { + output.SetCapacity(local_state.output_vector_size); } - PostgresFunctionGuard(ReleaseSysCache, tuple); - return ReplaceView(relid); + output.SetCardinality(local_state.output_vector_size); +} + +duckdb::unique_ptr +PostgresScanTableFunction::PostgresScanCardinality(duckdb::ClientContext &context, const duckdb::FunctionData *data) { + auto &bind_data = data->Cast(); + return duckdb::make_uniq(bind_data.cardinality, bind_data.cardinality); } } // namespace pgduckdb diff --git a/src/scan/postgres_seq_scan.cpp b/src/scan/postgres_seq_scan.cpp deleted file mode 100644 index 3a11e18d..00000000 --- a/src/scan/postgres_seq_scan.cpp +++ /dev/null @@ -1,115 +0,0 @@ -#include "duckdb.hpp" - -#include "pgduckdb/scan/postgres_seq_scan.hpp" -#include "pgduckdb/pgduckdb_types.hpp" -#include "pgduckdb/logger.hpp" -#include "pgduckdb/scan/heap_reader.hpp" -#include "pgduckdb/pg/relations.hpp" - -#include "pgduckdb/utility/cpp_only_file.hpp" // Must be last include. - -namespace pgduckdb { - -// -// PostgresSeqScanGlobalState -// - -PostgresSeqScanGlobalState::PostgresSeqScanGlobalState(Relation rel, duckdb::TableFunctionInitInput &input) - : m_global_state(duckdb::make_shared_ptr()), - m_heap_reader_global_state(duckdb::make_shared_ptr(rel)), m_rel(rel) { - m_global_state->InitGlobalState(input); - m_global_state->m_tuple_desc = RelationGetDescr(m_rel); - m_global_state->InitRelationMissingAttrs(m_global_state->m_tuple_desc); - pd_log(DEBUG2, "(DuckDB/PostgresSeqScanGlobalState) Running %" PRIu64 " threads -- ", (uint64_t)MaxThreads()); -} - -PostgresSeqScanGlobalState::~PostgresSeqScanGlobalState() { -} - -// -// PostgresSeqScanLocalState -// - -PostgresSeqScanLocalState::PostgresSeqScanLocalState(Relation rel, - duckdb::shared_ptr heap_reder_global_state, - duckdb::shared_ptr global_state) { - m_local_state = duckdb::make_shared_ptr(global_state.get()); - m_heap_table_reader = duckdb::make_uniq(rel, heap_reder_global_state, global_state, m_local_state); -} - -PostgresSeqScanLocalState::~PostgresSeqScanLocalState() { -} - -// -// PostgresSeqScanFunctionData -// - -PostgresSeqScanFunctionData::PostgresSeqScanFunctionData(Relation rel, uint64_t cardinality, Snapshot snapshot) - : m_rel(rel), m_cardinality(cardinality), m_snapshot(snapshot) { -} - -PostgresSeqScanFunctionData::~PostgresSeqScanFunctionData() { -} - -// -// PostgresSeqScanFunction -// - -PostgresSeqScanFunction::PostgresSeqScanFunction() - : TableFunction("postgres_seq_scan", {}, PostgresSeqScanFunc, nullptr, PostgresSeqScanInitGlobal, - PostgresSeqScanInitLocal) { - named_parameters["cardinality"] = duckdb::LogicalType::UBIGINT; - named_parameters["relid"] = duckdb::LogicalType::UINTEGER; - named_parameters["snapshot"] = duckdb::LogicalType::POINTER; - projection_pushdown = true; - filter_pushdown = true; - filter_prune = true; - cardinality = PostgresSeqScanCardinality; -} - -duckdb::unique_ptr -PostgresSeqScanFunction::PostgresSeqScanInitGlobal(duckdb::ClientContext &, duckdb::TableFunctionInitInput &input) { - auto &bind_data = input.bind_data->CastNoConst(); - auto global_state = duckdb::make_uniq(bind_data.m_rel, input); - global_state->m_global_state->m_snapshot = bind_data.m_snapshot; -#pragma GCC diagnostic push -#pragma GCC diagnostic ignored "-Wredundant-move" - return std::move(global_state); -#pragma GCC diagnostic pop -} - -duckdb::unique_ptr -PostgresSeqScanFunction::PostgresSeqScanInitLocal(duckdb::ExecutionContext &, duckdb::TableFunctionInitInput &, - duckdb::GlobalTableFunctionState *gstate) { - auto global_state = reinterpret_cast(gstate); - return duckdb::make_uniq(global_state->m_rel, global_state->m_heap_reader_global_state, - global_state->m_global_state); -} - -void -PostgresSeqScanFunction::PostgresSeqScanFunc(duckdb::ClientContext &, duckdb::TableFunctionInput &data, - duckdb::DataChunk &output) { - auto &local_state = data.local_state->Cast(); - - local_state.m_local_state->m_output_vector_size = 0; - - /* We have exhausted seq scan of heap table so we can return */ - if (local_state.m_local_state->m_exhausted_scan) { - output.SetCardinality(0); - return; - } - - auto hasTuple = local_state.m_heap_table_reader->ReadPageTuples(output); - - if (!hasTuple || !IsValidBlockNumber(local_state.m_heap_table_reader->GetCurrentBlockNumber())) { - local_state.m_local_state->m_exhausted_scan = true; - } -} - -duckdb::unique_ptr -PostgresSeqScanFunction::PostgresSeqScanCardinality(duckdb::ClientContext &, const duckdb::FunctionData *data) { - auto &bind_data = data->Cast(); - return duckdb::make_uniq(bind_data.m_cardinality, bind_data.m_cardinality); -} - -} // namespace pgduckdb diff --git a/src/scan/postgres_table_reader.cpp b/src/scan/postgres_table_reader.cpp new file mode 100644 index 00000000..7a1bb325 --- /dev/null +++ b/src/scan/postgres_table_reader.cpp @@ -0,0 +1,173 @@ +#include "pgduckdb/scan/postgres_table_reader.hpp" +#include "pgduckdb/pgduckdb_process_lock.hpp" + +extern "C" { +#include "postgres.h" +#include "miscadmin.h" +#include "access/htup_details.h" +#include "catalog/namespace.h" +#include "catalog/pg_class.h" +#include "optimizer/planmain.h" +#include "optimizer/planner.h" +#include "utils/builtins.h" +#include "utils/regproc.h" +#include "utils/rel.h" +#include "utils/snapmgr.h" +#include "utils/syscache.h" +#include "tcop/tcopprot.h" +#include "executor/execdesc.h" +#include "executor/executor.h" +#include "executor/execParallel.h" +#include "executor/tqueue.h" +#include "utils/wait_event.h" +#include "access/xact.h" +#include "pgduckdb/pgduckdb_guc.h" +} + +#include + +namespace pgduckdb { + +PostgresTableReader::PostgresTableReader(Cardinality cardinality, const char *table_scan_query) + : parallel_worker_readers(nullptr), nreaders(0), next_parallel_reader(0) { + + std::lock_guard lock(DuckdbProcessLock::GetLock()); + + pg_stack_base_t current_stack = set_stack_base(); + List *raw_parsetree_list; + raw_parsetree_list = pg_parse_query(table_scan_query); + List *query_list = + pg_analyze_and_rewrite_fixedparams(linitial_node(RawStmt, raw_parsetree_list), table_scan_query, NULL, 0, NULL); + PlannedStmt *planned_stmt = standard_planner(linitial_node(Query, query_list), table_scan_query, 0, NULL); + table_scan_query_desc = CreateQueryDesc(planned_stmt, table_scan_query, GetActiveSnapshot(), InvalidSnapshot, + None_Receiver, NULL, NULL, 0); + ExecutorStart(table_scan_query_desc, 0); + + table_scan_planstate = ExecInitNode(planned_stmt->planTree, table_scan_query_desc->estate, 0); + + table_scan_query_desc->planstate->plan->parallel_aware = true; + + int parallel_workers = ParalleWorkerNumber(cardinality); + + RESUME_CANCEL_INTERRUPTS(); + + ParallelContext *pcxt; + pei = ExecInitParallelPlan(table_scan_planstate, table_scan_query_desc->estate, NULL, parallel_workers, -1); + pcxt = pei->pcxt; + LaunchParallelWorkers(pcxt); + /* We save # workers launched for the benefit of EXPLAIN */ + nworkers_launched = pcxt->nworkers_launched; + + /* Set up tuple queue readers to read the results. */ + if (pcxt->nworkers_launched > 0) { + ExecParallelCreateReaders(pei); + /* Make a working array showing the active readers */ + nreaders = pcxt->nworkers_launched; + parallel_worker_readers = (void **)palloc(nreaders * sizeof(TupleQueueReader *)); + memcpy(parallel_worker_readers, pei->reader, nreaders * sizeof(TupleQueueReader *)); + } + + HOLD_CANCEL_INTERRUPTS(); + + slot = ExecInitExtraTupleSlot(table_scan_query_desc->estate, table_scan_planstate->ps_ResultTupleDesc, &TTSOpsMinimalTuple); + + restore_stack_base(current_stack); +} + +PostgresTableReader::~PostgresTableReader() { + std::lock_guard lock(DuckdbProcessLock::GetLock()); + + ExecEndNode(table_scan_planstate); + + if (pei != NULL) + ExecParallelFinish(pei); + + if (parallel_worker_readers) { + pfree(parallel_worker_readers); + } + + parallel_worker_readers = NULL; + + ExecParallelCleanup(pei); + + pei = NULL; + + ExecutorFinish(table_scan_query_desc); + ExecutorEnd(table_scan_query_desc); + FreeQueryDesc(table_scan_query_desc); +} + +int +PostgresTableReader::ParalleWorkerNumber(Cardinality cardinality) { + static const int base_log = 8; + int cardinality_log = std::log2(cardinality); + int base = cardinality_log / base_log; + return std::max(1, std::min(base, std::min(max_workers_per_postgres_scan, max_parallel_workers))); +} + +TupleTableSlot * +PostgresTableReader::GetNextTuple() { + MinimalTuple worker_minmal_tuple; + if (nreaders > 0) { + worker_minmal_tuple = GetNextWorkerTuple(); + if (HeapTupleIsValid(worker_minmal_tuple)) { + ExecStoreMinimalTuple(worker_minmal_tuple, slot, false); + return slot; + } + } else { + std::lock_guard lock(DuckdbProcessLock::GetLock()); + pg_stack_base_t current_stack = set_stack_base(); + table_scan_query_desc->estate->es_query_dsa = pei ? pei->area : NULL; + slot = ExecProcNode(table_scan_planstate); + restore_stack_base(current_stack); + table_scan_query_desc->estate->es_query_dsa = NULL; + if (!TupIsNull(slot)) { + return slot; + } + } + + return ExecClearTuple(slot); +} + +MinimalTuple +PostgresTableReader::GetNextWorkerTuple() { + int nvisited = 0; + for (;;) { + TupleQueueReader *reader; + MinimalTuple minimal_tuple; + bool readerdone; + + reader = (TupleQueueReader *) parallel_worker_readers[next_parallel_reader]; + minimal_tuple = TupleQueueReaderNext(reader, true, &readerdone); + + if (readerdone) { + --nreaders; + if (nreaders == 0) { + return NULL; + } + memmove(¶llel_worker_readers[next_parallel_reader], ¶llel_worker_readers[next_parallel_reader + 1], + sizeof(TupleQueueReader *) * (nreaders - next_parallel_reader)); + if (next_parallel_reader >= nreaders) { + next_parallel_reader = 0; + } + continue; + } + + if (minimal_tuple) { + return minimal_tuple; + } + + next_parallel_reader++; + if (next_parallel_reader >= nreaders) + next_parallel_reader = 0; + + nvisited++; + if (nvisited >= nreaders) { + (void)WaitLatch(MyLatch, WL_LATCH_SET | WL_EXIT_ON_PM_DEATH, 0, WAIT_EVENT_EXECUTE_GATHER); + ResetLatch(MyLatch); + nvisited = 0; + } + } +} + +} // namespace pgduckdb diff --git a/src/scan/postgres_view.cpp b/src/scan/postgres_view.cpp new file mode 100644 index 00000000..bf655560 --- /dev/null +++ b/src/scan/postgres_view.cpp @@ -0,0 +1,90 @@ +#include "duckdb/parser/parser.hpp" +#include "duckdb/parser/tableref/subqueryref.hpp" +#include "pgduckdb/pgduckdb_utils.hpp" + +extern "C" { +#include "postgres.h" +#include "access/htup_details.h" +#include "catalog/namespace.h" +#include "catalog/pg_class.h" +#include "optimizer/planmain.h" +#include "utils/builtins.h" +#include "utils/regproc.h" +#include "utils/syscache.h" +} + +namespace pgduckdb { + +static Oid +FindMatchingRelation(const duckdb::string &schema, const duckdb::string &table) { + List *name_list = NIL; + if (!schema.empty()) { + name_list = lappend(name_list, makeString(pstrdup(schema.c_str()))); + } + + name_list = lappend(name_list, makeString(pstrdup(table.c_str()))); + + RangeVar *table_range_var = makeRangeVarFromNameList(name_list); + return RangeVarGetRelid(table_range_var, AccessShareLock, true); +} + +const char * +pgduckdb_pg_get_viewdef(Oid view) { + auto oid = ObjectIdGetDatum(view); + Datum viewdef = DirectFunctionCall1(pg_get_viewdef, oid); + return text_to_cstring(DatumGetTextP(viewdef)); +} + +duckdb::unique_ptr +ReplaceView(Oid view) { + const auto view_definition = PostgresFunctionGuard(pgduckdb_pg_get_viewdef, view); + + if (!view_definition) { + throw duckdb::InvalidInputException("Could not retrieve view definition for Relation with relid: %u", view); + } + + duckdb::Parser parser; + parser.ParseQuery(view_definition); + auto &statements = parser.statements; + if (statements.size() != 1) { + throw duckdb::InvalidInputException("View definition contained more than 1 statement!"); + } + + if (statements[0]->type != duckdb::StatementType::SELECT_STATEMENT) { + throw duckdb::InvalidInputException("View definition (%s) did not contain a SELECT statement!", + view_definition); + } + + auto select = duckdb::unique_ptr_cast(std::move(statements[0])); + return duckdb::make_uniq(std::move(select)); +} + +duckdb::unique_ptr +PostgresViewScan(duckdb::ClientContext &context, duckdb::ReplacementScanInput &input, + duckdb::optional_ptr data) { + + auto &schema_name = input.schema_name; + auto &table_name = input.table_name; + + auto relid = PostgresFunctionGuard(FindMatchingRelation, schema_name, table_name); + if (relid == InvalidOid) { + return nullptr; + } + + auto tuple = PostgresFunctionGuard(SearchSysCache1, RELOID, ObjectIdGetDatum(relid)); + if (!HeapTupleIsValid(tuple)) { + elog(WARNING, "(PGDuckDB/PostgresReplacementScan) Cache lookup failed for relation %u", relid); + return nullptr; + } + + auto relForm = (Form_pg_class)GETSTRUCT(tuple); + if (relForm->relkind != RELKIND_VIEW) { + PostgresFunctionGuard(ReleaseSysCache, tuple); + return nullptr; + } + + PostgresFunctionGuard(ReleaseSysCache, tuple); + return ReplaceView(relid); +} + +} // namespace pgduckdb From 3bbeba110d2082ff51ddc7988045828ba3c8dc6d Mon Sep 17 00:00:00 2001 From: mkaruza Date: Wed, 4 Dec 2024 14:26:09 +0100 Subject: [PATCH 02/53] Simple comment for calculation of wanted number of workers --- src/scan/postgres_table_reader.cpp | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/scan/postgres_table_reader.cpp b/src/scan/postgres_table_reader.cpp index 7a1bb325..42511da6 100644 --- a/src/scan/postgres_table_reader.cpp +++ b/src/scan/postgres_table_reader.cpp @@ -97,6 +97,10 @@ PostgresTableReader::~PostgresTableReader() { FreeQueryDesc(table_scan_query_desc); } +/* + * This is simple calculation how much postgresql workers we will wanted to spawn for this + * postgres table scan. + */ int PostgresTableReader::ParalleWorkerNumber(Cardinality cardinality) { static const int base_log = 8; From 2a1d1896500e225657581ef1c9c367963bcfeaf2 Mon Sep 17 00:00:00 2001 From: mkaruza Date: Wed, 4 Dec 2024 14:32:33 +0100 Subject: [PATCH 03/53] Don't allow partitioned table scan for now --- src/pgduckdb_hooks.cpp | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/pgduckdb_hooks.cpp b/src/pgduckdb_hooks.cpp index c3144866..6936edbe 100644 --- a/src/pgduckdb_hooks.cpp +++ b/src/pgduckdb_hooks.cpp @@ -181,10 +181,10 @@ IsAllowedStatement(Query *query, bool throw_error = false) { /* * When accessing the partitioned table, we temporarily let PG handle it instead of DuckDB. */ - // if (ContainsPartitionedTable(query->rtable)) { - // elog(elevel, "DuckDB does not support querying PG partitioned table"); - // return false; - // } + if (ContainsPartitionedTable(query->rtable)) { + elog(elevel, "DuckDB does not support querying PG partitioned table"); + return false; + } /* Anything else is hopefully fine... */ return true; From ddef88741c528aebc6e79bbcd74072a4cb8054ca Mon Sep 17 00:00:00 2001 From: mkaruza Date: Fri, 6 Dec 2024 12:32:16 +0100 Subject: [PATCH 04/53] Added GlobalProcessLatch Added WaitLatch that can be used with multiple threads in process. --- include/pgduckdb/pgduckdb_process_latch.hpp | 37 +++++++++++++++ include/pgduckdb/scan/postgres_scan.hpp | 3 +- .../pgduckdb/scan/postgres_table_reader.hpp | 5 +- src/pgduckdb.cpp | 2 +- src/pgduckdb_duckdb.cpp | 1 - src/scan/postgres_scan.cpp | 12 ++--- src/scan/postgres_table_reader.cpp | 46 +++++++++---------- 7 files changed, 69 insertions(+), 37 deletions(-) create mode 100644 include/pgduckdb/pgduckdb_process_latch.hpp diff --git a/include/pgduckdb/pgduckdb_process_latch.hpp b/include/pgduckdb/pgduckdb_process_latch.hpp new file mode 100644 index 00000000..748c7689 --- /dev/null +++ b/include/pgduckdb/pgduckdb_process_latch.hpp @@ -0,0 +1,37 @@ +#pragma once + +extern "C" { +#include "postgres.h" +#include "miscadmin.h" +#include "storage/latch.h" +#include "utils/wait_event.h" +} + +#include +#include + +namespace pgduckdb { + +/* + * GlobalProcessLatch is used to wait on process postgres latch. First thread that enters will wait for latch while + * others will wait until this latch is released. + */ + +struct GlobalProcessLatch { +public: + static void WaitLatch() { + static std::condition_variable cv; + static std::mutex lock; + if (lock.try_lock()) { + ::WaitLatch(MyLatch, WL_LATCH_SET | WL_EXIT_ON_PM_DEATH, 0, PG_WAIT_EXTENSION); + ResetLatch(MyLatch); + cv.notify_all(); + lock.unlock(); + } else { + std::unique_lock lk(lock); + cv.wait(lk); + } + } +}; + +} // namespace pgduckdb diff --git a/include/pgduckdb/scan/postgres_scan.hpp b/include/pgduckdb/scan/postgres_scan.hpp index a16f5ab8..76067152 100644 --- a/include/pgduckdb/scan/postgres_scan.hpp +++ b/include/pgduckdb/scan/postgres_scan.hpp @@ -14,8 +14,7 @@ namespace pgduckdb { // Global State struct PostgresScanGlobalState : public duckdb::GlobalTableFunctionState { - explicit PostgresScanGlobalState(Snapshot, Relation rel, Cardinality cardinality, - duckdb::TableFunctionInitInput &input); + explicit PostgresScanGlobalState(Snapshot, Relation rel, duckdb::TableFunctionInitInput &input); ~PostgresScanGlobalState(); idx_t MaxThreads() const override { diff --git a/include/pgduckdb/scan/postgres_table_reader.hpp b/include/pgduckdb/scan/postgres_table_reader.hpp index 344f626f..a870f1b8 100644 --- a/include/pgduckdb/scan/postgres_table_reader.hpp +++ b/include/pgduckdb/scan/postgres_table_reader.hpp @@ -12,13 +12,13 @@ namespace pgduckdb { class PostgresTableReader { public: - PostgresTableReader(Cardinality cardinality, const char *table_scan_query); + PostgresTableReader(const char *table_scan_query); ~PostgresTableReader(); TupleTableSlot *GetNextTuple(); private: MinimalTuple GetNextWorkerTuple(); - int ParalleWorkerNumber(Cardinality cardinality); + int ParallelWorkerNumber(Cardinality cardinality); private: QueryDesc *table_scan_query_desc; @@ -29,6 +29,7 @@ class PostgresTableReader { int nworkers_launched; int nreaders; int next_parallel_reader; + bool entered_parallel_mode; }; } // namespace pgduckdb diff --git a/src/pgduckdb.cpp b/src/pgduckdb.cpp index a9aa97ae..dd4d0d6b 100644 --- a/src/pgduckdb.cpp +++ b/src/pgduckdb.cpp @@ -160,7 +160,7 @@ DuckdbInitGUC(void) { &duckdb_maximum_threads, -1, 1024, PGC_SUSET); DefineCustomVariable("duckdb.max_workers_per_postgres_scan", - "Maximum number of PostgreSQL threads used for a single Postgres scan", + "Maximum number of PostgreSQL workers used for a single Postgres scan", &max_workers_per_postgres_scan, 2, 8); DefineCustomVariable("duckdb.postgres_role", diff --git a/src/pgduckdb_duckdb.cpp b/src/pgduckdb_duckdb.cpp index c2a67c0a..622ca384 100644 --- a/src/pgduckdb_duckdb.cpp +++ b/src/pgduckdb_duckdb.cpp @@ -208,7 +208,6 @@ DuckDBManager::Initialize() { void DuckDBManager::LoadFunctions(duckdb::ClientContext &context) { - //auto &catalog = duckdb::Catalog::GetSystemCatalog(context); context.transaction.BeginTransaction(); auto &instance = *database->instance; duckdb::ExtensionUtil::RegisterType(instance, "UnsupportedPostgresType", duckdb::LogicalTypeId::VARCHAR); diff --git a/src/scan/postgres_scan.cpp b/src/scan/postgres_scan.cpp index a1947050..ea7f9c7d 100644 --- a/src/scan/postgres_scan.cpp +++ b/src/scan/postgres_scan.cpp @@ -118,10 +118,10 @@ PostgresScanGlobalState::ConstructTableScanQuery(duckdb::TableFunctionInitInput for (auto const &[attr_num, duckdb_scanned_index] : columns_to_scan) { auto filter = column_filters[duckdb_scanned_index]; if (filter) { - if (!first) { - scan_query << " AND "; - } else { + if (first) { scan_query << " WHERE "; + } else { + scan_query << " AND "; } first = false; scan_query << " ("; @@ -140,11 +140,11 @@ PostgresScanGlobalState::ConstructTableScanQuery(duckdb::TableFunctionInitInput // PostgresScanGlobalState // -PostgresScanGlobalState::PostgresScanGlobalState(Snapshot snapshot, Relation rel, Cardinality cardinality, duckdb::TableFunctionInitInput &input) +PostgresScanGlobalState::PostgresScanGlobalState(Snapshot snapshot, Relation rel, duckdb::TableFunctionInitInput &input) : snapshot(snapshot), rel(rel), table_tuple_desc(RelationGetDescr(rel)), count_tuples_only(false), total_row_count(0) { ConstructTableScanQuery(input); - table_reader_global_state = duckdb::make_shared_ptr(cardinality, scan_query.str().c_str()); + table_reader_global_state = duckdb::make_shared_ptr(scan_query.str().c_str()); pd_log(DEBUG2, "(DuckDB/PostgresSeqScanGlobalState) Running %" PRIu64 " threads -- ", (uint64_t)MaxThreads()); } @@ -192,7 +192,7 @@ duckdb::unique_ptr PostgresScanTableFunction::PostgresScanInitGlobal(duckdb::ClientContext &context, duckdb::TableFunctionInitInput &input) { auto &bind_data = input.bind_data->CastNoConst(); - auto global_state = duckdb::make_uniq(bind_data.snapshot, bind_data.rel, bind_data.cardinality, input); + auto global_state = duckdb::make_uniq(bind_data.snapshot, bind_data.rel, input); return std::move(global_state); } diff --git a/src/scan/postgres_table_reader.cpp b/src/scan/postgres_table_reader.cpp index 42511da6..5868791c 100644 --- a/src/scan/postgres_table_reader.cpp +++ b/src/scan/postgres_table_reader.cpp @@ -1,26 +1,19 @@ #include "pgduckdb/scan/postgres_table_reader.hpp" +#include "pgduckdb/pgduckdb_process_latch.hpp" #include "pgduckdb/pgduckdb_process_lock.hpp" extern "C" { #include "postgres.h" -#include "miscadmin.h" -#include "access/htup_details.h" -#include "catalog/namespace.h" -#include "catalog/pg_class.h" +#include "access/xact.h" +#include "executor/executor.h" +#include "executor/execParallel.h" +#include "executor/tqueue.h" #include "optimizer/planmain.h" #include "optimizer/planner.h" -#include "utils/builtins.h" -#include "utils/regproc.h" +#include "tcop/tcopprot.h" #include "utils/rel.h" #include "utils/snapmgr.h" -#include "utils/syscache.h" -#include "tcop/tcopprot.h" -#include "executor/execdesc.h" -#include "executor/executor.h" -#include "executor/execParallel.h" -#include "executor/tqueue.h" -#include "utils/wait_event.h" -#include "access/xact.h" + #include "pgduckdb/pgduckdb_guc.h" } @@ -28,14 +21,13 @@ extern "C" { namespace pgduckdb { -PostgresTableReader::PostgresTableReader(Cardinality cardinality, const char *table_scan_query) +PostgresTableReader::PostgresTableReader(const char *table_scan_query) : parallel_worker_readers(nullptr), nreaders(0), next_parallel_reader(0) { std::lock_guard lock(DuckdbProcessLock::GetLock()); pg_stack_base_t current_stack = set_stack_base(); - List *raw_parsetree_list; - raw_parsetree_list = pg_parse_query(table_scan_query); + List *raw_parsetree_list = pg_parse_query(table_scan_query); List *query_list = pg_analyze_and_rewrite_fixedparams(linitial_node(RawStmt, raw_parsetree_list), table_scan_query, NULL, 0, NULL); PlannedStmt *planned_stmt = standard_planner(linitial_node(Query, query_list), table_scan_query, 0, NULL); @@ -47,21 +39,22 @@ PostgresTableReader::PostgresTableReader(Cardinality cardinality, const char *ta table_scan_query_desc->planstate->plan->parallel_aware = true; - int parallel_workers = ParalleWorkerNumber(cardinality); - + int parallel_workers = ParallelWorkerNumber(planned_stmt->planTree->plan_rows); RESUME_CANCEL_INTERRUPTS(); + if (!IsInParallelMode()) { + EnterParallelMode(); + entered_parallel_mode = true; + } + ParallelContext *pcxt; pei = ExecInitParallelPlan(table_scan_planstate, table_scan_query_desc->estate, NULL, parallel_workers, -1); pcxt = pei->pcxt; LaunchParallelWorkers(pcxt); - /* We save # workers launched for the benefit of EXPLAIN */ nworkers_launched = pcxt->nworkers_launched; - /* Set up tuple queue readers to read the results. */ if (pcxt->nworkers_launched > 0) { ExecParallelCreateReaders(pei); - /* Make a working array showing the active readers */ nreaders = pcxt->nworkers_launched; parallel_worker_readers = (void **)palloc(nreaders * sizeof(TupleQueueReader *)); memcpy(parallel_worker_readers, pei->reader, nreaders * sizeof(TupleQueueReader *)); @@ -95,6 +88,10 @@ PostgresTableReader::~PostgresTableReader() { ExecutorFinish(table_scan_query_desc); ExecutorEnd(table_scan_query_desc); FreeQueryDesc(table_scan_query_desc); + + if (entered_parallel_mode) { + ExitParallelMode(); + } } /* @@ -102,7 +99,7 @@ PostgresTableReader::~PostgresTableReader() { * postgres table scan. */ int -PostgresTableReader::ParalleWorkerNumber(Cardinality cardinality) { +PostgresTableReader::ParallelWorkerNumber(Cardinality cardinality) { static const int base_log = 8; int cardinality_log = std::log2(cardinality); int base = cardinality_log / base_log; @@ -167,8 +164,7 @@ PostgresTableReader::GetNextWorkerTuple() { nvisited++; if (nvisited >= nreaders) { - (void)WaitLatch(MyLatch, WL_LATCH_SET | WL_EXIT_ON_PM_DEATH, 0, WAIT_EVENT_EXECUTE_GATHER); - ResetLatch(MyLatch); + GlobalProcessLatch::WaitLatch(); nvisited = 0; } } From 79a7a25c0f740ade134f6d1ed30817f60c29893b Mon Sep 17 00:00:00 2001 From: mkaruza Date: Fri, 6 Dec 2024 12:33:20 +0100 Subject: [PATCH 05/53] Renamed DuckdbProcessLock to GlobalProcessLock --- include/pgduckdb/logger.hpp | 20 ++++++++++---------- include/pgduckdb/pgduckdb_process_lock.hpp | 4 ++-- src/catalog/pgduckdb_table.cpp | 4 ++-- src/pgduckdb_detoast.cpp | 2 +- src/scan/postgres_table_reader.cpp | 6 +++--- 5 files changed, 18 insertions(+), 18 deletions(-) diff --git a/include/pgduckdb/logger.hpp b/include/pgduckdb/logger.hpp index 3ce6c5e0..e6a108fe 100644 --- a/include/pgduckdb/logger.hpp +++ b/include/pgduckdb/logger.hpp @@ -41,16 +41,16 @@ namespace pgduckdb { #define pd_prevent_errno_in_scope() #endif -#define pd_ereport_domain(elevel, domain, ...) \ - do { \ - pd_prevent_errno_in_scope(); \ - static_assert(elevel >= DEBUG5 && elevel <= WARNING_CLIENT_ONLY, "Invalid error level"); \ - if (message_level_is_interesting(elevel)) { \ - std::lock_guard lock(DuckdbProcessLock::GetLock()); \ - if (errstart(elevel, domain)) \ - __VA_ARGS__, errfinish(__FILE__, __LINE__, __func__); \ - } \ - } while (0) +#define pd_ereport_domain(elevel, domain, ...) \ + do { \ + pd_prevent_errno_in_scope(); \ + static_assert(elevel >= DEBUG5 && elevel <= WARNING_CLIENT_ONLY, "Invalid error level"); \ + if (message_level_is_interesting(elevel)) { \ + std::lock_guard lock(GlobalProcessLock::GetLock()); \ + if (errstart(elevel, domain)) \ + __VA_ARGS__, errfinish(__FILE__, __LINE__, __func__); \ + } \ + } while(0) #define TEXTDOMAIN NULL diff --git a/include/pgduckdb/pgduckdb_process_lock.hpp b/include/pgduckdb/pgduckdb_process_lock.hpp index 45947c98..dff227fe 100644 --- a/include/pgduckdb/pgduckdb_process_lock.hpp +++ b/include/pgduckdb/pgduckdb_process_lock.hpp @@ -5,11 +5,11 @@ namespace pgduckdb { /* - * DuckdbProcessLock is used to synchronize calls to PG functions that modify global variables. Examples + * GlobalProcessLock is used to synchronize calls to PG functions that modify global variables. Examples * for this synchronization are functions that read buffers/etc. This lock is shared between all threads and all * replacement scans. */ -struct DuckdbProcessLock { +struct GlobalProcessLock { public: static std::mutex & GetLock() { diff --git a/src/catalog/pgduckdb_table.cpp b/src/catalog/pgduckdb_table.cpp index 180fa51d..e204c83d 100644 --- a/src/catalog/pgduckdb_table.cpp +++ b/src/catalog/pgduckdb_table.cpp @@ -20,13 +20,13 @@ PostgresTable::PostgresTable(duckdb::Catalog &_catalog, duckdb::SchemaCatalogEnt } PostgresTable::~PostgresTable() { - std::lock_guard lock(DuckdbProcessLock::GetLock()); + std::lock_guard lock(GlobalProcessLock::GetLock()); CloseRelation(rel); } Relation PostgresTable::OpenRelation(Oid relid) { - std::lock_guard lock(DuckdbProcessLock::GetLock()); + std::lock_guard lock(GlobalProcessLock::GetLock()); return pgduckdb::OpenRelation(relid); } diff --git a/src/pgduckdb_detoast.cpp b/src/pgduckdb_detoast.cpp index 1597b8da..4f02c580 100644 --- a/src/pgduckdb_detoast.cpp +++ b/src/pgduckdb_detoast.cpp @@ -125,7 +125,7 @@ ToastFetchDatum(struct varlena *attr) { return result; } - std::lock_guard lock(DuckdbProcessLock::GetLock()); + std::lock_guard lock(GlobalProcessLock::GetLock()); if (!PostgresFunctionGuard(table_relation_fetch_toast_slice, toast_pointer, attrsize, result)) { duckdb_free(result); diff --git a/src/scan/postgres_table_reader.cpp b/src/scan/postgres_table_reader.cpp index 5868791c..4333e49e 100644 --- a/src/scan/postgres_table_reader.cpp +++ b/src/scan/postgres_table_reader.cpp @@ -24,7 +24,7 @@ namespace pgduckdb { PostgresTableReader::PostgresTableReader(const char *table_scan_query) : parallel_worker_readers(nullptr), nreaders(0), next_parallel_reader(0) { - std::lock_guard lock(DuckdbProcessLock::GetLock()); + std::lock_guard lock(GlobalProcessLock::GetLock()); pg_stack_base_t current_stack = set_stack_base(); List *raw_parsetree_list = pg_parse_query(table_scan_query); @@ -68,7 +68,7 @@ PostgresTableReader::PostgresTableReader(const char *table_scan_query) } PostgresTableReader::~PostgresTableReader() { - std::lock_guard lock(DuckdbProcessLock::GetLock()); + std::lock_guard lock(GlobalProcessLock::GetLock()); ExecEndNode(table_scan_planstate); @@ -116,7 +116,7 @@ PostgresTableReader::GetNextTuple() { return slot; } } else { - std::lock_guard lock(DuckdbProcessLock::GetLock()); + std::lock_guard lock(GlobalProcessLock::GetLock()); pg_stack_base_t current_stack = set_stack_base(); table_scan_query_desc->estate->es_query_dsa = pei ? pei->area : NULL; slot = ExecProcNode(table_scan_planstate); From d90aab576de361842fe7d4576311b7c91c3bb9dd Mon Sep 17 00:00:00 2001 From: mkaruza Date: Wed, 11 Dec 2024 13:45:56 +0100 Subject: [PATCH 06/53] Remove PostgresViewScan as it is not used anymore --- include/pgduckdb/scan/postgres_view.hpp | 13 ---- src/pgduckdb_duckdb.cpp | 4 +- src/scan/postgres_table_reader.cpp | 2 +- src/scan/postgres_view.cpp | 90 ------------------------- 4 files changed, 2 insertions(+), 107 deletions(-) delete mode 100644 include/pgduckdb/scan/postgres_view.hpp delete mode 100644 src/scan/postgres_view.cpp diff --git a/include/pgduckdb/scan/postgres_view.hpp b/include/pgduckdb/scan/postgres_view.hpp deleted file mode 100644 index 78cfec56..00000000 --- a/include/pgduckdb/scan/postgres_view.hpp +++ /dev/null @@ -1,13 +0,0 @@ -#pragma once - -#include "duckdb.hpp" - -#include "pgduckdb/utility/cpp_only_file.hpp" // Must be last include. - -namespace pgduckdb { - -duckdb::unique_ptr PostgresViewScan(duckdb::ClientContext &context, - duckdb::ReplacementScanInput &input, - duckdb::optional_ptr data); - -} diff --git a/src/pgduckdb_duckdb.cpp b/src/pgduckdb_duckdb.cpp index 622ca384..f80eadef 100644 --- a/src/pgduckdb_duckdb.cpp +++ b/src/pgduckdb_duckdb.cpp @@ -8,7 +8,6 @@ #include "pgduckdb/catalog/pgduckdb_storage.hpp" #include "pgduckdb/scan/postgres_scan.hpp" -#include "pgduckdb/scan/postgres_view.hpp" #include "pgduckdb/pg/transactions.hpp" #include "pgduckdb/pgduckdb_utils.hpp" @@ -135,8 +134,7 @@ DuckDBManager::Initialize() { duckdb::DBConfig config; config.SetOptionByName("custom_user_agent", "pg_duckdb"); config.SetOptionByName("extension_directory", CreateOrGetDirectoryPath("duckdb_extensions")); - // Transforms VIEWs into their view definition - config.replacement_scans.emplace_back(pgduckdb::PostgresViewScan); + SET_DUCKDB_OPTION(allow_unsigned_extensions); SET_DUCKDB_OPTION(enable_external_access); SET_DUCKDB_OPTION(autoinstall_known_extensions); diff --git a/src/scan/postgres_table_reader.cpp b/src/scan/postgres_table_reader.cpp index 4333e49e..8ad951d1 100644 --- a/src/scan/postgres_table_reader.cpp +++ b/src/scan/postgres_table_reader.cpp @@ -103,7 +103,7 @@ PostgresTableReader::ParallelWorkerNumber(Cardinality cardinality) { static const int base_log = 8; int cardinality_log = std::log2(cardinality); int base = cardinality_log / base_log; - return std::max(1, std::min(base, std::min(max_workers_per_postgres_scan, max_parallel_workers))); + return std::max(1, std::min(base, std::max(max_workers_per_postgres_scan, max_parallel_workers))); } TupleTableSlot * diff --git a/src/scan/postgres_view.cpp b/src/scan/postgres_view.cpp deleted file mode 100644 index bf655560..00000000 --- a/src/scan/postgres_view.cpp +++ /dev/null @@ -1,90 +0,0 @@ -#include "duckdb/parser/parser.hpp" -#include "duckdb/parser/tableref/subqueryref.hpp" -#include "pgduckdb/pgduckdb_utils.hpp" - -extern "C" { -#include "postgres.h" -#include "access/htup_details.h" -#include "catalog/namespace.h" -#include "catalog/pg_class.h" -#include "optimizer/planmain.h" -#include "utils/builtins.h" -#include "utils/regproc.h" -#include "utils/syscache.h" -} - -namespace pgduckdb { - -static Oid -FindMatchingRelation(const duckdb::string &schema, const duckdb::string &table) { - List *name_list = NIL; - if (!schema.empty()) { - name_list = lappend(name_list, makeString(pstrdup(schema.c_str()))); - } - - name_list = lappend(name_list, makeString(pstrdup(table.c_str()))); - - RangeVar *table_range_var = makeRangeVarFromNameList(name_list); - return RangeVarGetRelid(table_range_var, AccessShareLock, true); -} - -const char * -pgduckdb_pg_get_viewdef(Oid view) { - auto oid = ObjectIdGetDatum(view); - Datum viewdef = DirectFunctionCall1(pg_get_viewdef, oid); - return text_to_cstring(DatumGetTextP(viewdef)); -} - -duckdb::unique_ptr -ReplaceView(Oid view) { - const auto view_definition = PostgresFunctionGuard(pgduckdb_pg_get_viewdef, view); - - if (!view_definition) { - throw duckdb::InvalidInputException("Could not retrieve view definition for Relation with relid: %u", view); - } - - duckdb::Parser parser; - parser.ParseQuery(view_definition); - auto &statements = parser.statements; - if (statements.size() != 1) { - throw duckdb::InvalidInputException("View definition contained more than 1 statement!"); - } - - if (statements[0]->type != duckdb::StatementType::SELECT_STATEMENT) { - throw duckdb::InvalidInputException("View definition (%s) did not contain a SELECT statement!", - view_definition); - } - - auto select = duckdb::unique_ptr_cast(std::move(statements[0])); - return duckdb::make_uniq(std::move(select)); -} - -duckdb::unique_ptr -PostgresViewScan(duckdb::ClientContext &context, duckdb::ReplacementScanInput &input, - duckdb::optional_ptr data) { - - auto &schema_name = input.schema_name; - auto &table_name = input.table_name; - - auto relid = PostgresFunctionGuard(FindMatchingRelation, schema_name, table_name); - if (relid == InvalidOid) { - return nullptr; - } - - auto tuple = PostgresFunctionGuard(SearchSysCache1, RELOID, ObjectIdGetDatum(relid)); - if (!HeapTupleIsValid(tuple)) { - elog(WARNING, "(PGDuckDB/PostgresReplacementScan) Cache lookup failed for relation %u", relid); - return nullptr; - } - - auto relForm = (Form_pg_class)GETSTRUCT(tuple); - if (relForm->relkind != RELKIND_VIEW) { - PostgresFunctionGuard(ReleaseSysCache, tuple); - return nullptr; - } - - PostgresFunctionGuard(ReleaseSysCache, tuple); - return ReplaceView(relid); -} - -} // namespace pgduckdb From 86e32060aecdd8f6fd6a37c9e022edb0f140d6cc Mon Sep 17 00:00:00 2001 From: mkaruza Date: Fri, 13 Dec 2024 10:33:39 +0100 Subject: [PATCH 07/53] Protect pg function with guard. Handle TEMP tables. --- include/pgduckdb/pgduckdb_process_latch.hpp | 27 ++-- include/pgduckdb/pgduckdb_utils.hpp | 16 ++ include/pgduckdb/scan/postgres_scan.hpp | 3 + .../pgduckdb/scan/postgres_table_reader.hpp | 2 +- include/pgduckdb/utility/rename_ruleutils.h | 1 + src/pgduckdb_ruleutils.cpp | 1 + src/pgduckdb_types.cpp | 6 +- src/scan/postgres_scan.cpp | 27 +++- src/scan/postgres_table_reader.cpp | 137 ++++++++++-------- test/regression/expected/duckdb_recycle.out | 8 +- 10 files changed, 145 insertions(+), 83 deletions(-) diff --git a/include/pgduckdb/pgduckdb_process_latch.hpp b/include/pgduckdb/pgduckdb_process_latch.hpp index 748c7689..f5b07a73 100644 --- a/include/pgduckdb/pgduckdb_process_latch.hpp +++ b/include/pgduckdb/pgduckdb_process_latch.hpp @@ -1,15 +1,24 @@ #pragma once -extern "C" { -#include "postgres.h" -#include "miscadmin.h" -#include "storage/latch.h" -#include "utils/wait_event.h" -} - #include #include +extern "C" { +struct Latch; +extern struct Latch *MyLatch; +extern int WaitLatch(Latch *latch, int wakeEvents, long timeout, + uint32_t wait_event_info); +extern void ResetLatch(Latch *latch); + +/* Defined in storage/latch.h */ +#define WL_LATCH_SET (1 << 0) +#define WL_EXIT_ON_PM_DEATH (1 << 5) + +/* Defined in utils/wait_event.h */ +#define PG_WAIT_EXTENSION 0x07000000U + +} + namespace pgduckdb { /* @@ -19,11 +28,11 @@ namespace pgduckdb { struct GlobalProcessLatch { public: - static void WaitLatch() { + static void WaitGlobalLatch() { static std::condition_variable cv; static std::mutex lock; if (lock.try_lock()) { - ::WaitLatch(MyLatch, WL_LATCH_SET | WL_EXIT_ON_PM_DEATH, 0, PG_WAIT_EXTENSION); + WaitLatch(MyLatch, WL_LATCH_SET | WL_EXIT_ON_PM_DEATH, 0, PG_WAIT_EXTENSION); ResetLatch(MyLatch); cv.notify_all(); lock.unlock(); diff --git a/include/pgduckdb/pgduckdb_utils.hpp b/include/pgduckdb/pgduckdb_utils.hpp index 9787707e..4925ab93 100644 --- a/include/pgduckdb/pgduckdb_utils.hpp +++ b/include/pgduckdb/pgduckdb_utils.hpp @@ -17,12 +17,15 @@ struct ErrorContextCallback; struct MemoryContextData; typedef struct MemoryContextData *MemoryContext; +typedef char *pg_stack_base_t; extern sigjmp_buf *PG_exception_stack; extern MemoryContext CurrentMemoryContext; extern ErrorContextCallback *error_context_stack; extern ErrorData *CopyErrorData(); extern void FlushErrorState(); +extern pg_stack_base_t set_stack_base(); +extern void restore_stack_base(pg_stack_base_t base); } namespace pgduckdb { @@ -45,6 +48,19 @@ struct PgExceptionGuard { ErrorContextCallback *_save_context_stack; }; +/* + * DuckdbGlobalLock should be held before calling. + */ +struct PostgresScopedStackReset { + PostgresScopedStackReset() { + saved_current_stack = set_stack_base(); + } + ~PostgresScopedStackReset() { + restore_stack_base(saved_current_stack); + } + pg_stack_base_t saved_current_stack; +}; + /* * DuckdbGlobalLock should be held before calling. */ diff --git a/include/pgduckdb/scan/postgres_scan.hpp b/include/pgduckdb/scan/postgres_scan.hpp index 76067152..56e27635 100644 --- a/include/pgduckdb/scan/postgres_scan.hpp +++ b/include/pgduckdb/scan/postgres_scan.hpp @@ -22,6 +22,9 @@ struct PostgresScanGlobalState : public duckdb::GlobalTableFunctionState { } void ConstructTableScanQuery(duckdb::TableFunctionInitInput &input); +private: + std::string ConstructFullyQualifiedTableName(Relation rel); + public: Snapshot snapshot; Relation rel; diff --git a/include/pgduckdb/scan/postgres_table_reader.hpp b/include/pgduckdb/scan/postgres_table_reader.hpp index a870f1b8..a142ab2b 100644 --- a/include/pgduckdb/scan/postgres_table_reader.hpp +++ b/include/pgduckdb/scan/postgres_table_reader.hpp @@ -23,7 +23,7 @@ class PostgresTableReader { private: QueryDesc *table_scan_query_desc; PlanState *table_scan_planstate; - ParallelExecutorInfo *pei; + ParallelExecutorInfo *parallel_executor_info; void **parallel_worker_readers; TupleTableSlot *slot; int nworkers_launched; diff --git a/include/pgduckdb/utility/rename_ruleutils.h b/include/pgduckdb/utility/rename_ruleutils.h index d5cbf0e3..1ff77a8c 100644 --- a/include/pgduckdb/utility/rename_ruleutils.h +++ b/include/pgduckdb/utility/rename_ruleutils.h @@ -23,6 +23,7 @@ #define pg_get_statisticsobjdef_string pgduckdb_pg_get_statisticsobjdef_string #define get_list_partvalue_string pgduckdb_get_list_partvalue_string + /* * The following replaces all usages of generate_qualified_relation_name and * generate_relation_name with calls to the pgduckdb_relation_name function diff --git a/src/pgduckdb_ruleutils.cpp b/src/pgduckdb_ruleutils.cpp index 4622b061..5624b586 100644 --- a/src/pgduckdb_ruleutils.cpp +++ b/src/pgduckdb_ruleutils.cpp @@ -153,6 +153,7 @@ pgduckdb_relation_name(Oid relation_oid) { } const char *db_and_schema = pgduckdb_db_and_schema_string(postgres_schema_name, is_duckdb_table); + char *result = psprintf("%s.%s", db_and_schema, quote_identifier(relname)); ReleaseSysCache(tp); diff --git a/src/pgduckdb_types.cpp b/src/pgduckdb_types.cpp index 3fe1dc29..e1f144dc 100644 --- a/src/pgduckdb_types.cpp +++ b/src/pgduckdb_types.cpp @@ -1360,14 +1360,14 @@ InsertTupleIntoChunk(duckdb::DataChunk &output, PostgresScanLocalState &scan_loc } /* Write tuple columns in output vector. */ int duckdb_output_index = 0; - for (auto const &[duckdb_scanned_index, attr_num] : scan_global_state->output_columns) { + for (auto const &[_, attr_num] : scan_global_state->output_columns) { auto &result = output.data[duckdb_output_index]; - if (slot->tts_isnull[duckdb_scanned_index]) { + if (slot->tts_isnull[duckdb_output_index]) { auto &array_mask = duckdb::FlatVector::Validity(result); array_mask.SetInvalid(scan_local_state.output_vector_size); } else { auto attr = slot->tts_tupleDescriptor->attrs[attr_num - 1]; - ConvertPostgresToDuckValue(attr.atttypid, slot->tts_values[duckdb_scanned_index], result, + ConvertPostgresToDuckValue(attr.atttypid, slot->tts_values[duckdb_output_index], result, scan_local_state.output_vector_size); } duckdb_output_index++; diff --git a/src/scan/postgres_scan.cpp b/src/scan/postgres_scan.cpp index ea7f9c7d..0008b46d 100644 --- a/src/scan/postgres_scan.cpp +++ b/src/scan/postgres_scan.cpp @@ -30,6 +30,7 @@ extern "C" { #include "optimizer/planmain.h" #include "optimizer/planner.h" #include "utils/builtins.h" +#include "utils/lsyscache.h" #include "utils/regproc.h" #include "utils/rel.h" #include "utils/snapmgr.h" @@ -37,17 +38,24 @@ extern "C" { #include "tcop/tcopprot.h" #include "executor/execdesc.h" #include "executor/executor.h" + +#include "pgduckdb/vendor/pg_ruleutils.h" +#include "pgduckdb/pgduckdb_ruleutils.h" } #include "pgduckdb/pgduckdb_process_lock.hpp" namespace pgduckdb { +// +// PostgresScanGlobalState +// + void PostgresScanGlobalState::ConstructTableScanQuery(duckdb::TableFunctionInitInput &input) { /* SELECT COUNT(*) FROM */ if (input.column_ids.size() == 1 && input.column_ids[0] == UINT64_MAX) { - scan_query << "SELECT COUNT(*) FROM " << rel->rd_rel->relname.data; + scan_query << "SELECT COUNT(*) FROM " << ConstructFullyQualifiedTableName(rel); count_tuples_only = true; return; } @@ -108,10 +116,10 @@ PostgresScanGlobalState::ConstructTableScanQuery(duckdb::TableFunctionInitInput } first = false; auto attr = table_tuple_desc->attrs[attr_num - 1]; - scan_query << attr.attname.data; + scan_query << quote_identifier(attr.attname.data); } - scan_query << " FROM " << rel->rd_rel->relname.data; + scan_query << " FROM " << ConstructFullyQualifiedTableName(rel); first = true; @@ -126,7 +134,7 @@ PostgresScanGlobalState::ConstructTableScanQuery(duckdb::TableFunctionInitInput first = false; scan_query << " ("; auto attr = table_tuple_desc->attrs[attr_num - 1]; - scan_query << filter->ToString(attr.attname.data); + scan_query << filter->ToString(attr.attname.data).c_str(); scan_query << ") "; } } @@ -136,9 +144,11 @@ PostgresScanGlobalState::ConstructTableScanQuery(duckdb::TableFunctionInitInput elog(DEBUG1, "scan_query: %s", scan_query.str().c_str()); } -// -// PostgresScanGlobalState -// +std::string +PostgresScanGlobalState::ConstructFullyQualifiedTableName(Relation rel) { + return psprintf("%s.%s", quote_identifier(get_namespace_name_or_temp(get_rel_namespace(rel->rd_rel->oid))), + quote_identifier(get_rel_name(rel->rd_rel->oid))); +} PostgresScanGlobalState::PostgresScanGlobalState(Snapshot snapshot, Relation rel, duckdb::TableFunctionInitInput &input) : snapshot(snapshot), rel(rel), table_tuple_desc(RelationGetDescr(rel)), count_tuples_only(false), @@ -232,7 +242,8 @@ PostgresScanTableFunction::PostgresScanFunction(duckdb::ClientContext &context, local_state.exhausted_scan = true; } - if (local_state.global_state->count_tuples_only) { + /* Special case when we only use COUNT(*). We only update output capactity. */ + if (local_state.global_state->count_tuples_only && local_state.output_vector_size < STANDARD_VECTOR_SIZE) { output.SetCapacity(local_state.output_vector_size); } diff --git a/src/scan/postgres_table_reader.cpp b/src/scan/postgres_table_reader.cpp index 8ad951d1..3a42db33 100644 --- a/src/scan/postgres_table_reader.cpp +++ b/src/scan/postgres_table_reader.cpp @@ -1,9 +1,11 @@ #include "pgduckdb/scan/postgres_table_reader.hpp" #include "pgduckdb/pgduckdb_process_latch.hpp" #include "pgduckdb/pgduckdb_process_lock.hpp" +#include "pgduckdb/pgduckdb_utils.hpp" extern "C" { #include "postgres.h" +#include "miscadmin.h" #include "access/xact.h" #include "executor/executor.h" #include "executor/execParallel.h" @@ -13,6 +15,7 @@ extern "C" { #include "tcop/tcopprot.h" #include "utils/rel.h" #include "utils/snapmgr.h" +#include "utils/lsyscache.h" #include "pgduckdb/pgduckdb_guc.h" } @@ -22,82 +25,100 @@ extern "C" { namespace pgduckdb { PostgresTableReader::PostgresTableReader(const char *table_scan_query) - : parallel_worker_readers(nullptr), nreaders(0), next_parallel_reader(0) { - + : parallel_executor_info(nullptr), parallel_worker_readers(nullptr), nreaders(0), next_parallel_reader(0), + entered_parallel_mode(false) { std::lock_guard lock(GlobalProcessLock::GetLock()); + PostgresScopedStackReset scoped_stack_reset; + + List *raw_parsetree_list = PostgresFunctionGuard(pg_parse_query, table_scan_query); + Assert(list_length(raw_parsetree_list) == 1); + RawStmt *raw_parsetree = linitial_node(RawStmt, raw_parsetree_list); - pg_stack_base_t current_stack = set_stack_base(); - List *raw_parsetree_list = pg_parse_query(table_scan_query); List *query_list = - pg_analyze_and_rewrite_fixedparams(linitial_node(RawStmt, raw_parsetree_list), table_scan_query, NULL, 0, NULL); - PlannedStmt *planned_stmt = standard_planner(linitial_node(Query, query_list), table_scan_query, 0, NULL); - table_scan_query_desc = CreateQueryDesc(planned_stmt, table_scan_query, GetActiveSnapshot(), InvalidSnapshot, - None_Receiver, NULL, NULL, 0); - ExecutorStart(table_scan_query_desc, 0); + PostgresFunctionGuard(pg_analyze_and_rewrite_fixedparams, raw_parsetree, table_scan_query, nullptr, 0, nullptr); - table_scan_planstate = ExecInitNode(planned_stmt->planTree, table_scan_query_desc->estate, 0); + Assert(list_length(query_list) == 1); + Query *query = linitial_node(Query, query_list); - table_scan_query_desc->planstate->plan->parallel_aware = true; + Assert(list_length(query->rtable) == 1); + RangeTblEntry *rte = linitial_node(RangeTblEntry, query->rtable); - int parallel_workers = ParallelWorkerNumber(planned_stmt->planTree->plan_rows); - RESUME_CANCEL_INTERRUPTS(); + char persistence = get_rel_persistence(rte->relid); - if (!IsInParallelMode()) { - EnterParallelMode(); - entered_parallel_mode = true; - } + PlannedStmt *planned_stmt = PostgresFunctionGuard(standard_planner, query, table_scan_query, 0, nullptr); - ParallelContext *pcxt; - pei = ExecInitParallelPlan(table_scan_planstate, table_scan_query_desc->estate, NULL, parallel_workers, -1); - pcxt = pei->pcxt; - LaunchParallelWorkers(pcxt); - nworkers_launched = pcxt->nworkers_launched; - - if (pcxt->nworkers_launched > 0) { - ExecParallelCreateReaders(pei); - nreaders = pcxt->nworkers_launched; - parallel_worker_readers = (void **)palloc(nreaders * sizeof(TupleQueueReader *)); - memcpy(parallel_worker_readers, pei->reader, nreaders * sizeof(TupleQueueReader *)); - } + table_scan_query_desc = PostgresFunctionGuard(CreateQueryDesc, planned_stmt, table_scan_query, GetActiveSnapshot(), + InvalidSnapshot, None_Receiver, nullptr, nullptr, 0); + + PostgresFunctionGuard(ExecutorStart, table_scan_query_desc, 0); + + table_scan_planstate = + PostgresFunctionGuard(ExecInitNode, planned_stmt->planTree, table_scan_query_desc->estate, 0); + + /* Temp tables can be excuted with parallel workers */ + if (persistence != RELPERSISTENCE_TEMP) { + + table_scan_query_desc->planstate->plan->parallel_aware = true; + + int parallel_workers = ParallelWorkerNumber(planned_stmt->planTree->plan_rows); + + RESUME_CANCEL_INTERRUPTS(); - HOLD_CANCEL_INTERRUPTS(); + if (!IsInParallelMode()) { + EnterParallelMode(); + entered_parallel_mode = true; + } - slot = ExecInitExtraTupleSlot(table_scan_query_desc->estate, table_scan_planstate->ps_ResultTupleDesc, &TTSOpsMinimalTuple); + ParallelContext *pcxt; + parallel_executor_info = PostgresFunctionGuard(ExecInitParallelPlan, table_scan_planstate, + table_scan_query_desc->estate, nullptr, parallel_workers, -1); + pcxt = parallel_executor_info->pcxt; + PostgresFunctionGuard(LaunchParallelWorkers, pcxt); + nworkers_launched = pcxt->nworkers_launched; + + if (pcxt->nworkers_launched > 0) { + PostgresFunctionGuard(ExecParallelCreateReaders, parallel_executor_info); + nreaders = pcxt->nworkers_launched; + parallel_worker_readers = (void **)palloc(nreaders * sizeof(TupleQueueReader *)); + memcpy(parallel_worker_readers, parallel_executor_info->reader, nreaders * sizeof(TupleQueueReader *)); + } - restore_stack_base(current_stack); + HOLD_CANCEL_INTERRUPTS(); + } + + slot = PostgresFunctionGuard(ExecInitExtraTupleSlot, table_scan_query_desc->estate, + table_scan_planstate->ps_ResultTupleDesc, &TTSOpsMinimalTuple); } PostgresTableReader::~PostgresTableReader() { std::lock_guard lock(GlobalProcessLock::GetLock()); - ExecEndNode(table_scan_planstate); + PostgresScopedStackReset scoped_stack_reset; - if (pei != NULL) - ExecParallelFinish(pei); + PostgresFunctionGuard(ExecEndNode, table_scan_planstate); - if (parallel_worker_readers) { - pfree(parallel_worker_readers); + if (parallel_executor_info != NULL) { + PostgresFunctionGuard(ExecParallelFinish, parallel_executor_info); + PostgresFunctionGuard(ExecParallelCleanup, parallel_executor_info); } - parallel_worker_readers = NULL; + parallel_executor_info = NULL; - ExecParallelCleanup(pei); + if (parallel_worker_readers) { + PostgresFunctionGuard(pfree, parallel_worker_readers); + } - pei = NULL; + parallel_worker_readers = NULL; - ExecutorFinish(table_scan_query_desc); - ExecutorEnd(table_scan_query_desc); - FreeQueryDesc(table_scan_query_desc); + PostgresFunctionGuard(ExecutorFinish, table_scan_query_desc); + PostgresFunctionGuard(ExecutorEnd, table_scan_query_desc); + PostgresFunctionGuard(FreeQueryDesc, table_scan_query_desc); if (entered_parallel_mode) { ExitParallelMode(); } } -/* - * This is simple calculation how much postgresql workers we will wanted to spawn for this - * postgres table scan. - */ int PostgresTableReader::ParallelWorkerNumber(Cardinality cardinality) { static const int base_log = 8; @@ -109,25 +130,25 @@ PostgresTableReader::ParallelWorkerNumber(Cardinality cardinality) { TupleTableSlot * PostgresTableReader::GetNextTuple() { MinimalTuple worker_minmal_tuple; + TupleTableSlot *thread_scan_slot; if (nreaders > 0) { worker_minmal_tuple = GetNextWorkerTuple(); if (HeapTupleIsValid(worker_minmal_tuple)) { - ExecStoreMinimalTuple(worker_minmal_tuple, slot, false); + PostgresFunctionGuard(ExecStoreMinimalTuple, worker_minmal_tuple, slot, false); return slot; } } else { std::lock_guard lock(GlobalProcessLock::GetLock()); - pg_stack_base_t current_stack = set_stack_base(); - table_scan_query_desc->estate->es_query_dsa = pei ? pei->area : NULL; - slot = ExecProcNode(table_scan_planstate); - restore_stack_base(current_stack); + PostgresScopedStackReset scoped_stack_reset; + table_scan_query_desc->estate->es_query_dsa = parallel_executor_info ? parallel_executor_info->area : NULL; + thread_scan_slot = ExecProcNode(table_scan_planstate); table_scan_query_desc->estate->es_query_dsa = NULL; - if (!TupIsNull(slot)) { - return slot; + if (!TupIsNull(thread_scan_slot)) { + return thread_scan_slot; } } - return ExecClearTuple(slot); + return PostgresFunctionGuard(ExecClearTuple, slot); } MinimalTuple @@ -138,8 +159,8 @@ PostgresTableReader::GetNextWorkerTuple() { MinimalTuple minimal_tuple; bool readerdone; - reader = (TupleQueueReader *) parallel_worker_readers[next_parallel_reader]; - minimal_tuple = TupleQueueReaderNext(reader, true, &readerdone); + reader = (TupleQueueReader *)parallel_worker_readers[next_parallel_reader]; + minimal_tuple = PostgresFunctionGuard(TupleQueueReaderNext, reader, true, &readerdone); if (readerdone) { --nreaders; @@ -164,7 +185,7 @@ PostgresTableReader::GetNextWorkerTuple() { nvisited++; if (nvisited >= nreaders) { - GlobalProcessLatch::WaitLatch(); + GlobalProcessLatch::WaitGlobalLatch(); nvisited = 0; } } diff --git a/test/regression/expected/duckdb_recycle.out b/test/regression/expected/duckdb_recycle.out index fb9628a6..0a8521f1 100644 --- a/test/regression/expected/duckdb_recycle.out +++ b/test/regression/expected/duckdb_recycle.out @@ -13,10 +13,10 @@ EXPLAIN SELECT count(*) FROM ta; │ count_star() │ └─────────────┬─────────────┘ ┌─────────────┴─────────────┐ - │ POSTGRES_SEQ_SCAN │ + │ POSTGRES_SCAN │ │ ──────────────────── │ │ Function: │ - │ POSTGRES_SEQ_SCAN │ + │ POSTGRES_SCAN │ │ │ │ ~2550 Rows │ └───────────────────────────┘ @@ -38,10 +38,10 @@ EXPLAIN SELECT count(*) FROM ta; │ count_star() │ └─────────────┬─────────────┘ ┌─────────────┴─────────────┐ - │ POSTGRES_SEQ_SCAN │ + │ POSTGRES_SCAN │ │ ──────────────────── │ │ Function: │ - │ POSTGRES_SEQ_SCAN │ + │ POSTGRES_SCAN │ │ │ │ ~2550 Rows │ └───────────────────────────┘ From 41c9aea7dea9c1e80502f1b625549b63fad522d5 Mon Sep 17 00:00:00 2001 From: mkaruza Date: Fri, 13 Dec 2024 10:45:33 +0100 Subject: [PATCH 08/53] Keep compiler quite for unused variables --- include/pgduckdb/scan/postgres_scan.hpp | 2 +- src/catalog/pgduckdb_table.cpp | 9 +++++---- src/scan/postgres_scan.cpp | 21 +++++++++++---------- 3 files changed, 17 insertions(+), 15 deletions(-) diff --git a/include/pgduckdb/scan/postgres_scan.hpp b/include/pgduckdb/scan/postgres_scan.hpp index 56e27635..7c06ce4b 100644 --- a/include/pgduckdb/scan/postgres_scan.hpp +++ b/include/pgduckdb/scan/postgres_scan.hpp @@ -23,7 +23,7 @@ struct PostgresScanGlobalState : public duckdb::GlobalTableFunctionState { void ConstructTableScanQuery(duckdb::TableFunctionInitInput &input); private: - std::string ConstructFullyQualifiedTableName(Relation rel); + std::string ConstructFullyQualifiedTableName(); public: Snapshot snapshot; diff --git a/src/catalog/pgduckdb_table.cpp b/src/catalog/pgduckdb_table.cpp index e204c83d..6c348699 100644 --- a/src/catalog/pgduckdb_table.cpp +++ b/src/catalog/pgduckdb_table.cpp @@ -56,19 +56,20 @@ PostgresTable::GetTableCardinality(Relation rel) { } duckdb::unique_ptr -PostgresTable::GetStatistics(duckdb::ClientContext &context, duckdb::column_t column_id) { +PostgresTable::GetStatistics(__attribute__((unused)) duckdb::ClientContext &context, + __attribute__((unused)) duckdb::column_t column_id) { throw duckdb::NotImplementedException("GetStatistics not supported yet"); } duckdb::TableFunction -PostgresTable::GetScanFunction(duckdb::ClientContext &context, - duckdb::unique_ptr &bind_data) { +PostgresTable::GetScanFunction(__attribute__((unused)) duckdb::ClientContext &context, + duckdb::unique_ptr &bind_data) { bind_data = duckdb::make_uniq(rel, cardinality, snapshot); return PostgresScanTableFunction(); } duckdb::TableStorageInfo -PostgresTable::GetStorageInfo(duckdb::ClientContext &context) { +PostgresTable::GetStorageInfo(__attribute__((unused)) duckdb::ClientContext &context) { throw duckdb::NotImplementedException("GetStorageInfo not supported yet"); } diff --git a/src/scan/postgres_scan.cpp b/src/scan/postgres_scan.cpp index 0008b46d..0f666dde 100644 --- a/src/scan/postgres_scan.cpp +++ b/src/scan/postgres_scan.cpp @@ -55,7 +55,7 @@ void PostgresScanGlobalState::ConstructTableScanQuery(duckdb::TableFunctionInitInput &input) { /* SELECT COUNT(*) FROM */ if (input.column_ids.size() == 1 && input.column_ids[0] == UINT64_MAX) { - scan_query << "SELECT COUNT(*) FROM " << ConstructFullyQualifiedTableName(rel); + scan_query << "SELECT COUNT(*) FROM " << ConstructFullyQualifiedTableName(); count_tuples_only = true; return; } @@ -119,7 +119,7 @@ PostgresScanGlobalState::ConstructTableScanQuery(duckdb::TableFunctionInitInput scan_query << quote_identifier(attr.attname.data); } - scan_query << " FROM " << ConstructFullyQualifiedTableName(rel); + scan_query << " FROM " << ConstructFullyQualifiedTableName(); first = true; @@ -145,7 +145,7 @@ PostgresScanGlobalState::ConstructTableScanQuery(duckdb::TableFunctionInitInput } std::string -PostgresScanGlobalState::ConstructFullyQualifiedTableName(Relation rel) { +PostgresScanGlobalState::ConstructFullyQualifiedTableName() { return psprintf("%s.%s", quote_identifier(get_namespace_name_or_temp(get_rel_namespace(rel->rd_rel->oid))), quote_identifier(get_rel_name(rel->rd_rel->oid))); } @@ -199,24 +199,24 @@ PostgresScanTableFunction::PostgresScanTableFunction() } duckdb::unique_ptr -PostgresScanTableFunction::PostgresScanInitGlobal(duckdb::ClientContext &context, +PostgresScanTableFunction::PostgresScanInitGlobal(__attribute__((unused)) duckdb::ClientContext &context, duckdb::TableFunctionInitInput &input) { auto &bind_data = input.bind_data->CastNoConst(); auto global_state = duckdb::make_uniq(bind_data.snapshot, bind_data.rel, input); - return std::move(global_state); + return global_state; } duckdb::unique_ptr -PostgresScanTableFunction::PostgresScanInitLocal(duckdb::ExecutionContext &context, - duckdb::TableFunctionInitInput &input, +PostgresScanTableFunction::PostgresScanInitLocal(__attribute__((unused)) duckdb::ExecutionContext &context, + __attribute__((unused)) duckdb::TableFunctionInitInput &input, duckdb::GlobalTableFunctionState *gstate) { auto global_state = reinterpret_cast(gstate); return duckdb::make_uniq(global_state); } void -PostgresScanTableFunction::PostgresScanFunction(duckdb::ClientContext &context, duckdb::TableFunctionInput &data, - duckdb::DataChunk &output) { +PostgresScanTableFunction::PostgresScanFunction(__attribute__((unused)) duckdb::ClientContext &context, + duckdb::TableFunctionInput &data, duckdb::DataChunk &output) { auto &local_state = data.local_state->Cast(); local_state.output_vector_size = 0; @@ -251,7 +251,8 @@ PostgresScanTableFunction::PostgresScanFunction(duckdb::ClientContext &context, } duckdb::unique_ptr -PostgresScanTableFunction::PostgresScanCardinality(duckdb::ClientContext &context, const duckdb::FunctionData *data) { +PostgresScanTableFunction::PostgresScanCardinality(__attribute__((unused)) duckdb::ClientContext &context, + const duckdb::FunctionData *data) { auto &bind_data = data->Cast(); return duckdb::make_uniq(bind_data.cardinality, bind_data.cardinality); } From 437426f1f7723498532b0df9e366a221ced2a378 Mon Sep 17 00:00:00 2001 From: mkaruza Date: Fri, 13 Dec 2024 10:51:44 +0100 Subject: [PATCH 09/53] Minor header cleanup --- src/scan/postgres_scan.cpp | 34 ++-------------------------------- 1 file changed, 2 insertions(+), 32 deletions(-) diff --git a/src/scan/postgres_scan.cpp b/src/scan/postgres_scan.cpp index 0f666dde..3af7df43 100644 --- a/src/scan/postgres_scan.cpp +++ b/src/scan/postgres_scan.cpp @@ -1,50 +1,20 @@ -#include "duckdb/main/client_context.hpp" -#include "duckdb/function/replacement_scan.hpp" -#include "duckdb/parser/tableref/table_function_ref.hpp" -#include "duckdb/parser/parser.hpp" -#include "duckdb/parser/tableref/subqueryref.hpp" -#include "duckdb/parser/expression/function_expression.hpp" -#include "duckdb/parser/statement/select_statement.hpp" -#include "duckdb/parser/expression/constant_expression.hpp" -#include "duckdb/parser/expression/comparison_expression.hpp" -#include "duckdb/parser/expression/columnref_expression.hpp" -#include "duckdb/parser/qualified_name.hpp" -#include "duckdb/common/enums/statement_type.hpp" -#include "duckdb/common/enums/expression_type.hpp" - #include "pgduckdb/scan/postgres_scan.hpp" #include "pgduckdb/scan/postgres_table_reader.hpp" #include "pgduckdb/pgduckdb_types.hpp" #include "pgduckdb/pgduckdb_utils.hpp" #include "pgduckdb/pgduckdb_process_lock.hpp" -#include "pgduckdb/pgduckdb_types.hpp" -#include "pgduckdb/pgduckdb_utils.hpp" #include "pgduckdb/logger.hpp" extern "C" { #include "postgres.h" #include "access/htup_details.h" -#include "catalog/namespace.h" -#include "catalog/pg_class.h" -#include "optimizer/planmain.h" -#include "optimizer/planner.h" +#include "executor/tuptable.h" #include "utils/builtins.h" #include "utils/lsyscache.h" -#include "utils/regproc.h" #include "utils/rel.h" -#include "utils/snapmgr.h" -#include "utils/syscache.h" -#include "tcop/tcopprot.h" -#include "executor/execdesc.h" -#include "executor/executor.h" - -#include "pgduckdb/vendor/pg_ruleutils.h" -#include "pgduckdb/pgduckdb_ruleutils.h" } -#include "pgduckdb/pgduckdb_process_lock.hpp" - namespace pgduckdb { // @@ -141,7 +111,7 @@ PostgresScanGlobalState::ConstructTableScanQuery(duckdb::TableFunctionInitInput scan_query << ";"; - elog(DEBUG1, "scan_query: %s", scan_query.str().c_str()); + pd_log(DEBUG1, "scan_query: %s", scan_query.str().c_str()); } std::string From 22d473a5eb50c7c0143908d9c888910c0b4a14c0 Mon Sep 17 00:00:00 2001 From: mkaruza Date: Sat, 14 Dec 2024 11:27:10 +0100 Subject: [PATCH 10/53] Regression test for native postgres scan. Handle COUNT(*) in batches. --- .../pgduckdb/scan/postgres_table_reader.hpp | 5 +- src/scan/postgres_scan.cpp | 30 ++-- src/scan/postgres_table_reader.cpp | 65 ++++++- test/regression/expected/basic.out | 17 -- .../expected/scan_postgres_tables.out | 164 ++++++++++++++++++ test/regression/sql/basic.sql | 6 - test/regression/sql/scan_postgres_tables.sql | 46 +++++ 7 files changed, 290 insertions(+), 43 deletions(-) create mode 100644 test/regression/expected/scan_postgres_tables.out create mode 100644 test/regression/sql/scan_postgres_tables.sql diff --git a/include/pgduckdb/scan/postgres_table_reader.hpp b/include/pgduckdb/scan/postgres_table_reader.hpp index a142ab2b..cbd009b2 100644 --- a/include/pgduckdb/scan/postgres_table_reader.hpp +++ b/include/pgduckdb/scan/postgres_table_reader.hpp @@ -12,14 +12,15 @@ namespace pgduckdb { class PostgresTableReader { public: - PostgresTableReader(const char *table_scan_query); + PostgresTableReader(const char *table_scan_query, bool count_tuples_only); ~PostgresTableReader(); TupleTableSlot *GetNextTuple(); private: MinimalTuple GetNextWorkerTuple(); int ParallelWorkerNumber(Cardinality cardinality); - + std::string ExplainScanPlan(QueryDesc *query_desc); + bool MarkPlanParallelAware(Plan *plan); private: QueryDesc *table_scan_query_desc; PlanState *table_scan_planstate; diff --git a/src/scan/postgres_scan.cpp b/src/scan/postgres_scan.cpp index 3af7df43..ee7a58bf 100644 --- a/src/scan/postgres_scan.cpp +++ b/src/scan/postgres_scan.cpp @@ -108,10 +108,6 @@ PostgresScanGlobalState::ConstructTableScanQuery(duckdb::TableFunctionInitInput scan_query << ") "; } } - - scan_query << ";"; - - pd_log(DEBUG1, "scan_query: %s", scan_query.str().c_str()); } std::string @@ -124,7 +120,8 @@ PostgresScanGlobalState::PostgresScanGlobalState(Snapshot snapshot, Relation rel : snapshot(snapshot), rel(rel), table_tuple_desc(RelationGetDescr(rel)), count_tuples_only(false), total_row_count(0) { ConstructTableScanQuery(input); - table_reader_global_state = duckdb::make_shared_ptr(scan_query.str().c_str()); + table_reader_global_state = + duckdb::make_shared_ptr(scan_query.str().c_str(), count_tuples_only); pd_log(DEBUG2, "(DuckDB/PostgresSeqScanGlobalState) Running %" PRIu64 " threads -- ", (uint64_t)MaxThreads()); } @@ -183,20 +180,27 @@ PostgresScanTableFunction::PostgresScanInitLocal(__attribute__((unused)) duckdb: auto global_state = reinterpret_cast(gstate); return duckdb::make_uniq(global_state); } +void +SetOutputCardinality(duckdb::DataChunk &output, PostgresScanLocalState &local_state) { + idx_t output_cardinality = + local_state.output_vector_size <= STANDARD_VECTOR_SIZE ? local_state.output_vector_size : STANDARD_VECTOR_SIZE; + output.SetCardinality(output_cardinality); + local_state.output_vector_size -= output_cardinality; +} void PostgresScanTableFunction::PostgresScanFunction(__attribute__((unused)) duckdb::ClientContext &context, duckdb::TableFunctionInput &data, duckdb::DataChunk &output) { auto &local_state = data.local_state->Cast(); - local_state.output_vector_size = 0; - - /* We have exhausted seq scan of heap table so we can return */ + /* We have exhausted table scan */ if (local_state.exhausted_scan) { - output.SetCardinality(0); + SetOutputCardinality(output, local_state); return; } + local_state.output_vector_size = 0; + int i = 0; for (; i < STANDARD_VECTOR_SIZE; i++) { TupleTableSlot *slot = local_state.global_state->table_reader_global_state->GetNextTuple(); @@ -208,16 +212,12 @@ PostgresScanTableFunction::PostgresScanFunction(__attribute__((unused)) duckdb:: InsertTupleIntoChunk(output, local_state, slot); } + /* If we finish before reading complete vector means that scan was exhausted. */ if (i != STANDARD_VECTOR_SIZE) { local_state.exhausted_scan = true; } - /* Special case when we only use COUNT(*). We only update output capactity. */ - if (local_state.global_state->count_tuples_only && local_state.output_vector_size < STANDARD_VECTOR_SIZE) { - output.SetCapacity(local_state.output_vector_size); - } - - output.SetCardinality(local_state.output_vector_size); + SetOutputCardinality(output, local_state); } duckdb::unique_ptr diff --git a/src/scan/postgres_table_reader.cpp b/src/scan/postgres_table_reader.cpp index 3a42db33..e56610e8 100644 --- a/src/scan/postgres_table_reader.cpp +++ b/src/scan/postgres_table_reader.cpp @@ -7,6 +7,7 @@ extern "C" { #include "postgres.h" #include "miscadmin.h" #include "access/xact.h" +#include "commands/explain.h" #include "executor/executor.h" #include "executor/execParallel.h" #include "executor/tqueue.h" @@ -24,7 +25,7 @@ extern "C" { namespace pgduckdb { -PostgresTableReader::PostgresTableReader(const char *table_scan_query) +PostgresTableReader::PostgresTableReader(const char *table_scan_query, bool count_tuples_only) : parallel_executor_info(nullptr), parallel_worker_readers(nullptr), nreaders(0), next_parallel_reader(0), entered_parallel_mode(false) { std::lock_guard lock(GlobalProcessLock::GetLock()); @@ -55,10 +56,20 @@ PostgresTableReader::PostgresTableReader(const char *table_scan_query) table_scan_planstate = PostgresFunctionGuard(ExecInitNode, planned_stmt->planTree, table_scan_query_desc->estate, 0); - /* Temp tables can be excuted with parallel workers */ + bool marked_parallel_aware = false; + if (persistence != RELPERSISTENCE_TEMP) { + if (count_tuples_only) { + /* For count_tuples_only we will try to execute aggregate node on table scan */ + planned_stmt->planTree->parallel_aware = true; + marked_parallel_aware = MarkPlanParallelAware((Plan *)table_scan_query_desc->planstate->plan->lefttree); + } else { + marked_parallel_aware = MarkPlanParallelAware(table_scan_query_desc->planstate->plan); + } + } - table_scan_query_desc->planstate->plan->parallel_aware = true; + /* Temp tables can be excuted with parallel workers, and whole plan should be parallel aware */ + if (persistence != RELPERSISTENCE_TEMP && marked_parallel_aware) { int parallel_workers = ParallelWorkerNumber(planned_stmt->planTree->plan_rows); @@ -86,6 +97,10 @@ PostgresTableReader::PostgresTableReader(const char *table_scan_query) HOLD_CANCEL_INTERRUPTS(); } + elog(DEBUG1, "(PGDuckdDB/PostgresTableReader)\n\nQUERY: %s\nRUNNING: %s.\nEXECUTING: \n%s", table_scan_query, + !nreaders ? "IN PROCESS THREAD" : psprintf("ON %d PARALLEL WORKER(S)", nreaders), + ExplainScanPlan(table_scan_query_desc).c_str()); + slot = PostgresFunctionGuard(ExecInitExtraTupleSlot, table_scan_query_desc->estate, table_scan_planstate->ps_ResultTupleDesc, &TTSOpsMinimalTuple); } @@ -127,6 +142,50 @@ PostgresTableReader::ParallelWorkerNumber(Cardinality cardinality) { return std::max(1, std::min(base, std::max(max_workers_per_postgres_scan, max_parallel_workers))); } +std::string +PostgresTableReader::ExplainScanPlan(QueryDesc *query_desc) { + ExplainState *es = (ExplainState *)PostgresFunctionGuard(palloc0, sizeof(ExplainState)); + es->str = makeStringInfo(); + es->format = EXPLAIN_FORMAT_TEXT; + PostgresFunctionGuard(ExplainPrintPlan, es, query_desc); + // // Remove new line char from explain output + // es->str->data[es->str->len - 1] = 0; + std::string explain_scan(es->str->data); + return explain_scan; +} + +bool +PostgresTableReader::MarkPlanParallelAware(Plan *plan) { + switch (nodeTag(plan)) { + case T_SeqScan: + case T_IndexScan: + case T_IndexOnlyScan: { + plan->parallel_aware = true; + return true; + } + case T_BitmapHeapScan: { + plan->parallel_aware = true; + return MarkPlanParallelAware(plan->lefttree); + } + case T_BitmapIndexScan: { + ((BitmapIndexScan *)plan)->isshared = true; + return true; + } + case T_BitmapAnd: { + return MarkPlanParallelAware((Plan *)linitial(((BitmapAnd *)plan)->bitmapplans)); + } + case T_BitmapOr: { + ((BitmapOr *)plan)->isshared = true; + return MarkPlanParallelAware((Plan *)linitial(((BitmapOr *)plan)->bitmapplans)); + } + default: { + std::ostringstream oss; + oss << "Unknown postgres scan query plan node: " << nodeTag(plan); + throw duckdb::Exception(duckdb::ExceptionType::EXECUTOR, oss.str().c_str()); + } + } +} + TupleTableSlot * PostgresTableReader::GetNextTuple() { MinimalTuple worker_minmal_tuple; diff --git a/test/regression/expected/basic.out b/test/regression/expected/basic.out index 85e3391f..40e859a6 100644 --- a/test/regression/expected/basic.out +++ b/test/regression/expected/basic.out @@ -15,22 +15,6 @@ SELECT a, COUNT(*) FROM t WHERE a > 5 GROUP BY a ORDER BY a; 9 | 100 (4 rows) -SET duckdb.max_threads_per_postgres_scan to 4; -SELECT COUNT(*) FROM t; - count -------- - 1000 -(1 row) - -SELECT a, COUNT(*) FROM t WHERE a > 5 GROUP BY a ORDER BY a; - a | count ----+------- - 6 | 100 - 7 | 100 - 8 | 100 - 9 | 100 -(4 rows) - CREATE TABLE empty(a INT); SELECT COUNT(*) FROM empty; count @@ -38,7 +22,6 @@ SELECT COUNT(*) FROM empty; 0 (1 row) -SET duckdb.max_threads_per_postgres_scan TO default; SET client_min_messages TO default; DROP TABLE t; DROP TABLE empty; diff --git a/test/regression/expected/scan_postgres_tables.out b/test/regression/expected/scan_postgres_tables.out new file mode 100644 index 00000000..d98097c5 --- /dev/null +++ b/test/regression/expected/scan_postgres_tables.out @@ -0,0 +1,164 @@ +CREATE TABLE t1(a INT, b INT, c TEXT); +INSERT INTO t1 SELECT g, g % 100, 'scan_potgres_table_'||g from generate_series(1,100000) g; +SET client_min_messages TO DEBUG1; +-- COUNT(*) +SELECT COUNT(*) FROM t1; +DEBUG: (PGDuckdDB/PostgresTableReader) + +QUERY: SELECT COUNT(*) FROM public.t1 +RUNNING: ON 1 PARALLEL WORKER(S). +EXECUTING: +Parallel Aggregate + -> Parallel Seq Scan on t1 + + count +-------- + 100000 +(1 row) + +-- SEQ SCAN +SELECT COUNT(a) FROM t1 WHERE a < 10; +DEBUG: (PGDuckdDB/PostgresTableReader) + +QUERY: SELECT a FROM public.t1 WHERE (a<10 AND a IS NOT NULL) +RUNNING: ON 1 PARALLEL WORKER(S). +EXECUTING: +Parallel Seq Scan on t1 + Filter: ((a IS NOT NULL) AND (a < 10)) + + count +------- + 9 +(1 row) + +-- CREATE INDEX on t1 +SET client_min_messages TO DEFAULT; +CREATE INDEX ON t1(a); +SET client_min_messages TO DEBUG1; +-- BITMAP INDEX +SELECT COUNT(a) FROM t1 WHERE a < 10; +DEBUG: (PGDuckdDB/PostgresTableReader) + +QUERY: SELECT a FROM public.t1 WHERE (a<10 AND a IS NOT NULL) +RUNNING: ON 1 PARALLEL WORKER(S). +EXECUTING: +Parallel Bitmap Heap Scan on t1 + Recheck Cond: ((a < 10) AND (a IS NOT NULL)) + -> Bitmap Index Scan on t1_a_idx + Index Cond: ((a < 10) AND (a IS NOT NULL)) + + count +------- + 9 +(1 row) + +-- INDEXONLYSCAN +SET enable_bitmapscan TO false; +SELECT COUNT(a) FROM t1 WHERE a = 1; +DEBUG: (PGDuckdDB/PostgresTableReader) + +QUERY: SELECT a FROM public.t1 WHERE (a=1 AND a IS NOT NULL) +RUNNING: ON 1 PARALLEL WORKER(S). +EXECUTING: +Parallel Index Only Scan using t1_a_idx on t1 + Index Cond: ((a IS NOT NULL) AND (a = 1)) + + count +------- + 1 +(1 row) + +-- INDEXSCAN +SELECT COUNT(c) FROM t1 WHERE a = 1; +DEBUG: (PGDuckdDB/PostgresTableReader) + +QUERY: SELECT c FROM public.t1 WHERE (a=1 AND a IS NOT NULL) +RUNNING: ON 1 PARALLEL WORKER(S). +EXECUTING: +Parallel Index Scan using t1_a_idx on t1 + Index Cond: ((a IS NOT NULL) AND (a = 1)) + + count +------- + 1 +(1 row) + +-- TEMPORARY TABLES JOIN WITH HEAP TABLES +SET client_min_messages TO DEFAULT; +CREATE TEMP TABLE t2(a int); +INSERT INTO t2 VALUES (1), (2), (3); +SET client_min_messages TO DEBUG1; +SELECT t1.a, t2.a FROM t1, t2 WHERE t1.a = t2.a; +DEBUG: (PGDuckdDB/PostgresTableReader) + +QUERY: SELECT a FROM pg_temp.t2 +RUNNING: IN PROCESS THREAD. +EXECUTING: +Seq Scan on t2 + +DEBUG: (PGDuckdDB/PostgresTableReader) + +QUERY: SELECT a FROM public.t1 WHERE (a>=1 AND a<=3 AND a IS NOT NULL) +RUNNING: ON 1 PARALLEL WORKER(S). +EXECUTING: +Parallel Index Only Scan using t1_a_idx on t1 + Index Cond: ((a >= 1) AND (a <= 3) AND (a IS NOT NULL)) + + a | a +---+--- + 1 | 1 + 2 | 2 + 3 | 3 +(3 rows) + +-- JOIN WITH SAME TABLE (on WORKERS) +SELECT COUNT(*) FROM t1 AS t1_1, t1 AS t1_2 WHERE t1_1.a < 2 AND t1_2.a > 8; +DEBUG: (PGDuckdDB/PostgresTableReader) + +QUERY: SELECT a FROM public.t1 WHERE (a<2 AND a IS NOT NULL) +RUNNING: ON 1 PARALLEL WORKER(S). +EXECUTING: +Parallel Seq Scan on t1 + Filter: ((a IS NOT NULL) AND (a < 2)) + +DEBUG: (PGDuckdDB/PostgresTableReader) + +QUERY: SELECT a FROM public.t1 WHERE (a>8 AND a IS NOT NULL) +RUNNING: ON 1 PARALLEL WORKER(S). +EXECUTING: +Parallel Seq Scan on t1 + Filter: ((a IS NOT NULL) AND (a > 8)) + + count +------- + 99992 +(1 row) + +-- JOIN WITH SAME TABLE (in BACKEND PROCESS) +SET max_parallel_workers TO 0; +SELECT COUNT(*) FROM t1 AS t1_1, t1 AS t1_2 WHERE t1_1.a < 2 AND t1_2.a > 8; +DEBUG: (PGDuckdDB/PostgresTableReader) + +QUERY: SELECT a FROM public.t1 WHERE (a<2 AND a IS NOT NULL) +RUNNING: IN PROCESS THREAD. +EXECUTING: +Parallel Seq Scan on t1 + Filter: ((a IS NOT NULL) AND (a < 2)) + +DEBUG: (PGDuckdDB/PostgresTableReader) + +QUERY: SELECT a FROM public.t1 WHERE (a>8 AND a IS NOT NULL) +RUNNING: IN PROCESS THREAD. +EXECUTING: +Parallel Seq Scan on t1 + Filter: ((a IS NOT NULL) AND (a > 8)) + + count +------- + 99992 +(1 row) + +SET max_parallel_workers TO DEFAULT; +SET enable_bitmapscan TO DEFAULT; +SET client_min_messages TO DEFAULT; +DROP TABLE t1, t2; diff --git a/test/regression/sql/basic.sql b/test/regression/sql/basic.sql index 614773cd..aa69d83e 100644 --- a/test/regression/sql/basic.sql +++ b/test/regression/sql/basic.sql @@ -5,15 +5,9 @@ INSERT INTO t SELECT g % 10 from generate_series(1,1000) g; SELECT COUNT(*) FROM t; SELECT a, COUNT(*) FROM t WHERE a > 5 GROUP BY a ORDER BY a; -SET duckdb.max_threads_per_postgres_scan to 4; - -SELECT COUNT(*) FROM t; -SELECT a, COUNT(*) FROM t WHERE a > 5 GROUP BY a ORDER BY a; - CREATE TABLE empty(a INT); SELECT COUNT(*) FROM empty; -SET duckdb.max_threads_per_postgres_scan TO default; SET client_min_messages TO default; DROP TABLE t; diff --git a/test/regression/sql/scan_postgres_tables.sql b/test/regression/sql/scan_postgres_tables.sql new file mode 100644 index 00000000..faaa1d24 --- /dev/null +++ b/test/regression/sql/scan_postgres_tables.sql @@ -0,0 +1,46 @@ +CREATE TABLE t1(a INT, b INT, c TEXT); +INSERT INTO t1 SELECT g, g % 100, 'scan_potgres_table_'||g from generate_series(1,100000) g; + +SET client_min_messages TO DEBUG1; + +-- COUNT(*) +SELECT COUNT(*) FROM t1; + +-- SEQ SCAN +SELECT COUNT(a) FROM t1 WHERE a < 10; + +-- CREATE INDEX on t1 +SET client_min_messages TO DEFAULT; +CREATE INDEX ON t1(a); +SET client_min_messages TO DEBUG1; + +-- BITMAP INDEX +SELECT COUNT(a) FROM t1 WHERE a < 10; + +-- INDEXONLYSCAN +SET enable_bitmapscan TO false; +SELECT COUNT(a) FROM t1 WHERE a = 1; + +-- INDEXSCAN +SELECT COUNT(c) FROM t1 WHERE a = 1; + +-- TEMPORARY TABLES JOIN WITH HEAP TABLES +SET client_min_messages TO DEFAULT; +CREATE TEMP TABLE t2(a int); +INSERT INTO t2 VALUES (1), (2), (3); +SET client_min_messages TO DEBUG1; + +SELECT t1.a, t2.a FROM t1, t2 WHERE t1.a = t2.a; + +-- JOIN WITH SAME TABLE (on WORKERS) +SELECT COUNT(*) FROM t1 AS t1_1, t1 AS t1_2 WHERE t1_1.a < 2 AND t1_2.a > 8; + +-- JOIN WITH SAME TABLE (in BACKEND PROCESS) +SET max_parallel_workers TO 0; +SELECT COUNT(*) FROM t1 AS t1_1, t1 AS t1_2 WHERE t1_1.a < 2 AND t1_2.a > 8; +SET max_parallel_workers TO DEFAULT; + + +SET enable_bitmapscan TO DEFAULT; +SET client_min_messages TO DEFAULT; +DROP TABLE t1, t2; From 9d4454ffa5cc87b992c8af256b97aec08888b2ae Mon Sep 17 00:00:00 2001 From: mkaruza Date: Sat, 14 Dec 2024 11:38:28 +0100 Subject: [PATCH 11/53] Added small comment --- src/pgduckdb_types.cpp | 1 + 1 file changed, 1 insertion(+) diff --git a/src/pgduckdb_types.cpp b/src/pgduckdb_types.cpp index e1f144dc..0806382c 100644 --- a/src/pgduckdb_types.cpp +++ b/src/pgduckdb_types.cpp @@ -1354,6 +1354,7 @@ InsertTupleIntoChunk(duckdb::DataChunk &output, PostgresScanLocalState &scan_loc auto scan_global_state = scan_local_state.global_state; if (scan_global_state->count_tuples_only) { + /* COUNT(*) returned tuple will have only one value returned as first tuple element. */ scan_global_state->total_row_count += slot->tts_values[0]; scan_local_state.output_vector_size += slot->tts_values[0]; return; From d6848b5ce5cb610f4f2e44c70948fb45d41a0be3 Mon Sep 17 00:00:00 2001 From: mkaruza Date: Sat, 14 Dec 2024 11:41:49 +0100 Subject: [PATCH 12/53] Remove last part of view replacement scan. --- src/catalog/pgduckdb_transaction.cpp | 5 ----- 1 file changed, 5 deletions(-) diff --git a/src/catalog/pgduckdb_transaction.cpp b/src/catalog/pgduckdb_transaction.cpp index bb0f6e83..3c314c6a 100644 --- a/src/catalog/pgduckdb_transaction.cpp +++ b/src/catalog/pgduckdb_transaction.cpp @@ -44,11 +44,6 @@ SchemaItems::GetTable(const duckdb::string &entry_name) { } Relation rel = PostgresTable::OpenRelation(rel_oid); - if (IsRelView(rel)) { - // Let the replacement scan handle this, the ReplacementScan replaces the view with its view_definition, which - // will get bound again and hit a PostgresIndexTable / PostgresHeapTable. - return nullptr; - } duckdb::CreateTableInfo info; info.table = entry_name; From 8472d41e7810cdba4facbb191f110711fff64d52 Mon Sep 17 00:00:00 2001 From: mkaruza Date: Sat, 14 Dec 2024 11:45:42 +0100 Subject: [PATCH 13/53] Small output fix --- src/scan/postgres_scan.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/scan/postgres_scan.cpp b/src/scan/postgres_scan.cpp index ee7a58bf..d369e236 100644 --- a/src/scan/postgres_scan.cpp +++ b/src/scan/postgres_scan.cpp @@ -102,7 +102,7 @@ PostgresScanGlobalState::ConstructTableScanQuery(duckdb::TableFunctionInitInput scan_query << " AND "; } first = false; - scan_query << " ("; + scan_query << "("; auto attr = table_tuple_desc->attrs[attr_num - 1]; scan_query << filter->ToString(attr.attname.data).c_str(); scan_query << ") "; From e66cbbb5f4c5bd5de9b218b78f31197c3661968a Mon Sep 17 00:00:00 2001 From: mkaruza Date: Sat, 14 Dec 2024 11:51:04 +0100 Subject: [PATCH 14/53] Forgot to include regression test in schedule --- .../expected/scan_postgres_tables.out | 18 +++++++++--------- test/regression/schedule | 1 + 2 files changed, 10 insertions(+), 9 deletions(-) diff --git a/test/regression/expected/scan_postgres_tables.out b/test/regression/expected/scan_postgres_tables.out index d98097c5..65622544 100644 --- a/test/regression/expected/scan_postgres_tables.out +++ b/test/regression/expected/scan_postgres_tables.out @@ -20,7 +20,7 @@ Parallel Aggregate SELECT COUNT(a) FROM t1 WHERE a < 10; DEBUG: (PGDuckdDB/PostgresTableReader) -QUERY: SELECT a FROM public.t1 WHERE (a<10 AND a IS NOT NULL) +QUERY: SELECT a FROM public.t1 WHERE (a<10 AND a IS NOT NULL) RUNNING: ON 1 PARALLEL WORKER(S). EXECUTING: Parallel Seq Scan on t1 @@ -39,7 +39,7 @@ SET client_min_messages TO DEBUG1; SELECT COUNT(a) FROM t1 WHERE a < 10; DEBUG: (PGDuckdDB/PostgresTableReader) -QUERY: SELECT a FROM public.t1 WHERE (a<10 AND a IS NOT NULL) +QUERY: SELECT a FROM public.t1 WHERE (a<10 AND a IS NOT NULL) RUNNING: ON 1 PARALLEL WORKER(S). EXECUTING: Parallel Bitmap Heap Scan on t1 @@ -57,7 +57,7 @@ SET enable_bitmapscan TO false; SELECT COUNT(a) FROM t1 WHERE a = 1; DEBUG: (PGDuckdDB/PostgresTableReader) -QUERY: SELECT a FROM public.t1 WHERE (a=1 AND a IS NOT NULL) +QUERY: SELECT a FROM public.t1 WHERE (a=1 AND a IS NOT NULL) RUNNING: ON 1 PARALLEL WORKER(S). EXECUTING: Parallel Index Only Scan using t1_a_idx on t1 @@ -72,7 +72,7 @@ Parallel Index Only Scan using t1_a_idx on t1 SELECT COUNT(c) FROM t1 WHERE a = 1; DEBUG: (PGDuckdDB/PostgresTableReader) -QUERY: SELECT c FROM public.t1 WHERE (a=1 AND a IS NOT NULL) +QUERY: SELECT c FROM public.t1 WHERE (a=1 AND a IS NOT NULL) RUNNING: ON 1 PARALLEL WORKER(S). EXECUTING: Parallel Index Scan using t1_a_idx on t1 @@ -98,7 +98,7 @@ Seq Scan on t2 DEBUG: (PGDuckdDB/PostgresTableReader) -QUERY: SELECT a FROM public.t1 WHERE (a>=1 AND a<=3 AND a IS NOT NULL) +QUERY: SELECT a FROM public.t1 WHERE (a>=1 AND a<=3 AND a IS NOT NULL) RUNNING: ON 1 PARALLEL WORKER(S). EXECUTING: Parallel Index Only Scan using t1_a_idx on t1 @@ -115,7 +115,7 @@ Parallel Index Only Scan using t1_a_idx on t1 SELECT COUNT(*) FROM t1 AS t1_1, t1 AS t1_2 WHERE t1_1.a < 2 AND t1_2.a > 8; DEBUG: (PGDuckdDB/PostgresTableReader) -QUERY: SELECT a FROM public.t1 WHERE (a<2 AND a IS NOT NULL) +QUERY: SELECT a FROM public.t1 WHERE (a<2 AND a IS NOT NULL) RUNNING: ON 1 PARALLEL WORKER(S). EXECUTING: Parallel Seq Scan on t1 @@ -123,7 +123,7 @@ Parallel Seq Scan on t1 DEBUG: (PGDuckdDB/PostgresTableReader) -QUERY: SELECT a FROM public.t1 WHERE (a>8 AND a IS NOT NULL) +QUERY: SELECT a FROM public.t1 WHERE (a>8 AND a IS NOT NULL) RUNNING: ON 1 PARALLEL WORKER(S). EXECUTING: Parallel Seq Scan on t1 @@ -139,7 +139,7 @@ SET max_parallel_workers TO 0; SELECT COUNT(*) FROM t1 AS t1_1, t1 AS t1_2 WHERE t1_1.a < 2 AND t1_2.a > 8; DEBUG: (PGDuckdDB/PostgresTableReader) -QUERY: SELECT a FROM public.t1 WHERE (a<2 AND a IS NOT NULL) +QUERY: SELECT a FROM public.t1 WHERE (a<2 AND a IS NOT NULL) RUNNING: IN PROCESS THREAD. EXECUTING: Parallel Seq Scan on t1 @@ -147,7 +147,7 @@ Parallel Seq Scan on t1 DEBUG: (PGDuckdDB/PostgresTableReader) -QUERY: SELECT a FROM public.t1 WHERE (a>8 AND a IS NOT NULL) +QUERY: SELECT a FROM public.t1 WHERE (a>8 AND a IS NOT NULL) RUNNING: IN PROCESS THREAD. EXECUTING: Parallel Seq Scan on t1 diff --git a/test/regression/schedule b/test/regression/schedule index eec4208b..18d43899 100644 --- a/test/regression/schedule +++ b/test/regression/schedule @@ -28,3 +28,4 @@ test: prepare test: function test: timestamp_with_interval test: approx_count_distinct +test: scan_postgres_tables From e5e2a97489582958b1561bc596cdc4f94fe49cc1 Mon Sep 17 00:00:00 2001 From: mkaruza Date: Sat, 14 Dec 2024 13:59:00 +0100 Subject: [PATCH 15/53] Rename variable names so they are not shadowed --- include/pgduckdb/scan/postgres_scan.hpp | 1 - src/scan/postgres_scan.cpp | 12 ++++++------ src/scan/postgres_table_reader.cpp | 2 -- 3 files changed, 6 insertions(+), 9 deletions(-) diff --git a/include/pgduckdb/scan/postgres_scan.hpp b/include/pgduckdb/scan/postgres_scan.hpp index 7c06ce4b..71c41b11 100644 --- a/include/pgduckdb/scan/postgres_scan.hpp +++ b/include/pgduckdb/scan/postgres_scan.hpp @@ -29,7 +29,6 @@ struct PostgresScanGlobalState : public duckdb::GlobalTableFunctionState { Snapshot snapshot; Relation rel; TupleDesc table_tuple_desc; - std::mutex lock; // Lock for one replacement scan bool count_tuples_only; duckdb::vector> output_columns; std::atomic total_row_count; diff --git a/src/scan/postgres_scan.cpp b/src/scan/postgres_scan.cpp index d369e236..4c390d78 100644 --- a/src/scan/postgres_scan.cpp +++ b/src/scan/postgres_scan.cpp @@ -116,8 +116,8 @@ PostgresScanGlobalState::ConstructFullyQualifiedTableName() { quote_identifier(get_rel_name(rel->rd_rel->oid))); } -PostgresScanGlobalState::PostgresScanGlobalState(Snapshot snapshot, Relation rel, duckdb::TableFunctionInitInput &input) - : snapshot(snapshot), rel(rel), table_tuple_desc(RelationGetDescr(rel)), count_tuples_only(false), +PostgresScanGlobalState::PostgresScanGlobalState(Snapshot _snapshot, Relation _rel, duckdb::TableFunctionInitInput &input) + : snapshot(_snapshot), rel(_rel), table_tuple_desc(RelationGetDescr(rel)), count_tuples_only(false), total_row_count(0) { ConstructTableScanQuery(input); table_reader_global_state = @@ -132,8 +132,8 @@ PostgresScanGlobalState::~PostgresScanGlobalState() { // PostgresScanLocalState // -PostgresScanLocalState::PostgresScanLocalState(PostgresScanGlobalState *global_state) - : global_state(global_state), exhausted_scan(false) { +PostgresScanLocalState::PostgresScanLocalState(PostgresScanGlobalState *_global_state) + : global_state(_global_state), exhausted_scan(false) { } PostgresScanLocalState::~PostgresScanLocalState() { @@ -143,8 +143,8 @@ PostgresScanLocalState::~PostgresScanLocalState() { // PostgresSeqScanFunctionData // -PostgresScanFunctionData::PostgresScanFunctionData(Relation rel, uint64_t cardinality, Snapshot snapshot) - : rel(rel), cardinality(cardinality), snapshot(snapshot) { +PostgresScanFunctionData::PostgresScanFunctionData(Relation _rel, uint64_t _cardinality, Snapshot _snapshot) + : rel(_rel), cardinality(_cardinality), snapshot(_snapshot) { } PostgresScanFunctionData::~PostgresScanFunctionData() { diff --git a/src/scan/postgres_table_reader.cpp b/src/scan/postgres_table_reader.cpp index e56610e8..69a16e9a 100644 --- a/src/scan/postgres_table_reader.cpp +++ b/src/scan/postgres_table_reader.cpp @@ -148,8 +148,6 @@ PostgresTableReader::ExplainScanPlan(QueryDesc *query_desc) { es->str = makeStringInfo(); es->format = EXPLAIN_FORMAT_TEXT; PostgresFunctionGuard(ExplainPrintPlan, es, query_desc); - // // Remove new line char from explain output - // es->str->data[es->str->len - 1] = 0; std::string explain_scan(es->str->data); return explain_scan; } From 8d2151fcbfbd5eb4010e6c8330af0210a16ad561 Mon Sep 17 00:00:00 2001 From: mkaruza Date: Sat, 14 Dec 2024 15:13:28 +0100 Subject: [PATCH 16/53] Don't cancel interrupts if not needed. Python fix tests. --- src/scan/postgres_table_reader.cpp | 11 ++++++++--- test/pycheck/explain_test.py | 6 +++--- 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/src/scan/postgres_table_reader.cpp b/src/scan/postgres_table_reader.cpp index 69a16e9a..89a69fdd 100644 --- a/src/scan/postgres_table_reader.cpp +++ b/src/scan/postgres_table_reader.cpp @@ -62,7 +62,7 @@ PostgresTableReader::PostgresTableReader(const char *table_scan_query, bool coun if (count_tuples_only) { /* For count_tuples_only we will try to execute aggregate node on table scan */ planned_stmt->planTree->parallel_aware = true; - marked_parallel_aware = MarkPlanParallelAware((Plan *)table_scan_query_desc->planstate->plan->lefttree); + marked_parallel_aware = MarkPlanParallelAware((Plan *)table_scan_query_desc->planstate->plan->lefttree); } else { marked_parallel_aware = MarkPlanParallelAware(table_scan_query_desc->planstate->plan); } @@ -72,8 +72,11 @@ PostgresTableReader::PostgresTableReader(const char *table_scan_query, bool coun if (persistence != RELPERSISTENCE_TEMP && marked_parallel_aware) { int parallel_workers = ParallelWorkerNumber(planned_stmt->planTree->plan_rows); + bool interrupts_can_be_process = INTERRUPTS_CAN_BE_PROCESSED(); - RESUME_CANCEL_INTERRUPTS(); + if (!interrupts_can_be_process) { + RESUME_CANCEL_INTERRUPTS(); + } if (!IsInParallelMode()) { EnterParallelMode(); @@ -94,7 +97,9 @@ PostgresTableReader::PostgresTableReader(const char *table_scan_query, bool coun memcpy(parallel_worker_readers, parallel_executor_info->reader, nreaders * sizeof(TupleQueueReader *)); } - HOLD_CANCEL_INTERRUPTS(); + if (!interrupts_can_be_process) { + HOLD_CANCEL_INTERRUPTS(); + } } elog(DEBUG1, "(PGDuckdDB/PostgresTableReader)\n\nQUERY: %s\nRUNNING: %s.\nEXECUTING: \n%s", table_scan_query, diff --git a/test/pycheck/explain_test.py b/test/pycheck/explain_test.py index e84e440b..b9347e2b 100644 --- a/test/pycheck/explain_test.py +++ b/test/pycheck/explain_test.py @@ -59,21 +59,21 @@ def test_explain_ctas(cur: Cursor): cur.sql("CREATE TEMP TABLE heap1(id) AS SELECT 1") result = cur.sql("EXPLAIN CREATE TEMP TABLE heap2(id) AS SELECT * from heap1") plan = "\n".join(result) - assert "POSTGRES_SEQ_SCAN" in plan + assert "POSTGRES_SCAN" in plan assert "Total Time:" not in plan result = cur.sql( "EXPLAIN ANALYZE CREATE TEMP TABLE heap2(id) AS SELECT * from heap1" ) plan = "\n".join(result) - assert "POSTGRES_SEQ_SCAN" in plan + assert "POSTGRES_SCAN" in plan assert "Total Time:" in plan result = cur.sql( "EXPLAIN CREATE TEMP TABLE duckdb1(id) USING duckdb AS SELECT * from heap1" ) plan = "\n".join(result) - assert "POSTGRES_SEQ_SCAN" in plan + assert "POSTGRES_SCAN" in plan assert "Total Time:" not in plan # EXPLAIN ANALYZE is not supported for DuckDB CTAS (yet) From 05e1e1490600b5f4025640a4c5db5942c050fa2b Mon Sep 17 00:00:00 2001 From: mkaruza Date: Sat, 14 Dec 2024 17:08:18 +0100 Subject: [PATCH 17/53] Changes for Postgres 14 --- src/scan/postgres_table_reader.cpp | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/scan/postgres_table_reader.cpp b/src/scan/postgres_table_reader.cpp index 89a69fdd..0c88e341 100644 --- a/src/scan/postgres_table_reader.cpp +++ b/src/scan/postgres_table_reader.cpp @@ -35,8 +35,13 @@ PostgresTableReader::PostgresTableReader(const char *table_scan_query, bool coun Assert(list_length(raw_parsetree_list) == 1); RawStmt *raw_parsetree = linitial_node(RawStmt, raw_parsetree_list); +#if PG_VERSION_NUM >= 150000 List *query_list = PostgresFunctionGuard(pg_analyze_and_rewrite_fixedparams, raw_parsetree, table_scan_query, nullptr, 0, nullptr); +#else + List *query_list = + PostgresFunctionGuard(pg_analyze_and_rewrite, raw_parsetree, table_scan_query, nullptr, 0, nullptr); +#endif Assert(list_length(query_list) == 1); Query *query = linitial_node(Query, query_list); From 9acd30c5c757743401365663d36cb4da3c1759ee Mon Sep 17 00:00:00 2001 From: mkaruza Date: Mon, 16 Dec 2024 09:00:04 +0100 Subject: [PATCH 18/53] Remove unused attribute and instead don't name unused input argument --- src/catalog/pgduckdb_table.cpp | 8 +++----- src/scan/postgres_scan.cpp | 16 +++++++--------- 2 files changed, 10 insertions(+), 14 deletions(-) diff --git a/src/catalog/pgduckdb_table.cpp b/src/catalog/pgduckdb_table.cpp index 6c348699..61cc0619 100644 --- a/src/catalog/pgduckdb_table.cpp +++ b/src/catalog/pgduckdb_table.cpp @@ -56,20 +56,18 @@ PostgresTable::GetTableCardinality(Relation rel) { } duckdb::unique_ptr -PostgresTable::GetStatistics(__attribute__((unused)) duckdb::ClientContext &context, - __attribute__((unused)) duckdb::column_t column_id) { +PostgresTable::GetStatistics(duckdb::ClientContext &, duckdb::column_t) { throw duckdb::NotImplementedException("GetStatistics not supported yet"); } duckdb::TableFunction -PostgresTable::GetScanFunction(__attribute__((unused)) duckdb::ClientContext &context, - duckdb::unique_ptr &bind_data) { +PostgresTable::GetScanFunction(duckdb::ClientContext &, duckdb::unique_ptr &bind_data) { bind_data = duckdb::make_uniq(rel, cardinality, snapshot); return PostgresScanTableFunction(); } duckdb::TableStorageInfo -PostgresTable::GetStorageInfo(__attribute__((unused)) duckdb::ClientContext &context) { +PostgresTable::GetStorageInfo(duckdb::ClientContext &) { throw duckdb::NotImplementedException("GetStorageInfo not supported yet"); } diff --git a/src/scan/postgres_scan.cpp b/src/scan/postgres_scan.cpp index 4c390d78..2c058cba 100644 --- a/src/scan/postgres_scan.cpp +++ b/src/scan/postgres_scan.cpp @@ -116,7 +116,8 @@ PostgresScanGlobalState::ConstructFullyQualifiedTableName() { quote_identifier(get_rel_name(rel->rd_rel->oid))); } -PostgresScanGlobalState::PostgresScanGlobalState(Snapshot _snapshot, Relation _rel, duckdb::TableFunctionInitInput &input) +PostgresScanGlobalState::PostgresScanGlobalState(Snapshot _snapshot, Relation _rel, + duckdb::TableFunctionInitInput &input) : snapshot(_snapshot), rel(_rel), table_tuple_desc(RelationGetDescr(rel)), count_tuples_only(false), total_row_count(0) { ConstructTableScanQuery(input); @@ -166,16 +167,14 @@ PostgresScanTableFunction::PostgresScanTableFunction() } duckdb::unique_ptr -PostgresScanTableFunction::PostgresScanInitGlobal(__attribute__((unused)) duckdb::ClientContext &context, - duckdb::TableFunctionInitInput &input) { +PostgresScanTableFunction::PostgresScanInitGlobal(duckdb::ClientContext &, duckdb::TableFunctionInitInput &input) { auto &bind_data = input.bind_data->CastNoConst(); auto global_state = duckdb::make_uniq(bind_data.snapshot, bind_data.rel, input); return global_state; } duckdb::unique_ptr -PostgresScanTableFunction::PostgresScanInitLocal(__attribute__((unused)) duckdb::ExecutionContext &context, - __attribute__((unused)) duckdb::TableFunctionInitInput &input, +PostgresScanTableFunction::PostgresScanInitLocal(duckdb::ExecutionContext &, duckdb::TableFunctionInitInput &, duckdb::GlobalTableFunctionState *gstate) { auto global_state = reinterpret_cast(gstate); return duckdb::make_uniq(global_state); @@ -189,8 +188,8 @@ SetOutputCardinality(duckdb::DataChunk &output, PostgresScanLocalState &local_st } void -PostgresScanTableFunction::PostgresScanFunction(__attribute__((unused)) duckdb::ClientContext &context, - duckdb::TableFunctionInput &data, duckdb::DataChunk &output) { +PostgresScanTableFunction::PostgresScanFunction(duckdb::ClientContext &, duckdb::TableFunctionInput &data, + duckdb::DataChunk &output) { auto &local_state = data.local_state->Cast(); /* We have exhausted table scan */ @@ -221,8 +220,7 @@ PostgresScanTableFunction::PostgresScanFunction(__attribute__((unused)) duckdb:: } duckdb::unique_ptr -PostgresScanTableFunction::PostgresScanCardinality(__attribute__((unused)) duckdb::ClientContext &context, - const duckdb::FunctionData *data) { +PostgresScanTableFunction::PostgresScanCardinality(duckdb::ClientContext &, const duckdb::FunctionData *data) { auto &bind_data = data->Cast(); return duckdb::make_uniq(bind_data.cardinality, bind_data.cardinality); } From a33662d6e7e9420d4ee7c338c6130e05cd7d6e9d Mon Sep 17 00:00:00 2001 From: mkaruza Date: Tue, 17 Dec 2024 10:52:55 +0100 Subject: [PATCH 19/53] PR review changes. Cleanup scan, if possible, once the scan is exhausted. --- include/pgduckdb/pgduckdb_utils.hpp | 4 +- include/pgduckdb/scan/postgres_scan.hpp | 2 +- .../pgduckdb/scan/postgres_table_reader.hpp | 4 +- include/pgduckdb/utility/rename_ruleutils.h | 1 - src/pgduckdb_types.cpp | 2 +- src/scan/postgres_scan.cpp | 39 +++++++++---------- src/scan/postgres_table_reader.cpp | 33 +++++++++++----- 7 files changed, 48 insertions(+), 37 deletions(-) diff --git a/include/pgduckdb/pgduckdb_utils.hpp b/include/pgduckdb/pgduckdb_utils.hpp index 4925ab93..a6906c23 100644 --- a/include/pgduckdb/pgduckdb_utils.hpp +++ b/include/pgduckdb/pgduckdb_utils.hpp @@ -87,8 +87,8 @@ __PostgresFunctionGuard__(const char *func_name, FuncArgs... args) { throw duckdb::Exception(duckdb::ExceptionType::EXECUTOR, message); } -#define PostgresFunctionGuard(FUNC, ...) \ - pgduckdb::__PostgresFunctionGuard__(__func__, __VA_ARGS__) +#define PostgresFunctionGuard(FUNC, ...) \ + pgduckdb::__PostgresFunctionGuard__(__func__, ##__VA_ARGS__) duckdb::unique_ptr DuckDBQueryOrThrow(duckdb::ClientContext &context, const std::string &query); diff --git a/include/pgduckdb/scan/postgres_scan.hpp b/include/pgduckdb/scan/postgres_scan.hpp index 71c41b11..ee4af2f2 100644 --- a/include/pgduckdb/scan/postgres_scan.hpp +++ b/include/pgduckdb/scan/postgres_scan.hpp @@ -30,7 +30,7 @@ struct PostgresScanGlobalState : public duckdb::GlobalTableFunctionState { Relation rel; TupleDesc table_tuple_desc; bool count_tuples_only; - duckdb::vector> output_columns; + duckdb::vector output_columns; std::atomic total_row_count; std::ostringstream scan_query; duckdb::shared_ptr table_reader_global_state; diff --git a/include/pgduckdb/scan/postgres_table_reader.hpp b/include/pgduckdb/scan/postgres_table_reader.hpp index cbd009b2..db0b2359 100644 --- a/include/pgduckdb/scan/postgres_table_reader.hpp +++ b/include/pgduckdb/scan/postgres_table_reader.hpp @@ -15,11 +15,11 @@ class PostgresTableReader { PostgresTableReader(const char *table_scan_query, bool count_tuples_only); ~PostgresTableReader(); TupleTableSlot *GetNextTuple(); - + void PostgresTableReaderCleanup(); private: MinimalTuple GetNextWorkerTuple(); int ParallelWorkerNumber(Cardinality cardinality); - std::string ExplainScanPlan(QueryDesc *query_desc); + const char * ExplainScanPlan(QueryDesc *query_desc); bool MarkPlanParallelAware(Plan *plan); private: QueryDesc *table_scan_query_desc; diff --git a/include/pgduckdb/utility/rename_ruleutils.h b/include/pgduckdb/utility/rename_ruleutils.h index 1ff77a8c..d5cbf0e3 100644 --- a/include/pgduckdb/utility/rename_ruleutils.h +++ b/include/pgduckdb/utility/rename_ruleutils.h @@ -23,7 +23,6 @@ #define pg_get_statisticsobjdef_string pgduckdb_pg_get_statisticsobjdef_string #define get_list_partvalue_string pgduckdb_get_list_partvalue_string - /* * The following replaces all usages of generate_qualified_relation_name and * generate_relation_name with calls to the pgduckdb_relation_name function diff --git a/src/pgduckdb_types.cpp b/src/pgduckdb_types.cpp index 0806382c..038d523f 100644 --- a/src/pgduckdb_types.cpp +++ b/src/pgduckdb_types.cpp @@ -1361,7 +1361,7 @@ InsertTupleIntoChunk(duckdb::DataChunk &output, PostgresScanLocalState &scan_loc } /* Write tuple columns in output vector. */ int duckdb_output_index = 0; - for (auto const &[_, attr_num] : scan_global_state->output_columns) { + for (auto const &attr_num : scan_global_state->output_columns) { auto &result = output.data[duckdb_output_index]; if (slot->tts_isnull[duckdb_output_index]) { auto &array_mask = duckdb::FlatVector::Validity(result); diff --git a/src/scan/postgres_scan.cpp b/src/scan/postgres_scan.cpp index 2c058cba..31bb158e 100644 --- a/src/scan/postgres_scan.cpp +++ b/src/scan/postgres_scan.cpp @@ -68,19 +68,18 @@ PostgresScanGlobalState::ConstructTableScanQuery(duckdb::TableFunctionInitInput */ if (input.CanRemoveFilterColumns()) { for (const auto &projection_id : input.projection_ids) { - output_columns.emplace_back(projection_id, input.column_ids[projection_id] + 1); + output_columns.emplace_back(input.column_ids[projection_id] + 1); } } else { - duckdb::idx_t output_index = 0; for (const auto &column_id : input.column_ids) { - output_columns.emplace_back(output_index++, column_id + 1); + output_columns.emplace_back(column_id + 1); } } scan_query << "SELECT "; bool first = true; - for (auto const &[duckdb_scanned_index, attr_num] : output_columns) { + for (auto const &attr_num : output_columns) { if (!first) { scan_query << ", "; } @@ -95,18 +94,22 @@ PostgresScanGlobalState::ConstructTableScanQuery(duckdb::TableFunctionInitInput for (auto const &[attr_num, duckdb_scanned_index] : columns_to_scan) { auto filter = column_filters[duckdb_scanned_index]; - if (filter) { - if (first) { - scan_query << " WHERE "; - } else { - scan_query << " AND "; - } - first = false; - scan_query << "("; - auto attr = table_tuple_desc->attrs[attr_num - 1]; - scan_query << filter->ToString(attr.attname.data).c_str(); - scan_query << ") "; + + if (!filter) { + continue; } + + if (first) { + scan_query << " WHERE "; + } else { + scan_query << " AND "; + } + + first = false; + scan_query << "("; + auto attr = table_tuple_desc->attrs[attr_num - 1]; + scan_query << filter->ToString(attr.attname.data).c_str(); + scan_query << ") "; } } @@ -204,6 +207,7 @@ PostgresScanTableFunction::PostgresScanFunction(duckdb::ClientContext &, duckdb: for (; i < STANDARD_VECTOR_SIZE; i++) { TupleTableSlot *slot = local_state.global_state->table_reader_global_state->GetNextTuple(); if (TupIsNull(slot)) { + local_state.global_state->table_reader_global_state->PostgresTableReaderCleanup(); local_state.exhausted_scan = true; break; } @@ -211,11 +215,6 @@ PostgresScanTableFunction::PostgresScanFunction(duckdb::ClientContext &, duckdb: InsertTupleIntoChunk(output, local_state, slot); } - /* If we finish before reading complete vector means that scan was exhausted. */ - if (i != STANDARD_VECTOR_SIZE) { - local_state.exhausted_scan = true; - } - SetOutputCardinality(output, local_state); } diff --git a/src/scan/postgres_table_reader.cpp b/src/scan/postgres_table_reader.cpp index 0c88e341..e4b99e62 100644 --- a/src/scan/postgres_table_reader.cpp +++ b/src/scan/postgres_table_reader.cpp @@ -109,13 +109,19 @@ PostgresTableReader::PostgresTableReader(const char *table_scan_query, bool coun elog(DEBUG1, "(PGDuckdDB/PostgresTableReader)\n\nQUERY: %s\nRUNNING: %s.\nEXECUTING: \n%s", table_scan_query, !nreaders ? "IN PROCESS THREAD" : psprintf("ON %d PARALLEL WORKER(S)", nreaders), - ExplainScanPlan(table_scan_query_desc).c_str()); + ExplainScanPlan(table_scan_query_desc)); slot = PostgresFunctionGuard(ExecInitExtraTupleSlot, table_scan_query_desc->estate, table_scan_planstate->ps_ResultTupleDesc, &TTSOpsMinimalTuple); } PostgresTableReader::~PostgresTableReader() { + if (table_scan_query_desc) { + PostgresTableReaderCleanup(); + } +} + +void PostgresTableReader::PostgresTableReaderCleanup() { std::lock_guard lock(GlobalProcessLock::GetLock()); PostgresScopedStackReset scoped_stack_reset; @@ -127,13 +133,13 @@ PostgresTableReader::~PostgresTableReader() { PostgresFunctionGuard(ExecParallelCleanup, parallel_executor_info); } - parallel_executor_info = NULL; + parallel_executor_info = nullptr; if (parallel_worker_readers) { PostgresFunctionGuard(pfree, parallel_worker_readers); } - parallel_worker_readers = NULL; + parallel_worker_readers = nullptr; PostgresFunctionGuard(ExecutorFinish, table_scan_query_desc); PostgresFunctionGuard(ExecutorEnd, table_scan_query_desc); @@ -142,6 +148,8 @@ PostgresTableReader::~PostgresTableReader() { if (entered_parallel_mode) { ExitParallelMode(); } + + table_scan_query_desc = nullptr; } int @@ -152,16 +160,21 @@ PostgresTableReader::ParallelWorkerNumber(Cardinality cardinality) { return std::max(1, std::min(base, std::max(max_workers_per_postgres_scan, max_parallel_workers))); } -std::string -PostgresTableReader::ExplainScanPlan(QueryDesc *query_desc) { - ExplainState *es = (ExplainState *)PostgresFunctionGuard(palloc0, sizeof(ExplainState)); +const char * +ExplainScanPlan_Unsafe(QueryDesc *query_desc) { + ExplainState *es = (ExplainState *)palloc0(sizeof(ExplainState)); es->str = makeStringInfo(); es->format = EXPLAIN_FORMAT_TEXT; - PostgresFunctionGuard(ExplainPrintPlan, es, query_desc); - std::string explain_scan(es->str->data); - return explain_scan; + ExplainPrintPlan(es, query_desc); + return es->str->data; } +const char * +PostgresTableReader::ExplainScanPlan(QueryDesc *query_desc) { + return PostgresFunctionGuard(ExplainScanPlan_Unsafe, query_desc); +} + + bool PostgresTableReader::MarkPlanParallelAware(Plan *plan) { switch (nodeTag(plan)) { @@ -208,7 +221,7 @@ PostgresTableReader::GetNextTuple() { std::lock_guard lock(GlobalProcessLock::GetLock()); PostgresScopedStackReset scoped_stack_reset; table_scan_query_desc->estate->es_query_dsa = parallel_executor_info ? parallel_executor_info->area : NULL; - thread_scan_slot = ExecProcNode(table_scan_planstate); + thread_scan_slot = PostgresFunctionGuard(ExecProcNode, table_scan_planstate); table_scan_query_desc->estate->es_query_dsa = NULL; if (!TupIsNull(thread_scan_slot)) { return thread_scan_slot; From 79cc554a7a0de21b3af92bea61fc2f0dc4e39bac Mon Sep 17 00:00:00 2001 From: mkaruza Date: Wed, 18 Dec 2024 11:36:07 +0100 Subject: [PATCH 20/53] Allow FDW tables and paritioned tables --- include/pgduckdb/catalog/pgduckdb_table.hpp | 1 - include/pgduckdb/pg/relations.hpp | 4 +- include/pgduckdb/pgduckdb_process_latch.hpp | 1 - .../pgduckdb/scan/postgres_table_reader.hpp | 3 + src/catalog/pgduckdb_table.cpp | 9 -- src/catalog/pgduckdb_transaction.cpp | 3 +- src/pg/relations.cpp | 19 ++-- src/pgduckdb_hooks.cpp | 25 ++++- src/scan/postgres_table_reader.cpp | 48 ++++++++-- .../expected/scan_postgres_tables.out | 93 ++++++++++++++++++- test/regression/sql/scan_postgres_tables.sql | 19 +++- 11 files changed, 186 insertions(+), 39 deletions(-) diff --git a/include/pgduckdb/catalog/pgduckdb_table.hpp b/include/pgduckdb/catalog/pgduckdb_table.hpp index 8157d6e8..521f559e 100644 --- a/include/pgduckdb/catalog/pgduckdb_table.hpp +++ b/include/pgduckdb/catalog/pgduckdb_table.hpp @@ -26,7 +26,6 @@ class PostgresTable : public duckdb::TableCatalogEntry { public: static Relation OpenRelation(Oid relid); static void SetTableInfo(duckdb::CreateTableInfo &info, Relation rel); - static Cardinality GetTableCardinality(Relation rel); protected: Relation rel; diff --git a/include/pgduckdb/pg/relations.hpp b/include/pgduckdb/pg/relations.hpp index e4fd62a7..099e39fe 100644 --- a/include/pgduckdb/pg/relations.hpp +++ b/include/pgduckdb/pg/relations.hpp @@ -18,7 +18,7 @@ const char *GetAttName(const Form_pg_attribute); Form_pg_attribute GetAttr(const TupleDesc tupleDesc, int i); -void EstimateRelSize(Relation rel, int32_t *attr_widths, BlockNumber *pages, double *tuples, double *allvisfrac); +double EstimateRelSize(Relation rel); Oid GetRelidFromSchemaAndTable(const char *, const char *); @@ -26,6 +26,4 @@ bool IsValidOid(Oid); bool IsValidBlockNumber(BlockNumber); -bool IsRelView(Relation); - } // namespace pgduckdb diff --git a/include/pgduckdb/pgduckdb_process_latch.hpp b/include/pgduckdb/pgduckdb_process_latch.hpp index f5b07a73..cd6210c8 100644 --- a/include/pgduckdb/pgduckdb_process_latch.hpp +++ b/include/pgduckdb/pgduckdb_process_latch.hpp @@ -16,7 +16,6 @@ extern void ResetLatch(Latch *latch); /* Defined in utils/wait_event.h */ #define PG_WAIT_EXTENSION 0x07000000U - } namespace pgduckdb { diff --git a/include/pgduckdb/scan/postgres_table_reader.hpp b/include/pgduckdb/scan/postgres_table_reader.hpp index db0b2359..76dfbd67 100644 --- a/include/pgduckdb/scan/postgres_table_reader.hpp +++ b/include/pgduckdb/scan/postgres_table_reader.hpp @@ -16,11 +16,14 @@ class PostgresTableReader { ~PostgresTableReader(); TupleTableSlot *GetNextTuple(); void PostgresTableReaderCleanup(); + private: MinimalTuple GetNextWorkerTuple(); int ParallelWorkerNumber(Cardinality cardinality); const char * ExplainScanPlan(QueryDesc *query_desc); + bool CanTableScanRunInParallel(Plan *plan); bool MarkPlanParallelAware(Plan *plan); + private: QueryDesc *table_scan_query_desc; PlanState *table_scan_planstate; diff --git a/src/catalog/pgduckdb_table.cpp b/src/catalog/pgduckdb_table.cpp index 61cc0619..05778d6a 100644 --- a/src/catalog/pgduckdb_table.cpp +++ b/src/catalog/pgduckdb_table.cpp @@ -46,15 +46,6 @@ PostgresTable::SetTableInfo(duckdb::CreateTableInfo &info, Relation rel) { } } -Cardinality -PostgresTable::GetTableCardinality(Relation rel) { - Cardinality cardinality; - BlockNumber n_pages; - double allvisfrac; - EstimateRelSize(rel, NULL, &n_pages, &cardinality, &allvisfrac); - return cardinality; -} - duckdb::unique_ptr PostgresTable::GetStatistics(duckdb::ClientContext &, duckdb::column_t) { throw duckdb::NotImplementedException("GetStatistics not supported yet"); diff --git a/src/catalog/pgduckdb_transaction.cpp b/src/catalog/pgduckdb_transaction.cpp index 3c314c6a..93fc0f91 100644 --- a/src/catalog/pgduckdb_transaction.cpp +++ b/src/catalog/pgduckdb_transaction.cpp @@ -39,6 +39,7 @@ SchemaItems::GetTable(const duckdb::string &entry_name) { } Oid rel_oid = GetRelidFromSchemaAndTable(name.c_str(), entry_name.c_str()); + if (!IsValidOid(rel_oid)) { return nullptr; // Table could not be found } @@ -49,7 +50,7 @@ SchemaItems::GetTable(const duckdb::string &entry_name) { info.table = entry_name; PostgresTable::SetTableInfo(info, rel); - auto cardinality = PostgresTable::GetTableCardinality(rel); + auto cardinality = EstimateRelSize(rel); tables.emplace(entry_name, duckdb::make_uniq(schema->catalog, *schema, info, rel, cardinality, schema->snapshot)); return tables[entry_name].get(); diff --git a/src/pg/relations.cpp b/src/pg/relations.cpp index 1bb1b6be..dbc10d04 100644 --- a/src/pg/relations.cpp +++ b/src/pg/relations.cpp @@ -67,9 +67,17 @@ CloseRelation(Relation rel) { CurrentResourceOwner = saveResourceOwner; } -void -EstimateRelSize(Relation rel, int32_t *attr_widths, BlockNumber *pages, double *tuples, double *allvisfrac) { - PostgresFunctionGuard(estimate_rel_size, rel, attr_widths, pages, tuples, allvisfrac); +double +EstimateRelSize(Relation rel) { + Cardinality cardinality = 0; + + if (RELKIND_HAS_TABLE_AM(rel->rd_rel->relkind) || rel->rd_rel->relkind == RELKIND_INDEX) { + BlockNumber pages; + double allvisfrac; + PostgresFunctionGuard(estimate_rel_size, rel, nullptr, &pages, &cardinality, &allvisfrac); + } + + return cardinality; } Oid @@ -91,11 +99,6 @@ IsValidOid(Oid oid) { return oid != InvalidOid; } -bool -IsRelView(Relation rel) { - return rel->rd_rel->relkind == RELKIND_VIEW; -} - bool IsValidBlockNumber(BlockNumber block_number) { return block_number != InvalidBlockNumber; diff --git a/src/pgduckdb_hooks.cpp b/src/pgduckdb_hooks.cpp index 6936edbe..53eb8f46 100644 --- a/src/pgduckdb_hooks.cpp +++ b/src/pgduckdb_hooks.cpp @@ -61,18 +61,34 @@ ContainsCatalogTable(List *rtes) { } static bool -ContainsPartitionedTable(List *rtes) { +ContainsAllowedTableType(List *rtes, int elevel) { foreach_node(RangeTblEntry, rte, rtes) { if (rte->rtekind == RTE_SUBQUERY) { /* Check whether any table in the subquery is a partitioned table */ - if (ContainsPartitionedTable(rte->subquery->rtable)) { + if (ContainsAllowedTableType(rte->subquery->rtable, elevel)) { return true; } } - if (rte->relkind == RELKIND_PARTITIONED_TABLE) { + /* Allow functions */ + if (rte->rtekind == RTE_FUNCTION) { return true; } + + switch (rte->relkind) { + case RELKIND_RELATION: + case RELKIND_INDEX: + case RELKIND_TOASTVALUE: + case RELKIND_VIEW: + case RELKIND_MATVIEW: + case RELKIND_FOREIGN_TABLE: + case RELKIND_PARTITIONED_TABLE: + case RELKIND_PARTITIONED_INDEX: + return true; + default: + elog(elevel, "DuckDB does not support querying table type '%c'", rte->relkind); + return false; + } } return false; } @@ -181,8 +197,7 @@ IsAllowedStatement(Query *query, bool throw_error = false) { /* * When accessing the partitioned table, we temporarily let PG handle it instead of DuckDB. */ - if (ContainsPartitionedTable(query->rtable)) { - elog(elevel, "DuckDB does not support querying PG partitioned table"); + if (!ContainsAllowedTableType(query->rtable, elevel)) { return false; } diff --git a/src/scan/postgres_table_reader.cpp b/src/scan/postgres_table_reader.cpp index e4b99e62..c5e8bb09 100644 --- a/src/scan/postgres_table_reader.cpp +++ b/src/scan/postgres_table_reader.cpp @@ -21,6 +21,8 @@ extern "C" { #include "pgduckdb/pgduckdb_guc.h" } +#include "pgduckdb/vendor/pg_list.hpp" + #include namespace pgduckdb { @@ -61,20 +63,19 @@ PostgresTableReader::PostgresTableReader(const char *table_scan_query, bool coun table_scan_planstate = PostgresFunctionGuard(ExecInitNode, planned_stmt->planTree, table_scan_query_desc->estate, 0); - bool marked_parallel_aware = false; + bool run_scan_with_parallel_workers = persistence != RELPERSISTENCE_TEMP; + run_scan_with_parallel_workers &= CanTableScanRunInParallel(table_scan_query_desc->planstate->plan); + + /* Temp tables can be excuted with parallel workers, and whole plan should be parallel aware */ + if (run_scan_with_parallel_workers) { - if (persistence != RELPERSISTENCE_TEMP) { if (count_tuples_only) { /* For count_tuples_only we will try to execute aggregate node on table scan */ planned_stmt->planTree->parallel_aware = true; - marked_parallel_aware = MarkPlanParallelAware((Plan *)table_scan_query_desc->planstate->plan->lefttree); + MarkPlanParallelAware((Plan *)table_scan_query_desc->planstate->plan->lefttree); } else { - marked_parallel_aware = MarkPlanParallelAware(table_scan_query_desc->planstate->plan); + MarkPlanParallelAware(table_scan_query_desc->planstate->plan); } - } - - /* Temp tables can be excuted with parallel workers, and whole plan should be parallel aware */ - if (persistence != RELPERSISTENCE_TEMP && marked_parallel_aware) { int parallel_workers = ParallelWorkerNumber(planned_stmt->planTree->plan_rows); bool interrupts_can_be_process = INTERRUPTS_CAN_BE_PROCESSED(); @@ -121,7 +122,8 @@ PostgresTableReader::~PostgresTableReader() { } } -void PostgresTableReader::PostgresTableReaderCleanup() { +void +PostgresTableReader::PostgresTableReaderCleanup() { std::lock_guard lock(GlobalProcessLock::GetLock()); PostgresScopedStackReset scoped_stack_reset; @@ -174,6 +176,30 @@ PostgresTableReader::ExplainScanPlan(QueryDesc *query_desc) { return PostgresFunctionGuard(ExplainScanPlan_Unsafe, query_desc); } +bool +PostgresTableReader::CanTableScanRunInParallel(Plan *plan) { + switch (nodeTag(plan)) { + case T_SeqScan: + case T_IndexScan: + case T_IndexOnlyScan: + case T_BitmapHeapScan: + return true; + case T_Append: { + ListCell *l; + foreach (l, ((Append *)plan)->appendplans) { + if (!CanTableScanRunInParallel((Plan *)lfirst(l))) { + return false; + } + } + return true; + } + /* This is special case for COUNT(*) */ + case T_Agg: + return true; + default: + return false; + } +} bool PostgresTableReader::MarkPlanParallelAware(Plan *plan) { @@ -184,6 +210,10 @@ PostgresTableReader::MarkPlanParallelAware(Plan *plan) { plan->parallel_aware = true; return true; } + case T_Append: { + plan->parallel_aware = true; + return true; + } case T_BitmapHeapScan: { plan->parallel_aware = true; return MarkPlanParallelAware(plan->lefttree); diff --git a/test/regression/expected/scan_postgres_tables.out b/test/regression/expected/scan_postgres_tables.out index 65622544..918912ee 100644 --- a/test/regression/expected/scan_postgres_tables.out +++ b/test/regression/expected/scan_postgres_tables.out @@ -159,6 +159,97 @@ Parallel Seq Scan on t1 (1 row) SET max_parallel_workers TO DEFAULT; +-- PARTITIONED TABLE +SET client_min_messages TO DEFAULT; +CREATE TABLE partitioned_table(a int, b INT, c text) PARTITION BY RANGE (a); +CREATE TABLE partition_1 PARTITION OF partitioned_table FOR VALUES FROM (0) TO (50); +CREATE TABLE partition_2 PARTITION OF partitioned_table FOR VALUES FROM (50) TO (100); +INSERT INTO partitioned_table SELECT g % 100, g, 'abcde_'||g from generate_series(1,100000) g; +CREATE INDEX ON partitioned_table(b); +SET client_min_messages TO DEBUG1; +SELECT COUNT(*) FROM partitioned_table WHERE a < 25; +DEBUG: (PGDuckdDB/PostgresTableReader) + +QUERY: SELECT a FROM public.partitioned_table WHERE (a<25 AND a IS NOT NULL) +RUNNING: ON 1 PARALLEL WORKER(S). +EXECUTING: +Parallel Seq Scan on partition_1 partitioned_table + Filter: ((a IS NOT NULL) AND (a < 25)) + + count +------- + 25000 +(1 row) + +SELECT COUNT(*) FROM partitioned_table WHERE a < 75; +DEBUG: (PGDuckdDB/PostgresTableReader) + +QUERY: SELECT a FROM public.partitioned_table WHERE (a<75 AND a IS NOT NULL) +RUNNING: ON 1 PARALLEL WORKER(S). +EXECUTING: +Parallel Append + -> Seq Scan on partition_1 partitioned_table_1 + Filter: ((a IS NOT NULL) AND (a < 75)) + -> Seq Scan on partition_2 partitioned_table_2 + Filter: ((a IS NOT NULL) AND (a < 75)) + + count +------- + 75000 +(1 row) + +SELECT COUNT(*) FROM partitioned_table WHERE a < 25 OR a > 75; +DEBUG: (PGDuckdDB/PostgresTableReader) + +QUERY: SELECT a FROM public.partitioned_table +RUNNING: ON 2 PARALLEL WORKER(S). +EXECUTING: +Parallel Append + -> Seq Scan on partition_1 partitioned_table_1 + -> Seq Scan on partition_2 partitioned_table_2 + + count +------- + 49000 +(1 row) + +SELECT COUNT(*) FROM partitioned_table WHERE a < 25 AND b = 1; +DEBUG: (PGDuckdDB/PostgresTableReader) + +QUERY: SELECT a, b FROM public.partitioned_table WHERE (a<25 AND a IS NOT NULL) AND (b=1 AND b IS NOT NULL) +RUNNING: ON 1 PARALLEL WORKER(S). +EXECUTING: +Parallel Index Scan using partition_1_b_idx on partition_1 partitioned_table + Index Cond: ((b IS NOT NULL) AND (b = 1)) + Filter: ((a IS NOT NULL) AND (a < 25)) + + count +------- + 1 +(1 row) + +SELECT COUNT(*) FROM partitioned_table, t2 WHERE partitioned_table.a = t2.a AND partitioned_table.a < 2; +DEBUG: (PGDuckdDB/PostgresTableReader) + +QUERY: SELECT a FROM public.partitioned_table WHERE (a<2 AND a IS NOT NULL) +RUNNING: ON 1 PARALLEL WORKER(S). +EXECUTING: +Parallel Seq Scan on partition_1 partitioned_table + Filter: ((a IS NOT NULL) AND (a < 2)) + +DEBUG: (PGDuckdDB/PostgresTableReader) + +QUERY: SELECT a FROM pg_temp.t2 WHERE (a<2 AND a IS NOT NULL AND a>=0 AND a<=1 AND a IS NOT NULL) +RUNNING: IN PROCESS THREAD. +EXECUTING: +Seq Scan on t2 + Filter: ((a IS NOT NULL) AND (a IS NOT NULL) AND (a < 2) AND (a >= 0) AND (a <= 1)) + + count +------- + 1000 +(1 row) + SET enable_bitmapscan TO DEFAULT; SET client_min_messages TO DEFAULT; -DROP TABLE t1, t2; +DROP TABLE t1, t2, partitioned_table; diff --git a/test/regression/sql/scan_postgres_tables.sql b/test/regression/sql/scan_postgres_tables.sql index faaa1d24..efff9649 100644 --- a/test/regression/sql/scan_postgres_tables.sql +++ b/test/regression/sql/scan_postgres_tables.sql @@ -41,6 +41,23 @@ SELECT COUNT(*) FROM t1 AS t1_1, t1 AS t1_2 WHERE t1_1.a < 2 AND t1_2.a > 8; SET max_parallel_workers TO DEFAULT; +-- PARTITIONED TABLE + +SET client_min_messages TO DEFAULT; +CREATE TABLE partitioned_table(a int, b INT, c text) PARTITION BY RANGE (a); +CREATE TABLE partition_1 PARTITION OF partitioned_table FOR VALUES FROM (0) TO (50); +CREATE TABLE partition_2 PARTITION OF partitioned_table FOR VALUES FROM (50) TO (100); +INSERT INTO partitioned_table SELECT g % 100, g, 'abcde_'||g from generate_series(1,100000) g; +CREATE INDEX ON partitioned_table(b); +SET client_min_messages TO DEBUG1; + +SELECT COUNT(*) FROM partitioned_table WHERE a < 25; +SELECT COUNT(*) FROM partitioned_table WHERE a < 75; +SELECT COUNT(*) FROM partitioned_table WHERE a < 25 OR a > 75; +SELECT COUNT(*) FROM partitioned_table WHERE a < 25 AND b = 1; +SELECT COUNT(*) FROM partitioned_table, t2 WHERE partitioned_table.a = t2.a AND partitioned_table.a < 2; + + SET enable_bitmapscan TO DEFAULT; SET client_min_messages TO DEFAULT; -DROP TABLE t1, t2; +DROP TABLE t1, t2, partitioned_table; From c850250030e22af33a7f0df2eb9c4b22f5ff008a Mon Sep 17 00:00:00 2001 From: Yves Date: Wed, 18 Dec 2024 14:30:59 +0100 Subject: [PATCH 21/53] Fix Mac build --- include/pgduckdb/pgduckdb_types.hpp | 4 ++-- include/pgduckdb/scan/postgres_scan.hpp | 2 +- src/scan/postgres_scan.cpp | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/include/pgduckdb/pgduckdb_types.hpp b/include/pgduckdb/pgduckdb_types.hpp index a47cd3c6..ebc7622d 100644 --- a/include/pgduckdb/pgduckdb_types.hpp +++ b/include/pgduckdb/pgduckdb_types.hpp @@ -7,8 +7,8 @@ namespace pgduckdb { -class PostgresScanGlobalState; -class PostgresScanLocalState; +struct PostgresScanGlobalState; +struct PostgresScanLocalState; // DuckDB has date starting from 1/1/1970 while PG starts from 1/1/2000 constexpr int32_t PGDUCKDB_DUCK_DATE_OFFSET = 10957; diff --git a/include/pgduckdb/scan/postgres_scan.hpp b/include/pgduckdb/scan/postgres_scan.hpp index ee4af2f2..a48f5a6d 100644 --- a/include/pgduckdb/scan/postgres_scan.hpp +++ b/include/pgduckdb/scan/postgres_scan.hpp @@ -45,7 +45,7 @@ struct PostgresScanLocalState : public duckdb::LocalTableFunctionState { public: PostgresScanGlobalState *global_state; - int output_vector_size; + size_t output_vector_size; bool exhausted_scan; }; diff --git a/src/scan/postgres_scan.cpp b/src/scan/postgres_scan.cpp index 31bb158e..0152fb61 100644 --- a/src/scan/postgres_scan.cpp +++ b/src/scan/postgres_scan.cpp @@ -203,7 +203,7 @@ PostgresScanTableFunction::PostgresScanFunction(duckdb::ClientContext &, duckdb: local_state.output_vector_size = 0; - int i = 0; + size_t i = 0; for (; i < STANDARD_VECTOR_SIZE; i++) { TupleTableSlot *slot = local_state.global_state->table_reader_global_state->GetNextTuple(); if (TupIsNull(slot)) { From fed5f370c878109b721044542c2f5f1aa02ff3af Mon Sep 17 00:00:00 2001 From: Yves Date: Wed, 18 Dec 2024 14:35:51 +0100 Subject: [PATCH 22/53] Quote column name --- src/scan/postgres_scan.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/scan/postgres_scan.cpp b/src/scan/postgres_scan.cpp index 0152fb61..edf3797f 100644 --- a/src/scan/postgres_scan.cpp +++ b/src/scan/postgres_scan.cpp @@ -108,7 +108,8 @@ PostgresScanGlobalState::ConstructTableScanQuery(duckdb::TableFunctionInitInput first = false; scan_query << "("; auto attr = table_tuple_desc->attrs[attr_num - 1]; - scan_query << filter->ToString(attr.attname.data).c_str(); + auto col = quote_identifier(attr.attname.data); + scan_query << filter->ToString(col).c_str(); scan_query << ") "; } } From 076f47a099faa9cf14ed8f1c0ecf2bc8041b6bba Mon Sep 17 00:00:00 2001 From: Yves Date: Wed, 18 Dec 2024 14:43:20 +0100 Subject: [PATCH 23/53] Vendor `RELKIND_HAS_TABLE_AM` for PG14 --- src/pg/relations.cpp | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/src/pg/relations.cpp b/src/pg/relations.cpp index dbc10d04..a0279435 100644 --- a/src/pg/relations.cpp +++ b/src/pg/relations.cpp @@ -17,6 +17,20 @@ namespace pgduckdb { #undef RelationGetDescr +#if PG_VERSION_NUM < 150000 +/* + * Relation kinds with a table access method (rd_tableam). Although sequences + * use the heap table AM, they are enough of a special case in most uses that + * they are not included here. Likewise, partitioned tables can have an access + * method defined so that their partitions can inherit it, but they do not set + * rd_tableam; hence, this is handled specially outside of this macro. + */ +#define RELKIND_HAS_TABLE_AM(relkind) \ + ((relkind) == RELKIND_RELATION || \ + (relkind) == RELKIND_TOASTVALUE || \ + (relkind) == RELKIND_MATVIEW) +#endif + TupleDesc RelationGetDescr(Relation rel) { return rel->rd_att; From 0e42b9ee7253d4c26a3c5b0170b3be9426c2d024 Mon Sep 17 00:00:00 2001 From: mkaruza Date: Thu, 19 Dec 2024 11:44:35 +0100 Subject: [PATCH 24/53] Calculate number of parallel workers on scan for count_tuples_only --- src/scan/postgres_table_reader.cpp | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/scan/postgres_table_reader.cpp b/src/scan/postgres_table_reader.cpp index c5e8bb09..c8256d3c 100644 --- a/src/scan/postgres_table_reader.cpp +++ b/src/scan/postgres_table_reader.cpp @@ -69,15 +69,17 @@ PostgresTableReader::PostgresTableReader(const char *table_scan_query, bool coun /* Temp tables can be excuted with parallel workers, and whole plan should be parallel aware */ if (run_scan_with_parallel_workers) { + int parallel_workers = 0; if (count_tuples_only) { /* For count_tuples_only we will try to execute aggregate node on table scan */ planned_stmt->planTree->parallel_aware = true; MarkPlanParallelAware((Plan *)table_scan_query_desc->planstate->plan->lefttree); + parallel_workers = ParallelWorkerNumber(planned_stmt->planTree->lefttree->plan_rows); } else { MarkPlanParallelAware(table_scan_query_desc->planstate->plan); + parallel_workers = ParallelWorkerNumber(planned_stmt->planTree->plan_rows); } - int parallel_workers = ParallelWorkerNumber(planned_stmt->planTree->plan_rows); bool interrupts_can_be_process = INTERRUPTS_CAN_BE_PROCESSED(); if (!interrupts_can_be_process) { From c92901dde54c128e8379f5bb284be7aeefdac894 Mon Sep 17 00:00:00 2001 From: mkaruza Date: Thu, 19 Dec 2024 11:51:27 +0100 Subject: [PATCH 25/53] Fixed regression test for previous commit --- test/regression/expected/scan_postgres_tables.out | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/regression/expected/scan_postgres_tables.out b/test/regression/expected/scan_postgres_tables.out index 918912ee..b6cca102 100644 --- a/test/regression/expected/scan_postgres_tables.out +++ b/test/regression/expected/scan_postgres_tables.out @@ -6,7 +6,7 @@ SELECT COUNT(*) FROM t1; DEBUG: (PGDuckdDB/PostgresTableReader) QUERY: SELECT COUNT(*) FROM public.t1 -RUNNING: ON 1 PARALLEL WORKER(S). +RUNNING: ON 2 PARALLEL WORKER(S). EXECUTING: Parallel Aggregate -> Parallel Seq Scan on t1 From 06c53753ff69ae80a2a982b2ce0a33dea26dc882 Mon Sep 17 00:00:00 2001 From: mkaruza Date: Fri, 20 Dec 2024 10:23:54 +0100 Subject: [PATCH 26/53] Refactor logic for WaitGlobalLatch --- include/pgduckdb/pgduckdb_process_latch.hpp | 62 +++++++++++++++++---- src/scan/postgres_table_reader.cpp | 7 ++- 2 files changed, 57 insertions(+), 12 deletions(-) diff --git a/include/pgduckdb/pgduckdb_process_latch.hpp b/include/pgduckdb/pgduckdb_process_latch.hpp index cd6210c8..14854086 100644 --- a/include/pgduckdb/pgduckdb_process_latch.hpp +++ b/include/pgduckdb/pgduckdb_process_latch.hpp @@ -2,17 +2,17 @@ #include #include +#include extern "C" { struct Latch; extern struct Latch *MyLatch; -extern int WaitLatch(Latch *latch, int wakeEvents, long timeout, - uint32_t wait_event_info); +extern int WaitLatch(Latch *latch, int wakeEvents, long timeout, uint32_t wait_event_info); extern void ResetLatch(Latch *latch); /* Defined in storage/latch.h */ -#define WL_LATCH_SET (1 << 0) -#define WL_EXIT_ON_PM_DEATH (1 << 5) +#define WL_LATCH_SET (1 << 0) +#define WL_EXIT_ON_PM_DEATH (1 << 5) /* Defined in utils/wait_event.h */ #define PG_WAIT_EXTENSION 0x07000000U @@ -27,17 +27,57 @@ namespace pgduckdb { struct GlobalProcessLatch { public: - static void WaitGlobalLatch() { - static std::condition_variable cv; + static void + WaitGlobalLatch() { static std::mutex lock; - if (lock.try_lock()) { + static std::condition_variable cv; + static std::atomic threads_waiting_for_latch_event = 0; + static bool latch_released = false; + static std::atomic latch_event_thread_done = false; + + // We exit if "latch" thread is waiting to finish + if (!latch_event_thread_done.load()) { + return; + } + + if (!threads_waiting_for_latch_event.fetch_add(1)) { + // Lets start waiting for latch on this thread (we are consider this as "latch" thread) + std::unique_lock lk(lock); + lock.lock(); + WaitLatch(MyLatch, WL_LATCH_SET | WL_EXIT_ON_PM_DEATH, 0, PG_WAIT_EXTENSION); ResetLatch(MyLatch); - cv.notify_all(); - lock.unlock(); + + // Use for notify waiting threads to exit + latch_released = true; + + // All new threads arriving on latch waiting will return immediately + latch_event_thread_done = true; + + // Notify one thread waiting + cv.notify_one(); + lk.unlock(); + + // Wait until we are only thread left + cv.wait(lk, [] { return threads_waiting_for_latch_event = 1; }); + threads_waiting_for_latch_event--; + + // Reset variables + latch_released = false; + latch_event_thread_done = false; } else { - std::unique_lock lk(lock); - cv.wait(lk); + std::unique_lock lk(lock); + + // Wait for "latch" thread to notify + cv.wait(lk, [] { return latch_released; }); + + // We are done with this threads + threads_waiting_for_latch_event--; + + lk.unlock(); + + // Notify another thread (either waiting thread or "latch" thread) + cv.notify_one(); } } }; diff --git a/src/scan/postgres_table_reader.cpp b/src/scan/postgres_table_reader.cpp index c8256d3c..50b7ece8 100644 --- a/src/scan/postgres_table_reader.cpp +++ b/src/scan/postgres_table_reader.cpp @@ -272,7 +272,12 @@ PostgresTableReader::GetNextWorkerTuple() { bool readerdone; reader = (TupleQueueReader *)parallel_worker_readers[next_parallel_reader]; - minimal_tuple = PostgresFunctionGuard(TupleQueueReaderNext, reader, true, &readerdone); + + { + // We need to take global lock for `TupleQueueReaderNext` call + std::lock_guard lock(GlobalProcessLock::GetLock()); + minimal_tuple = PostgresFunctionGuard(TupleQueueReaderNext, reader, true, &readerdone); + } if (readerdone) { --nreaders; From 5aaa154661953c25509c559e4d0b5027bb77bf9b Mon Sep 17 00:00:00 2001 From: Jelte Fennema-Nio Date: Fri, 20 Dec 2024 14:13:14 +0100 Subject: [PATCH 27/53] Simplify latch waiting logic --- include/pgduckdb/pg/latch.hpp | 7 ++ include/pgduckdb/pgduckdb_process_latch.hpp | 85 ------------------- .../pgduckdb/scan/postgres_table_reader.hpp | 1 + src/pg/latch.cpp | 48 +++++++++++ src/scan/postgres_table_reader.cpp | 8 +- 5 files changed, 61 insertions(+), 88 deletions(-) create mode 100644 include/pgduckdb/pg/latch.hpp delete mode 100644 include/pgduckdb/pgduckdb_process_latch.hpp create mode 100644 src/pg/latch.cpp diff --git a/include/pgduckdb/pg/latch.hpp b/include/pgduckdb/pg/latch.hpp new file mode 100644 index 00000000..a78d80b4 --- /dev/null +++ b/include/pgduckdb/pg/latch.hpp @@ -0,0 +1,7 @@ +#pragma once + +#include + +namespace pgduckdb::pg { +void WaitMyLatch(uint64_t &last_known_latch_update_count); +} // namespace pgduckdb::pg diff --git a/include/pgduckdb/pgduckdb_process_latch.hpp b/include/pgduckdb/pgduckdb_process_latch.hpp deleted file mode 100644 index 14854086..00000000 --- a/include/pgduckdb/pgduckdb_process_latch.hpp +++ /dev/null @@ -1,85 +0,0 @@ -#pragma once - -#include -#include -#include - -extern "C" { -struct Latch; -extern struct Latch *MyLatch; -extern int WaitLatch(Latch *latch, int wakeEvents, long timeout, uint32_t wait_event_info); -extern void ResetLatch(Latch *latch); - -/* Defined in storage/latch.h */ -#define WL_LATCH_SET (1 << 0) -#define WL_EXIT_ON_PM_DEATH (1 << 5) - -/* Defined in utils/wait_event.h */ -#define PG_WAIT_EXTENSION 0x07000000U -} - -namespace pgduckdb { - -/* - * GlobalProcessLatch is used to wait on process postgres latch. First thread that enters will wait for latch while - * others will wait until this latch is released. - */ - -struct GlobalProcessLatch { -public: - static void - WaitGlobalLatch() { - static std::mutex lock; - static std::condition_variable cv; - static std::atomic threads_waiting_for_latch_event = 0; - static bool latch_released = false; - static std::atomic latch_event_thread_done = false; - - // We exit if "latch" thread is waiting to finish - if (!latch_event_thread_done.load()) { - return; - } - - if (!threads_waiting_for_latch_event.fetch_add(1)) { - // Lets start waiting for latch on this thread (we are consider this as "latch" thread) - std::unique_lock lk(lock); - lock.lock(); - - WaitLatch(MyLatch, WL_LATCH_SET | WL_EXIT_ON_PM_DEATH, 0, PG_WAIT_EXTENSION); - ResetLatch(MyLatch); - - // Use for notify waiting threads to exit - latch_released = true; - - // All new threads arriving on latch waiting will return immediately - latch_event_thread_done = true; - - // Notify one thread waiting - cv.notify_one(); - lk.unlock(); - - // Wait until we are only thread left - cv.wait(lk, [] { return threads_waiting_for_latch_event = 1; }); - threads_waiting_for_latch_event--; - - // Reset variables - latch_released = false; - latch_event_thread_done = false; - } else { - std::unique_lock lk(lock); - - // Wait for "latch" thread to notify - cv.wait(lk, [] { return latch_released; }); - - // We are done with this threads - threads_waiting_for_latch_event--; - - lk.unlock(); - - // Notify another thread (either waiting thread or "latch" thread) - cv.notify_one(); - } - } -}; - -} // namespace pgduckdb diff --git a/include/pgduckdb/scan/postgres_table_reader.hpp b/include/pgduckdb/scan/postgres_table_reader.hpp index 76dfbd67..c1fab87d 100644 --- a/include/pgduckdb/scan/postgres_table_reader.hpp +++ b/include/pgduckdb/scan/postgres_table_reader.hpp @@ -34,6 +34,7 @@ class PostgresTableReader { int nreaders; int next_parallel_reader; bool entered_parallel_mode; + uint64_t last_known_latch_update_count; }; } // namespace pgduckdb diff --git a/src/pg/latch.cpp b/src/pg/latch.cpp new file mode 100644 index 00000000..0afdac17 --- /dev/null +++ b/src/pg/latch.cpp @@ -0,0 +1,48 @@ + +#include "pgduckdb/pgduckdb_utils.hpp" + +extern "C" { +#include "postgres.h" + +#include "miscadmin.h" +#include "storage/latch.h" +#include "utils/wait_event.h" +} + +namespace pgduckdb::pg { + +/* + * WaitGlobalLatch is used to wait on process postgres its MyLatch. To make + * sure one thread doesn't miss any updates on the latch, which were received by + * another thread. We use a global counter to keep track of the number of times + * the latch was updated. If the counter is different from the last known + * counter, we know another thread consumed an update that might have been + * intended for us, so we simply return without waiting. + * + * NOTE: We need to take the GlobalProcessLock because WaitLatch can throw + * errors, which modifies global variables. So to avoid dataraces conditions, + * we need take the global lock. This also obviously modifies the MyLatch + * global, so even if WaitLatch did not throw errors. We'd still need to take + * the global lock for that, or at the very least we'd need to take a specific + * lock for MyLatch and ensure that no other code modifies MyLatch. + */ +void +WaitMyLatch(uint64_t &last_known_latch_update_count) { + // A 64bit uint takes 584 years to overflow when it's updated every + // nanosecond. That's more than enough for any practical purposes. If + // someone keeps a Postgres connection open for more than 584 years, + // they deserve what they get. + static uint64_t latch_update_count = 0; + + std::lock_guard lock(GlobalProcessLock::GetLock()); + if (last_known_latch_update_count == latch_update_count) { + PostgresFunctionGuard(WaitLatch, MyLatch, WL_LATCH_SET | WL_EXIT_ON_PM_DEATH, 0, PG_WAIT_EXTENSION); + /* Note: No need to use PostgresFunctionGuard here, because ResetLatch + * is a trivial function */ + ResetLatch(MyLatch); + latch_update_count++; + } + last_known_latch_update_count = latch_update_count; +} + +} // namespace pgduckdb::pg diff --git a/src/scan/postgres_table_reader.cpp b/src/scan/postgres_table_reader.cpp index 50b7ece8..1aa0ba64 100644 --- a/src/scan/postgres_table_reader.cpp +++ b/src/scan/postgres_table_reader.cpp @@ -1,7 +1,7 @@ #include "pgduckdb/scan/postgres_table_reader.hpp" -#include "pgduckdb/pgduckdb_process_latch.hpp" #include "pgduckdb/pgduckdb_process_lock.hpp" #include "pgduckdb/pgduckdb_utils.hpp" +#include "pgduckdb/pg/latch.hpp" extern "C" { #include "postgres.h" @@ -29,7 +29,8 @@ namespace pgduckdb { PostgresTableReader::PostgresTableReader(const char *table_scan_query, bool count_tuples_only) : parallel_executor_info(nullptr), parallel_worker_readers(nullptr), nreaders(0), next_parallel_reader(0), - entered_parallel_mode(false) { + entered_parallel_mode(false), last_known_latch_update_count(0) { + std::lock_guard lock(GlobalProcessLock::GetLock()); PostgresScopedStackReset scoped_stack_reset; @@ -302,7 +303,8 @@ PostgresTableReader::GetNextWorkerTuple() { nvisited++; if (nvisited >= nreaders) { - GlobalProcessLatch::WaitGlobalLatch(); + std::lock_guard lock(GlobalProcessLock::GetLock()); + pgduckdb::pg::WaitMyLatch(last_known_latch_update_count); nvisited = 0; } } From 5da2802907b82bfbbf303a7ac6cae6ce62400987 Mon Sep 17 00:00:00 2001 From: Jelte Fennema-Nio Date: Fri, 20 Dec 2024 14:47:48 +0100 Subject: [PATCH 28/53] Add comment to PostgresScopedStackReset --- include/pgduckdb/pgduckdb_utils.hpp | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/include/pgduckdb/pgduckdb_utils.hpp b/include/pgduckdb/pgduckdb_utils.hpp index a6906c23..192865ae 100644 --- a/include/pgduckdb/pgduckdb_utils.hpp +++ b/include/pgduckdb/pgduckdb_utils.hpp @@ -49,7 +49,21 @@ struct PgExceptionGuard { }; /* - * DuckdbGlobalLock should be held before calling. + * PostgresScopedStackReset is a RAII class that saves the current stack base + * and restores it on destruction. When calling certain Postgres C functions + * from other threads than the main thread this is necessary to avoid Postgres + * throwing an error running out of stack space. In codepaths that postgres + * expects to be called recursively it checks if the stack size is still within + * the limit set by max_stack_depth. It does so by comparing the current stack + * pointer to the pointer it saved when starting the process. But since + * different threads have different stacks, this check will fail basically + * automatically if the thread is not the main thread. This class is a + * workaround for this problem, by configuring a new stack base matching the + * current location of the stack. This does mean that the stack might grow + * higher than, but for our use case this shouldn't matter anyway because we + * don't expect any recursive functions to be called. And even if we did expect + * that, the default max_stack_depth is conservative enough to handle this small + * bit of extra stack space. */ struct PostgresScopedStackReset { PostgresScopedStackReset() { From 7c5ab0db78b4126f50abdebe585c4cff9c9f7da5 Mon Sep 17 00:00:00 2001 From: Jelte Fennema-Nio Date: Fri, 20 Dec 2024 14:57:59 +0100 Subject: [PATCH 29/53] clang-format off around vendored code --- src/pg/relations.cpp | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/pg/relations.cpp b/src/pg/relations.cpp index a0279435..caa21c18 100644 --- a/src/pg/relations.cpp +++ b/src/pg/relations.cpp @@ -18,6 +18,8 @@ namespace pgduckdb { #undef RelationGetDescr #if PG_VERSION_NUM < 150000 +// clang-format off + /* * Relation kinds with a table access method (rd_tableam). Although sequences * use the heap table AM, they are enough of a special case in most uses that @@ -29,6 +31,8 @@ namespace pgduckdb { ((relkind) == RELKIND_RELATION || \ (relkind) == RELKIND_TOASTVALUE || \ (relkind) == RELKIND_MATVIEW) + +// clang-format on #endif TupleDesc From fe6b73b29f4697983a8f9f618268a8515af87c41 Mon Sep 17 00:00:00 2001 From: Jelte Fennema-Nio Date: Fri, 20 Dec 2024 15:34:00 +0100 Subject: [PATCH 30/53] Move ConstructFullyQualifiedTableName to relations.cpp --- include/pgduckdb/pg/relations.hpp | 2 ++ include/pgduckdb/scan/postgres_scan.hpp | 3 --- src/pg/relations.cpp | 22 ++++++++++++++++++++++ src/pgduckdb_hooks.cpp | 17 ----------------- src/scan/postgres_scan.cpp | 11 +++-------- 5 files changed, 27 insertions(+), 28 deletions(-) diff --git a/include/pgduckdb/pg/relations.hpp b/include/pgduckdb/pg/relations.hpp index 099e39fe..9c07c233 100644 --- a/include/pgduckdb/pg/relations.hpp +++ b/include/pgduckdb/pg/relations.hpp @@ -26,4 +26,6 @@ bool IsValidOid(Oid); bool IsValidBlockNumber(BlockNumber); +char *GenerateQualifiedRelationName(Relation rel); + } // namespace pgduckdb diff --git a/include/pgduckdb/scan/postgres_scan.hpp b/include/pgduckdb/scan/postgres_scan.hpp index a48f5a6d..5b45136d 100644 --- a/include/pgduckdb/scan/postgres_scan.hpp +++ b/include/pgduckdb/scan/postgres_scan.hpp @@ -22,9 +22,6 @@ struct PostgresScanGlobalState : public duckdb::GlobalTableFunctionState { } void ConstructTableScanQuery(duckdb::TableFunctionInitInput &input); -private: - std::string ConstructFullyQualifiedTableName(); - public: Snapshot snapshot; Relation rel; diff --git a/src/pg/relations.cpp b/src/pg/relations.cpp index caa21c18..3524d3b0 100644 --- a/src/pg/relations.cpp +++ b/src/pg/relations.cpp @@ -8,6 +8,8 @@ extern "C" { #include "access/relation.h" // relation_open and relation_close #include "catalog/namespace.h" // makeRangeVarFromNameList, RangeVarGetRelid #include "optimizer/plancat.h" // estimate_rel_size +#include "utils/builtins.h" +#include "utils/lsyscache.h" #include "utils/rel.h" #include "utils/resowner.h" // CurrentResourceOwner and TopTransactionResourceOwner #include "utils/syscache.h" // RELOID @@ -122,4 +124,24 @@ IsValidBlockNumber(BlockNumber block_number) { return block_number != InvalidBlockNumber; } +/* + * generate_qualified_relation_name + * Compute the name to display for a relation specified by OID + * + * As above, but unconditionally schema-qualify the name. + */ +static char * +GenerateQualifiedRelationName_Unsafe(Relation rel) { + char *nspname = get_namespace_name_or_temp(rel->rd_rel->relnamespace); + if (!nspname) + elog(ERROR, "cache lookup failed for namespace %u", rel->rd_rel->relnamespace); + + return quote_qualified_identifier(nspname, NameStr(rel->rd_rel->relname)); +} + +char * +GenerateQualifiedRelationName(Relation rel) { + return PostgresFunctionGuard(GenerateQualifiedRelationName_Unsafe, rel); +} + } // namespace pgduckdb diff --git a/src/pgduckdb_hooks.cpp b/src/pgduckdb_hooks.cpp index 53eb8f46..44147c11 100644 --- a/src/pgduckdb_hooks.cpp +++ b/src/pgduckdb_hooks.cpp @@ -184,23 +184,6 @@ IsAllowedStatement(Query *query, bool throw_error = false) { return false; } - /* - * If any table is from pg_catalog, we don't want to use DuckDB. This is - * because DuckDB has its own pg_catalog tables that contain different data - * then Postgres its pg_catalog tables. - */ - if (ContainsCatalogTable(query->rtable)) { - elog(elevel, "DuckDB does not support querying PG catalog tables"); - return false; - } - - /* - * When accessing the partitioned table, we temporarily let PG handle it instead of DuckDB. - */ - if (!ContainsAllowedTableType(query->rtable, elevel)) { - return false; - } - /* Anything else is hopefully fine... */ return true; } diff --git a/src/scan/postgres_scan.cpp b/src/scan/postgres_scan.cpp index edf3797f..6ee80ca0 100644 --- a/src/scan/postgres_scan.cpp +++ b/src/scan/postgres_scan.cpp @@ -2,6 +2,7 @@ #include "pgduckdb/scan/postgres_table_reader.hpp" #include "pgduckdb/pgduckdb_types.hpp" #include "pgduckdb/pgduckdb_utils.hpp" +#include "pgduckdb/pg/relations.hpp" #include "pgduckdb/pgduckdb_process_lock.hpp" #include "pgduckdb/logger.hpp" @@ -25,7 +26,7 @@ void PostgresScanGlobalState::ConstructTableScanQuery(duckdb::TableFunctionInitInput &input) { /* SELECT COUNT(*) FROM */ if (input.column_ids.size() == 1 && input.column_ids[0] == UINT64_MAX) { - scan_query << "SELECT COUNT(*) FROM " << ConstructFullyQualifiedTableName(); + scan_query << "SELECT COUNT(*) FROM " << pgduckdb::GenerateQualifiedRelationName(rel); count_tuples_only = true; return; } @@ -88,7 +89,7 @@ PostgresScanGlobalState::ConstructTableScanQuery(duckdb::TableFunctionInitInput scan_query << quote_identifier(attr.attname.data); } - scan_query << " FROM " << ConstructFullyQualifiedTableName(); + scan_query << " FROM " << GenerateQualifiedRelationName(rel); first = true; @@ -114,12 +115,6 @@ PostgresScanGlobalState::ConstructTableScanQuery(duckdb::TableFunctionInitInput } } -std::string -PostgresScanGlobalState::ConstructFullyQualifiedTableName() { - return psprintf("%s.%s", quote_identifier(get_namespace_name_or_temp(get_rel_namespace(rel->rd_rel->oid))), - quote_identifier(get_rel_name(rel->rd_rel->oid))); -} - PostgresScanGlobalState::PostgresScanGlobalState(Snapshot _snapshot, Relation _rel, duckdb::TableFunctionInitInput &input) : snapshot(_snapshot), rel(_rel), table_tuple_desc(RelationGetDescr(rel)), count_tuples_only(false), From 1db3d34dab0eeed6bc975eee2f8f4667276c2199 Mon Sep 17 00:00:00 2001 From: Jelte Fennema-Nio Date: Fri, 20 Dec 2024 15:34:11 +0100 Subject: [PATCH 31/53] Remove accidental lock --- src/pg/latch.cpp | 1 - 1 file changed, 1 deletion(-) diff --git a/src/pg/latch.cpp b/src/pg/latch.cpp index 0afdac17..ba870bab 100644 --- a/src/pg/latch.cpp +++ b/src/pg/latch.cpp @@ -34,7 +34,6 @@ WaitMyLatch(uint64_t &last_known_latch_update_count) { // they deserve what they get. static uint64_t latch_update_count = 0; - std::lock_guard lock(GlobalProcessLock::GetLock()); if (last_known_latch_update_count == latch_update_count) { PostgresFunctionGuard(WaitLatch, MyLatch, WL_LATCH_SET | WL_EXIT_ON_PM_DEATH, 0, PG_WAIT_EXTENSION); /* Note: No need to use PostgresFunctionGuard here, because ResetLatch From bad8e19e0e538f9075e72431ac3ae0b728b3394f Mon Sep 17 00:00:00 2001 From: Jelte Fennema-Nio Date: Fri, 20 Dec 2024 15:42:04 +0100 Subject: [PATCH 32/53] Move QuoteIdentifier to relations.cpp --- include/pgduckdb/pg/relations.hpp | 1 + src/pg/relations.cpp | 5 +++++ src/scan/postgres_scan.cpp | 7 ++----- 3 files changed, 8 insertions(+), 5 deletions(-) diff --git a/include/pgduckdb/pg/relations.hpp b/include/pgduckdb/pg/relations.hpp index 9c07c233..4876d825 100644 --- a/include/pgduckdb/pg/relations.hpp +++ b/include/pgduckdb/pg/relations.hpp @@ -27,5 +27,6 @@ bool IsValidOid(Oid); bool IsValidBlockNumber(BlockNumber); char *GenerateQualifiedRelationName(Relation rel); +const char *QuoteIdentifier(const char *ident); } // namespace pgduckdb diff --git a/src/pg/relations.cpp b/src/pg/relations.cpp index 3524d3b0..ad3ca05c 100644 --- a/src/pg/relations.cpp +++ b/src/pg/relations.cpp @@ -144,4 +144,9 @@ GenerateQualifiedRelationName(Relation rel) { return PostgresFunctionGuard(GenerateQualifiedRelationName_Unsafe, rel); } +const char * +QuoteIdentifier(const char *ident) { + return PostgresFunctionGuard(quote_identifier, ident); +} + } // namespace pgduckdb diff --git a/src/scan/postgres_scan.cpp b/src/scan/postgres_scan.cpp index 6ee80ca0..9f1b3161 100644 --- a/src/scan/postgres_scan.cpp +++ b/src/scan/postgres_scan.cpp @@ -9,10 +9,7 @@ extern "C" { #include "postgres.h" -#include "access/htup_details.h" #include "executor/tuptable.h" -#include "utils/builtins.h" -#include "utils/lsyscache.h" #include "utils/rel.h" } @@ -86,7 +83,7 @@ PostgresScanGlobalState::ConstructTableScanQuery(duckdb::TableFunctionInitInput } first = false; auto attr = table_tuple_desc->attrs[attr_num - 1]; - scan_query << quote_identifier(attr.attname.data); + scan_query << pgduckdb::QuoteIdentifier(attr.attname.data); } scan_query << " FROM " << GenerateQualifiedRelationName(rel); @@ -109,7 +106,7 @@ PostgresScanGlobalState::ConstructTableScanQuery(duckdb::TableFunctionInitInput first = false; scan_query << "("; auto attr = table_tuple_desc->attrs[attr_num - 1]; - auto col = quote_identifier(attr.attname.data); + auto col = pgduckdb::QuoteIdentifier(attr.attname.data); scan_query << filter->ToString(col).c_str(); scan_query << ") "; } From 534c608ad9e26d2ec37b6a5b5cefcb94d49b3afd Mon Sep 17 00:00:00 2001 From: Jelte Fennema-Nio Date: Fri, 20 Dec 2024 15:42:12 +0100 Subject: [PATCH 33/53] Undo accidental removal --- src/pgduckdb_hooks.cpp | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/src/pgduckdb_hooks.cpp b/src/pgduckdb_hooks.cpp index 44147c11..46f2bf5d 100644 --- a/src/pgduckdb_hooks.cpp +++ b/src/pgduckdb_hooks.cpp @@ -184,6 +184,24 @@ IsAllowedStatement(Query *query, bool throw_error = false) { return false; } + /* + * If any table is from pg_catalog, we don't want to use DuckDB. This is + * because DuckDB has its own pg_catalog tables that contain different data + * then Postgres its pg_catalog tables. + */ + if (ContainsCatalogTable(query->rtable)) { + elog(elevel, "DuckDB does not support querying PG catalog tables"); + return false; + } + + /* + * Check that we're only accessing table types that we know we have support + * for. + */ + if (!ContainsAllowedTableType(query->rtable, elevel)) { + return false; + } + /* Anything else is hopefully fine... */ return true; } From 6767f962441bed87427cdb4450cbd9c664671d6e Mon Sep 17 00:00:00 2001 From: Jelte Fennema-Nio Date: Fri, 20 Dec 2024 16:01:16 +0100 Subject: [PATCH 34/53] Make postgres_scan.cpp c++ only --- include/pgduckdb/pg/relations.hpp | 4 ++++ src/pg/relations.cpp | 15 +++++++++++++-- src/scan/postgres_scan.cpp | 18 ++++++------------ 3 files changed, 23 insertions(+), 14 deletions(-) diff --git a/include/pgduckdb/pg/relations.hpp b/include/pgduckdb/pg/relations.hpp index 4876d825..cafdd3cc 100644 --- a/include/pgduckdb/pg/relations.hpp +++ b/include/pgduckdb/pg/relations.hpp @@ -18,6 +18,10 @@ const char *GetAttName(const Form_pg_attribute); Form_pg_attribute GetAttr(const TupleDesc tupleDesc, int i); +bool TupleIsNull(TupleTableSlot *slot); + +void SlotGetAllAttrs(TupleTableSlot *slot); + double EstimateRelSize(Relation rel); Oid GetRelidFromSchemaAndTable(const char *, const char *); diff --git a/src/pg/relations.cpp b/src/pg/relations.cpp index ad3ca05c..5afb2cb4 100644 --- a/src/pg/relations.cpp +++ b/src/pg/relations.cpp @@ -11,8 +11,9 @@ extern "C" { #include "utils/builtins.h" #include "utils/lsyscache.h" #include "utils/rel.h" -#include "utils/resowner.h" // CurrentResourceOwner and TopTransactionResourceOwner -#include "utils/syscache.h" // RELOID +#include "utils/resowner.h" // CurrentResourceOwner and TopTransactionResourceOwner +#include "executor/tuptable.h" // TupIsNull +#include "utils/syscache.h" // RELOID } namespace pgduckdb { @@ -57,6 +58,16 @@ GetAttr(const TupleDesc tupleDesc, int i) { return &tupleDesc->attrs[i]; } +bool +TupleIsNull(TupleTableSlot *slot) { + return TupIsNull(slot); +} + +void +SlotGetAllAttrs(TupleTableSlot *slot) { + PostgresFunctionGuard(slot_getallattrs, slot); +} + Relation OpenRelation(Oid relationId) { /* diff --git a/src/scan/postgres_scan.cpp b/src/scan/postgres_scan.cpp index 9f1b3161..fcc664b2 100644 --- a/src/scan/postgres_scan.cpp +++ b/src/scan/postgres_scan.cpp @@ -7,12 +7,6 @@ #include "pgduckdb/pgduckdb_process_lock.hpp" #include "pgduckdb/logger.hpp" -extern "C" { -#include "postgres.h" -#include "executor/tuptable.h" -#include "utils/rel.h" -} - namespace pgduckdb { // @@ -82,8 +76,8 @@ PostgresScanGlobalState::ConstructTableScanQuery(duckdb::TableFunctionInitInput scan_query << ", "; } first = false; - auto attr = table_tuple_desc->attrs[attr_num - 1]; - scan_query << pgduckdb::QuoteIdentifier(attr.attname.data); + auto attr = GetAttr(table_tuple_desc, attr_num - 1); + scan_query << pgduckdb::QuoteIdentifier(GetAttName(attr)); } scan_query << " FROM " << GenerateQualifiedRelationName(rel); @@ -105,8 +99,8 @@ PostgresScanGlobalState::ConstructTableScanQuery(duckdb::TableFunctionInitInput first = false; scan_query << "("; - auto attr = table_tuple_desc->attrs[attr_num - 1]; - auto col = pgduckdb::QuoteIdentifier(attr.attname.data); + auto attr = GetAttr(table_tuple_desc, attr_num - 1); + auto col = pgduckdb::QuoteIdentifier(GetAttName(attr)); scan_query << filter->ToString(col).c_str(); scan_query << ") "; } @@ -199,12 +193,12 @@ PostgresScanTableFunction::PostgresScanFunction(duckdb::ClientContext &, duckdb: size_t i = 0; for (; i < STANDARD_VECTOR_SIZE; i++) { TupleTableSlot *slot = local_state.global_state->table_reader_global_state->GetNextTuple(); - if (TupIsNull(slot)) { + if (pgduckdb::TupleIsNull(slot)) { local_state.global_state->table_reader_global_state->PostgresTableReaderCleanup(); local_state.exhausted_scan = true; break; } - slot_getallattrs(slot); + SlotGetAllAttrs(slot); InsertTupleIntoChunk(output, local_state, slot); } From 45c5ab4b9f0507aa777d670735227e4ae07067ee Mon Sep 17 00:00:00 2001 From: Jelte Fennema-Nio Date: Fri, 20 Dec 2024 16:09:33 +0100 Subject: [PATCH 35/53] Take GlobalProcessLock for the duration of a chunk At the very least SlotGetAllAttrs should should be called while holding the lock. --- src/scan/postgres_scan.cpp | 1 + src/scan/postgres_table_reader.cpp | 6 +----- 2 files changed, 2 insertions(+), 5 deletions(-) diff --git a/src/scan/postgres_scan.cpp b/src/scan/postgres_scan.cpp index fcc664b2..417aa7a8 100644 --- a/src/scan/postgres_scan.cpp +++ b/src/scan/postgres_scan.cpp @@ -191,6 +191,7 @@ PostgresScanTableFunction::PostgresScanFunction(duckdb::ClientContext &, duckdb: local_state.output_vector_size = 0; size_t i = 0; + std::lock_guard lock(GlobalProcessLock::GetLock()); for (; i < STANDARD_VECTOR_SIZE; i++) { TupleTableSlot *slot = local_state.global_state->table_reader_global_state->GetNextTuple(); if (pgduckdb::TupleIsNull(slot)) { diff --git a/src/scan/postgres_table_reader.cpp b/src/scan/postgres_table_reader.cpp index 1aa0ba64..db7e77c3 100644 --- a/src/scan/postgres_table_reader.cpp +++ b/src/scan/postgres_table_reader.cpp @@ -121,14 +121,13 @@ PostgresTableReader::PostgresTableReader(const char *table_scan_query, bool coun PostgresTableReader::~PostgresTableReader() { if (table_scan_query_desc) { + std::lock_guard lock(GlobalProcessLock::GetLock()); PostgresTableReaderCleanup(); } } void PostgresTableReader::PostgresTableReaderCleanup() { - std::lock_guard lock(GlobalProcessLock::GetLock()); - PostgresScopedStackReset scoped_stack_reset; PostgresFunctionGuard(ExecEndNode, table_scan_planstate); @@ -251,7 +250,6 @@ PostgresTableReader::GetNextTuple() { return slot; } } else { - std::lock_guard lock(GlobalProcessLock::GetLock()); PostgresScopedStackReset scoped_stack_reset; table_scan_query_desc->estate->es_query_dsa = parallel_executor_info ? parallel_executor_info->area : NULL; thread_scan_slot = PostgresFunctionGuard(ExecProcNode, table_scan_planstate); @@ -276,7 +274,6 @@ PostgresTableReader::GetNextWorkerTuple() { { // We need to take global lock for `TupleQueueReaderNext` call - std::lock_guard lock(GlobalProcessLock::GetLock()); minimal_tuple = PostgresFunctionGuard(TupleQueueReaderNext, reader, true, &readerdone); } @@ -303,7 +300,6 @@ PostgresTableReader::GetNextWorkerTuple() { nvisited++; if (nvisited >= nreaders) { - std::lock_guard lock(GlobalProcessLock::GetLock()); pgduckdb::pg::WaitMyLatch(last_known_latch_update_count); nvisited = 0; } From 77df0e1837c2018079a134c0177d0a1e203d6629 Mon Sep 17 00:00:00 2001 From: mkaruza Date: Mon, 23 Dec 2024 09:40:07 +0100 Subject: [PATCH 36/53] GlobalLock held for duration of populating output vector If GlobaclLock is held for duration of populating output vector we don't need any special WaitLatch wrappers and we can call directly this function. --- include/pgduckdb/pg/latch.hpp | 7 --- .../pgduckdb/scan/postgres_table_reader.hpp | 1 - src/pg/latch.cpp | 47 ------------------- src/scan/postgres_table_reader.cpp | 22 +++++---- 4 files changed, 14 insertions(+), 63 deletions(-) delete mode 100644 include/pgduckdb/pg/latch.hpp delete mode 100644 src/pg/latch.cpp diff --git a/include/pgduckdb/pg/latch.hpp b/include/pgduckdb/pg/latch.hpp deleted file mode 100644 index a78d80b4..00000000 --- a/include/pgduckdb/pg/latch.hpp +++ /dev/null @@ -1,7 +0,0 @@ -#pragma once - -#include - -namespace pgduckdb::pg { -void WaitMyLatch(uint64_t &last_known_latch_update_count); -} // namespace pgduckdb::pg diff --git a/include/pgduckdb/scan/postgres_table_reader.hpp b/include/pgduckdb/scan/postgres_table_reader.hpp index c1fab87d..76dfbd67 100644 --- a/include/pgduckdb/scan/postgres_table_reader.hpp +++ b/include/pgduckdb/scan/postgres_table_reader.hpp @@ -34,7 +34,6 @@ class PostgresTableReader { int nreaders; int next_parallel_reader; bool entered_parallel_mode; - uint64_t last_known_latch_update_count; }; } // namespace pgduckdb diff --git a/src/pg/latch.cpp b/src/pg/latch.cpp deleted file mode 100644 index ba870bab..00000000 --- a/src/pg/latch.cpp +++ /dev/null @@ -1,47 +0,0 @@ - -#include "pgduckdb/pgduckdb_utils.hpp" - -extern "C" { -#include "postgres.h" - -#include "miscadmin.h" -#include "storage/latch.h" -#include "utils/wait_event.h" -} - -namespace pgduckdb::pg { - -/* - * WaitGlobalLatch is used to wait on process postgres its MyLatch. To make - * sure one thread doesn't miss any updates on the latch, which were received by - * another thread. We use a global counter to keep track of the number of times - * the latch was updated. If the counter is different from the last known - * counter, we know another thread consumed an update that might have been - * intended for us, so we simply return without waiting. - * - * NOTE: We need to take the GlobalProcessLock because WaitLatch can throw - * errors, which modifies global variables. So to avoid dataraces conditions, - * we need take the global lock. This also obviously modifies the MyLatch - * global, so even if WaitLatch did not throw errors. We'd still need to take - * the global lock for that, or at the very least we'd need to take a specific - * lock for MyLatch and ensure that no other code modifies MyLatch. - */ -void -WaitMyLatch(uint64_t &last_known_latch_update_count) { - // A 64bit uint takes 584 years to overflow when it's updated every - // nanosecond. That's more than enough for any practical purposes. If - // someone keeps a Postgres connection open for more than 584 years, - // they deserve what they get. - static uint64_t latch_update_count = 0; - - if (last_known_latch_update_count == latch_update_count) { - PostgresFunctionGuard(WaitLatch, MyLatch, WL_LATCH_SET | WL_EXIT_ON_PM_DEATH, 0, PG_WAIT_EXTENSION); - /* Note: No need to use PostgresFunctionGuard here, because ResetLatch - * is a trivial function */ - ResetLatch(MyLatch); - latch_update_count++; - } - last_known_latch_update_count = latch_update_count; -} - -} // namespace pgduckdb::pg diff --git a/src/scan/postgres_table_reader.cpp b/src/scan/postgres_table_reader.cpp index db7e77c3..8970973d 100644 --- a/src/scan/postgres_table_reader.cpp +++ b/src/scan/postgres_table_reader.cpp @@ -1,7 +1,6 @@ #include "pgduckdb/scan/postgres_table_reader.hpp" #include "pgduckdb/pgduckdb_process_lock.hpp" #include "pgduckdb/pgduckdb_utils.hpp" -#include "pgduckdb/pg/latch.hpp" extern "C" { #include "postgres.h" @@ -14,9 +13,10 @@ extern "C" { #include "optimizer/planmain.h" #include "optimizer/planner.h" #include "tcop/tcopprot.h" +#include "utils/lsyscache.h" #include "utils/rel.h" #include "utils/snapmgr.h" -#include "utils/lsyscache.h" +#include "utils/wait_event.h" #include "pgduckdb/pgduckdb_guc.h" } @@ -29,7 +29,7 @@ namespace pgduckdb { PostgresTableReader::PostgresTableReader(const char *table_scan_query, bool count_tuples_only) : parallel_executor_info(nullptr), parallel_worker_readers(nullptr), nreaders(0), next_parallel_reader(0), - entered_parallel_mode(false), last_known_latch_update_count(0) { + entered_parallel_mode(false) { std::lock_guard lock(GlobalProcessLock::GetLock()); PostgresScopedStackReset scoped_stack_reset; @@ -239,6 +239,9 @@ PostgresTableReader::MarkPlanParallelAware(Plan *plan) { } } +/* + * GlobalProcessLock should be held before calling this. + */ TupleTableSlot * PostgresTableReader::GetNextTuple() { MinimalTuple worker_minmal_tuple; @@ -272,10 +275,7 @@ PostgresTableReader::GetNextWorkerTuple() { reader = (TupleQueueReader *)parallel_worker_readers[next_parallel_reader]; - { - // We need to take global lock for `TupleQueueReaderNext` call - minimal_tuple = PostgresFunctionGuard(TupleQueueReaderNext, reader, true, &readerdone); - } + minimal_tuple = PostgresFunctionGuard(TupleQueueReaderNext, reader, true, &readerdone); if (readerdone) { --nreaders; @@ -300,7 +300,13 @@ PostgresTableReader::GetNextWorkerTuple() { nvisited++; if (nvisited >= nreaders) { - pgduckdb::pg::WaitMyLatch(last_known_latch_update_count); + /* + * It should be safe to make this call because function calling GetNextTuple() and transitively + * GetNextWorkerTuple() should held GlobalProcesLock. + */ + PostgresFunctionGuard(WaitLatch, MyLatch, WL_LATCH_SET | WL_EXIT_ON_PM_DEATH, 0, PG_WAIT_EXTENSION); + /* No need to use PostgresFunctionGuard here, because ResetLatch is a trivial function */ + ResetLatch(MyLatch); nvisited = 0; } } From 9eed6af3cbdf02996795cf9a5a8f7cb0c351b7ce Mon Sep 17 00:00:00 2001 From: mkaruza Date: Mon, 23 Dec 2024 09:45:42 +0100 Subject: [PATCH 37/53] Use returned tuple slot attr information for populating output vector --- src/pgduckdb_types.cpp | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/pgduckdb_types.cpp b/src/pgduckdb_types.cpp index 038d523f..1b72cc9a 100644 --- a/src/pgduckdb_types.cpp +++ b/src/pgduckdb_types.cpp @@ -1360,18 +1360,17 @@ InsertTupleIntoChunk(duckdb::DataChunk &output, PostgresScanLocalState &scan_loc return; } /* Write tuple columns in output vector. */ - int duckdb_output_index = 0; - for (auto const &attr_num : scan_global_state->output_columns) { + for (int duckdb_output_index = 0; duckdb_output_index < slot->tts_tupleDescriptor->natts; duckdb_output_index++) { auto &result = output.data[duckdb_output_index]; if (slot->tts_isnull[duckdb_output_index]) { auto &array_mask = duckdb::FlatVector::Validity(result); array_mask.SetInvalid(scan_local_state.output_vector_size); } else { - auto attr = slot->tts_tupleDescriptor->attrs[attr_num - 1]; + /* Use ruturned tuple slot attr information. */ + auto attr = slot->tts_tupleDescriptor->attrs[duckdb_output_index]; ConvertPostgresToDuckValue(attr.atttypid, slot->tts_values[duckdb_output_index], result, scan_local_state.output_vector_size); } - duckdb_output_index++; } scan_local_state.output_vector_size++; From abb086f35c1cbda608f46e4a5c4fdf238c98e68e Mon Sep 17 00:00:00 2001 From: mkaruza Date: Mon, 23 Dec 2024 09:48:51 +0100 Subject: [PATCH 38/53] Rename GUC variable according to PR suggestion --- include/pgduckdb/pgduckdb_guc.h | 2 +- src/pgduckdb.cpp | 4 ++-- src/scan/postgres_table_reader.cpp | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/include/pgduckdb/pgduckdb_guc.h b/include/pgduckdb/pgduckdb_guc.h index 35a8369a..629de5ac 100644 --- a/include/pgduckdb/pgduckdb_guc.h +++ b/include/pgduckdb/pgduckdb_guc.h @@ -8,7 +8,7 @@ extern bool duckdb_enable_external_access; extern bool duckdb_allow_unsigned_extensions; extern bool duckdb_autoinstall_known_extensions; extern bool duckdb_autoload_known_extensions; -extern int max_workers_per_postgres_scan; +extern int duckdb_max_workers_per_postgres_scan; extern char *duckdb_motherduck_postgres_database; extern int duckdb_motherduck_enabled; extern char *duckdb_motherduck_token; diff --git a/src/pgduckdb.cpp b/src/pgduckdb.cpp index dd4d0d6b..0718d126 100644 --- a/src/pgduckdb.cpp +++ b/src/pgduckdb.cpp @@ -14,7 +14,7 @@ extern "C" { static void DuckdbInitGUC(void); bool duckdb_force_execution = false; -int max_workers_per_postgres_scan = 2; +int duckdb_max_workers_per_postgres_scan = 2; int duckdb_motherduck_enabled = MotherDuckEnabled::MOTHERDUCK_AUTO; char *duckdb_motherduck_token = strdup(""); char *duckdb_motherduck_postgres_database = strdup("postgres"); @@ -161,7 +161,7 @@ DuckdbInitGUC(void) { DefineCustomVariable("duckdb.max_workers_per_postgres_scan", "Maximum number of PostgreSQL workers used for a single Postgres scan", - &max_workers_per_postgres_scan, 2, 8); + &duckdb_max_workers_per_postgres_scan, 2, 8); DefineCustomVariable("duckdb.postgres_role", "Which postgres role should be allowed to use DuckDB execution, use the secrets and create " diff --git a/src/scan/postgres_table_reader.cpp b/src/scan/postgres_table_reader.cpp index 8970973d..2443296b 100644 --- a/src/scan/postgres_table_reader.cpp +++ b/src/scan/postgres_table_reader.cpp @@ -161,7 +161,7 @@ PostgresTableReader::ParallelWorkerNumber(Cardinality cardinality) { static const int base_log = 8; int cardinality_log = std::log2(cardinality); int base = cardinality_log / base_log; - return std::max(1, std::min(base, std::max(max_workers_per_postgres_scan, max_parallel_workers))); + return std::max(1, std::min(base, std::max(duckdb_max_workers_per_postgres_scan, max_parallel_workers))); } const char * From aa42e61e68829b3b72024d66dcdb64ac50ea7245 Mon Sep 17 00:00:00 2001 From: mkaruza Date: Mon, 23 Dec 2024 10:21:34 +0100 Subject: [PATCH 39/53] Remove check for allowed table types We should now support all types with native postgres scan --- src/pgduckdb_hooks.cpp | 41 ----------------------------------------- 1 file changed, 41 deletions(-) diff --git a/src/pgduckdb_hooks.cpp b/src/pgduckdb_hooks.cpp index 46f2bf5d..92a6801c 100644 --- a/src/pgduckdb_hooks.cpp +++ b/src/pgduckdb_hooks.cpp @@ -60,39 +60,6 @@ ContainsCatalogTable(List *rtes) { return false; } -static bool -ContainsAllowedTableType(List *rtes, int elevel) { - foreach_node(RangeTblEntry, rte, rtes) { - if (rte->rtekind == RTE_SUBQUERY) { - /* Check whether any table in the subquery is a partitioned table */ - if (ContainsAllowedTableType(rte->subquery->rtable, elevel)) { - return true; - } - } - - /* Allow functions */ - if (rte->rtekind == RTE_FUNCTION) { - return true; - } - - switch (rte->relkind) { - case RELKIND_RELATION: - case RELKIND_INDEX: - case RELKIND_TOASTVALUE: - case RELKIND_VIEW: - case RELKIND_MATVIEW: - case RELKIND_FOREIGN_TABLE: - case RELKIND_PARTITIONED_TABLE: - case RELKIND_PARTITIONED_INDEX: - return true; - default: - elog(elevel, "DuckDB does not support querying table type '%c'", rte->relkind); - return false; - } - } - return false; -} - static bool IsDuckdbTable(Oid relid) { if (relid == InvalidOid) { @@ -194,14 +161,6 @@ IsAllowedStatement(Query *query, bool throw_error = false) { return false; } - /* - * Check that we're only accessing table types that we know we have support - * for. - */ - if (!ContainsAllowedTableType(query->rtable, elevel)) { - return false; - } - /* Anything else is hopefully fine... */ return true; } From 60f2028e5ffa51680939212918b9686d891e28ca Mon Sep 17 00:00:00 2001 From: Jelte Fennema-Nio Date: Thu, 2 Jan 2025 13:15:28 +0100 Subject: [PATCH 40/53] Fix typo --- src/pgduckdb_types.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pgduckdb_types.cpp b/src/pgduckdb_types.cpp index 1b72cc9a..2597dece 100644 --- a/src/pgduckdb_types.cpp +++ b/src/pgduckdb_types.cpp @@ -1366,7 +1366,7 @@ InsertTupleIntoChunk(duckdb::DataChunk &output, PostgresScanLocalState &scan_loc auto &array_mask = duckdb::FlatVector::Validity(result); array_mask.SetInvalid(scan_local_state.output_vector_size); } else { - /* Use ruturned tuple slot attr information. */ + /* Use returned tuple slot attr information. */ auto attr = slot->tts_tupleDescriptor->attrs[duckdb_output_index]; ConvertPostgresToDuckValue(attr.atttypid, slot->tts_values[duckdb_output_index], result, scan_local_state.output_vector_size); From 5d560df8fcfd96ad7ee63863081e89f3a3fbbf0e Mon Sep 17 00:00:00 2001 From: mkaruza Date: Tue, 7 Jan 2025 12:08:08 +0100 Subject: [PATCH 41/53] Simplify max_workers_per_postgres_scan * Setting this variable to `0` disables parallelization * Cardinality of table less than 65536 use only single parallel process * Higher cardinality will try to use `max_workers_per_postgres_scan` parallel processes with upper limit of `max_parallel_workers` --- src/pgduckdb.cpp | 3 ++- src/scan/postgres_table_reader.cpp | 22 ++++++++++++++++++---- 2 files changed, 20 insertions(+), 5 deletions(-) diff --git a/src/pgduckdb.cpp b/src/pgduckdb.cpp index 0718d126..763c6969 100644 --- a/src/pgduckdb.cpp +++ b/src/pgduckdb.cpp @@ -4,6 +4,7 @@ extern "C" { #include "postgres.h" #include "miscadmin.h" #include "utils/guc.h" +#include "postmaster/bgworker_internals.h" } #include "pgduckdb/pgduckdb.h" @@ -161,7 +162,7 @@ DuckdbInitGUC(void) { DefineCustomVariable("duckdb.max_workers_per_postgres_scan", "Maximum number of PostgreSQL workers used for a single Postgres scan", - &duckdb_max_workers_per_postgres_scan, 2, 8); + &duckdb_max_workers_per_postgres_scan, 0, MAX_PARALLEL_WORKER_LIMIT); DefineCustomVariable("duckdb.postgres_role", "Which postgres role should be allowed to use DuckDB execution, use the secrets and create " diff --git a/src/scan/postgres_table_reader.cpp b/src/scan/postgres_table_reader.cpp index 2443296b..e5741562 100644 --- a/src/scan/postgres_table_reader.cpp +++ b/src/scan/postgres_table_reader.cpp @@ -156,12 +156,26 @@ PostgresTableReader::PostgresTableReaderCleanup() { table_scan_query_desc = nullptr; } + +/* + * Logic is straightforward, if `duckdb_max_workers_per_postgres_scan` is set to 0 we don't want any + * parallelization. For cardinality less equal than 2^16 we only try to run one parallel process. When cardinality + * is bigger than we should spawn numer of parallel processes set by `duckdb_max_workers_per_postgres_scan` but + * not bigger than `max_parallel_workers`. + */ + int PostgresTableReader::ParallelWorkerNumber(Cardinality cardinality) { - static const int base_log = 8; - int cardinality_log = std::log2(cardinality); - int base = cardinality_log / base_log; - return std::max(1, std::min(base, std::max(duckdb_max_workers_per_postgres_scan, max_parallel_workers))); + static const int cardinality_threshold = 1 << 16; + /* No parallel scan wanted */ + if (!duckdb_max_workers_per_postgres_scan) { + return 0; + } + /* */ + if (cardinality <= cardinality_threshold) { + return 1; + } + return std::min(duckdb_max_workers_per_postgres_scan, max_parallel_workers); } const char * From 94f60af467149624cccf783224351e7c8020aad4 Mon Sep 17 00:00:00 2001 From: mkaruza Date: Tue, 7 Jan 2025 16:29:37 +0100 Subject: [PATCH 42/53] Update empty comment --- src/scan/postgres_table_reader.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/scan/postgres_table_reader.cpp b/src/scan/postgres_table_reader.cpp index e5741562..583d4f0e 100644 --- a/src/scan/postgres_table_reader.cpp +++ b/src/scan/postgres_table_reader.cpp @@ -167,11 +167,11 @@ PostgresTableReader::PostgresTableReaderCleanup() { int PostgresTableReader::ParallelWorkerNumber(Cardinality cardinality) { static const int cardinality_threshold = 1 << 16; - /* No parallel scan wanted */ + /* No parallel worker scan wanted */ if (!duckdb_max_workers_per_postgres_scan) { return 0; } - /* */ + /* Use only one worker when scan is done on low cardinality */ if (cardinality <= cardinality_threshold) { return 1; } From 68bb7a6922873321d245d7533530ab5309ff6508 Mon Sep 17 00:00:00 2001 From: Jelte Fennema-Nio Date: Tue, 7 Jan 2025 18:00:41 +0100 Subject: [PATCH 43/53] Handle cleanup correctly in case of stopped parallel workers --- src/scan/postgres_table_reader.cpp | 34 ++++++++++++++++++------------ 1 file changed, 20 insertions(+), 14 deletions(-) diff --git a/src/scan/postgres_table_reader.cpp b/src/scan/postgres_table_reader.cpp index 583d4f0e..38d999d3 100644 --- a/src/scan/postgres_table_reader.cpp +++ b/src/scan/postgres_table_reader.cpp @@ -130,33 +130,39 @@ void PostgresTableReader::PostgresTableReaderCleanup() { PostgresScopedStackReset scoped_stack_reset; - PostgresFunctionGuard(ExecEndNode, table_scan_planstate); + if (table_scan_planstate) { + auto tmp = table_scan_planstate; + table_scan_planstate = nullptr; + PostgresFunctionGuard(ExecEndNode, tmp); + } if (parallel_executor_info != NULL) { - PostgresFunctionGuard(ExecParallelFinish, parallel_executor_info); - PostgresFunctionGuard(ExecParallelCleanup, parallel_executor_info); + auto tmp = parallel_executor_info; + parallel_executor_info = nullptr; + PostgresFunctionGuard(ExecParallelFinish, tmp); + PostgresFunctionGuard(ExecParallelCleanup, tmp); } - parallel_executor_info = nullptr; - if (parallel_worker_readers) { - PostgresFunctionGuard(pfree, parallel_worker_readers); + auto tmp = parallel_worker_readers; + parallel_worker_readers = nullptr; + PostgresFunctionGuard(pfree, tmp); } - parallel_worker_readers = nullptr; - - PostgresFunctionGuard(ExecutorFinish, table_scan_query_desc); - PostgresFunctionGuard(ExecutorEnd, table_scan_query_desc); - PostgresFunctionGuard(FreeQueryDesc, table_scan_query_desc); + if (table_scan_query_desc) { + auto tmp = table_scan_query_desc; + table_scan_query_desc = nullptr; + PostgresFunctionGuard(ExecutorFinish, tmp); + PostgresFunctionGuard(ExecutorEnd, tmp); + PostgresFunctionGuard(FreeQueryDesc, tmp); + } if (entered_parallel_mode) { + entered_parallel_mode = false; ExitParallelMode(); } - - table_scan_query_desc = nullptr; } - /* * Logic is straightforward, if `duckdb_max_workers_per_postgres_scan` is set to 0 we don't want any * parallelization. For cardinality less equal than 2^16 we only try to run one parallel process. When cardinality From e0135e8423335cd6fef23b9f003e918c8ce9d893 Mon Sep 17 00:00:00 2001 From: Jelte Fennema-Nio Date: Tue, 7 Jan 2025 18:03:36 +0100 Subject: [PATCH 44/53] Fix formatting --- include/pgduckdb/logger.hpp | 20 +++++++++---------- include/pgduckdb/pgduckdb_utils.hpp | 2 +- .../pgduckdb/scan/postgres_table_reader.hpp | 2 +- 3 files changed, 12 insertions(+), 12 deletions(-) diff --git a/include/pgduckdb/logger.hpp b/include/pgduckdb/logger.hpp index e6a108fe..67ae9ce8 100644 --- a/include/pgduckdb/logger.hpp +++ b/include/pgduckdb/logger.hpp @@ -41,16 +41,16 @@ namespace pgduckdb { #define pd_prevent_errno_in_scope() #endif -#define pd_ereport_domain(elevel, domain, ...) \ - do { \ - pd_prevent_errno_in_scope(); \ - static_assert(elevel >= DEBUG5 && elevel <= WARNING_CLIENT_ONLY, "Invalid error level"); \ - if (message_level_is_interesting(elevel)) { \ - std::lock_guard lock(GlobalProcessLock::GetLock()); \ - if (errstart(elevel, domain)) \ - __VA_ARGS__, errfinish(__FILE__, __LINE__, __func__); \ - } \ - } while(0) +#define pd_ereport_domain(elevel, domain, ...) \ + do { \ + pd_prevent_errno_in_scope(); \ + static_assert(elevel >= DEBUG5 && elevel <= WARNING_CLIENT_ONLY, "Invalid error level"); \ + if (message_level_is_interesting(elevel)) { \ + std::lock_guard lock(GlobalProcessLock::GetLock()); \ + if (errstart(elevel, domain)) \ + __VA_ARGS__, errfinish(__FILE__, __LINE__, __func__); \ + } \ + } while (0) #define TEXTDOMAIN NULL diff --git a/include/pgduckdb/pgduckdb_utils.hpp b/include/pgduckdb/pgduckdb_utils.hpp index 192865ae..dfe74c1a 100644 --- a/include/pgduckdb/pgduckdb_utils.hpp +++ b/include/pgduckdb/pgduckdb_utils.hpp @@ -101,7 +101,7 @@ __PostgresFunctionGuard__(const char *func_name, FuncArgs... args) { throw duckdb::Exception(duckdb::ExceptionType::EXECUTOR, message); } -#define PostgresFunctionGuard(FUNC, ...) \ +#define PostgresFunctionGuard(FUNC, ...) \ pgduckdb::__PostgresFunctionGuard__(__func__, ##__VA_ARGS__) duckdb::unique_ptr DuckDBQueryOrThrow(duckdb::ClientContext &context, const std::string &query); diff --git a/include/pgduckdb/scan/postgres_table_reader.hpp b/include/pgduckdb/scan/postgres_table_reader.hpp index 76dfbd67..f4424bf1 100644 --- a/include/pgduckdb/scan/postgres_table_reader.hpp +++ b/include/pgduckdb/scan/postgres_table_reader.hpp @@ -20,7 +20,7 @@ class PostgresTableReader { private: MinimalTuple GetNextWorkerTuple(); int ParallelWorkerNumber(Cardinality cardinality); - const char * ExplainScanPlan(QueryDesc *query_desc); + const char *ExplainScanPlan(QueryDesc *query_desc); bool CanTableScanRunInParallel(Plan *plan); bool MarkPlanParallelAware(Plan *plan); From 0b6c8ec88f7a629e257ff220ea4a2cdd56d9c0c3 Mon Sep 17 00:00:00 2001 From: Jelte Fennema-Nio Date: Thu, 9 Jan 2025 17:38:17 +0100 Subject: [PATCH 45/53] Simplify PostgresTableReader cleanup logic --- .../pgduckdb/scan/postgres_table_reader.hpp | 1 + src/scan/postgres_table_reader.cpp | 31 +++++++++---------- 2 files changed, 16 insertions(+), 16 deletions(-) diff --git a/include/pgduckdb/scan/postgres_table_reader.hpp b/include/pgduckdb/scan/postgres_table_reader.hpp index f4424bf1..88998188 100644 --- a/include/pgduckdb/scan/postgres_table_reader.hpp +++ b/include/pgduckdb/scan/postgres_table_reader.hpp @@ -34,6 +34,7 @@ class PostgresTableReader { int nreaders; int next_parallel_reader; bool entered_parallel_mode; + bool cleaned_up; }; } // namespace pgduckdb diff --git a/src/scan/postgres_table_reader.cpp b/src/scan/postgres_table_reader.cpp index 38d999d3..ff84fac5 100644 --- a/src/scan/postgres_table_reader.cpp +++ b/src/scan/postgres_table_reader.cpp @@ -29,7 +29,7 @@ namespace pgduckdb { PostgresTableReader::PostgresTableReader(const char *table_scan_query, bool count_tuples_only) : parallel_executor_info(nullptr), parallel_worker_readers(nullptr), nreaders(0), next_parallel_reader(0), - entered_parallel_mode(false) { + entered_parallel_mode(false), cleaned_up(false) { std::lock_guard lock(GlobalProcessLock::GetLock()); PostgresScopedStackReset scoped_stack_reset; @@ -120,46 +120,45 @@ PostgresTableReader::PostgresTableReader(const char *table_scan_query, bool coun } PostgresTableReader::~PostgresTableReader() { - if (table_scan_query_desc) { - std::lock_guard lock(GlobalProcessLock::GetLock()); - PostgresTableReaderCleanup(); + if (cleaned_up) { + return; } + std::lock_guard lock(GlobalProcessLock::GetLock()); + PostgresTableReaderCleanup(); } void PostgresTableReader::PostgresTableReaderCleanup() { + D_ASSERT(!cleaned_up); + cleaned_up = true; PostgresScopedStackReset scoped_stack_reset; if (table_scan_planstate) { - auto tmp = table_scan_planstate; + PostgresFunctionGuard(ExecEndNode, table_scan_planstate); table_scan_planstate = nullptr; - PostgresFunctionGuard(ExecEndNode, tmp); } if (parallel_executor_info != NULL) { - auto tmp = parallel_executor_info; + PostgresFunctionGuard(ExecParallelFinish, parallel_executor_info); + PostgresFunctionGuard(ExecParallelCleanup, parallel_executor_info); parallel_executor_info = nullptr; - PostgresFunctionGuard(ExecParallelFinish, tmp); - PostgresFunctionGuard(ExecParallelCleanup, tmp); } if (parallel_worker_readers) { - auto tmp = parallel_worker_readers; + PostgresFunctionGuard(pfree, parallel_worker_readers); parallel_worker_readers = nullptr; - PostgresFunctionGuard(pfree, tmp); } if (table_scan_query_desc) { - auto tmp = table_scan_query_desc; + PostgresFunctionGuard(ExecutorFinish, table_scan_query_desc); + PostgresFunctionGuard(ExecutorEnd, table_scan_query_desc); + PostgresFunctionGuard(FreeQueryDesc, table_scan_query_desc); table_scan_query_desc = nullptr; - PostgresFunctionGuard(ExecutorFinish, tmp); - PostgresFunctionGuard(ExecutorEnd, tmp); - PostgresFunctionGuard(FreeQueryDesc, tmp); } if (entered_parallel_mode) { - entered_parallel_mode = false; ExitParallelMode(); + entered_parallel_mode = false; } } From 431e5fb8792ee4b2adaa93a86c8f421514b55fbc Mon Sep 17 00:00:00 2001 From: Jelte Fennema-Nio Date: Thu, 9 Jan 2025 17:43:32 +0100 Subject: [PATCH 46/53] Add the BLOB alias for bytea, which duckdb uses --- sql/pg_duckdb--0.2.0--0.3.0.sql | 3 +++ 1 file changed, 3 insertions(+) diff --git a/sql/pg_duckdb--0.2.0--0.3.0.sql b/sql/pg_duckdb--0.2.0--0.3.0.sql index f64f71b7..a23bb628 100644 --- a/sql/pg_duckdb--0.2.0--0.3.0.sql +++ b/sql/pg_duckdb--0.2.0--0.3.0.sql @@ -14,3 +14,6 @@ CREATE AGGREGATE @extschema@.approx_count_distinct(anyelement) stype = bigint, initcond = 0 ); + +CREATE DOMAIN pg_catalog.blob AS bytea; +COMMENT ON DOMAIN pg_catalog.blob IS 'The DuckDB BLOB alias for BYTEA'; From d69c2228afe2245f76b183a46b1a698ceaf0d1e4 Mon Sep 17 00:00:00 2001 From: Jelte Fennema-Nio Date: Fri, 10 Jan 2025 11:24:30 +0100 Subject: [PATCH 47/53] Fix null deref if execution finished before cancel is processed --- src/pgduckdb_node.cpp | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/pgduckdb_node.cpp b/src/pgduckdb_node.cpp index 355c2c89..37e89079 100644 --- a/src/pgduckdb_node.cpp +++ b/src/pgduckdb_node.cpp @@ -137,8 +137,12 @@ ExecuteQuery(DuckdbScanState *state) { } duckdb::PendingExecutionResult execution_result; - do { + while (true) { execution_result = pending->ExecuteTask(); + if (duckdb::PendingQueryResult::IsResultReady(execution_result)) { + break; + } + if (QueryCancelPending) { // Send an interrupt connection->Interrupt(); @@ -150,7 +154,7 @@ ExecuteQuery(DuckdbScanState *state) { ProcessInterrupts(); throw duckdb::Exception(duckdb::ExceptionType::EXECUTOR, "Query cancelled"); } - } while (!duckdb::PendingQueryResult::IsResultReady(execution_result)); + } if (execution_result == duckdb::PendingExecutionResult::EXECUTION_ERROR) { return pending->ThrowError(); From a7274dde680e7ba2ffb28d0221fb49234cae53bc Mon Sep 17 00:00:00 2001 From: Yves Date: Fri, 10 Jan 2025 14:44:49 +0100 Subject: [PATCH 48/53] Return table name when stringifying TF --- include/pgduckdb/pg/relations.hpp | 2 ++ include/pgduckdb/scan/postgres_scan.hpp | 1 + src/pg/relations.cpp | 5 +++++ src/scan/postgres_scan.cpp | 6 ++++++ 4 files changed, 14 insertions(+) diff --git a/include/pgduckdb/pg/relations.hpp b/include/pgduckdb/pg/relations.hpp index cafdd3cc..25d15782 100644 --- a/include/pgduckdb/pg/relations.hpp +++ b/include/pgduckdb/pg/relations.hpp @@ -33,4 +33,6 @@ bool IsValidBlockNumber(BlockNumber); char *GenerateQualifiedRelationName(Relation rel); const char *QuoteIdentifier(const char *ident); +const char * GetRelationName(Relation rel); + } // namespace pgduckdb diff --git a/include/pgduckdb/scan/postgres_scan.hpp b/include/pgduckdb/scan/postgres_scan.hpp index 5b45136d..3f4d0171 100644 --- a/include/pgduckdb/scan/postgres_scan.hpp +++ b/include/pgduckdb/scan/postgres_scan.hpp @@ -76,6 +76,7 @@ struct PostgresScanTableFunction : public duckdb::TableFunction { duckdb::DataChunk &output); static duckdb::unique_ptr PostgresScanCardinality(duckdb::ClientContext &context, const duckdb::FunctionData *data); + static std::string ToString(const duckdb::FunctionData *bind_data); }; } // namespace pgduckdb diff --git a/src/pg/relations.cpp b/src/pg/relations.cpp index 5afb2cb4..ed0a8c01 100644 --- a/src/pg/relations.cpp +++ b/src/pg/relations.cpp @@ -160,4 +160,9 @@ QuoteIdentifier(const char *ident) { return PostgresFunctionGuard(quote_identifier, ident); } +const char * +GetRelationName(Relation rel) { + return RelationGetRelationName(rel); +} + } // namespace pgduckdb diff --git a/src/scan/postgres_scan.cpp b/src/scan/postgres_scan.cpp index 417aa7a8..8d2914e5 100644 --- a/src/scan/postgres_scan.cpp +++ b/src/scan/postgres_scan.cpp @@ -154,6 +154,12 @@ PostgresScanTableFunction::PostgresScanTableFunction() filter_pushdown = true; filter_prune = true; cardinality = PostgresScanCardinality; + to_string = ToString; +} + +std::string PostgresScanTableFunction::ToString(const duckdb::FunctionData *data) { + auto &bind_data = data->Cast(); + return GetRelationName(bind_data.rel); } duckdb::unique_ptr From 6e9a2e92fffd3f36b33dae73e7c49532989df8f5 Mon Sep 17 00:00:00 2001 From: Yves Date: Fri, 10 Jan 2025 14:44:56 +0100 Subject: [PATCH 49/53] Cleanup --- src/pgduckdb_node.cpp | 7 +++---- src/scan/postgres_scan.cpp | 3 +-- 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/src/pgduckdb_node.cpp b/src/pgduckdb_node.cpp index 37e89079..9c796ba6 100644 --- a/src/pgduckdb_node.cpp +++ b/src/pgduckdb_node.cpp @@ -258,10 +258,9 @@ Duckdb_ExplainCustomScan_Cpp(CustomScanState *node, ExplainState *es) { chunk = duckdb_scan_state->query_results->Fetch(); } while (chunk && chunk->size() > 0); - std::string explain_output = "\n\n"; - explain_output += value; - explain_output += "\n"; - ExplainPropertyText("DuckDB Execution Plan", explain_output.c_str(), es); + std::ostringstream explain_output; + explain_output << "\n\n" << value << "\n"; + ExplainPropertyText("DuckDB Execution Plan", explain_output.str().c_str(), es); } void diff --git a/src/scan/postgres_scan.cpp b/src/scan/postgres_scan.cpp index 8d2914e5..8aa36539 100644 --- a/src/scan/postgres_scan.cpp +++ b/src/scan/postgres_scan.cpp @@ -165,8 +165,7 @@ std::string PostgresScanTableFunction::ToString(const duckdb::FunctionData *data duckdb::unique_ptr PostgresScanTableFunction::PostgresScanInitGlobal(duckdb::ClientContext &, duckdb::TableFunctionInitInput &input) { auto &bind_data = input.bind_data->CastNoConst(); - auto global_state = duckdb::make_uniq(bind_data.snapshot, bind_data.rel, input); - return global_state; + return duckdb::make_uniq(bind_data.snapshot, bind_data.rel, input); } duckdb::unique_ptr From c5122db826fc6a430e0c5f2e353e31d6b8aba530 Mon Sep 17 00:00:00 2001 From: Yves Date: Fri, 10 Jan 2025 14:47:42 +0100 Subject: [PATCH 50/53] Format --- include/pgduckdb/pg/relations.hpp | 2 +- src/scan/postgres_scan.cpp | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/include/pgduckdb/pg/relations.hpp b/include/pgduckdb/pg/relations.hpp index 25d15782..8a81fbc9 100644 --- a/include/pgduckdb/pg/relations.hpp +++ b/include/pgduckdb/pg/relations.hpp @@ -33,6 +33,6 @@ bool IsValidBlockNumber(BlockNumber); char *GenerateQualifiedRelationName(Relation rel); const char *QuoteIdentifier(const char *ident); -const char * GetRelationName(Relation rel); +const char *GetRelationName(Relation rel); } // namespace pgduckdb diff --git a/src/scan/postgres_scan.cpp b/src/scan/postgres_scan.cpp index 8aa36539..18c39274 100644 --- a/src/scan/postgres_scan.cpp +++ b/src/scan/postgres_scan.cpp @@ -157,7 +157,8 @@ PostgresScanTableFunction::PostgresScanTableFunction() to_string = ToString; } -std::string PostgresScanTableFunction::ToString(const duckdb::FunctionData *data) { +std::string +PostgresScanTableFunction::ToString(const duckdb::FunctionData *data) { auto &bind_data = data->Cast(); return GetRelationName(bind_data.rel); } From 28838673da4c883294e30a76c33a9ee71ad426b8 Mon Sep 17 00:00:00 2001 From: Yves Date: Fri, 10 Jan 2025 14:49:30 +0100 Subject: [PATCH 51/53] Update duckdb_recycle.out --- test/regression/expected/duckdb_recycle.out | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/test/regression/expected/duckdb_recycle.out b/test/regression/expected/duckdb_recycle.out index 0a8521f1..af491a2b 100644 --- a/test/regression/expected/duckdb_recycle.out +++ b/test/regression/expected/duckdb_recycle.out @@ -15,8 +15,7 @@ EXPLAIN SELECT count(*) FROM ta; ┌─────────────┴─────────────┐ │ POSTGRES_SCAN │ │ ──────────────────── │ - │ Function: │ - │ POSTGRES_SCAN │ + │ ta │ │ │ │ ~2550 Rows │ └───────────────────────────┘ @@ -40,8 +39,7 @@ EXPLAIN SELECT count(*) FROM ta; ┌─────────────┴─────────────┐ │ POSTGRES_SCAN │ │ ──────────────────── │ - │ Function: │ - │ POSTGRES_SCAN │ + │ ta │ │ │ │ ~2550 Rows │ └───────────────────────────┘ From e79a99a2058c5f044e8300126b033b556df052c0 Mon Sep 17 00:00:00 2001 From: Yves Date: Fri, 10 Jan 2025 14:53:05 +0100 Subject: [PATCH 52/53] Update tests --- test/regression/expected/duckdb_recycle.out | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/regression/expected/duckdb_recycle.out b/test/regression/expected/duckdb_recycle.out index af491a2b..d2dbe70b 100644 --- a/test/regression/expected/duckdb_recycle.out +++ b/test/regression/expected/duckdb_recycle.out @@ -21,7 +21,7 @@ EXPLAIN SELECT count(*) FROM ta; └───────────────────────────┘ -(19 rows) +(18 rows) CALL duckdb.recycle_ddb(); EXPLAIN SELECT count(*) FROM ta; @@ -45,7 +45,7 @@ EXPLAIN SELECT count(*) FROM ta; └───────────────────────────┘ -(19 rows) +(18 rows) -- Not allowed in a transaction BEGIN; From 7908a9ad46e0ef32b2cff35dc28dad6b9f3c3f33 Mon Sep 17 00:00:00 2001 From: Yves Date: Fri, 10 Jan 2025 15:03:52 +0100 Subject: [PATCH 53/53] Keep `POSTGRES_SCAN` in the output --- src/scan/postgres_scan.cpp | 4 +++- test/regression/expected/duckdb_recycle.out | 4 ++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/src/scan/postgres_scan.cpp b/src/scan/postgres_scan.cpp index 18c39274..11a85306 100644 --- a/src/scan/postgres_scan.cpp +++ b/src/scan/postgres_scan.cpp @@ -160,7 +160,9 @@ PostgresScanTableFunction::PostgresScanTableFunction() std::string PostgresScanTableFunction::ToString(const duckdb::FunctionData *data) { auto &bind_data = data->Cast(); - return GetRelationName(bind_data.rel); + std::ostringstream oss; + oss << "(POSTGRES_SCAN) " << GetRelationName(bind_data.rel); + return oss.str(); } duckdb::unique_ptr diff --git a/test/regression/expected/duckdb_recycle.out b/test/regression/expected/duckdb_recycle.out index d2dbe70b..723cf523 100644 --- a/test/regression/expected/duckdb_recycle.out +++ b/test/regression/expected/duckdb_recycle.out @@ -15,7 +15,7 @@ EXPLAIN SELECT count(*) FROM ta; ┌─────────────┴─────────────┐ │ POSTGRES_SCAN │ │ ──────────────────── │ - │ ta │ + │ (POSTGRES_SCAN) ta │ │ │ │ ~2550 Rows │ └───────────────────────────┘ @@ -39,7 +39,7 @@ EXPLAIN SELECT count(*) FROM ta; ┌─────────────┴─────────────┐ │ POSTGRES_SCAN │ │ ──────────────────── │ - │ ta │ + │ (POSTGRES_SCAN) ta │ │ │ │ ~2550 Rows │ └───────────────────────────┘