Add nquad file type, add scan multiple rdf files

kuzudb · Dec 1, 2023 · 5da59c2 · 5da59c2
1 parent 3ab6bf1
commit 5da59c2
Show file tree

Hide file tree

Showing 17 changed files with 1,257 additions and 70 deletions.
diff --git a/dataset/rdf/spb1k/copy.cypher b/dataset/rdf/spb1k/copy.cypher
@@ -0,0 +1,4 @@
+COPY spb_resource_t FROM "dataset/rdf/spb1k/*.nq";
+COPY spb_literal_t FROM glob("dataset/rdf/spb1k/*.nq");
+COPY spb_resource_triples_t FROM glob("dataset/rdf/spb1k/*.nq");
+COPY spb_literal_triples_t FROM glob("dataset/rdf/spb1k/*.nq");
diff --git a/dataset/rdf/spb1k/generatedCreativeWorks-000001.nq b/dataset/rdf/spb1k/generatedCreativeWorks-000001.nq
diff --git a/dataset/rdf/spb1k/generatedCreativeWorks-000002.nq b/dataset/rdf/spb1k/generatedCreativeWorks-000002.nq
diff --git a/dataset/rdf/spb1k/generatedCreativeWorks-000003.nq b/dataset/rdf/spb1k/generatedCreativeWorks-000003.nq
diff --git a/dataset/rdf/spb1k/generatedCreativeWorks-000004.nq b/dataset/rdf/spb1k/generatedCreativeWorks-000004.nq
diff --git a/dataset/rdf/spb1k/schema.cypher b/dataset/rdf/spb1k/schema.cypher
@@ -0,0 +1 @@
+CREATE RDF GRAPH spb;
diff --git a/src/binder/bind/bind_copy.cpp b/src/binder/bind/bind_copy.cpp
@@ -88,15 +88,21 @@ std::unique_ptr<BoundStatement> Binder::bindCopyFromClause(const Statement& stat
     }
     switch (tableSchema->tableType) {
     case TableType::NODE:
-        if (readerConfig->fileType == FileType::TURTLE) {
+        switch (readerConfig->fileType) {
+        case FileType::TURTLE:
+        case FileType::NQUADS: {
             return bindCopyRdfNodeFrom(statement, std::move(readerConfig), tableSchema);
-        } else {
+        }
+        default:
             return bindCopyNodeFrom(statement, std::move(readerConfig), tableSchema);
         }
     case TableType::REL: {
-        if (readerConfig->fileType == FileType::TURTLE) {
+        switch (readerConfig->fileType) {
+        case FileType::TURTLE:
+        case FileType::NQUADS: {
             return bindCopyRdfRelFrom(statement, std::move(readerConfig), tableSchema);
-        } else {
+        }
+        default:
             return bindCopyRelFrom(statement, std::move(readerConfig), tableSchema);
         }
     }

diff --git a/src/binder/binder.cpp b/src/binder/binder.cpp
@@ -217,19 +217,20 @@ function::TableFunction* Binder::getScanFunction(FileType fileType, const Reader
     inputTypes.push_back(&stringType);
     auto functions = catalog.getBuiltInFunctions();
     switch (fileType) {
-    case common::FileType::PARQUET: {
+    case FileType::PARQUET: {
         func = functions->matchScalarFunction(READ_PARQUET_FUNC_NAME, inputTypes);
     } break;
-    case common::FileType::NPY: {
+    case FileType::NPY: {
         func = functions->matchScalarFunction(READ_NPY_FUNC_NAME, inputTypes);
     } break;
-    case common::FileType::CSV: {
+    case FileType::CSV: {
         auto csvConfig = reinterpret_cast<CSVReaderConfig*>(config.extraConfig.get());
         func = functions->matchScalarFunction(
             csvConfig->parallel ? READ_CSV_PARALLEL_FUNC_NAME : READ_CSV_SERIAL_FUNC_NAME,
             inputTypes);
     } break;
-    case common::FileType::TURTLE: {
+    case FileType::NQUADS:
+    case FileType::TURTLE: {
         func = functions->matchScalarFunction(READ_RDF_FUNC_NAME, inputTypes);
     } break;
     default:

diff --git a/src/common/copier_config/reader_config.cpp b/src/common/copier_config/reader_config.cpp
@@ -19,6 +19,9 @@ FileType FileTypeUtils::getFileTypeFromExtension(std::string_view extension) {
     if (extension == ".ttl") {
         return FileType::TURTLE;
     }
+    if (extension == ".nq") {
+        return FileType::NQUADS;
+    }
     throw CopyException(std::string("Unsupported file type ").append(extension));
 }
 
@@ -39,6 +42,9 @@ std::string FileTypeUtils::toString(FileType fileType) {
     case FileType::TURTLE: {
         return "TURTLE";
     }
+    case FileType::NQUADS: {
+        return "NQUADS";
+    }
     default: {
         KU_UNREACHABLE;
     }

diff --git a/src/include/common/copier_config/reader_config.h b/src/include/common/copier_config/reader_config.h
@@ -7,7 +7,14 @@
 namespace kuzu {
 namespace common {
 
-enum class FileType : uint8_t { UNKNOWN = 0, CSV = 1, PARQUET = 2, NPY = 3, TURTLE = 4 };
+enum class FileType : uint8_t {
+    UNKNOWN = 0,
+    CSV = 1,
+    PARQUET = 2,
+    NPY = 3,
+    TURTLE = 4, // Terse triples http://www.w3.org/TR/turtle
+    NQUADS = 5  // Line-based quads http://www.w3.org/TR/n-quads/
+};
 
 struct FileTypeUtils {
     static FileType getFileTypeFromExtension(std::string_view extension);

diff --git a/src/include/processor/operator/persistent/reader/csv/serial_csv_reader.h b/src/include/processor/operator/persistent/reader/csv/serial_csv_reader.h
@@ -24,7 +24,7 @@ class SerialCSVReader final : public BaseCSVReader {
 };
 
 struct SerialCSVScanSharedState final : public function::ScanSharedState {
-    explicit SerialCSVScanSharedState(
+    SerialCSVScanSharedState(
         common::ReaderConfig readerConfig, uint64_t numRows, uint64_t numColumns)
         : ScanSharedState{std::move(readerConfig), numRows}, numColumns{numColumns} {
         initReader();

diff --git a/src/include/processor/operator/persistent/reader/rdf/rdf_reader.h b/src/include/processor/operator/persistent/reader/rdf/rdf_reader.h
@@ -6,39 +6,49 @@
 #include "function/table_functions.h"
 #include "function/table_functions/bind_data.h"
 #include "function/table_functions/bind_input.h"
+#include "function/table_functions/scan_functions.h"
 #include "serd.h"
 
 namespace kuzu {
 namespace processor {
 
-class RDFReader {
+class RdfReader {
 public:
-    RDFReader(std::string filePath, const common::RdfReaderConfig& config);
+    RdfReader(
+        std::string filePath, common::FileType fileType, const common::RdfReaderConfig& config)
+        : filePath{std::move(filePath)}, fileType{fileType}, mode{config.mode}, index{config.index},
+          reader{nullptr}, rowOffset{0}, vectorSize{0}, sVector{nullptr}, pVector{nullptr},
+          oVector{nullptr}, status{SERD_SUCCESS} {}
 
-    ~RDFReader();
+    ~RdfReader();
+
+    inline void initReader() { initReader(prefixHandle, statementHandle); }
+    inline void initCountReader() { initReader(nullptr, countStatementHandle); }
 
     common::offset_t read(common::DataChunk* dataChunkToRead);
     common::offset_t countLine();
 
 private:
     static SerdStatus errorHandle(void* handle, const SerdError* error);
-    static SerdStatus readerStatementSink(void* handle, SerdStatementFlags flags,
-        const SerdNode* graph, const SerdNode* subject, const SerdNode* predicate,
-        const SerdNode* object, const SerdNode* object_datatype, const SerdNode* object_lang);
-    static SerdStatus prefixSink(void* handle, const SerdNode* name, const SerdNode* uri);
+    static SerdStatus statementHandle(void* handle, SerdStatementFlags flags, const SerdNode* graph,
+        const SerdNode* subject, const SerdNode* predicate, const SerdNode* object,
+        const SerdNode* object_datatype, const SerdNode* object_lang);
+    static SerdStatus prefixHandle(void* handle, const SerdNode* name, const SerdNode* uri);
 
-    static SerdStatus counterStatementSink(void* handle, SerdStatementFlags flags,
+    static SerdStatus countStatementHandle(void* handle, SerdStatementFlags flags,
         const SerdNode* graph, const SerdNode* subject, const SerdNode* predicate,
         const SerdNode* object, const SerdNode* object_datatype, const SerdNode* object_lang);
 
+    void initReader(const SerdPrefixSink prefixSink_, const SerdStatementSink statementSink_);
+
 private:
     const std::string filePath;
+    common::FileType fileType;
     common::RdfReaderMode mode;
     storage::PrimaryKeyIndex* index;
 
     FILE* fp;
     SerdReader* reader;
-    SerdReader* counter;
 
     // TODO(Xiyang): use prefix to expand CURIE.
     const char* currentPrefix;
@@ -51,8 +61,17 @@ class RDFReader {
     common::ValueVector* oVector; // object
 };
 
-struct RdfScanLocalState final : public function::TableFuncLocalState {
-    std::unique_ptr<RDFReader> reader;
+struct RdfScanSharedState final : public function::ScanSharedState {
+    std::unique_ptr<RdfReader> reader;
+
+    RdfScanSharedState(common::ReaderConfig readerConfig, uint64_t numRows)
+        : ScanSharedState{readerConfig, numRows} {
+        initReader();
+    }
+
+    void read(common::DataChunk& dataChunk);
+
+    void initReader();
 };
 
 struct RdfScan {

diff --git a/src/planner/plan/plan_copy.cpp b/src/planner/plan/plan_copy.cpp
@@ -36,6 +36,7 @@ static void appendPartitioner(BoundCopyFromInfo* copyFromInfo, LogicalPlan& plan
     payloads.push_back(copyFromInfo->fileScanInfo->offset);
     // TODO(Xiyang): Merge TURTLE case with other data types.
     switch (fileType) {
+    case FileType::NQUADS:
     case FileType::TURTLE: {
         auto extraInfo = reinterpret_cast<ExtraBoundCopyRdfRelInfo*>(copyFromInfo->extraInfo.get());
         infos.push_back(std::make_unique<LogicalPartitionerInfo>(
@@ -83,7 +84,7 @@ std::unique_ptr<LogicalPlan> Planner::planCopyFrom(const BoundStatement& stateme
     QueryPlanner::appendScanFile(copyFromInfo->fileScanInfo.get(), *plan);
     auto tableType = copyFromInfo->tableSchema->tableType;
     if (tableType == TableType::REL) {
-        if (fileType != FileType::TURTLE) {
+        if (fileType != FileType::TURTLE && fileType != FileType::NQUADS) {
             auto extraInfo = dynamic_cast<ExtraBoundCopyRelInfo*>(copyFromInfo->extraInfo.get());
             std::vector<std::unique_ptr<LogicalIndexScanNodeInfo>> infos;
             auto srcNodeTableSchema = dynamic_cast<NodeTableSchema*>(extraInfo->srcTableSchema);

diff --git a/src/processor/map/map_copy_from.cpp b/src/processor/map/map_copy_from.cpp
@@ -130,9 +130,10 @@ std::unique_ptr<PhysicalOperator> PlanMapper::mapCopyNodeFrom(LogicalOperator* l
     auto readerConfig =
         reinterpret_cast<function::ScanBindData*>(copyFromInfo->fileScanInfo->bindData.get())
             ->config;
-    if (readerConfig.fileType == FileType::TURTLE &&
-        reinterpret_cast<RdfReaderConfig*>(readerConfig.extraConfig.get())->mode ==
-            RdfReaderMode::RESOURCE) {
+    bool isRdfFile =
+        readerConfig.fileType == FileType::TURTLE || readerConfig.fileType == FileType::NQUADS;
+    if (isRdfFile && reinterpret_cast<RdfReaderConfig*>(readerConfig.extraConfig.get())->mode ==
+                         RdfReaderMode::RESOURCE) {
         copyNode = std::make_unique<CopyRdfResource>(sharedState, std::move(info),
             std::make_unique<ResultSetDescriptor>(copyFrom->getSchema()), std::move(prevOperator),
             getOperatorID(), copyFrom->getExpressionsForPrinting());