Implemented progress for in memory RDF scan (#3208)

I have read and agree to the CLA of the Kuzu repository.
kuzudb · Apr 4, 2024 · 37de692 · 37de692
1 parent fa0ef79
commit 37de692
Showing 1 changed file with 41 additions and 4 deletions.
diff --git a/src/processor/operator/persistent/reader/rdf/rdf_scan.cpp b/src/processor/operator/persistent/reader/rdf/rdf_scan.cpp
@@ -203,6 +203,39 @@ static std::unique_ptr<TableFuncSharedState> RdfAllTripleScanInitSharedState(
         bindData->config.copy(), std::move(rdfConfig), bindData->store);
 }
 
+static double RdfResourceInMemScanProgressFunc(TableFuncSharedState* sharedState) {
+    auto rdfSharedState =
+        ku_dynamic_cast<TableFuncSharedState*, RdfInMemScanSharedState*>(sharedState);
+    uint64_t rtSize =
+        ku_dynamic_cast<RdfStore*, TripleStore*>(rdfSharedState->store.get())->rtStore.size();
+    if (rtSize == 0) {
+        return 0.0;
+    }
+    return static_cast<double>(rdfSharedState->rtCursor) / rtSize;
+}
+
+static double RdfLiteralInMemScanProgressFunc(TableFuncSharedState* sharedState) {
+    auto rdfSharedState =
+        ku_dynamic_cast<TableFuncSharedState*, RdfInMemScanSharedState*>(sharedState);
+    uint64_t ltSize =
+        ku_dynamic_cast<RdfStore*, TripleStore*>(rdfSharedState->store.get())->ltStore.size();
+    if (ltSize == 0) {
+        return 0.0;
+    }
+    return static_cast<double>(rdfSharedState->ltCursor) / ltSize;
+}
+
+static double RdfResourceTripleInMemScanProgressFunc(TableFuncSharedState* sharedState) {
+    auto rdfSharedState =
+        ku_dynamic_cast<TableFuncSharedState*, RdfInMemScanSharedState*>(sharedState);
+    TripleStore* store = ku_dynamic_cast<RdfStore*, TripleStore*>(rdfSharedState->store.get());
+    uint64_t size = store->rtStore.size() + store->ltStore.size();
+    if (size == 0) {
+        return 0.0;
+    }
+    return static_cast<double>(rdfSharedState->rtCursor + rdfSharedState->ltCursor) / size;
+}
+
 function_set RdfResourceScan::getFunctionSet() {
     function_set functionSet;
     auto func = std::make_unique<TableFunction>(name, scanTableFunc, nullptr,
@@ -247,31 +280,35 @@ function::function_set RdfAllTripleScan::getFunctionSet() {
 function_set RdfResourceInMemScan::getFunctionSet() {
     function_set functionSet;
     auto func = std::make_unique<TableFunction>(name, RdfResourceInMemScanTableFunc, nullptr,
-        inMemScanInitSharedState, initLocalState, std::vector<LogicalTypeID>{});
+        inMemScanInitSharedState, initLocalState, RdfResourceInMemScanProgressFunc,
+        std::vector<LogicalTypeID>{});
     functionSet.push_back(std::move(func));
     return functionSet;
 }
 
 function_set RdfLiteralInMemScan::getFunctionSet() {
     function_set functionSet;
     auto func = std::make_unique<TableFunction>(name, RdfLiteralInMemScanTableFunc, nullptr,
-        inMemScanInitSharedState, initLocalState, std::vector<LogicalTypeID>{});
+        inMemScanInitSharedState, initLocalState, RdfLiteralInMemScanProgressFunc,
+        std::vector<LogicalTypeID>{});
     functionSet.push_back(std::move(func));
     return functionSet;
 }
 
 function_set RdfResourceTripleInMemScan::getFunctionSet() {
     function_set functionSet;
     auto func = std::make_unique<TableFunction>(name, RdfResourceTripleInMemScanTableFunc, nullptr,
-        inMemScanInitSharedState, initLocalState, std::vector<LogicalTypeID>{});
+        inMemScanInitSharedState, initLocalState, RdfResourceTripleInMemScanProgressFunc,
+        std::vector<LogicalTypeID>{});
     functionSet.push_back(std::move(func));
     return functionSet;
 }
 
 function_set RdfLiteralTripleInMemScan::getFunctionSet() {
     function_set functionSet;
     auto func = std::make_unique<TableFunction>(name, RdfLiteralTripleInMemScanTableFunc, nullptr,
-        inMemScanInitSharedState, initLocalState, std::vector<LogicalTypeID>{});
+        inMemScanInitSharedState, initLocalState, RdfLiteralInMemScanProgressFunc,
+        std::vector<LogicalTypeID>{});
     functionSet.push_back(std::move(func));
     return functionSet;
 }