feat: log number of row groups

Print to external log file the number of fetched row groups.
rimarin · Feb 22, 2024 · 5164a1c · 5164a1c
1 parent 6eb7745
commit 5164a1c
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 5 deletions.
diff --git a/extension/parquet/include/parquet_reader.hpp b/extension/parquet/include/parquet_reader.hpp
@@ -64,6 +64,7 @@ struct ParquetReaderScanState {
 	bool prefetch_mode = false;
 	bool current_group_prefetched = false;
 	set<string> fetchedFiles;
+	unordered_map<string, set<int64_t>> fetchedRowGroups;
 };
 
 struct ParquetColumnDefinition {

diff --git a/extension/parquet/parquet_extension.cpp b/extension/parquet/parquet_extension.cpp
@@ -109,6 +109,7 @@ struct ParquetReadGlobalState : public GlobalTableFunctionState {
 	vector<column_t> column_ids;
 	TableFilterSet *filters;
 	set<string> fetchedFiles;
+	unordered_map<string, set<int64_t>> fetchedRowsGroups;
 
 	idx_t MaxThreads() const override {
 		return max_threads;
@@ -615,6 +616,9 @@ class ParquetScanFunction {
 			}
 
 			gstate.fetchedFiles.insert(data.scan_state.fetchedFiles.begin(), data.scan_state.fetchedFiles.end());
+			for (const auto &fileToRowsGroups : data.scan_state.fetchedRowGroups){
+				gstate.fetchedRowsGroups[fileToRowsGroups.first].insert(fileToRowsGroups.second.begin(), fileToRowsGroups.second.end());
+			}
 
 			bind_data.chunk_count++;
 			if (output.size() > 0) {
@@ -629,6 +633,15 @@ class ParquetScanFunction {
 		numPartitionsFile.open("partitions.log", std::fstream::out);
 		numPartitionsFile << gstate.fetchedFiles.size() << "\n";
 		numPartitionsFile.close();
+		// Log the number of fetched row groups
+		uint64_t totalRowGroups = 0;
+		for (const auto &fileToRowsGroups : gstate.fetchedRowsGroups){
+			totalRowGroups += fileToRowsGroups.second.size();
+		}
+		std::ofstream numRowGroupsFile;
+		numRowGroupsFile.open("row_groups.log", std::fstream::out);
+		numRowGroupsFile << totalRowGroups << "\n";
+		numRowGroupsFile.close();
 	}
 
 	static unique_ptr<NodeStatistics> ParquetCardinality(ClientContext &context, const FunctionData *bind_data) {

diff --git a/extension/parquet/parquet_reader.cpp b/extension/parquet/parquet_reader.cpp
@@ -910,19 +910,24 @@ bool ParquetReader::ScanInternal(ParquetReaderScanState &state, DataChunk &resul
 		}
 
 		uint64_t to_scan_compressed_bytes = 0;
-		bool columnSkipped = false;
+		bool rowGroupSkipped = false;
 		for (idx_t col_idx = 0; col_idx < reader_data.column_ids.size(); col_idx++) {
-			columnSkipped |= PrepareRowGroupBuffer(state, col_idx);
-			// True: row group is relevant for the query filters
-			// False: row group was pruned thanks to columnar statistics
+			// True: column chunk was pruned thanks to columnar statistics
+			// False: row group is relevant for the query filters
+			bool columnChunkSkipped = PrepareRowGroupBuffer(state, col_idx);
+			rowGroupSkipped |= columnChunkSkipped;
 
 			auto file_col_idx = reader_data.column_ids[col_idx];
 
 			auto &root_reader = state.root_reader->Cast<StructColumnReader>();
 			to_scan_compressed_bytes += root_reader.GetChildReader(file_col_idx)->TotalCompressedSize();
 		}
 
-		if (!columnSkipped){
+		if (!rowGroupSkipped){
+			// Keep track of the fetched row group
+			auto fetchedRowGroup = state.group_idx_list[state.current_group];
+			state.fetchedRowGroups[file_name].emplace(fetchedRowGroup);
+			// Keep track of the fetched file
 			state.fetchedFiles.emplace(file_name);
 		}