From e4ac9e93e0bd27428ed55eeb3a02ac8c54f10dd0 Mon Sep 17 00:00:00 2001 From: mkaruza Date: Wed, 17 Jul 2024 12:30:58 +0200 Subject: [PATCH] WIP IndexOnlyScan --- include/pgduckdb/catalog/pgduckdb_table.hpp | 10 +- include/pgduckdb/pgduckdb_types.hpp | 4 + include/pgduckdb/scan/postgres_index_scan.hpp | 18 +- src/catalog/pgduckdb_table.cpp | 95 ++++++++-- src/catalog/pgduckdb_transaction.cpp | 30 ++-- src/pgduckdb_types.cpp | 63 +++++++ src/scan/index_scan_utils.cpp | 92 +++++++++- src/scan/postgres_index_scan.cpp | 168 ++++++++++-------- src/scan/postgres_seq_scan.cpp | 2 +- test/regression/expected/index_scan.out | 0 test/regression/schedule | 1 + test/regression/sql/index_scan.sql | 19 ++ 12 files changed, 392 insertions(+), 110 deletions(-) create mode 100644 test/regression/expected/index_scan.out create mode 100644 test/regression/sql/index_scan.sql diff --git a/include/pgduckdb/catalog/pgduckdb_table.hpp b/include/pgduckdb/catalog/pgduckdb_table.hpp index 70966383..12ec1a43 100644 --- a/include/pgduckdb/catalog/pgduckdb_table.hpp +++ b/include/pgduckdb/catalog/pgduckdb_table.hpp @@ -26,12 +26,16 @@ class PostgresTable : public TableCatalogEntry { } public: - static bool PopulateColumns(CreateTableInfo &info, Oid relid, Snapshot snapshot); + static bool PopulateHeapColumns(CreateTableInfo &info, Oid relid, Snapshot snapshot); + static bool PopulateIndexColumns(CreateTableInfo &info, Oid relid, Path *path, bool indexonly); protected: PostgresTable(Catalog &catalog, SchemaCatalogEntry &schema, CreateTableInfo &info, Cardinality cardinality, Snapshot snapshot); +private: + static TupleDesc ExecTypeFromTLWithNames(List *target_list, TupleDesc tuple_desc); + protected: Cardinality cardinality; Snapshot snapshot; @@ -55,7 +59,7 @@ class PostgresHeapTable : public PostgresTable { class PostgresIndexTable : public PostgresTable { public: PostgresIndexTable(Catalog &catalog, SchemaCatalogEntry &schema, CreateTableInfo &info, Cardinality cardinality, - Snapshot snapshot, Path *path, PlannerInfo *planner_info); + Snapshot snapshot, bool is_indexonly_scan, Path *path, PlannerInfo *planner_info, Oid oid); public: // -- Table API -- @@ -66,6 +70,8 @@ class PostgresIndexTable : public PostgresTable { private: Path *path; PlannerInfo *planner_info; + bool indexonly_scan; + Oid oid; }; } // namespace duckdb diff --git a/include/pgduckdb/pgduckdb_types.hpp b/include/pgduckdb/pgduckdb_types.hpp index d444bef3..8b41b103 100644 --- a/include/pgduckdb/pgduckdb_types.hpp +++ b/include/pgduckdb/pgduckdb_types.hpp @@ -24,5 +24,9 @@ void ConvertPostgresToDuckValue(Datum value, duckdb::Vector &result, idx_t offse bool ConvertDuckToPostgresValue(TupleTableSlot *slot, duckdb::Value &value, idx_t col); void InsertTupleIntoChunk(duckdb::DataChunk &output, duckdb::shared_ptr scan_global_state, duckdb::shared_ptr scan_local_state, HeapTupleData *tuple); +void InsertTupleValuesIntoChunk(duckdb::DataChunk &output, + duckdb::shared_ptr scan_global_state, + duckdb::shared_ptr scan_local_state, Datum *values, + bool *nulls); } // namespace pgduckdb diff --git a/include/pgduckdb/scan/postgres_index_scan.hpp b/include/pgduckdb/scan/postgres_index_scan.hpp index b4985ba2..cec31dcf 100644 --- a/include/pgduckdb/scan/postgres_index_scan.hpp +++ b/include/pgduckdb/scan/postgres_index_scan.hpp @@ -17,8 +17,8 @@ namespace pgduckdb { // Global State struct PostgresIndexScanGlobalState : public duckdb::GlobalTableFunctionState { - explicit PostgresIndexScanGlobalState(IndexScanState *index_scan_state, Relation relation, - duckdb::TableFunctionInitInput &input); + explicit PostgresIndexScanGlobalState(IndexOptInfo *index, IndexScanState *index_scan_state, bool indexonly, + Relation relation, duckdb::TableFunctionInitInput &input); ~PostgresIndexScanGlobalState(); idx_t MaxThreads() const override { @@ -27,7 +27,9 @@ struct PostgresIndexScanGlobalState : public duckdb::GlobalTableFunctionState { public: duckdb::shared_ptr m_global_state; + IndexOptInfo *m_index; IndexScanState *m_index_scan; + bool m_indexonly; Relation m_relation; }; @@ -35,7 +37,7 @@ struct PostgresIndexScanGlobalState : public duckdb::GlobalTableFunctionState { struct PostgresIndexScanLocalState : public duckdb::LocalTableFunctionState { public: - PostgresIndexScanLocalState(IndexScanDesc index_scan_desc, Relation relation); + PostgresIndexScanLocalState(IndexScanDesc index_scan_desc, TupleDesc desc, Relation relation); ~PostgresIndexScanLocalState() override; public: @@ -43,18 +45,20 @@ struct PostgresIndexScanLocalState : public duckdb::LocalTableFunctionState { IndexScanDesc m_index_scan_desc; Relation m_relation; TupleTableSlot *m_slot; + TupleTableSlot *m_index_only_slot; }; // PostgresIndexScanFunctionData struct PostgresIndexScanFunctionData : public duckdb::TableFunctionData { public: - PostgresIndexScanFunctionData(uint64_t cardinality, Path *path, PlannerInfo *planner_info, Oid relation_oid, - Snapshot snapshot); + PostgresIndexScanFunctionData(uint64_t cardinality, bool indexonly, Path *path, PlannerInfo *planner_info, + Oid relation_oid, Snapshot Snapshot); ~PostgresIndexScanFunctionData() override; public: uint64_t m_cardinality; + bool m_indexonly; Path *m_path; PlannerInfo *m_planner_info; Snapshot m_snapshot; @@ -68,10 +72,6 @@ struct PostgresIndexScanFunction : public duckdb::TableFunction { PostgresIndexScanFunction(); public: - static duckdb::unique_ptr - PostgresIndexScanBind(duckdb::ClientContext &context, duckdb::TableFunctionBindInput &input, - duckdb::vector &return_types, duckdb::vector &names); - static duckdb::unique_ptr PostgresIndexScanInitGlobal(duckdb::ClientContext &context, duckdb::TableFunctionInitInput &input); diff --git a/src/catalog/pgduckdb_table.cpp b/src/catalog/pgduckdb_table.cpp index 15be3008..2b14b39a 100644 --- a/src/catalog/pgduckdb_table.cpp +++ b/src/catalog/pgduckdb_table.cpp @@ -21,6 +21,14 @@ extern "C" { #include "utils/syscache.h" #include "access/htup_details.h" #include "parser/parsetree.h" +#include "postgres.h" +#include "executor/nodeIndexscan.h" +#include "nodes/pathnodes.h" +#include "nodes/execnodes.h" +#include "nodes/makefuncs.h" +#include "nodes/nodeFuncs.h" +#include "parser/parsetree.h" +#include "utils/rel.h" } namespace duckdb { @@ -31,23 +39,84 @@ PostgresTable::PostgresTable(Catalog &catalog, SchemaCatalogEntry &schema, Creat } bool -PostgresTable::PopulateColumns(CreateTableInfo &info, Oid relid, Snapshot snapshot) { +PostgresTable::PopulateHeapColumns(CreateTableInfo &info, Oid relid, Snapshot snapshot) { auto rel = RelationIdGetRelation(relid); - auto tupleDesc = RelationGetDescr(rel); + auto tuple_desc = RelationGetDescr(rel); + + if (!tuple_desc) { + elog(WARNING, "Failed to get tuple descriptor for relation with OID %u", relid); + RelationClose(rel); + return false; + } + + for (int i = 0; i < tuple_desc->natts; i++) { + Form_pg_attribute attr = &tuple_desc->attrs[i]; + auto col_name = duckdb::string(NameStr(attr->attname)); + auto duck_type = pgduckdb::ConvertPostgresToDuckColumnType(attr); + info.columns.AddColumn(duckdb::ColumnDefinition(col_name, duck_type)); + /* Log column name and type */ + elog(DEBUG2, "(DuckDB/PopulateHeapColumns) Column name: %s, Type: %s --", col_name.c_str(), + duck_type.ToString().c_str()); + } + + RelationClose(rel); + return true; +} + +/* + * Generate tuple descriptor from target_list with column names matching + * relation original tuple descriptor. + */ +TupleDesc +PostgresTable::ExecTypeFromTLWithNames(List *target_list, TupleDesc tuple_desc) { + TupleDesc type_info; + ListCell *l; + int len; + int cur_resno = 1; + + len = ExecTargetListLength(target_list); + type_info = CreateTemplateTupleDesc(len); + + foreach (l, target_list) { + TargetEntry *tle = (TargetEntry *)lfirst(l); + const Var *var = (Var *)tle->expr; + + TupleDescInitEntry(type_info, cur_resno, tuple_desc->attrs[var->varattno - 1].attname.data, + exprType((Node *)tle->expr), exprTypmod((Node *)tle->expr), 0); + + TupleDescInitEntryCollation(type_info, cur_resno, exprCollation((Node *)tle->expr)); + cur_resno++; + } + + return type_info; +} + +bool +PostgresTable::PopulateIndexColumns(CreateTableInfo &info, Oid relid, Path *path, bool indexonly) { + auto rel = RelationIdGetRelation(relid); + auto tuple_desc = RelationGetDescr(rel); + + if (indexonly) { + tuple_desc = ExecTypeFromTLWithNames(((IndexPath *)path)->indexinfo->indextlist, RelationGetDescr(rel)); + elog(DEBUG2, "(PostgresTable::PopulateIndexColumns) IndexOnlyScan selected"); + } else { + tuple_desc = RelationGetDescr(rel); + elog(DEBUG2, "(PostgresTable::PopulateIndexColumns) IndexScan selected"); + } - if (!tupleDesc) { - elog(ERROR, "Failed to get tuple descriptor for relation with OID %u", relid); + if (!tuple_desc) { + elog(WARNING, "Failed to get tuple descriptor for relation with OID %u", relid); RelationClose(rel); return false; } - for (int i = 0; i < tupleDesc->natts; i++) { - Form_pg_attribute attr = &tupleDesc->attrs[i]; + for (int i = 0; i < tuple_desc->natts; i++) { + Form_pg_attribute attr = &tuple_desc->attrs[i]; auto col_name = duckdb::string(NameStr(attr->attname)); auto duck_type = pgduckdb::ConvertPostgresToDuckColumnType(attr); info.columns.AddColumn(duckdb::ColumnDefinition(col_name, duck_type)); /* Log column name and type */ - elog(DEBUG3, "-- (DuckDB/PostgresHeapBind) Column name: %s, Type: %s --", col_name.c_str(), + elog(DEBUG2, "-- (DuckDB/PopulateIndexColumns) Column name: %s, Type: %s --", col_name.c_str(), duck_type.ToString().c_str()); } @@ -85,9 +154,10 @@ PostgresHeapTable::GetStorageInfo(ClientContext &context) { //===--------------------------------------------------------------------===// PostgresIndexTable::PostgresIndexTable(Catalog &catalog, SchemaCatalogEntry &schema, CreateTableInfo &info, - Cardinality cardinality, Snapshot snapshot, Path *path, - PlannerInfo *planner_info) - : PostgresTable(catalog, schema, info, cardinality, snapshot), path(path), planner_info(planner_info) { + Cardinality cardinality, Snapshot snapshot, bool is_indexonly_scan, Path *path, + PlannerInfo *planner_info, Oid oid) + : PostgresTable(catalog, schema, info, cardinality, snapshot), path(path), planner_info(planner_info), + indexonly_scan(is_indexonly_scan), oid(oid) { } unique_ptr @@ -97,9 +167,8 @@ PostgresIndexTable::GetStatistics(ClientContext &context, column_t column_id) { TableFunction PostgresIndexTable::GetScanFunction(ClientContext &context, unique_ptr &bind_data) { - RangeTblEntry *rte = planner_rt_fetch(path->parent->relid, planner_info); - bind_data = duckdb::make_uniq(cardinality, path, planner_info, rte->relid, - snapshot); + bind_data = duckdb::make_uniq(cardinality, indexonly_scan, path, + planner_info, oid, snapshot); return pgduckdb::PostgresIndexScanFunction(); } diff --git a/src/catalog/pgduckdb_transaction.cpp b/src/catalog/pgduckdb_transaction.cpp index 67258ac4..64c8dfd3 100644 --- a/src/catalog/pgduckdb_transaction.cpp +++ b/src/catalog/pgduckdb_transaction.cpp @@ -41,8 +41,10 @@ FindMatchingRelEntry(Oid relid, PlannerInfo *planner_info) { int i = 1; RelOptInfo *node = nullptr; for (; i < planner_info->simple_rel_array_size; i++) { - if (planner_info->simple_rte_array[i]->rtekind == RTE_SUBQUERY && planner_info->simple_rel_array[i]) { + if (planner_info->simple_rte_array[i]->rtekind == RTE_SUBQUERY && planner_info->simple_rel_array[i] + && planner_info->simple_rel_array[i]->subroot) { node = FindMatchingRelEntry(relid, planner_info->simple_rel_array[i]->subroot); + //node = FindMatchingRelEntry(relid, planner_info->simple_rte_array[i]->subquery); if (node) { return node; } @@ -106,12 +108,20 @@ SchemaItems::GetTable(const string &entry_name, PlannerInfo *planner_info) { ReleaseSysCache(tuple); Path *node_path = nullptr; + RelOptInfo *node = nullptr; if (planner_info) { - auto node = FindMatchingRelEntry(rel_oid, planner_info); - if (node) { - node_path = get_cheapest_fractional_path(node, 0.0); + node = FindMatchingRelEntry(rel_oid, planner_info); + ListCell *lc; + /* We should prefer IndexOnlyScan */ + foreach (lc, node->pathlist) { + Path *p = (Path *)lfirst(lc); + if (p->pathtype == T_IndexOnlyScan) { + node_path = p; + } } + if (node && (node_path == nullptr)) + node_path = get_cheapest_fractional_path(node, 0.0); } unique_ptr table; @@ -119,14 +129,14 @@ SchemaItems::GetTable(const string &entry_name, PlannerInfo *planner_info) { info.table = entry_name; Cardinality cardinality = node_path ? node_path->rows : 1; if (IsIndexScan(node_path)) { - RangeTblEntry *rte = planner_rt_fetch(node_path->parent->relid, planner_info); - rel_oid = rte->relid; - if (!PostgresTable::PopulateColumns(info, rel_oid, snapshot)) { + auto is_indexonly_scan = node_path->pathtype == T_IndexOnlyScan; + if (!PostgresTable::PopulateIndexColumns(info, rel_oid, node_path, is_indexonly_scan)) { return nullptr; } - table = make_uniq(catalog, *schema, info, cardinality, snapshot, node_path, planner_info); + table = make_uniq(catalog, *schema, info, cardinality, snapshot, is_indexonly_scan, + node_path, planner_info, rel_oid); } else { - if (!PostgresTable::PopulateColumns(info, rel_oid, snapshot)) { + if (!PostgresTable::PopulateHeapColumns(info, rel_oid, snapshot)) { return nullptr; } table = make_uniq(catalog, *schema, info, cardinality, snapshot, rel_oid); @@ -175,5 +185,3 @@ PostgresTransaction::GetCatalogEntry(CatalogType type, const string &schema, con } } // namespace duckdb - -// namespace duckdb diff --git a/src/pgduckdb_types.cpp b/src/pgduckdb_types.cpp index f9d89dfb..3f145e49 100644 --- a/src/pgduckdb_types.cpp +++ b/src/pgduckdb_types.cpp @@ -1084,4 +1084,67 @@ InsertTupleIntoChunk(duckdb::DataChunk &output, duckdb::shared_ptr scan_global_state, + duckdb::shared_ptr scan_local_state, Datum *values, bool *nulls) { + + if (scan_global_state->m_count_tuples_only) { + scan_local_state->m_output_vector_size++; + return; + } + + bool valid_tuple = true; + + /* 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 &[columnIdx, valueIdx] : scan_global_state->m_read_columns_ids) { + if (scan_global_state->m_filters && + (scan_global_state->m_filters->filters.find(valueIdx) != scan_global_state->m_filters->filters.end())) { + auto &filter = scan_global_state->m_filters->filters[valueIdx]; + valid_tuple = ApplyValueFilter(*filter, values[valueIdx], nulls[valueIdx], + scan_global_state->m_tuple_desc->attrs[columnIdx].atttypid); + } + + if (!valid_tuple) { + break; + } + } + + /* Write tuple columns in output vector. */ + for (idx_t idx = 0; valid_tuple && idx < scan_global_state->m_output_columns_ids.size(); idx++) { + auto &result = output.data[idx]; + if (nulls[idx]) { + auto &array_mask = duckdb::FlatVector::Validity(result); + array_mask.SetInvalid(scan_local_state->m_output_vector_size); + } else { + idx_t output_column_idx = + scan_global_state->m_read_columns_ids[scan_global_state->m_output_columns_ids[idx]]; + if (scan_global_state->m_tuple_desc->attrs[scan_global_state->m_output_columns_ids[idx]].attlen == -1) { + bool should_free = false; + values[output_column_idx] = + DetoastPostgresDatum(reinterpret_cast(values[output_column_idx]), &should_free); + ConvertPostgresToDuckValue(values[output_column_idx], result, scan_local_state->m_output_vector_size); + if (should_free) { + duckdb_free(reinterpret_cast(values[output_column_idx])); + } + } else { + ConvertPostgresToDuckValue(values[output_column_idx], result, scan_local_state->m_output_vector_size); + } + } + } + + if (valid_tuple) { + scan_local_state->m_output_vector_size++; + scan_global_state->m_total_row_count++; + } + + if (!(scan_global_state->m_total_row_count % 100000)) { + elog(WARNING, "JOE %d", scan_global_state->m_total_row_count.load()); + } +} + } // namespace pgduckdb diff --git a/src/scan/index_scan_utils.cpp b/src/scan/index_scan_utils.cpp index 228f65d4..b2a10df9 100644 --- a/src/scan/index_scan_utils.cpp +++ b/src/scan/index_scan_utils.cpp @@ -1,7 +1,97 @@ +/* + * + */ + +extern "C"{ +#include "postgres.h" +#include "optimizer/paramassign.h" +#include "optimizer/placeholder.h" +} + #include "pgduckdb/scan/index_scan_utils.hpp" + namespace pgduckdb { +static Node * +replace_nestloop_params_mutator(Node *node, PlannerInfo *root); + +/* + * replace_nestloop_params + * Replace outer-relation Vars and PlaceHolderVars in the given expression + * with nestloop Params + * + * All Vars and PlaceHolderVars belonging to the relation(s) identified by + * root->curOuterRels are replaced by Params, and entries are added to + * root->curOuterParams if not already present. + */ +static Node * +replace_nestloop_params(PlannerInfo *root, Node *expr) +{ + /* No setup needed for tree walk, so away we go */ + return replace_nestloop_params_mutator(expr, root); +} + +static Node * +replace_nestloop_params_mutator(Node *node, PlannerInfo *root) +{ + if (node == NULL) + return NULL; + if (IsA(node, Var)) + { + Var *var = (Var *) node; + + /* Upper-level Vars should be long gone at this point */ + Assert(var->varlevelsup == 0); + /* If not to be replaced, we can just return the Var unmodified */ + if (IS_SPECIAL_VARNO(var->varno) || + !bms_is_member(var->varno, root->curOuterRels)) + return node; + /* Replace the Var with a nestloop Param */ + return (Node *) replace_nestloop_param_var(root, var); + } + if (IsA(node, PlaceHolderVar)) + { + PlaceHolderVar *phv = (PlaceHolderVar *) node; + + /* Upper-level PlaceHolderVars should be long gone at this point */ + Assert(phv->phlevelsup == 0); + + /* Check whether we need to replace the PHV */ + if (!bms_is_subset(find_placeholder_info(root, phv)->ph_eval_at, + root->curOuterRels)) + { + /* + * We can't replace the whole PHV, but we might still need to + * replace Vars or PHVs within its expression, in case it ends up + * actually getting evaluated here. (It might get evaluated in + * this plan node, or some child node; in the latter case we don't + * really need to process the expression here, but we haven't got + * enough info to tell if that's the case.) Flat-copy the PHV + * node and then recurse on its expression. + * + * Note that after doing this, we might have different + * representations of the contents of the same PHV in different + * parts of the plan tree. This is OK because equal() will just + * match on phid/phlevelsup, so setrefs.c will still recognize an + * upper-level reference to a lower-level copy of the same PHV. + */ + PlaceHolderVar *newphv = makeNode(PlaceHolderVar); + + memcpy(newphv, phv, sizeof(PlaceHolderVar)); + newphv->phexpr = (Expr *) + replace_nestloop_params_mutator((Node *) phv->phexpr, + root); + return (Node *) newphv; + } + /* Replace the PlaceHolderVar with a nestloop Param */ + return (Node *) replace_nestloop_param_placeholdervar(root, phv); + } + return expression_tree_mutator(node, + replace_nestloop_params_mutator, + (void *) root); +} + /* * fix_indexqual_operand * Convert an indexqual expression to a Var referencing the index column. @@ -81,7 +171,7 @@ FixIndexQualClause(PlannerInfo *root, IndexOptInfo *index, int indexcol, Node *c * This also makes a copy of the clause, so it's safe to modify it * in-place below. */ - // clause = replace_nestloop_params(root, clause); + clause = replace_nestloop_params(root, clause); if (IsA(clause, OpExpr)) { OpExpr *op = (OpExpr *)clause; diff --git a/src/scan/postgres_index_scan.cpp b/src/scan/postgres_index_scan.cpp index 729a319c..d83e45f9 100644 --- a/src/scan/postgres_index_scan.cpp +++ b/src/scan/postgres_index_scan.cpp @@ -11,9 +11,11 @@ extern "C" { #include "utils/rel.h" } +#include "pgduckdb/pgduckdb_utils.hpp" #include "pgduckdb/scan/index_scan_utils.hpp" #include "pgduckdb/scan/postgres_index_scan.hpp" #include "pgduckdb/pgduckdb_types.hpp" +#include "pgduckdb/pgduckdb_process_lock.hpp" #include "pgduckdb/vendor/pg_list.hpp" namespace pgduckdb { @@ -22,12 +24,18 @@ namespace pgduckdb { // PostgresIndexScanGlobalState // -PostgresIndexScanGlobalState::PostgresIndexScanGlobalState(IndexScanState *indexScanState, Relation relation, +PostgresIndexScanGlobalState::PostgresIndexScanGlobalState(IndexOptInfo *index, IndexScanState *index_scan_state, + bool indexonly, Relation relation, duckdb::TableFunctionInitInput &input) - : m_global_state(duckdb::make_shared_ptr()), m_index_scan(indexScanState), - m_relation(relation) { + : m_global_state(duckdb::make_shared_ptr()), m_index(index), + m_index_scan(index_scan_state), m_indexonly(indexonly), m_relation(relation) { m_global_state->InitGlobalState(input); - m_global_state->m_tuple_desc = RelationGetDescr(m_relation); + + if (m_indexonly) { + m_global_state->m_tuple_desc = ExecTypeFromTL(index->indextlist); + } else { + m_global_state->m_tuple_desc = RelationGetDescr(m_relation); + } } PostgresIndexScanGlobalState::~PostgresIndexScanGlobalState() { @@ -39,10 +47,11 @@ PostgresIndexScanGlobalState::~PostgresIndexScanGlobalState() { // PostgresIndexScanLocalState // -PostgresIndexScanLocalState::PostgresIndexScanLocalState(IndexScanDesc index_scan_desc, Relation relation) +PostgresIndexScanLocalState::PostgresIndexScanLocalState(IndexScanDesc index_scan_desc, TupleDesc desc, + Relation relation) : m_local_state(duckdb::make_shared_ptr()), m_index_scan_desc(index_scan_desc), m_relation(relation) { - m_slot = MakeTupleTableSlot(CreateTupleDescCopy(RelationGetDescr(m_relation)), &TTSOpsBufferHeapTuple); + m_slot = MakeTupleTableSlot(CreateTupleDescCopy(desc), &TTSOpsBufferHeapTuple); } PostgresIndexScanLocalState::~PostgresIndexScanLocalState() { @@ -53,11 +62,11 @@ PostgresIndexScanLocalState::~PostgresIndexScanLocalState() { // PostgresIndexScanFunctionData // -PostgresIndexScanFunctionData::PostgresIndexScanFunctionData(uint64_t cardinality, Path *path, +PostgresIndexScanFunctionData::PostgresIndexScanFunctionData(uint64_t cardinality, bool indexonly, Path *path, PlannerInfo *planner_info, Oid relation_oid, Snapshot snapshot) - : m_cardinality(cardinality), m_path(path), m_planner_info(planner_info), m_snapshot(snapshot), - m_relation_oid(relation_oid) { + : m_cardinality(cardinality), m_indexonly(indexonly), m_path(path), m_planner_info(planner_info), + m_snapshot(snapshot), m_relation_oid(relation_oid) { } PostgresIndexScanFunctionData::~PostgresIndexScanFunctionData() { @@ -68,9 +77,10 @@ PostgresIndexScanFunctionData::~PostgresIndexScanFunctionData() { // PostgresIndexScanFunction::PostgresIndexScanFunction() - : TableFunction("postgres_index_scan", {}, PostgresIndexScanFunc, PostgresIndexScanBind, - PostgresIndexScanInitGlobal, PostgresIndexScanInitLocal) { + : TableFunction("postgres_index_scan", {}, PostgresIndexScanFunc, nullptr, PostgresIndexScanInitGlobal, + PostgresIndexScanInitLocal) { named_parameters["cardinality"] = duckdb::LogicalType::UBIGINT; + named_parameters["indexonly"] = duckdb::LogicalType::BOOLEAN; named_parameters["path"] = duckdb::LogicalType::POINTER; named_parameters["planner_info"] = duckdb::LogicalType::POINTER; named_parameters["snapshot"] = duckdb::LogicalType::POINTER; @@ -80,51 +90,21 @@ PostgresIndexScanFunction::PostgresIndexScanFunction() cardinality = PostgresIndexScanCardinality; } -duckdb::unique_ptr -PostgresIndexScanFunction::PostgresIndexScanBind(duckdb::ClientContext &context, duckdb::TableFunctionBindInput &input, - duckdb::vector &return_types, - duckdb::vector &names) { - auto cardinality = input.named_parameters["cardinality"].GetValue(); - auto path = (reinterpret_cast(input.named_parameters["path"].GetPointer())); - auto planner_info = (reinterpret_cast(input.named_parameters["planner_info"].GetPointer())); - auto snapshot = (reinterpret_cast(input.named_parameters["snapshot"].GetPointer())); - - RangeTblEntry *rte = planner_rt_fetch(path->parent->relid, planner_info); - - auto rel = RelationIdGetRelation(rte->relid); - auto relation_descr = RelationGetDescr(rel); - - if (!relation_descr) { - elog(WARNING, "(PGDuckDB/PostgresIndexScanBind) Failed to get tuple descriptor for relation with OID %u", - rel->rd_id); - return nullptr; - } - - for (int i = 0; i < relation_descr->natts; i++) { - Form_pg_attribute attr = &relation_descr->attrs[i]; - auto col_name = duckdb::string(NameStr(attr->attname)); - auto duck_type = ConvertPostgresToDuckColumnType(attr); - return_types.push_back(duck_type); - names.push_back(col_name); - /* Log column name and type */ - elog(DEBUG2, "(PGDuckDB/PostgresIndexScanBind) Column name: %s, Type: %s --", col_name.c_str(), - duck_type.ToString().c_str()); - } - - RelationClose(rel); - - return duckdb::make_uniq(cardinality, path, planner_info, rte->relid, snapshot); -} - duckdb::unique_ptr PostgresIndexScanFunction::PostgresIndexScanInitGlobal(duckdb::ClientContext &context, duckdb::TableFunctionInitInput &input) { auto &bind_data = input.bind_data->CastNoConst(); + auto old_stack_base = set_stack_base(); + IndexScanState *index_state = makeNode(IndexScanState); IndexPath *index_path = (IndexPath *)bind_data.m_path; + bool indexonly = bind_data.m_indexonly; + + elog(WARNING, "--%d", index_path->indexinfo->indexoid); - index_state->iss_RelationDesc = index_open(index_path->indexinfo->indexoid, AccessShareLock); + index_state->iss_RelationDesc = + PostgresFunctionGuard(index_open, index_path->indexinfo->indexoid, AccessShareLock); index_state->iss_RuntimeKeysReady = false; index_state->iss_RuntimeKeys = NULL; index_state->iss_NumRuntimeKeys = 0; @@ -132,22 +112,31 @@ PostgresIndexScanFunction::PostgresIndexScanInitGlobal(duckdb::ClientContext &co List *stripped_clause_list = NIL; IndexOptInfo *index = index_path->indexinfo; - foreach_node(IndexClause, iclause, index_path->indexclauses) { - int indexcol = iclause->indexcol; - foreach_node(RestrictInfo, rinfo, iclause->indexquals) { - Node *clause = (Node *)rinfo->clause; + ListCell *lc; + + foreach(lc, index_path->indexclauses) + { + IndexClause *iclause = lfirst_node(IndexClause, lc); + int indexcol = iclause->indexcol; + ListCell *lc2; + + foreach(lc2, iclause->indexquals) + { + RestrictInfo *rinfo = lfirst_node(RestrictInfo, lc2); + Node *clause = (Node *) rinfo->clause; clause = FixIndexQualClause(bind_data.m_planner_info, index, indexcol, clause, iclause->indexcols); stripped_clause_list = lappend(stripped_clause_list, clause); } } - ExecIndexBuildScanKeys((PlanState *)index_state, index_state->iss_RelationDesc, stripped_clause_list, false, - &index_state->iss_ScanKeys, &index_state->iss_NumScanKeys, &index_state->iss_RuntimeKeys, - &index_state->iss_NumRuntimeKeys, NULL, /* no ArrayKeys */ - NULL); + PostgresFunctionGuard(ExecIndexBuildScanKeys, (PlanState *)index_state, index_state->iss_RelationDesc, + stripped_clause_list, false, &index_state->iss_ScanKeys, &index_state->iss_NumScanKeys, + &index_state->iss_RuntimeKeys, &index_state->iss_NumRuntimeKeys, nullptr, nullptr); - return duckdb::make_uniq(index_state, RelationIdGetRelation(bind_data.m_relation_oid), - input); + restore_stack_base(old_stack_base); + + return duckdb::make_uniq(index, index_state, indexonly, + RelationIdGetRelation(bind_data.m_relation_oid), input); } duckdb::unique_ptr @@ -157,15 +146,22 @@ PostgresIndexScanFunction::PostgresIndexScanInitLocal(duckdb::ExecutionContext & auto &bind_data = input.bind_data->CastNoConst(); auto global_state = reinterpret_cast(gstate); - IndexScanDesc scandesc = - index_beginscan(global_state->m_relation, global_state->m_index_scan->iss_RelationDesc, bind_data.m_snapshot, - global_state->m_index_scan->iss_NumScanKeys, global_state->m_index_scan->iss_NumOrderByKeys); + IndexScanDesc scandesc = PostgresFunctionGuard( + index_beginscan, global_state->m_relation, global_state->m_index_scan->iss_RelationDesc, bind_data.m_snapshot, + global_state->m_index_scan->iss_NumScanKeys, global_state->m_index_scan->iss_NumOrderByKeys); - if (global_state->m_index_scan->iss_NumRuntimeKeys == 0 || global_state->m_index_scan->iss_RuntimeKeysReady) - index_rescan(scandesc, global_state->m_index_scan->iss_ScanKeys, global_state->m_index_scan->iss_NumScanKeys, - global_state->m_index_scan->iss_OrderByKeys, global_state->m_index_scan->iss_NumOrderByKeys); + if (global_state->m_indexonly) { + scandesc->xs_want_itup = true; + } - return duckdb::make_uniq(scandesc, global_state->m_relation); + if (global_state->m_index_scan->iss_NumRuntimeKeys == 0 || global_state->m_index_scan->iss_RuntimeKeysReady) { + PostgresFunctionGuard(index_rescan, scandesc, global_state->m_index_scan->iss_ScanKeys, + global_state->m_index_scan->iss_NumScanKeys, global_state->m_index_scan->iss_OrderByKeys, + global_state->m_index_scan->iss_NumOrderByKeys); + } + + return duckdb::make_uniq(scandesc, global_state->m_global_state->m_tuple_desc, + global_state->m_relation); } void @@ -173,7 +169,7 @@ PostgresIndexScanFunction::PostgresIndexScanFunc(duckdb::ClientContext &context, duckdb::DataChunk &output) { auto &local_state = data.local_state->Cast(); auto &global_state = data.global_state->Cast(); - bool next_index_tuple = false; + ItemPointer next_index_tid = nullptr; local_state.m_local_state->m_output_vector_size = 0; @@ -182,16 +178,42 @@ PostgresIndexScanFunction::PostgresIndexScanFunc(duckdb::ClientContext &context, return; } - while (local_state.m_local_state->m_output_vector_size < STANDARD_VECTOR_SIZE && - (next_index_tuple = - index_getnext_slot(local_state.m_index_scan_desc, ForwardScanDirection, local_state.m_slot))) { - bool should_free; - auto tuple = ExecFetchSlotHeapTuple(local_state.m_slot, false, &should_free); - InsertTupleIntoChunk(output, global_state.m_global_state, local_state.m_local_state, tuple); + while (local_state.m_local_state->m_output_vector_size < STANDARD_VECTOR_SIZE) { + + DuckdbProcessLock::GetLock().lock(); + next_index_tid = index_getnext_tid(local_state.m_index_scan_desc, ForwardScanDirection); + DuckdbProcessLock::GetLock().unlock(); + + /* No more index tuples to read */ + if (!next_index_tid) { + break; + } + + if (!global_state.m_indexonly) { + DuckdbProcessLock::GetLock().lock(); + auto found = index_fetch_heap(local_state.m_index_scan_desc, local_state.m_slot); + DuckdbProcessLock::GetLock().unlock(); + /* Tuple not found */ + if (!found) { + continue; + } + InsertTupleIntoChunk(output, global_state.m_global_state, local_state.m_local_state, + reinterpret_cast(local_state.m_slot)->base.tuple); + } else { + if (local_state.m_index_scan_desc->xs_hitup) { + InsertTupleIntoChunk(output, global_state.m_global_state, local_state.m_local_state, + local_state.m_index_scan_desc->xs_hitup); + } else if (local_state.m_index_scan_desc->xs_itup) { + index_deform_tuple(local_state.m_index_scan_desc->xs_itup, local_state.m_index_scan_desc->xs_itupdesc, + local_state.m_slot->tts_values, local_state.m_slot->tts_isnull); + InsertTupleValuesIntoChunk(output, global_state.m_global_state, local_state.m_local_state, + local_state.m_slot->tts_values, local_state.m_slot->tts_isnull); + } + } ExecClearTuple(local_state.m_slot); } - if (!next_index_tuple) { + if (!next_index_tid) { local_state.m_local_state->m_exhausted_scan = true; } diff --git a/src/scan/postgres_seq_scan.cpp b/src/scan/postgres_seq_scan.cpp index 0026839e..ac100baf 100644 --- a/src/scan/postgres_seq_scan.cpp +++ b/src/scan/postgres_seq_scan.cpp @@ -15,7 +15,7 @@ PostgresSeqScanGlobalState::PostgresSeqScanGlobalState(Relation relation, duckdb m_heap_reader_global_state(duckdb::make_shared_ptr(relation)), m_relation(relation) { m_global_state->InitGlobalState(input); m_global_state->m_tuple_desc = RelationGetDescr(m_relation); - elog(DEBUG2, "-- (DuckDB/PostgresReplacementScanGlobalState) Running %" PRIu64 " threads -- ", + elog(DEBUG2, "(DuckDB/PostgresReplacementScanGlobalState) Running %" PRIu64 " threads -- ", (uint64_t)MaxThreads()); } diff --git a/test/regression/expected/index_scan.out b/test/regression/expected/index_scan.out new file mode 100644 index 00000000..e69de29b diff --git a/test/regression/schedule b/test/regression/schedule index 360ad831..537e813b 100644 --- a/test/regression/schedule +++ b/test/regression/schedule @@ -13,3 +13,4 @@ test: cte test: create_table_as test: standard_conforming_strings test: query_filter +test: index_scan diff --git a/test/regression/sql/index_scan.sql b/test/regression/sql/index_scan.sql new file mode 100644 index 00000000..a74853a1 --- /dev/null +++ b/test/regression/sql/index_scan.sql @@ -0,0 +1,19 @@ +CREATE TABLE t(a INT PRIMARY KEY, b TEXT, c FLOAT); + +INSERT INTO t SELECT g, 'abcde_'||g, g % 1000 FROM generate_series(1,1000000) g; + +SET client_min_messages to 'DEBUG2'; + +SELECT b FROM t WHERE a = 10; + +CREATE INDEX t_idx_mixed on t(a, c); + +SET duckdb.execution TO FALSE; +EXPLAIN SELECT a,c, count(*) FROM t WHERE a < 1000 GROUP BY a,c ORDER BY a LIMIT 10; +SET duckdb.execution TO TRUE; + +SELECT a,c, count(*) FROM t WHERE a = 10 GROUP BY a,c ORDER BY a LIMIT 10; + +SET client_min_messages to default; + +DROP TABLE t;