-
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
1 parent
abca5e6
commit 00ad3ae
Showing
36 changed files
with
2,404 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
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
98 changes: 98 additions & 0 deletions
98
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,98 @@ | ||
#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 { | ||
// static constexpr uint64_t BYTES_PER_ROW = 1024; | ||
static constexpr uint64_t rowGroupSize = common::StorageConstants::NODE_GROUP_SIZE; | ||
// static constexpr uint64_t rowGroupSizeInBytes = BYTES_PER_ROW * rowGroupSize; | ||
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 { | ||
inline void init(CopyToParquetInfo* info, storage::MemoryManager* mm, ResultSet* resultSet) { | ||
ft = std::make_unique<FactorizedTable>(mm, std::move(info->tableSchema)); | ||
vectorsToAppend.reserve(info->dataPoses.size()); | ||
for (auto& pos : info->dataPoses) { | ||
vectorsToAppend.push_back(resultSet->getValueVector(pos).get()); | ||
} | ||
this->mm = mm; | ||
} | ||
|
||
inline void append() { ft->append(vectorsToAppend); } | ||
|
||
std::unique_ptr<FactorizedTable> ft; | ||
std::vector<common::ValueVector*> vectorsToAppend; | ||
storage::MemoryManager* mm; | ||
}; | ||
|
||
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)} {} | ||
|
||
void initLocalStateInternal(ResultSet* resultSet, ExecutionContext* context) final; | ||
|
||
void executeInternal(ExecutionContext* context) final; | ||
|
||
inline void finalize(ExecutionContext* executionContext) final { | ||
sharedState->writer->finalize(); | ||
} | ||
|
||
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 |
Oops, something went wrong.