-
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 #2177 from kuzudb/parquet-writer
Parquet writer
- Loading branch information
Showing
37 changed files
with
2,449 additions
and
417 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
84 changes: 84 additions & 0 deletions
84
src/include/processor/operator/persistent/copy_to_parquet.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 "parquet/parquet_types.h" | ||
#include "processor/operator/persistent/writer/parquet/parquet_writer.h" | ||
#include "processor/operator/sink.h" | ||
#include "processor/result/factorized_table.h" | ||
#include "processor/result/result_set.h" | ||
|
||
namespace kuzu { | ||
namespace processor { | ||
|
||
class CopyToParquetSharedState { | ||
public: | ||
std::unique_ptr<ParquetWriter> writer; | ||
}; | ||
|
||
struct CopyToParquetInfo { | ||
kuzu_parquet::format::CompressionCodec::type codec = | ||
kuzu_parquet::format::CompressionCodec::SNAPPY; | ||
std::unique_ptr<FactorizedTableSchema> tableSchema; | ||
std::vector<std::unique_ptr<common::LogicalType>> types; | ||
std::vector<std::string> names; | ||
std::vector<DataPos> dataPoses; | ||
std::string fileName; | ||
|
||
CopyToParquetInfo(std::unique_ptr<FactorizedTableSchema> tableSchema, | ||
std::vector<std::unique_ptr<common::LogicalType>> types, std::vector<std::string> names, | ||
std::vector<DataPos> dataPoses, std::string fileName) | ||
: tableSchema{std::move(tableSchema)}, types{std::move(types)}, names{std::move(names)}, | ||
dataPoses{std::move(dataPoses)}, fileName{std::move(fileName)} {} | ||
|
||
std::vector<std::unique_ptr<common::LogicalType>> copyTypes(); | ||
|
||
std::unique_ptr<CopyToParquetInfo> copy() { | ||
return std::make_unique<CopyToParquetInfo>( | ||
tableSchema->copy(), copyTypes(), names, dataPoses, fileName); | ||
} | ||
}; | ||
|
||
struct CopyToParquetLocalState { | ||
std::unique_ptr<FactorizedTable> ft; | ||
std::vector<common::ValueVector*> vectorsToAppend; | ||
storage::MemoryManager* mm; | ||
|
||
void init(CopyToParquetInfo* info, storage::MemoryManager* mm, ResultSet* resultSet); | ||
|
||
inline void append() { ft->append(vectorsToAppend); } | ||
}; | ||
|
||
class CopyToParquet : public Sink { | ||
public: | ||
CopyToParquet(std::unique_ptr<ResultSetDescriptor> resultSetDescriptor, | ||
std::unique_ptr<CopyToParquetInfo> info, | ||
std::shared_ptr<CopyToParquetSharedState> sharedState, | ||
std::unique_ptr<PhysicalOperator> child, uint32_t id, const std::string& paramsString) | ||
: Sink{std::move(resultSetDescriptor), PhysicalOperatorType::COPY_TO_PARQUET, | ||
std::move(child), id, paramsString}, | ||
info{std::move(info)}, localState{std::make_unique<CopyToParquetLocalState>()}, | ||
sharedState{std::move(sharedState)} {} | ||
|
||
inline void finalize(ExecutionContext* executionContext) final { | ||
sharedState->writer->finalize(); | ||
} | ||
void initLocalStateInternal(ResultSet* resultSet, ExecutionContext* context) final; | ||
void executeInternal(ExecutionContext* context) final; | ||
std::unique_ptr<PhysicalOperator> clone() final { | ||
return std::make_unique<CopyToParquet>(resultSetDescriptor->copy(), info->copy(), | ||
sharedState, children[0]->clone(), id, paramsString); | ||
} | ||
|
||
private: | ||
void initGlobalStateInternal(ExecutionContext* context) final { | ||
sharedState->writer = std::make_unique<ParquetWriter>( | ||
info->fileName, info->copyTypes(), info->names, info->codec, context->memoryManager); | ||
} | ||
|
||
private: | ||
std::shared_ptr<CopyToParquetSharedState> sharedState; | ||
std::unique_ptr<CopyToParquetLocalState> localState; | ||
std::unique_ptr<CopyToParquetInfo> info; | ||
}; | ||
|
||
} // namespace processor | ||
} // namespace kuzu |
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
58 changes: 0 additions & 58 deletions
58
src/include/processor/operator/persistent/parquet_file_writer.h
This file was deleted.
Oops, something went wrong.
1 change: 1 addition & 0 deletions
1
src/include/processor/operator/persistent/reader/parquet/thrift_tools.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
94 changes: 94 additions & 0 deletions
94
src/include/processor/operator/persistent/writer/parquet/basic_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,94 @@ | ||
#pragma once | ||
|
||
#include "parquet/parquet_types.h" | ||
#include "processor/operator/persistent/writer/parquet/column_writer.h" | ||
|
||
namespace kuzu { | ||
namespace processor { | ||
|
||
class BasicColumnWriterState : public ColumnWriterState { | ||
public: | ||
BasicColumnWriterState(kuzu_parquet::format::RowGroup& rowGroup, uint64_t colIdx) | ||
: rowGroup{rowGroup}, colIdx{colIdx} { | ||
pageInfo.emplace_back(); | ||
} | ||
|
||
kuzu_parquet::format::RowGroup& rowGroup; | ||
uint64_t colIdx; | ||
std::vector<PageInformation> pageInfo; | ||
std::vector<PageWriteInformation> writeInfo; | ||
std::unique_ptr<ColumnWriterStatistics> statsState; | ||
uint64_t currentPage = 0; | ||
}; | ||
|
||
class BasicColumnWriter : public ColumnWriter { | ||
public: | ||
BasicColumnWriter(ParquetWriter& writer, uint64_t schemaIdx, | ||
std::vector<std::string> schemaPath, uint64_t maxRepeat, uint64_t maxDefine, | ||
bool canHaveNulls) | ||
: ColumnWriter( | ||
writer, schemaIdx, std::move(schemaPath), maxRepeat, maxDefine, canHaveNulls) {} | ||
|
||
public: | ||
std::unique_ptr<ColumnWriterState> initializeWriteState( | ||
kuzu_parquet::format::RowGroup& rowGroup) override; | ||
void prepare(ColumnWriterState& state, ColumnWriterState* parent, common::ValueVector* vector, | ||
uint64_t count) override; | ||
void beginWrite(ColumnWriterState& state) override; | ||
void write(ColumnWriterState& state, common::ValueVector* vector, uint64_t count) override; | ||
void finalizeWrite(ColumnWriterState& state) override; | ||
|
||
protected: | ||
void writeLevels(BufferedSerializer& bufferedSerializer, const std::vector<uint16_t>& levels, | ||
uint64_t maxValue, uint64_t startOffset, uint64_t count); | ||
|
||
virtual kuzu_parquet::format::Encoding::type getEncoding(BasicColumnWriterState& state) { | ||
return kuzu_parquet::format::Encoding::PLAIN; | ||
} | ||
|
||
void nextPage(BasicColumnWriterState& state); | ||
void flushPage(BasicColumnWriterState& state); | ||
|
||
// Initializes the state used to track statistics during writing. Only used for scalar types. | ||
virtual std::unique_ptr<ColumnWriterStatistics> initializeStatsState() { | ||
return std::make_unique<ColumnWriterStatistics>(); | ||
} | ||
|
||
// Initialize the writer for a specific page. Only used for scalar types. | ||
virtual std::unique_ptr<ColumnWriterPageState> initializePageState( | ||
BasicColumnWriterState& state) { | ||
return nullptr; | ||
} | ||
|
||
// Flushes the writer for a specific page. Only used for scalar types. | ||
virtual void flushPageState( | ||
BufferedSerializer& bufferedSerializer, ColumnWriterPageState* state) {} | ||
|
||
// Retrieves the row size of a vector at the specified location. Only used for scalar types. | ||
virtual uint64_t getRowSize( | ||
common::ValueVector* vector, uint64_t index, BasicColumnWriterState& state) { | ||
throw common::NotImplementedException{"BasicColumnWriter::getRowSize"}; | ||
} | ||
// Writes a (subset of a) vector to the specified serializer. Only used for scalar types. | ||
virtual void writeVector(BufferedSerializer& bufferedSerializer, ColumnWriterStatistics* stats, | ||
ColumnWriterPageState* pageState, common::ValueVector* vector, uint64_t chunkStart, | ||
uint64_t chunkEnd) = 0; | ||
|
||
virtual bool hasDictionary(BasicColumnWriterState& writerState) { return false; } | ||
// The number of elements in the dictionary. | ||
virtual uint64_t dictionarySize(BasicColumnWriterState& writerState) { | ||
throw common::NotImplementedException{"BasicColumnWriter::dictionarySize"}; | ||
} | ||
void writeDictionary(BasicColumnWriterState& state, | ||
std::unique_ptr<BufferedSerializer> bufferedSerializer, uint64_t rowCount); | ||
virtual void flushDictionary(BasicColumnWriterState& state, ColumnWriterStatistics* stats) { | ||
throw common::NotImplementedException{"BasicColumnWriter::flushDictionary"}; | ||
} | ||
|
||
void setParquetStatistics( | ||
BasicColumnWriterState& state, kuzu_parquet::format::ColumnChunk& column); | ||
void registerToRowGroup(kuzu_parquet::format::RowGroup& rowGroup); | ||
}; | ||
|
||
} // namespace processor | ||
} // namespace kuzu |
67 changes: 67 additions & 0 deletions
67
src/include/processor/operator/persistent/writer/parquet/boolean_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,67 @@ | ||
#pragma once | ||
|
||
#include "processor/operator/persistent/writer/parquet/basic_column_writer.h" | ||
|
||
namespace kuzu { | ||
namespace processor { | ||
|
||
class BooleanStatisticsState : public ColumnWriterStatistics { | ||
public: | ||
BooleanStatisticsState() : min{true}, max{false} {} | ||
|
||
bool min; | ||
bool max; | ||
|
||
public: | ||
bool hasStats() { return !(min && !max); } | ||
|
||
std::string getMin() override { return getMinValue(); } | ||
std::string getMax() override { return getMaxValue(); } | ||
std::string getMinValue() override { | ||
return hasStats() ? std::string(reinterpret_cast<const char*>(&min), sizeof(bool)) : | ||
std::string(); | ||
} | ||
std::string getMaxValue() override { | ||
return hasStats() ? std::string(reinterpret_cast<const char*>(&max), sizeof(bool)) : | ||
std::string(); | ||
} | ||
}; | ||
|
||
class BooleanWriterPageState : public ColumnWriterPageState { | ||
public: | ||
uint8_t byte = 0; | ||
uint8_t bytePos = 0; | ||
}; | ||
|
||
class BooleanColumnWriter : public BasicColumnWriter { | ||
public: | ||
BooleanColumnWriter(ParquetWriter& writer, uint64_t schemaIdx, | ||
std::vector<std::string> schemaPath, uint64_t maxRepeat, uint64_t maxDefine, | ||
bool canHaveNulls) | ||
: BasicColumnWriter( | ||
writer, schemaIdx, std::move(schemaPath), maxRepeat, maxDefine, canHaveNulls) {} | ||
|
||
inline std::unique_ptr<ColumnWriterStatistics> initializeStatsState() override { | ||
return std::make_unique<BooleanStatisticsState>(); | ||
} | ||
|
||
inline uint64_t getRowSize( | ||
common::ValueVector* vector, uint64_t index, BasicColumnWriterState& state) override { | ||
return sizeof(bool); | ||
} | ||
|
||
inline std::unique_ptr<ColumnWriterPageState> initializePageState( | ||
BasicColumnWriterState& state) override { | ||
return std::make_unique<BooleanWriterPageState>(); | ||
} | ||
|
||
void writeVector(BufferedSerializer& bufferedSerializer, | ||
ColumnWriterStatistics* writerStatistics, ColumnWriterPageState* writerPageState, | ||
common::ValueVector* vector, uint64_t chunkStart, uint64_t chunkEnd) override; | ||
|
||
void flushPageState( | ||
BufferedSerializer& temp_writer, ColumnWriterPageState* writerPageState) override; | ||
}; | ||
|
||
} // namespace processor | ||
} // namespace kuzu |
Oops, something went wrong.