-
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
1,035 additions
and
26 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 |
---|---|---|
@@ -0,0 +1,9 @@ | ||
cmake_minimum_required(VERSION 3.11) | ||
|
||
project(examples VERSION 0.0.3.1 LANGUAGES CXX) | ||
|
||
include_directories(../src/include) | ||
link_directories(../build/release/src) | ||
|
||
add_executable(example main.cpp) | ||
target_link_libraries(example PRIVATE 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,86 @@ | ||
#include "main/database.h" | ||
#include "main/connection.h" | ||
#include "main/prepared_statement.h" | ||
#include "binder/bound_statement_result.h" | ||
#include "planner/logical_plan/logical_plan.h" | ||
|
||
#include <iostream> | ||
|
||
using namespace kuzu::main; | ||
using namespace kuzu::common; | ||
|
||
void testPyG(Connection* conn, const std::string& srcTableName, | ||
const std::string& relName, const std::string& dstTableName, size_t queryBatchSize) { | ||
// Get the number of nodes in the src table for batching. | ||
auto numNodesQuery = "MATCH (a:{}) RETURN count(*)"; | ||
auto numNodesQueryWithParams = StringUtils::string_format(numNodesQuery, srcTableName); | ||
std::cout << "Execute query: " << numNodesQueryWithParams << std::endl; | ||
auto numNodesResult = conn->query(numNodesQueryWithParams); | ||
std::cout << "Done execute query: " << numNodesQueryWithParams << std::endl; | ||
|
||
if (!numNodesResult->isSuccess()) { | ||
throw std::runtime_error(numNodesResult->getErrorMessage()); | ||
} | ||
auto numNodes = numNodesResult->getNext()->getValue(0)->getValue<int64_t>(); | ||
int64_t batches = numNodes / queryBatchSize; | ||
if (numNodes % queryBatchSize != 0) { | ||
batches += 1; | ||
} | ||
|
||
// Get total number of edges for allocating the buffer. | ||
auto countQuery = "MATCH (a:{})-[:{}]->(b:{}) RETURN count(*)"; | ||
auto countQueryWithParams = | ||
StringUtils::string_format(countQuery, srcTableName, relName, dstTableName); | ||
std::cout << "Execute query: " << countQueryWithParams << std::endl; | ||
auto countResult = conn->query(countQueryWithParams); | ||
std::cout << "Done execute query: " << countQueryWithParams << std::endl; | ||
|
||
if (!countResult->isSuccess()) { | ||
throw std::runtime_error(countResult->getErrorMessage()); | ||
} | ||
uint64_t count = countResult->getNext()->getValue(0)->getValue<int64_t>(); | ||
std::cout << "Allocate buffer." << std::endl; | ||
uint64_t bufferSize = count * 2 * sizeof(int64_t); | ||
auto* buffer = (int64_t*)malloc(bufferSize); | ||
memset(buffer, 0, bufferSize); | ||
std::cout << "Done allocating buffer." << std::endl; | ||
// Run queries in batch to fetch edges. | ||
auto queryString = "MATCH (a:{})-[:{}]->(b:{}) WHERE offset(id(a)) >= $s AND offset(id(a)) < " | ||
"$e RETURN offset(id(a)), offset(id(b))"; | ||
auto query = StringUtils::string_format(queryString, srcTableName, relName, dstTableName); | ||
std::cout << "Prepare query: " << query << std::endl; | ||
auto preparedStatement = conn->prepare(query); | ||
uint64_t i = 0; | ||
for (int64_t batch = 0; batch < batches; ++batch) { | ||
int64_t start = batch * queryBatchSize; | ||
int64_t end = (batch + 1) * queryBatchSize; | ||
end = end > numNodes ? numNodes : end; | ||
std::unordered_map<std::string, std::shared_ptr<Value>> parameters; | ||
parameters["s"] = std::make_shared<Value>(start); | ||
parameters["e"] = std::make_shared<Value>(end); | ||
std::cout << "Execute query with params " << start << ", " << end << std::endl; | ||
auto result = conn->executeWithParams(preparedStatement.get(), parameters); | ||
std::cout << "Done execute query with params " << start << ", " << end << std::endl; | ||
if (!result->isSuccess()) { | ||
throw std::runtime_error(result->getErrorMessage()); | ||
} | ||
while (result->hasNext()) { | ||
auto row = result->getNext(); | ||
buffer[i] = row->getValue(0)->getValue<int64_t>(); | ||
buffer[i + count] = row->getValue(1)->getValue<int64_t>(); | ||
i += 1; | ||
} | ||
} | ||
} | ||
|
||
int main() { | ||
std::string dbPath = "/home/lc/Developer/mag100m-kuzu"; | ||
// std::string dbPath = "/home/kuzu/Developer/mag100m-kuzu"; | ||
kuzu::main::SystemConfig systemConfig; | ||
systemConfig.bufferPoolSize = (std::uint64_t)10 * 1024 * 1024 * 1024; | ||
auto database = std::make_unique<Database>(dbPath, systemConfig); | ||
auto conn = std::make_unique<Connection>(database.get()); | ||
testPyG(conn.get(), "Paper", "Cites", "Paper", 1000000); | ||
std::cout << "Done processing! Now we wait" << std::endl; | ||
std::this_thread::sleep_for(std::chrono::seconds(1000)); | ||
} |
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
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,37 @@ | ||
#pragma once | ||
|
||
#include <queue> | ||
|
||
#include "common/types/types_include.h" | ||
#include "database.h" | ||
|
||
namespace kuzu { | ||
namespace storage { | ||
class Column; | ||
} | ||
|
||
namespace main { | ||
|
||
class StorageDriver { | ||
public: | ||
explicit StorageDriver(Database* database); | ||
|
||
~StorageDriver(); | ||
|
||
void scan(const std::string& nodeName, const std::string& propertyName, | ||
common::offset_t* offsets, size_t size, uint8_t* result, size_t numThreads); | ||
|
||
uint64_t getNumNodes(const std::string& nodeName); | ||
uint64_t getNumRels(const std::string& relName); | ||
|
||
private: | ||
void scanColumn( | ||
storage::Column* column, common::offset_t* offsets, size_t size, uint8_t* result); | ||
|
||
private: | ||
catalog::Catalog* catalog; | ||
storage::StorageManager* storageManager; | ||
}; | ||
|
||
} // namespace main | ||
} // 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
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,64 @@ | ||
#include "main/storage_driver.h" | ||
|
||
#include "storage/storage_manager.h" | ||
|
||
using namespace kuzu::common; | ||
|
||
namespace kuzu { | ||
namespace main { | ||
|
||
StorageDriver::StorageDriver(kuzu::main::Database* database) | ||
: catalog{database->catalog.get()}, storageManager{database->storageManager.get()} {} | ||
|
||
StorageDriver::~StorageDriver() = default; | ||
|
||
void StorageDriver::scan(const std::string& nodeName, const std::string& propertyName, | ||
common::offset_t* offsets, size_t size, uint8_t* result, size_t numThreads) { | ||
// Resolve files to read from | ||
auto catalogContent = catalog->getReadOnlyVersion(); | ||
auto nodeTableID = catalogContent->getTableID(nodeName); | ||
auto propertyID = catalogContent->getTableSchema(nodeTableID)->getPropertyID(propertyName); | ||
auto nodeTable = storageManager->getNodesStore().getNodeTable(nodeTableID); | ||
auto column = nodeTable->getPropertyColumn(propertyID); | ||
|
||
auto current_buffer = result; | ||
std::vector<std::thread> threads; | ||
auto numElementsPerThread = size / numThreads + 1; | ||
auto sizeLeft = size; | ||
while (sizeLeft > 0) { | ||
uint64_t sizeToRead = std::min(numElementsPerThread, sizeLeft); | ||
threads.emplace_back( | ||
&StorageDriver::scanColumn, this, column, offsets, sizeToRead, current_buffer); | ||
offsets += sizeToRead; | ||
current_buffer += sizeToRead * column->elementSize; | ||
sizeLeft -= sizeToRead; | ||
} | ||
for (auto& thread : threads) { | ||
thread.join(); | ||
} | ||
} | ||
|
||
uint64_t StorageDriver::getNumNodes(const std::string& nodeName) { | ||
auto catalogContent = catalog->getReadOnlyVersion(); | ||
auto nodeTableID = catalogContent->getTableID(nodeName); | ||
auto nodeStatistics = storageManager->getNodesStore() | ||
.getNodesStatisticsAndDeletedIDs() | ||
.getNodeStatisticsAndDeletedIDs(nodeTableID); | ||
return nodeStatistics->getNumTuples(); | ||
} | ||
|
||
uint64_t StorageDriver::getNumRels(const std::string& relName) { | ||
auto catalogContent = catalog->getReadOnlyVersion(); | ||
auto relTableID = catalogContent->getTableID(relName); | ||
auto relStatistics = | ||
storageManager->getRelsStore().getRelsStatistics().getRelStatistics(relTableID); | ||
return relStatistics->getNumTuples(); | ||
} | ||
|
||
void StorageDriver::scanColumn( | ||
storage::Column* column, common::offset_t* offsets, size_t size, uint8_t* result) { | ||
column->scan(offsets, size, result); | ||
} | ||
|
||
} // namespace main | ||
} // 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
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,28 @@ | ||
#include "main/storage_driver.h" | ||
#include "main_test_helper/main_test_helper.h" | ||
|
||
using namespace kuzu::testing; | ||
using namespace kuzu::common; | ||
|
||
TEST_F(ApiTest, StorageDriverScan) { | ||
auto storageDriver = std::make_unique<StorageDriver>(database.get()); | ||
auto size = 6; | ||
auto nodeOffsetsBuffer = std::make_unique<uint8_t[]>(sizeof(offset_t) * size); | ||
auto nodeOffsets = (offset_t*)nodeOffsetsBuffer.get(); | ||
nodeOffsets[0] = 7; | ||
nodeOffsets[1] = 0; | ||
nodeOffsets[2] = 3; | ||
nodeOffsets[3] = 1; | ||
nodeOffsets[4] = 2; | ||
nodeOffsets[5] = 6; | ||
auto result = std::make_unique<uint8_t[]>(sizeof(int64_t) * size); | ||
auto resultBuffer = (uint8_t*)result.get(); | ||
storageDriver->scan("person", "ID", nodeOffsets, size, resultBuffer, 3); | ||
auto ids = (int64_t*)resultBuffer; | ||
ASSERT_EQ(ids[0], 10); | ||
ASSERT_EQ(ids[1], 0); | ||
ASSERT_EQ(ids[2], 5); | ||
ASSERT_EQ(ids[3], 2); | ||
ASSERT_EQ(ids[4], 3); | ||
ASSERT_EQ(ids[5], 9); | ||
} |
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
Oops, something went wrong.