-
Notifications
You must be signed in to change notification settings - Fork 86
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #1834 from kuzudb/clean-copy
Clean up duplicated morsel and sharedState between copy node and rel
- Loading branch information
Showing
20 changed files
with
392 additions
and
491 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,116 @@ | ||
| #pragma once | ||
|
|
||
| #include "storage/copier/table_copy_utils.h" | ||
|
|
||
| namespace kuzu { | ||
| namespace storage { | ||
|
|
||
| class ReadFileMorsel { | ||
| public: | ||
| ReadFileMorsel(common::row_idx_t rowIdx, common::block_idx_t blockIdx, | ||
| common::row_idx_t numRows, std::string filePath, common::row_idx_t rowIdxInFile) | ||
| : rowIdx{rowIdx}, blockIdx{blockIdx}, numRows{numRows}, filePath{std::move(filePath)}, | ||
| rowIdxInFile{rowIdxInFile} {}; | ||
| virtual ~ReadFileMorsel() = default; | ||
|
|
||
| public: | ||
| // When reading from multiple files, rowIdx is accumulated. | ||
| common::row_idx_t rowIdx; | ||
| common::block_idx_t blockIdx; | ||
| common::row_idx_t numRows; | ||
| std::string filePath; | ||
| // Row idx in the current file. Equal to `rowIdx` when reading from only a single file. | ||
| common::row_idx_t rowIdxInFile; | ||
| }; | ||
|
|
||
| // For CSV file, we need to read in streaming mode, so we need to read one batch at a time. | ||
| class ReadCSVMorsel : public ReadFileMorsel { | ||
| public: | ||
| ReadCSVMorsel(common::offset_t startRowIdx, std::string filePath, | ||
| common::row_idx_t rowIdxInFile, std::shared_ptr<arrow::RecordBatch> recordBatch) | ||
| : ReadFileMorsel{startRowIdx, common::INVALID_BLOCK_IDX, common::INVALID_ROW_IDX, | ||
| std::move(filePath), rowIdxInFile}, | ||
| recordBatch{std::move(recordBatch)} {} | ||
|
|
||
| std::shared_ptr<arrow::RecordBatch> recordBatch; | ||
| }; | ||
|
|
||
| class ReadNPYMorsel : public ReadFileMorsel { | ||
| public: | ||
| ReadNPYMorsel(common::row_idx_t rowIdx, common::block_idx_t blockIdx, common::row_idx_t numRows, | ||
| std::string filePath, common::vector_idx_t curFileIdx, common::row_idx_t rowIdxInFile) | ||
| : ReadFileMorsel{rowIdx, blockIdx, numRows, std::move(filePath), rowIdxInFile}, | ||
| columnIdx{curFileIdx} {} | ||
|
|
||
| inline common::vector_idx_t getColumnIdx() const { return columnIdx; } | ||
|
|
||
| private: | ||
| common::vector_idx_t columnIdx; | ||
| }; | ||
|
|
||
| class ReadFileSharedState { | ||
| public: | ||
| ReadFileSharedState(std::vector<std::string> filePaths, common::CSVReaderConfig csvReaderConfig, | ||
| catalog::TableSchema* tableSchema) | ||
| : numRows{0}, tableSchema{tableSchema}, filePaths{std::move(filePaths)}, | ||
| csvReaderConfig{csvReaderConfig}, currRowIdx{0}, currBlockIdx{0}, currFileIdx{0}, | ||
| currRowIdxInCurrFile{1} {}; | ||
| virtual ~ReadFileSharedState() = default; | ||
|
|
||
| virtual void countNumRows() = 0; | ||
| virtual std::unique_ptr<ReadFileMorsel> getMorsel() = 0; | ||
|
|
||
| public: | ||
| std::mutex mtx; | ||
| common::row_idx_t numRows; | ||
| catalog::TableSchema* tableSchema; | ||
| std::vector<std::string> filePaths; | ||
| common::CSVReaderConfig csvReaderConfig; | ||
| std::unordered_map<std::string, FileBlockInfo> fileBlockInfos; | ||
| common::row_idx_t currRowIdx; | ||
| common::block_idx_t currBlockIdx; | ||
| common::vector_idx_t currFileIdx; | ||
| common::row_idx_t currRowIdxInCurrFile; | ||
| }; | ||
|
|
||
| // For CSV file, we need to read in streaming mode, so we need to keep the reader in the shared | ||
| // state. | ||
| class ReadCSVSharedState : public ReadFileSharedState { | ||
| public: | ||
| ReadCSVSharedState(std::vector<std::string> filePaths, common::CSVReaderConfig csvReaderConfig, | ||
| catalog::TableSchema* tableSchema) | ||
| : ReadFileSharedState{std::move(filePaths), csvReaderConfig, tableSchema} {}; | ||
|
|
||
| void countNumRows() final; | ||
| std::unique_ptr<ReadFileMorsel> getMorsel() final; | ||
|
|
||
| private: | ||
| std::shared_ptr<arrow::csv::StreamingReader> reader; | ||
| }; | ||
|
|
||
| class ReadParquetSharedState : public ReadFileSharedState { | ||
| public: | ||
| explicit ReadParquetSharedState(std::vector<std::string> filePaths, | ||
| common::CSVReaderConfig csvReaderConfig, catalog::TableSchema* tableSchema) | ||
| : ReadFileSharedState{std::move(filePaths), csvReaderConfig, tableSchema} {} | ||
|
|
||
| private: | ||
| void countNumRows() override; | ||
|
|
||
| std::unique_ptr<storage::ReadFileMorsel> getMorsel() override; | ||
| }; | ||
|
|
||
| class ReadNPYSharedState : public ReadFileSharedState { | ||
| public: | ||
| ReadNPYSharedState(std::vector<std::string> filePaths, common::CSVReaderConfig csvReaderConfig, | ||
| catalog::NodeTableSchema* tableSchema) | ||
| : ReadFileSharedState{std::move(filePaths), csvReaderConfig, tableSchema} {} | ||
|
|
||
| std::unique_ptr<storage::ReadFileMorsel> getMorsel() final; | ||
|
|
||
| private: | ||
| void countNumRows() final; | ||
| }; | ||
|
|
||
| } // namespace storage | ||
| } // namespace kuzu |
Oops, something went wrong.