feat: log number of rows

Print to external log file the number of fetched rows.
rimarin · Feb 22, 2024 · 8e1d3f1 · 8e1d3f1
1 parent 5164a1c
commit 8e1d3f1
Show file tree

Hide file tree

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

diff --git a/extension/parquet/parquet_extension.cpp b/extension/parquet/parquet_extension.cpp
@@ -109,7 +109,8 @@ struct ParquetReadGlobalState : public GlobalTableFunctionState {
 	vector<column_t> column_ids;
 	TableFilterSet *filters;
 	set<string> fetchedFiles;
-	unordered_map<string, set<int64_t>> fetchedRowsGroups;
+	unordered_map<string, set<uint64_t>> fetchedRowsGroups;
+	unordered_map<string, uint64_t> fetchedGroupsToRows;
 
 	idx_t MaxThreads() const override {
 		return max_threads;
@@ -618,6 +619,14 @@ 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());
+				auto fileRowGroups = gstate.fetchedRowsGroups[data.reader->file_name];
+				for (const auto &rowGroupIdx : fileRowGroups){
+					auto rowGroupsMetadata = data.reader->metadata->metadata->row_groups;
+					if (rowGroupIdx < rowGroupsMetadata.size()){
+						auto rowGroupMetadata = rowGroupsMetadata[rowGroupIdx];
+						gstate.fetchedGroupsToRows[data.reader->file_name + to_string(rowGroupIdx)] = rowGroupMetadata.num_rows;
+					}
+				}
 			}
 
 			bind_data.chunk_count++;
@@ -633,6 +642,15 @@ class ParquetScanFunction {
 		numPartitionsFile.open("partitions.log", std::fstream::out);
 		numPartitionsFile << gstate.fetchedFiles.size() << "\n";
 		numPartitionsFile.close();
+		// Log the number of rows
+		uint64_t totalRows = 0;
+		for (const auto &rowGroupToRows : gstate.fetchedGroupsToRows){
+			totalRows += rowGroupToRows.second;
+		}
+		std::ofstream numRowsFile;
+		numRowsFile.open("rows.log", std::fstream::out);
+		numRowsFile << totalRows << "\n";
+		numRowsFile.close();
 		// Log the number of fetched row groups
 		uint64_t totalRowGroups = 0;
 		for (const auto &fileToRowsGroups : gstate.fetchedRowsGroups){

diff --git a/extension/parquet/parquet_reader.cpp b/extension/parquet/parquet_reader.cpp
@@ -925,7 +925,7 @@ bool ParquetReader::ScanInternal(ParquetReaderScanState &state, DataChunk &resul
 
 		if (!rowGroupSkipped){
 			// Keep track of the fetched row group
-			auto fetchedRowGroup = state.group_idx_list[state.current_group];
+			uint64_t 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);