-
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.
Support Parquet filetype on COPY TO (#1893)
- Loading branch information
Showing
21 changed files
with
1,910 additions
and
65 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
19 changes: 10 additions & 9 deletions
19
src/include/processor/operator/persistent/csv_file_writer.h
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,27 @@ | ||
#pragma once | ||
|
||
#include "common/copier_config/copier_config.h" | ||
#include "processor/result/result_set.h" | ||
|
||
namespace kuzu { | ||
namespace processor { | ||
|
||
class FileWriter { | ||
public: | ||
FileWriter(std::string filePath, std::vector<std::string> columnNames, | ||
std::vector<common::LogicalType> columnTypes) | ||
: filePath{filePath}, columnNames{columnNames}, columnTypes{columnTypes} {} | ||
virtual ~FileWriter(){}; | ||
virtual void init() = 0; | ||
virtual void openFile() = 0; | ||
virtual void closeFile() = 0; | ||
virtual void writeValues(std::vector<common::ValueVector*>& outputVectors) = 0; | ||
|
||
protected: | ||
std::string filePath; | ||
std::vector<std::string> columnNames; | ||
std::vector<common::LogicalType> columnTypes; | ||
}; | ||
|
||
} // namespace processor | ||
} // namespace kuzu |
76 changes: 76 additions & 0 deletions
76
src/include/processor/operator/persistent/parquet_column_writer.h
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,76 @@ | ||
#pragma once | ||
|
||
#include "common/types/date_t.h" | ||
#include "common/types/timestamp_t.h" | ||
#include "processor/operator/persistent/file_writer.h" | ||
#include <parquet/api/writer.h> | ||
|
||
namespace kuzu { | ||
namespace processor { | ||
|
||
// Kuzu column represents a single column in Kuzu, which may contains multiple Parquet columns | ||
// writeColumn will be called for each Kuzu column. | ||
class ParquetColumnWriter { | ||
public: | ||
ParquetColumnWriter(int totalColumns, std::vector<int> maxDefinitionLevels) | ||
: maxDefinitionLevels(maxDefinitionLevels), totalColumns(totalColumns){}; | ||
|
||
void writeColumn( | ||
int column, common::ValueVector* vector, parquet::RowGroupWriter* rowGroupWriter); | ||
|
||
int64_t estimatedRowBytes; | ||
|
||
private: | ||
void nextParquetColumn(common::LogicalTypeID logicalTypeID); | ||
void writePrimitiveValue(common::LogicalTypeID logicalTypeID, uint8_t* value, | ||
int16_t definitionLevel = 0, int16_t repetitionLevel = 0); | ||
|
||
// ParquetBatch contains the information needed by low-level arrow API writeBatch: | ||
// WriteBatch(int64_t num_values, const int16_t* def_levels, | ||
// const int16_t* rep_levels, const T* values | ||
struct ParquetColumn { | ||
common::LogicalTypeID logicalTypeID; | ||
std::vector<int16_t> repetitionLevels; | ||
std::vector<int16_t> definitionLevels; | ||
std::vector<uint8_t*> values; | ||
}; | ||
|
||
void addToParquetColumns(uint8_t* value, common::ValueVector* vector, | ||
std::map<std::string, ParquetColumn>& parquetColumns, int currentElementIdx = 0, | ||
int parentElementIdx = 0, int depth = 0, std::string parentStructFieldName = ""); | ||
|
||
void extractList(const common::list_entry_t& list, const common::ValueVector* vector, | ||
std::map<std::string, ParquetColumn>& parquetColumns, int currentElementIdx = 0, | ||
int parentElementIdx = 0, int depth = 0, std::string parentStructFieldName = ""); | ||
|
||
void extractStruct(const common::struct_entry_t& val, const common::ValueVector* vector, | ||
std::map<std::string, ParquetColumn>& parquetColumns, int currentElementIdx = 0, | ||
int parentElementIdx = 0, int depth = 0, std::string parentStructFieldName = ""); | ||
|
||
void extractNested(uint8_t* value, const common::ValueVector* vector, | ||
std::map<std::string, ParquetColumn>& parquetColumns, int currentElementIdx = 0, | ||
int parentElementIdx = 0, int depth = 0, std::string parentStructFieldName = ""); | ||
|
||
inline void initNewColumn() { isListStarting = true; } | ||
|
||
void writeLittleEndianUint32(uint8_t* buffer, uint32_t value); | ||
|
||
// Extract dremel encoding levels | ||
int getRepetitionLevel(int currentElementIdx, int parentElementIdx, int depth); | ||
|
||
std::vector<int> maxDefinitionLevels; | ||
|
||
// Properties for nested lists and structs | ||
bool isListStarting; | ||
|
||
int currentColumn = 0; | ||
int currentParquetColumn = 0; | ||
int totalColumns = 0; | ||
|
||
// define the writers | ||
parquet::RowGroupWriter* rowGroupWriter; | ||
parquet::ColumnWriter* columnWriter; | ||
}; | ||
|
||
} // namespace processor | ||
} // namespace kuzu |
58 changes: 58 additions & 0 deletions
58
src/include/processor/operator/persistent/parquet_file_writer.h
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,58 @@ | ||
#pragma once | ||
|
||
#include "arrow/io/file.h" | ||
#include "processor/operator/persistent/file_writer.h" | ||
#include "processor/operator/persistent/parquet_column_writer.h" | ||
#include <parquet/api/writer.h> | ||
|
||
namespace kuzu { | ||
namespace processor { | ||
|
||
// ParquetFileWriter performs the following: | ||
// openFile: opens the file and create an arrow::io::FileOutputStream object | ||
// init: | ||
// - generate the schema | ||
// - calculate the max definition levels and number of primitive nodes. The | ||
// definition levels are used for nested types. | ||
// - initialize parquetColumnWriter | ||
// writeValues : take a vector of ValueVector and pass to parquetColumnWriter | ||
// More information about the encoding: | ||
// https://parquet.apache.org/docs/file-format/data-pages/encodings | ||
class ParquetFileWriter : public FileWriter { | ||
public: | ||
using FileWriter::FileWriter; | ||
void openFile() final; | ||
void init() final; | ||
void closeFile() final; | ||
void writeValues(std::vector<common::ValueVector*>& outputVectors) final; | ||
|
||
private: | ||
static std::shared_ptr<parquet::schema::Node> createParquetSchemaNode(int& parquetColumnsCount, | ||
std::string& columnName, const common::LogicalType& logicalType, | ||
parquet::Repetition::type repetition = parquet::Repetition::REQUIRED, int length = -1); | ||
|
||
static std::shared_ptr<parquet::schema::Node> createNestedNode( | ||
int& parquetColumnsCount, std::string& columnName, const common::LogicalType& logicalType); | ||
|
||
static std::shared_ptr<parquet::schema::Node> createPrimitiveNode(int& parquetColumnsCount, | ||
std::string& columnName, const common::LogicalType& logicalType, | ||
parquet::Repetition::type repetition, int length); | ||
|
||
void writeValue(common::LogicalTypeID type, void* value); | ||
|
||
void flush(); | ||
|
||
void generateSchema( | ||
std::shared_ptr<parquet::schema::GroupNode>& schema, int& parquetColumnsCount); | ||
|
||
std::shared_ptr<ParquetColumnWriter> parquetColumnWriter; | ||
|
||
parquet::RowGroupWriter* rowGroupWriter; | ||
|
||
std::shared_ptr<parquet::ParquetFileWriter> fileWriter; | ||
|
||
std::shared_ptr<arrow::io::FileOutputStream> outFile; | ||
}; | ||
|
||
} // namespace processor | ||
} // namespace kuzu |
Oops, something went wrong.