Skip to content

Commit

Permalink
impl py secondary index
Browse files Browse the repository at this point in the history
  • Loading branch information
shivam-880 committed Nov 28, 2024
1 parent fdc31d3 commit fbebc01
Show file tree
Hide file tree
Showing 4 changed files with 165 additions and 37 deletions.
27 changes: 21 additions & 6 deletions raphtory/src/python/graph/edge.rs
Original file line number Diff line number Diff line change
Expand Up @@ -385,18 +385,33 @@ impl PyMutableEdge {
/// This function allows for the addition of property updates to an edge within the graph. The updates are time-stamped, meaning they are applied at the specified time.
///
/// Parameters:
/// t (TimeInput): The timestamp at which the updates should be applied.
/// properties (PropInput, optional): A dictionary of properties to update.
/// layer (str, optional): The layer you want these properties to be added on to.
#[pyo3(signature = (t, properties=None, layer=None))]
/// t (TimeInput): The timestamp at which the updates should be applied.
/// properties (PropInput, optional): A dictionary of properties to update.
/// layer (str, optional): The layer you want these properties to be added on to.
/// secondary_index (int, optional): The optional integer which will be used as a secondary index
///
/// Returns:
/// None: This function does not return a value, if the operation is successful.
///
/// Raises:
/// GraphError: If the operation fails.
#[pyo3(signature = (t, properties=None, layer=None, secondary_index=None))]
fn add_updates(
&self,
t: PyTime,
properties: Option<HashMap<String, Prop>>,
layer: Option<&str>,
secondary_index: Option<usize>,
) -> Result<(), GraphError> {
self.edge
.add_updates(t, properties.unwrap_or_default(), layer)
match secondary_index {
None => self
.edge
.add_updates(t, properties.unwrap_or_default(), layer),
Some(secondary_index) => {
self.edge
.add_updates((t, secondary_index), properties.unwrap_or_default(), layer)
}
}
}

/// Mark the edge as deleted at the specified time.
Expand Down
77 changes: 63 additions & 14 deletions raphtory/src/python/graph/graph.rs
Original file line number Diff line number Diff line change
Expand Up @@ -140,7 +140,7 @@ impl PyGraphEncoder {
#[pymethods]
impl PyGraph {
#[new]
#[pyo3(signature=(num_shards=None))]
#[pyo3(signature = (num_shards = None))]
pub fn py_new(num_shards: Option<usize>) -> (Self, PyGraphView) {
let graph = match num_shards {
None => Graph::new(),
Expand Down Expand Up @@ -177,22 +177,35 @@ impl PyGraph {
/// id (str|int): The id of the node.
/// properties (PropInput, optional): The properties of the node.
/// node_type (str, optional): The optional string which will be used as a node type
/// secondary_index (int, optional): The optional integer which will be used as a secondary index
///
/// Returns:
/// MutableNode: The added node.
///
/// Raises:
/// GraphError: If the operation fails.
#[pyo3(signature = (timestamp, id, properties = None, node_type = None))]
#[pyo3(
signature = (timestamp, id, properties = None, node_type = None, secondary_index = None)
)]
pub fn add_node(
&self,
timestamp: PyTime,
id: GID,
properties: Option<HashMap<String, Prop>>,
node_type: Option<&str>,
secondary_index: Option<usize>,
) -> Result<NodeView<Graph, Graph>, GraphError> {
self.graph
.add_node(timestamp, id, properties.unwrap_or_default(), node_type)
match secondary_index {
None => self
.graph
.add_node(timestamp, id, properties.unwrap_or_default(), node_type),
Some(secondary_index) => self.graph.add_node(
(timestamp, secondary_index),
id,
properties.unwrap_or_default(),
node_type,
),
}
}

/// Creates a new node with the given id and properties to the graph. It fails if the node already exists.
Expand All @@ -202,41 +215,61 @@ impl PyGraph {
/// id (str|int): The id of the node.
/// properties (PropInput, optional): The properties of the node.
/// node_type (str, optional): The optional string which will be used as a node type
/// secondary_index (int, optional): The optional integer which will be used as a secondary index
///
/// Returns:
/// MutableNode: The created node.
///
/// Raises:
/// GraphError: If the operation fails.
#[pyo3(signature = (timestamp, id, properties = None, node_type = None))]
#[pyo3(signature = (timestamp, id, properties = None, node_type = None, secondary_index = None))]
pub fn create_node(
&self,
timestamp: PyTime,
id: GID,
properties: Option<HashMap<String, Prop>>,
node_type: Option<&str>,
secondary_index: Option<usize>,
) -> Result<NodeView<Graph, Graph>, GraphError> {
self.graph
.create_node(timestamp, id, properties.unwrap_or_default(), node_type)
match secondary_index {
None => {
self.graph
.create_node(timestamp, id, properties.unwrap_or_default(), node_type)
}
Some(secondary_index) => self.graph.create_node(
(timestamp, secondary_index),
id,
properties.unwrap_or_default(),
node_type,
),
}
}

/// Adds properties to the graph.
///
/// Arguments:
/// timestamp (TimeInput): The timestamp of the temporal property.
/// properties (PropInput): The temporal properties of the graph.
/// secondary_index (int, optional): The optional integer which will be used as a secondary index
///
/// Returns:
/// None: This function does not return a value, if the operation is successful.
///
/// Raises:
/// GraphError: If the operation fails.
#[pyo3(signature = (timestamp, properties, secondary_index = None))]
pub fn add_property(
&self,
timestamp: PyTime,
properties: HashMap<String, Prop>,
secondary_index: Option<usize>,
) -> Result<(), GraphError> {
self.graph.add_properties(timestamp, properties)
match secondary_index {
None => self.graph.add_properties(timestamp, properties),
Some(secondary_index) => self
.graph
.add_properties((timestamp, secondary_index), properties),
}
}

/// Adds static properties to the graph.
Expand Down Expand Up @@ -281,23 +314,35 @@ impl PyGraph {
/// dst (str|int): The id of the destination node.
/// properties (PropInput, optional): The properties of the edge, as a dict of string and properties.
/// layer (str, optional): The layer of the edge.
/// secondary_index (int, optional): The optional integer which will be used as a secondary index
///
/// Returns:
/// MutableEdge: The added edge.
///
/// Raises:
/// GraphError: If the operation fails.
#[pyo3(signature = (timestamp, src, dst, properties = None, layer = None))]
#[pyo3(signature = (timestamp, src, dst, properties = None, layer = None, secondary_index = None))]
pub fn add_edge(
&self,
timestamp: PyTime,
src: GID,
dst: GID,
properties: Option<HashMap<String, Prop>>,
layer: Option<&str>,
secondary_index: Option<usize>,
) -> Result<EdgeView<Graph, Graph>, GraphError> {
self.graph
.add_edge(timestamp, src, dst, properties.unwrap_or_default(), layer)
match secondary_index {
None => self
.graph
.add_edge(timestamp, src, dst, properties.unwrap_or_default(), layer),
Some(secondary_index) => self.graph.add_edge(
(timestamp, secondary_index),
src,
dst,
properties.unwrap_or_default(),
layer,
),
}
}

/// Import a single node into the graph.
Expand Down Expand Up @@ -557,7 +602,7 @@ impl PyGraph {
/// Raises:
/// GraphError: If the operation fails.
#[pyo3(
signature = (df,time, id, node_type = None, node_type_col = None, properties = None, constant_properties = None, shared_constant_properties = None)
signature = (df, time, id, node_type = None, node_type_col = None, properties = None, constant_properties = None, shared_constant_properties = None)
)]
fn load_nodes_from_pandas<'py>(
&self,
Expand Down Expand Up @@ -744,7 +789,9 @@ impl PyGraph {
///
/// Raises:
/// GraphError: If the operation fails.
#[pyo3(signature = (df, id, node_type=None, node_type_col=None, constant_properties = None, shared_constant_properties = None))]
#[pyo3(
signature = (df, id, node_type = None, node_type_col = None, constant_properties = None, shared_constant_properties = None)
)]
fn load_node_props_from_pandas(
&self,
df: &Bound<PyAny>,
Expand Down Expand Up @@ -781,7 +828,9 @@ impl PyGraph {
///
/// Raises:
/// GraphError: If the operation fails.
#[pyo3(signature = (parquet_path, id, node_type=None,node_type_col=None, constant_properties = None, shared_constant_properties = None))]
#[pyo3(
signature = (parquet_path, id, node_type = None, node_type_col = None, constant_properties = None, shared_constant_properties = None)
)]
fn load_node_props_from_parquet(
&self,
parquet_path: PathBuf,
Expand Down
77 changes: 64 additions & 13 deletions raphtory/src/python/graph/graph_with_deletions.rs
Original file line number Diff line number Diff line change
Expand Up @@ -116,22 +116,33 @@ impl PyPersistentGraph {
/// id (str | int): The id of the node.
/// properties (dict): The properties of the node.
/// node_type (str) : The optional string which will be used as a node type
/// secondary_index (int, optional): The optional integer which will be used as a secondary index
///
/// Returns:
/// None: This function does not return a value, if the operation is successful.
///
/// Raises:
/// GraphError: If the operation fails.
#[pyo3(signature = (timestamp, id, properties = None, node_type = None))]
#[pyo3(signature = (timestamp, id, properties = None, node_type = None, secondary_index = None))]
pub fn add_node(
&self,
timestamp: PyTime,
id: GID,
properties: Option<HashMap<String, Prop>>,
node_type: Option<&str>,
secondary_index: Option<usize>,
) -> Result<NodeView<PersistentGraph>, GraphError> {
self.graph
.add_node(timestamp, id, properties.unwrap_or_default(), node_type)
match secondary_index {
None => self
.graph
.add_node(timestamp, id, properties.unwrap_or_default(), node_type),
Some(secondary_index) => self.graph.add_node(
(timestamp, secondary_index),
id,
properties.unwrap_or_default(),
node_type,
),
}
}

/// Creates a new node with the given id and properties to the graph. It fails if the node already exists.
Expand All @@ -141,41 +152,61 @@ impl PyPersistentGraph {
/// id (str | int): The id of the node.
/// properties (dict): The properties of the node.
/// node_type (str) : The optional string which will be used as a node type
/// secondary_index (int, optional): The optional integer which will be used as a secondary index
///
/// Returns:
/// MutableNode
///
/// Raises:
/// GraphError: If the operation fails.
#[pyo3(signature = (timestamp, id, properties = None, node_type = None))]
#[pyo3(signature = (timestamp, id, properties = None, node_type = None, secondary_index = None))]
pub fn create_node(
&self,
timestamp: PyTime,
id: GID,
properties: Option<HashMap<String, Prop>>,
node_type: Option<&str>,
secondary_index: Option<usize>,
) -> Result<NodeView<PersistentGraph>, GraphError> {
self.graph
.create_node(timestamp, id, properties.unwrap_or_default(), node_type)
match secondary_index {
None => {
self.graph
.create_node(timestamp, id, properties.unwrap_or_default(), node_type)
}
Some(secondary_index) => self.graph.create_node(
(timestamp, secondary_index),
id,
properties.unwrap_or_default(),
node_type,
),
}
}

/// Adds properties to the graph.
///
/// Arguments:
/// timestamp (TimeInput): The timestamp of the temporal property.
/// properties (dict): The temporal properties of the graph.
/// secondary_index (int, optional): The optional integer which will be used as a secondary index
///
/// Returns:
/// None: This function does not return a value, if the operation is successful.
///
/// Raises:
/// GraphError: If the operation fails.
pub fn add_property(
#[pyo3(signature = (timestamp, properties, secondary_index = None))]
pub fn add_properties(
&self,
timestamp: PyTime,
properties: HashMap<String, Prop>,
secondary_index: Option<usize>,
) -> Result<(), GraphError> {
self.graph.add_properties(timestamp, properties)
match secondary_index {
None => self.graph.add_properties(timestamp, properties),
Some(secondary_index) => self
.graph
.add_properties((timestamp, secondary_index), properties),
}
}

/// Adds static properties to the graph.
Expand Down Expand Up @@ -220,23 +251,35 @@ impl PyPersistentGraph {
/// dst (str | int): The id of the destination node.
/// properties (dict): The properties of the edge, as a dict of string and properties
/// layer (str): The layer of the edge.
/// secondary_index (int, optional): The optional integer which will be used as a secondary index
///
/// Returns:
/// None: This function does not return a value, if the operation is successful.
///
/// Raises:
/// GraphError: If the operation fails.
#[pyo3(signature = (timestamp, src, dst, properties = None, layer = None))]
#[pyo3(signature = (timestamp, src, dst, properties = None, layer = None, secondary_index = None))]
pub fn add_edge(
&self,
timestamp: PyTime,
src: GID,
dst: GID,
properties: Option<HashMap<String, Prop>>,
layer: Option<&str>,
secondary_index: Option<usize>,
) -> Result<EdgeView<PersistentGraph, PersistentGraph>, GraphError> {
self.graph
.add_edge(timestamp, src, dst, properties.unwrap_or_default(), layer)
match secondary_index {
None => self
.graph
.add_edge(timestamp, src, dst, properties.unwrap_or_default(), layer),
Some(secondary_index) => self.graph.add_edge(
(timestamp, secondary_index),
src,
dst,
properties.unwrap_or_default(),
layer,
),
}
}

/// Deletes an edge given the timestamp, src and dst nodes and layer (optional)
Expand All @@ -246,21 +289,29 @@ impl PyPersistentGraph {
/// src (str | int): The id of the source node.
/// dst (str | int): The id of the destination node.
/// layer (str): The layer of the edge. (optional)
/// secondary_index (int, optional): The optional integer which will be used as a secondary index
///
/// Returns:
/// The deleted edge
///
/// Raises:
/// GraphError: If the operation fails.
#[pyo3(signature = (timestamp, src, dst, layer=None))]
#[pyo3(signature = (timestamp, src, dst, layer=None, secondary_index = None))]
pub fn delete_edge(
&self,
timestamp: PyTime,
src: GID,
dst: GID,
layer: Option<&str>,
secondary_index: Option<usize>,
) -> Result<EdgeView<PersistentGraph>, GraphError> {
self.graph.delete_edge(timestamp, src, dst, layer)
match secondary_index {
None => self.graph.delete_edge(timestamp, src, dst, layer),
Some(secondary_index) => {
self.graph
.delete_edge((timestamp, secondary_index), src, dst, layer)
}
}
}

//FIXME: This is reimplemented here to get mutable views. If we switch the underlying graph to enum dispatch, this won't be necessary!
Expand Down
Loading

0 comments on commit fbebc01

Please sign in to comment.