Feature Storage complete with tests, almost ready for V0.2

Author: KrishSingaria

.gitignore

@@ -5,8 +5,10 @@ benchmark/dataset
 
 *.exe
 *.gl
+*.gd
 *.zip
 *.npz
 benchmark/*.csv
 
-*_pycache__/
+*_pycache__/
+.pytest_cache/

graphzero/__init__.py

@@ -1,4 +1,4 @@
-from .graphzero import Graph, convert_csv_to_gl
+from .graphzero import Graph, convert_csv_to_gl, convert_csv_to_gd, DataType, FeatureStore
 
 # Metadata
 __version__ = "0.1.2"

src/FeatureStore.hpp

@@ -0,0 +1,80 @@
+#ifndef FEATURE_STORE_H
+#define FEATURE_STORE_H
+#include <cstdint>
+#include <stdexcept>
+#include "MemoryMap.hpp"
+#include <span>
+
+enum class DataType : int32_t { 
+    INT32   = 0, 
+    INT64   = 1, 
+    FLOAT32 = 2, 
+    FLOAT64 = 3
+};
+
+#pragma pack(push,1) // to complier to use 1-byte alignment(no padding)
+struct FeatureHeader
+{
+    char magic[8];          // 8bytes "GZDATA26"
+    uint32_t flags;         // 4bytes
+    DataType dtype;         // 4bytes
+    uint64_t num_nodes;     // 8bytes
+    uint64_t feature_dim;   // 8bytes
+}; // 32 bytes aligned 
+#pragma pack(pop) // restore default compiler alignment
+
+class FeatureStore {
+private:
+    MemoryMap* fileMap;
+    struct FeatureHeader header;
+    char* data_ptr;
+public:
+    int64_t num_nodes;
+    int64_t feature_dim;
+    std::string filename;
+
+    FeatureStore(const char* filename);
+    ~FeatureStore();
+
+    DataType get_dtype() const;
+    char* get_data_ptr() const; // 1byte pointer
+
+    template <typename T>
+    std::span<T> get_data(int64_t nodeId);
+};
+
+inline FeatureStore::FeatureStore(const char* filename){
+    this->filename = std::string(filename);
+    this->fileMap = new MemoryMap(filename);
+    this->header = reinterpret_cast<FeatureHeader*>(this->fileMap->get_data())[0];
+    if (std::string_view(header.magic, 8) != "GZDATA26") {
+        throw std::runtime_error("Corrupted or invalid .gd file! Magic bytes mismatch.");
+    }
+    this->num_nodes = this->header.num_nodes;
+    this->feature_dim = this->header.feature_dim;
+    this->data_ptr = reinterpret_cast<char*>(this->fileMap->get_data()) + sizeof(FeatureHeader);
+}
+
+inline FeatureStore::~FeatureStore(){
+    this->data_ptr = nullptr;
+    delete this->fileMap;
+}
+
+inline DataType FeatureStore::get_dtype() const{
+    return this->header.dtype;
+}
+
+inline char* FeatureStore::get_data_ptr() const {
+    return this->data_ptr;
+}
+
+template <typename T>
+inline std::span<T> FeatureStore::get_data(int64_t nodeId){
+    if(nodeId >= num_nodes){
+        throw std::runtime_error("NodeId is greater than number of nodes.");
+    }
+    int64_t offset = nodeId * feature_dim; // start of data of nodeId
+    T* p = reinterpret_cast<T*>(this->data_ptr) + offset;
+    return std::span<T>(p,this->feature_dim);
+}
+#endif

src/bindings.cpp

@@ -5,6 +5,8 @@
 #include <nanobind/stl/string.h>
 #include "Graphzero.hpp"
 #include "csrFilegen.hpp"
+#include "featureFilegen.hpp"
+#include "FeatureStore.hpp"
 #include <vector>
 namespace nb = nanobind;
 
@@ -23,9 +25,9 @@ It holds the mmap / zero-copy memory.
 )doc",
             nb::arg("filename")
         ) 
-        .def_rw("num_nodes",&Graphzero::num_nodes)
-        .def_rw("num_edges",&Graphzero::num_edges)
-        .def_rw("has_weights",&Graphzero::has_weights)
+        .def_ro("num_nodes",&Graphzero::num_nodes)
+        .def_ro("num_edges",&Graphzero::num_edges)
+        .def_ro("has_weights",&Graphzero::has_weights)
 
 
         .def("get_degree", [](Graphzero &self, int64_t node_id) {
@@ -215,9 +217,159 @@ R"doc(Performs uniform random neighbour sampling for a node.
 
     // convert csv to gl 
     m.def("convert_csv_to_gl", &convert_csv,
-        "Convert a CSV edge list to GraphZero binary format (.gl)",
+R"doc(Convert a CSV edge list to GraphZero binary format (.gl))doc",
         nb::arg("csv_path"), 
         nb::arg("out_path"), 
         nb::arg("directed") = false,
         nb::call_guard<nb::gil_scoped_release>());
+        
+    // Feature store class
+    
+    // Bind the DataType Enum so Python can pass it to the converter
+    nb::enum_<DataType>(m, "DataType")
+        .value("INT32", DataType::INT32)
+        .value("INT64", DataType::INT64)
+        .value("FLOAT32", DataType::FLOAT32)
+        .value("FLOAT64", DataType::FLOAT64)
+        .export_values();
+
+    nb::class_<FeatureStore>(m, "FeatureStore")
+        .def(nb::init<const char*>(), // constructor
+R"doc(Data Class contains the Datafile and its relevant functions and methods.
+It holds the mmap / zero-copy memory.
+Args:
+    filename (str): either absolute path or relative path (depends on the current working directory).
+Returns:
+    FeatureStorage instance.
+)doc",
+            nb::arg("filename")
+        ) 
+        .def_ro("num_nodes",&FeatureStore::num_nodes)
+        .def_ro("feature_dim",&FeatureStore::feature_dim)
+        .def("get_data",[](FeatureStore &self,int64_t nodeId) -> nb::object { 
+            switch(self.get_dtype()){
+                case DataType::INT32: {                    
+                    auto data = self.get_data<int32_t>(nodeId);
+                    auto arr = nb::ndarray<nb::numpy, int32_t, nb::shape<1>>(
+                        const_cast<int32_t*>(data.data()), // pointer to data
+                        { data.size() }                    // shape
+                    );
+                    return nb::cast(arr);
+                }
+                case DataType::INT64: {
+                    
+                    auto data = self.get_data<int64_t>(nodeId);
+                    auto arr = nb::ndarray<nb::numpy, int64_t, nb::shape<1>>(
+                        const_cast<int64_t*>(data.data()), // pointer to data
+                        { data.size() }                    // shape
+                    );
+                    return nb::cast(arr);
+                }
+                case DataType::FLOAT32: {  
+                    auto data = self.get_data<float>(nodeId);
+                    auto arr = nb::ndarray<nb::numpy, float, nb::shape<1>>(
+                        const_cast<float*>(data.data()),    // pointer to data
+                        { data.size() }                    // shape
+                    );
+                    return nb::cast(arr);
+                }
+                case DataType::FLOAT64: {
+                    auto data = self.get_data<double>(nodeId);
+                    auto arr = nb::ndarray<nb::numpy, double, nb::shape<1>>(
+                        const_cast<double*>(data.data()), // pointer to data
+                        { data.size() }                    // shape
+                    );
+                    return nb::cast(arr);
+                }
+                default: throw std::runtime_error("Unsupported data type");
+            }
+
+            // Return a zero-copy view into the underlying data buffer and keep
+        },
+            nb::keep_alive<0,1>(),
+R"doc(Returns the data of a node.
+Args:
+    node_id (int)
+Returns:
+    1-D numpy ndarray of data for nodeId.
+)doc",
+            nb::arg("node_id")
+        )
+        .def("get_tensor",[](FeatureStore &self) -> nb::object {
+            char* data = self.get_data_ptr();
+            size_t n = self.num_nodes;
+            size_t f = self.feature_dim;
+
+            switch(self.get_dtype()){
+                case DataType::INT32: {
+                    int32_t* ptr = reinterpret_cast<int32_t*>(data);
+
+                    auto arr = nb::ndarray<nb::numpy, int32_t, nb::shape<2>>(
+                        ptr, // pointer to data
+                        {n,f } // shape
+                    );
+                    return nb::cast(arr);
+                }
+                case DataType::INT64: {
+                    int64_t* ptr = reinterpret_cast<int64_t*>(data);
+
+                    auto arr = nb::ndarray<nb::numpy, int64_t, nb::shape<2>>(
+                        ptr, // pointer to data
+                        {n,f } // shape
+                    );
+                    return nb::cast(arr);
+                }
+                case DataType::FLOAT32: {
+                    float* ptr = reinterpret_cast<float*>(data);
+
+                    auto arr = nb::ndarray<nb::numpy, float, nb::shape<2>>(
+                        ptr, // pointer to data
+                        {n,f } // shape
+                    );
+                    return nb::cast(arr);
+                }
+                case DataType::FLOAT64: {
+                    double* ptr = reinterpret_cast<double*>(data);
+
+                    auto arr = nb::ndarray<nb::numpy, double, nb::shape<2>>(
+                        ptr, // pointer to data
+                        {n,f } // shape
+                    );
+                    return nb::cast(arr);
+                }
+                default: throw std::runtime_error("Unsupported data type");
+            }
+        },
+            nb::keep_alive<0,1>(),
+R"doc(Returns the entire Data (tensor).
+Returns:
+    NxF data of given dtype.
+)doc"
+        )
+        
+        // serialization (Pack)
+        .def("__getstate__", [](const FeatureStore &d){
+
+            return nb::make_tuple(d.filename); // only filename required to rebuild the object
+        })
+        // deserialization (unpack)
+        .def("__setstate__",[](nb::tuple &t){
+            
+            if (t.size() != 1) 
+                throw std::runtime_error("Invalid state!");
+            
+            std::string filename = nb::cast<std::string>(t[0]);
+
+            // create new c++ object using the filename
+            return new FeatureStore(filename.c_str());
+        })
+        ;
+        
+    //feature store data
+    m.def("convert_csv_to_gd", &convert_csv_to_binary,
+R"doc(Convert a CSV data to GraphZero data format (.gd).)doc",
+        nb::arg("csv_path"), 
+        nb::arg("out_path"), 
+        nb::arg("dtype"),
+        nb::call_guard<nb::gil_scoped_release>());
 }