read parquet metadata

kuzudb · Apr 21, 2023 · 3a41372 · 3a41372
1 parent a0a51d5
commit 3a41372
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 8 deletions.
diff --git a/src/storage/copier/node_copier.cpp b/src/storage/copier/node_copier.cpp
@@ -215,8 +215,6 @@ void ParquetNodeCopier<T>::executeInternal(std::unique_ptr<NodeCopyMorsel> morse
     TableCopyExecutor::throwCopyExceptionIfNotOK(
         reader->RowGroup(morsel->blockIdx)->ReadTable(&table));
     arrow::TableBatchReader batchReader(*table);
-    // TODO(GUODONG): We assume here each time the table only contains one record batch. Needs to
-    // verify if this always holds true.
     std::shared_ptr<arrow::RecordBatch> recordBatch;
     TableCopyExecutor::throwCopyExceptionIfNotOK(batchReader.ReadNext(&recordBatch));
     std::vector<std::unique_ptr<InMemColumnChunk>> columnChunks(this->columns.size());

diff --git a/src/storage/copier/table_copy_executor.cpp b/src/storage/copier/table_copy_executor.cpp
@@ -82,17 +82,16 @@ void TableCopyExecutor::countNumLinesParquet(const std::vector<std::string>& fil
     numRows = 0;
     for (auto& filePath : filePaths) {
         std::unique_ptr<parquet::arrow::FileReader> reader = createParquetReader(filePath);
-        uint64_t numBlocks = reader->num_row_groups();
-        std::vector<uint64_t> numLinesPerBlock;
-        std::shared_ptr<arrow::Table> table;
+        auto metadata = reader->parquet_reader()->metadata();
+        uint64_t numBlocks = metadata->num_row_groups();
+        std::vector<uint64_t> numLinesPerBlock(numBlocks);
         auto startNodeOffset = numRows;
         for (auto blockIdx = 0; blockIdx < numBlocks; ++blockIdx) {
-            throwCopyExceptionIfNotOK(reader->RowGroup(blockIdx)->ReadTable(&table));
-            numLinesPerBlock.push_back(table->num_rows());
-            numRows += table->num_rows();
+            numLinesPerBlock[blockIdx] = metadata->RowGroup(blockIdx)->num_rows();
         }
         fileBlockInfos.emplace(
             filePath, FileBlockInfo{startNodeOffset, numBlocks, numLinesPerBlock});
+        numRows += metadata->num_rows();
     }
 }