Query COUNT(*)

* COUNT(*) doesn't require any columns to be retrieved so we only count tuples that pass visibility without fetching.
duckdb · Apr 27, 2024 · ebe5014 · ebe5014
1 parent 28e46d3
commit ebe5014
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 2 deletions.
diff --git 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<duckdb::column_t, duckdb::idx_t> m_columns;
 	duckdb::map<duckdb::idx_t, duckdb::column_t> m_projections;
 	duckdb::TableFilterSet *m_filters = nullptr;

diff --git a/src/quack_heap_seq_scan.cpp b/src/quack_heap_seq_scan.cpp
@@ -60,13 +60,21 @@ PostgresHeapSeqScan::InitParallelScanState(const duckdb::vector<duckdb::column_t
                                            const duckdb::vector<duckdb::idx_t> &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
@@ -239,6 +239,11 @@ InsertTupleIntoChunk(duckdb::DataChunk &output, PostgresHeapSeqScanThreadInfo &t
                      PostgresHeapSeqParallelScanState &parallelScanState) {
 	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());