-
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.
- Loading branch information
Showing
29 changed files
with
3,218 additions
and
2,381 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
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1 @@ | ||
COPY T FROM "dataset/copy-test/list-of-struct/test.csv" (HEADER=true); |
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 @@ | ||
create node table T (id INT64, fields STRUCT(a INT64, b INT64)[], PRIMARY KEY (id)); |
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,4 @@ | ||
id,fields | ||
0,"[{a:1,b:2},{a:3,b:4}]" | ||
1,"[{a:5,b:6},{a:7,b:8},{a:9,b:10}]" | ||
2,"[{a:11,b:12}]" |
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
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
31 changes: 31 additions & 0 deletions
31
src/include/processor/operator/copy_to/csv_parquet_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,31 @@ | ||
#pragma once | ||
|
||
#include "common/copier_config/copier_config.h" | ||
#include "processor/result/result_set.h" | ||
|
||
namespace kuzu { | ||
namespace processor { | ||
|
||
class CSVParquetWriter { | ||
public: | ||
virtual ~CSVParquetWriter(){}; | ||
virtual void openFile(const std::string& filePath) = 0; | ||
virtual void init() = 0; | ||
virtual void closeFile() = 0; | ||
virtual void writeValues(std::vector<common::ValueVector*>& outputVectors) = 0; | ||
|
||
inline void setColumns(const std::vector<std::string>& columnNames, | ||
const std::vector<common::LogicalType>& columnTypes) { | ||
this->columnNames = columnNames; | ||
this->columnTypes = columnTypes; | ||
} | ||
inline std::vector<std::string>& getColumnNames() { return columnNames; } | ||
inline std::vector<common::LogicalType>& getColumnTypes() { return columnTypes; } | ||
|
||
private: | ||
std::vector<std::string> columnNames; | ||
std::vector<common::LogicalType> columnTypes; | ||
}; | ||
|
||
} // namespace processor | ||
} // namespace kuzu |
19 changes: 10 additions & 9 deletions
19
...cessor/operator/copy_to/csv_file_writer.h → ...e/processor/operator/copy_to/csv_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
84 changes: 84 additions & 0 deletions
84
src/include/processor/operator/copy_to/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,84 @@ | ||
#pragma once | ||
|
||
#include "processor/operator/copy_to/csv_parquet_writer.h" | ||
#include <parquet/api/writer.h> | ||
|
||
namespace kuzu { | ||
namespace processor { | ||
|
||
// Additional computed information for each column | ||
// that will be used to write values | ||
// Each element in the vector is related to a column | ||
// TODO (Rui): as other information is no longer needed, we can remove this struct | ||
struct ParquetColumnDescriptor { | ||
int maxDefinitionLevel = 1; | ||
}; | ||
|
||
// 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(std::unordered_map<int, ParquetColumnDescriptor> columnDescriptors) | ||
: columnDescriptors(columnDescriptors){}; | ||
|
||
void writeColumn( | ||
int kuzuColumn, common::ValueVector* vector, parquet::RowGroupWriter* rowWriter); | ||
|
||
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::unordered_map<int, ParquetColumnDescriptor> columnDescriptors; | ||
|
||
// Properties for nested lists and structs | ||
bool isListStarting; | ||
|
||
int currentKuzuColumn; | ||
|
||
// define the writers | ||
parquet::Int64Writer* int64Writer; | ||
parquet::ByteArrayWriter* byteArrayWriter; | ||
parquet::FixedLenByteArrayWriter* fixedLenByteArrayWriter; | ||
parquet::BoolWriter* boolWriter; | ||
parquet::RowGroupWriter* rowWriter; | ||
parquet::Int32Writer* int32Writer; | ||
parquet::DoubleWriter* doubleWriter; | ||
parquet::FloatWriter* floatWriter; | ||
}; | ||
|
||
} // namespace processor | ||
} // namespace kuzu |
Oops, something went wrong.