-
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 #1193 from kuzudb/arrow
Export query result (fixed sized values only) to arrow
- Loading branch information
Showing
38 changed files
with
734 additions
and
148 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
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,3 +1,4 @@ | ||
add_subdirectory(arrow) | ||
add_subdirectory(csv_reader) | ||
add_subdirectory(data_chunk) | ||
add_subdirectory(task_system) | ||
|
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,8 @@ | ||
add_library(kuzu_common_arrow | ||
OBJECT | ||
arrow_row_batch.cpp | ||
arrow_converter.cpp) | ||
|
||
set(ALL_OBJECT_FILES | ||
${ALL_OBJECT_FILES} $<TARGET_OBJECTS:kuzu_common_arrow> | ||
PARENT_SCOPE) |
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,100 @@ | ||
#include "common/arrow/arrow_converter.h" | ||
|
||
#include "common/arrow/arrow_row_batch.h" | ||
|
||
namespace kuzu { | ||
namespace common { | ||
|
||
struct ArrowSchemaHolder { | ||
vector<ArrowSchema> children; | ||
vector<ArrowSchema*> childrenPtrs; | ||
vector<unique_ptr<char[]>> ownedTypeNames; | ||
}; | ||
|
||
static void releaseArrowSchema(ArrowSchema* schema) { | ||
if (!schema || schema->release) { | ||
return; | ||
} | ||
schema->release = nullptr; | ||
auto holder = static_cast<ArrowSchemaHolder*>(schema->private_data); | ||
delete holder; | ||
} | ||
|
||
static void initializeChild(ArrowSchema& child, const string& name = "") { | ||
//! Child is cleaned up by parent | ||
child.private_data = nullptr; | ||
child.release = releaseArrowSchema; | ||
|
||
//! Store the child schema | ||
child.flags = ARROW_FLAG_NULLABLE; | ||
child.name = name.c_str(); | ||
child.n_children = 0; | ||
child.children = nullptr; | ||
child.metadata = nullptr; | ||
child.dictionary = nullptr; | ||
} | ||
|
||
static void setArrowFormat(ArrowSchema& child, const DataType& type) { | ||
switch (type.typeID) { | ||
case DataTypeID::BOOL: { | ||
child.format = "b"; | ||
} break; | ||
case DataTypeID::INT64: { | ||
child.format = "l"; | ||
} break; | ||
case DataTypeID::DOUBLE: { | ||
child.format = "g"; | ||
} break; | ||
case DataTypeID::DATE: { | ||
child.format = "tdD"; | ||
} break; | ||
case DataTypeID::TIMESTAMP: { | ||
child.format = "tsu:"; | ||
} break; | ||
case DataTypeID::INTERVAL: { | ||
child.format = "tDm"; | ||
} break; | ||
default: | ||
throw InternalException("Unsupported Arrow type " + Types::dataTypeToString(type)); | ||
} | ||
} | ||
|
||
void ArrowConverter::toArrowSchema( | ||
ArrowSchema* outSchema, std::vector<DataType>& types, std::vector<std::string>& names) { | ||
assert(outSchema && types.size() == names.size()); | ||
auto rootHolder = make_unique<ArrowSchemaHolder>(); | ||
|
||
auto columnCount = (int64_t)names.size(); | ||
rootHolder->children.resize(columnCount); | ||
rootHolder->childrenPtrs.resize(columnCount); | ||
for (auto i = 0u; i < columnCount; i++) { | ||
rootHolder->childrenPtrs[i] = &rootHolder->children[i]; | ||
} | ||
outSchema->children = rootHolder->childrenPtrs.data(); | ||
outSchema->n_children = columnCount; | ||
|
||
outSchema->format = "+s"; // struct apparently | ||
outSchema->flags = 0; | ||
outSchema->metadata = nullptr; | ||
outSchema->name = "kuzu_query_result"; | ||
outSchema->dictionary = nullptr; | ||
|
||
for (auto i = 0u; i < columnCount; i++) { | ||
auto& child = rootHolder->children[i]; | ||
initializeChild(child, names[i]); | ||
setArrowFormat(child, types[i]); | ||
} | ||
|
||
outSchema->private_data = rootHolder.release(); | ||
outSchema->release = releaseArrowSchema; | ||
} | ||
|
||
void ArrowConverter::toArrowArray( | ||
main::QueryResult& queryResult, ArrowArray* outArray, std::int64_t chunkSize) { | ||
auto types = queryResult.getColumnDataTypes(); | ||
auto rowBatch = make_unique<ArrowRowBatch>(types, chunkSize); | ||
*outArray = rowBatch->append(queryResult, chunkSize); | ||
} | ||
|
||
} // namespace common | ||
} // 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
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,147 @@ | ||
#include "common/arrow/arrow_row_batch.h" | ||
|
||
ArrowRowBatch::ArrowRowBatch(std::vector<DataType> types, std::int64_t capacity) | ||
: types{std::move(types)} { | ||
auto numVectors = this->types.size(); | ||
vectors.resize(numVectors); | ||
for (auto i = 0u; i < numVectors; i++) { | ||
vectors[i] = make_unique<ArrowVector>(); | ||
initializeArrowVector(vectors[i].get(), this->types[i], capacity); | ||
} | ||
} | ||
|
||
void ArrowRowBatch::initializeArrowVector( | ||
ArrowVector* vector, const DataType& type, std::int64_t capacity) { | ||
auto numBytesForValidity = getNumBytesForBits(capacity); | ||
vector->validity.resize(numBytesForValidity, 0xFF); | ||
std::int64_t numBytesForData = 0; | ||
if (type.typeID == BOOL) { | ||
numBytesForData = getNumBytesForBits(capacity); | ||
} else { | ||
numBytesForData = Types::getDataTypeSize(type) * capacity; | ||
} | ||
vector->data.reserve(numBytesForData); | ||
} | ||
|
||
static void getBitPosition(std::int64_t pos, std::int64_t& bytePos, std::int64_t& bitOffset) { | ||
bytePos = pos / 8; | ||
bitOffset = pos % 8; | ||
} | ||
|
||
static void setBitToZero(std::uint8_t* data, std::int64_t pos) { | ||
std::int64_t bytePos, bitOffset; | ||
getBitPosition(pos, bytePos, bitOffset); | ||
data[bytePos] &= ~((std::uint64_t)1 << bitOffset); | ||
} | ||
|
||
static void setBitToOne(std::uint8_t* data, std::int64_t pos) { | ||
std::int64_t bytePos, bitOffset; | ||
getBitPosition(pos, bytePos, bitOffset); | ||
data[bytePos] |= ((std::uint64_t)1 << bitOffset); | ||
} | ||
|
||
void ArrowRowBatch::setValue(ArrowVector* vector, Value* value, std::int64_t pos) { | ||
if (value->isNull_) { | ||
setBitToZero(vector->validity.data(), pos); | ||
vector->numNulls++; | ||
return; | ||
} | ||
switch (value->dataType.typeID) { | ||
case BOOL: { | ||
if (value->val.booleanVal) { | ||
setBitToOne(vector->data.data(), pos); | ||
} else { | ||
setBitToZero(vector->data.data(), pos); | ||
} | ||
} break; | ||
case INT64: { | ||
((std::int64_t*)vector->data.data())[pos] = value->val.int64Val; | ||
} break; | ||
case DOUBLE: { | ||
((std::double_t*)vector->data.data())[pos] = value->val.doubleVal; | ||
} break; | ||
case DATE: { | ||
((date_t*)vector->data.data())[pos] = value->val.dateVal; | ||
} break; | ||
case TIMESTAMP: { | ||
((timestamp_t*)vector->data.data())[pos] = value->val.timestampVal; | ||
} break; | ||
case INTERVAL: { | ||
((interval_t*)vector->data.data())[pos] = value->val.intervalVal; | ||
} break; | ||
default: { | ||
throw RuntimeException("Data type " + Types::dataTypeToString(value->dataType) + | ||
" is not supported for arrow exportation."); | ||
} | ||
} | ||
} | ||
|
||
static void releaseArrowVector(ArrowArray* array) { | ||
if (!array || !array->release) { | ||
return; | ||
} | ||
array->release = nullptr; | ||
auto holder = static_cast<ArrowVector*>(array->private_data); | ||
delete holder; | ||
} | ||
|
||
ArrowArray* ArrowRowBatch::finalizeArrowChild(const DataType& type, ArrowVector& vector) { | ||
auto result = make_unique<ArrowArray>(); | ||
|
||
result->private_data = nullptr; | ||
result->release = releaseArrowVector; | ||
result->n_children = 0; | ||
result->offset = 0; | ||
result->dictionary = nullptr; | ||
result->buffers = vector.buffers.data(); | ||
result->null_count = vector.numNulls; | ||
result->length = vector.numValues; | ||
result->buffers[0] = vector.validity.data(); | ||
result->buffers[1] = vector.data.data(); | ||
result->n_buffers = 2; | ||
|
||
vector.array = std::move(result); | ||
return vector.array.get(); | ||
} | ||
|
||
ArrowArray ArrowRowBatch::finalize() { | ||
auto rootHolder = make_unique<ArrowVector>(); | ||
ArrowArray result; | ||
rootHolder->childPointers.resize(types.size()); | ||
result.children = rootHolder->childPointers.data(); | ||
result.n_children = (std::int64_t)types.size(); | ||
result.length = numRows; | ||
result.n_buffers = 1; | ||
result.buffers = rootHolder->buffers.data(); // no actual buffer | ||
result.offset = 0; | ||
result.null_count = 0; | ||
result.dictionary = nullptr; | ||
rootHolder->childData = std::move(vectors); | ||
for (auto i = 0u; i < rootHolder->childData.size(); i++) { | ||
rootHolder->childPointers[i] = finalizeArrowChild(types[i], *rootHolder->childData[i]); | ||
} | ||
|
||
result.private_data = rootHolder.release(); | ||
result.release = releaseArrowVector; | ||
return result; | ||
} | ||
|
||
ArrowArray ArrowRowBatch::append(main::QueryResult& queryResult, std::int64_t chunkSize) { | ||
std::int64_t numTuples = 0; | ||
auto numColumns = queryResult.getColumnNames().size(); | ||
while (numTuples < chunkSize) { | ||
if (!queryResult.hasNext()) { | ||
break; | ||
} | ||
auto tuple = queryResult.getNext(); | ||
for (auto i = 0u; i < numColumns; i++) { | ||
setValue(vectors[i].get(), tuple->getValue(i), numTuples); | ||
} | ||
numTuples++; | ||
} | ||
for (auto i = 0u; i < numColumns; i++) { | ||
vectors[i]->numValues = numTuples; | ||
} | ||
numRows = numTuples; | ||
return finalize(); | ||
} |
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,56 @@ | ||
#pragma once | ||
|
||
// The Arrow C data interface. | ||
// https://arrow.apache.org/docs/format/CDataInterface.html | ||
|
||
#include <stdint.h> | ||
|
||
#ifdef __cplusplus | ||
extern "C" { | ||
#endif | ||
|
||
#ifndef ARROW_C_DATA_INTERFACE | ||
#define ARROW_C_DATA_INTERFACE | ||
|
||
#define ARROW_FLAG_DICTIONARY_ORDERED 1 | ||
#define ARROW_FLAG_NULLABLE 2 | ||
#define ARROW_FLAG_MAP_KEYS_SORTED 4 | ||
|
||
struct ArrowSchema { | ||
// Array type description | ||
const char* format; | ||
const char* name; | ||
const char* metadata; | ||
int64_t flags; | ||
int64_t n_children; | ||
struct ArrowSchema** children; | ||
struct ArrowSchema* dictionary; | ||
|
||
// Release callback | ||
void (*release)(struct ArrowSchema*); | ||
// Opaque producer-specific data | ||
void* private_data; | ||
}; | ||
|
||
struct ArrowArray { | ||
// Array data description | ||
int64_t length; | ||
int64_t null_count; | ||
int64_t offset; | ||
int64_t n_buffers; | ||
int64_t n_children; | ||
const void** buffers; | ||
struct ArrowArray** children; | ||
struct ArrowArray* dictionary; | ||
|
||
// Release callback | ||
void (*release)(struct ArrowArray*); | ||
// Opaque producer-specific data | ||
void* private_data; | ||
}; | ||
|
||
#endif // ARROW_C_DATA_INTERFACE | ||
|
||
#ifdef __cplusplus | ||
} | ||
#endif |
Oops, something went wrong.