diff --git a/include/quack/quack_heap_seq_scan.hpp b/include/quack/quack_heap_seq_scan.hpp index a209702d..4021ef33 100644 --- a/include/quack/quack_heap_seq_scan.hpp +++ b/include/quack/quack_heap_seq_scan.hpp @@ -39,8 +39,8 @@ class PostgresHeapSeqParallelScanState { public: PostgresHeapSeqParallelScanState() - : m_nblocks(InvalidBlockNumber), m_last_assigned_block_number(InvalidBlockNumber), m_total_row_count(0), - m_last_prefetch_block(0), m_strategy(nullptr) { + : m_nblocks(InvalidBlockNumber), m_last_assigned_block_number(InvalidBlockNumber), m_count_tuple_only(false), + m_total_row_count(0), m_last_prefetch_block(0), m_strategy(nullptr) { } ~PostgresHeapSeqParallelScanState() { if (m_strategy) @@ -51,6 +51,7 @@ class PostgresHeapSeqParallelScanState { std::mutex m_lock; BlockNumber m_nblocks; BlockNumber m_last_assigned_block_number; + bool m_count_tuple_only; duckdb::map m_columns; duckdb::map m_projections; duckdb::TableFilterSet *m_filters = nullptr; diff --git a/src/quack_heap_seq_scan.cpp b/src/quack_heap_seq_scan.cpp index dd79f8f7..562498da 100644 --- a/src/quack_heap_seq_scan.cpp +++ b/src/quack_heap_seq_scan.cpp @@ -60,13 +60,21 @@ PostgresHeapSeqScan::InitParallelScanState(const duckdb::vector &projections, duckdb::TableFilterSet *filters) { m_parallel_scan_state.m_nblocks = RelationGetNumberOfBlocks(m_rel); + + if (columns.size() == 1 && columns[0] == UINT64_MAX) { + m_parallel_scan_state.m_count_tuple_only = true; + return; + } + /* We need ordered columns ids for tuple fetch */ for (duckdb::idx_t i = 0; i < columns.size(); i++) { m_parallel_scan_state.m_columns[columns[i]] = i; } + for (duckdb::idx_t i = 0; i < columns.size(); i++) { m_parallel_scan_state.m_projections[projections[i]] = columns[i]; } + //m_parallel_scan_state.PrefetchNextRelationPages(m_rel); m_parallel_scan_state.m_filters = filters; } diff --git a/src/quack_types.cpp b/src/quack_types.cpp index f8190284..d6c8fc48 100644 --- a/src/quack_types.cpp +++ b/src/quack_types.cpp @@ -239,6 +239,11 @@ InsertTupleIntoChunk(duckdb::DataChunk &output, PostgresHeapSeqScanThreadInfo &t PostgresHeapSeqParallelScanState ¶llelScanState) { HeapTupleReadState heapTupleReadState = {}; + if (parallelScanState.m_count_tuple_only) { + threadScanInfo.m_output_vector_size++; + return; + } + Datum *values = (Datum *)duckdb_malloc(sizeof(Datum) * parallelScanState.m_columns.size()); bool *nulls = (bool *)duckdb_malloc(sizeof(bool) * parallelScanState.m_columns.size());