-
Notifications
You must be signed in to change notification settings - Fork 86
Commit
- Loading branch information
There are no files selected for viewing
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 was deleted.
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>(); | ||
Check warning on line 54 in src/include/processor/operator/persistent/writer/parquet/basic_column_writer.h
|
||
} | ||
|
||
// 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"}; | ||
Check warning on line 80 in src/include/processor/operator/persistent/writer/parquet/basic_column_writer.h
|
||
} | ||
void writeDictionary(BasicColumnWriterState& state, | ||
std::unique_ptr<BufferedSerializer> bufferedSerializer, uint64_t rowCount); | ||
virtual void flushDictionary(BasicColumnWriterState& state, ColumnWriterStatistics* stats) { | ||
throw common::NotImplementedException{"BasicColumnWriter::flushDictionary"}; | ||
Check warning on line 85 in src/include/processor/operator/persistent/writer/parquet/basic_column_writer.h
|
||
} | ||
|
||
void setParquetStatistics( | ||
BasicColumnWriterState& state, kuzu_parquet::format::ColumnChunk& column); | ||
void registerToRowGroup(kuzu_parquet::format::RowGroup& rowGroup); | ||
}; | ||
|
||
} // namespace processor | ||
} // namespace kuzu |
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 |