From a4aa4f985a8c652c69ef9c30285965a2ff13437d Mon Sep 17 00:00:00 2001 From: Lucas Jeub Date: Mon, 8 Jan 2024 15:17:08 +0100 Subject: [PATCH] add shrink to python and gql --- python/tests/test_graphdb.py | 28 +++++------ raphtory-graphql/src/model/graph/edge.rs | 12 +++++ raphtory-graphql/src/model/graph/graph.rs | 12 +++++ raphtory-graphql/src/model/graph/node.rs | 12 +++++ raphtory/src/db/api/view/time.rs | 16 +++---- raphtory/src/python/types/macros/timeops.rs | 52 ++++++++++++++++----- 6 files changed, 98 insertions(+), 34 deletions(-) diff --git a/python/tests/test_graphdb.py b/python/tests/test_graphdb.py index f3a52a5afc..9656127493 100644 --- a/python/tests/test_graphdb.py +++ b/python/tests/test_graphdb.py @@ -1026,14 +1026,14 @@ def test_all_degrees_window(): view = g.before(5) v = view.node(2) assert v.window(0, 4).in_degree() == 3 - assert v.window(start=2).in_degree() == 2 - assert v.window(end=3).in_degree() == 2 + assert v.after(1).in_degree() == 2 + assert v.before(3).in_degree() == 2 assert v.window(0, 4).out_degree() == 1 - assert v.window(start=2).out_degree() == 1 - assert v.window(end=3).out_degree() == 1 + assert v.after(1).out_degree() == 1 + assert v.before(end=3).out_degree() == 1 assert v.window(0, 4).degree() == 3 - assert v.window(start=2).degree() == 2 - assert v.window(end=3).degree() == 2 + assert v.after(1).degree() == 2 + assert v.before(end=3).degree() == 2 def test_all_edge_window(): @@ -1049,24 +1049,24 @@ def test_all_edge_window(): view = g.before(5) v = view.node(2) assert sorted(v.window(0, 4).in_edges.src.id) == [1, 3, 4] - assert sorted(v.window(end=4).in_edges.src.id) == [1, 3, 4] - assert sorted(v.window(start=2).in_edges.src.id) == [3, 4] + assert sorted(v.before(end=4).in_edges.src.id) == [1, 3, 4] + assert sorted(v.after(start=1).in_edges.src.id) == [3, 4] assert sorted(v.window(0, 4).out_edges.dst.id) == [3] - assert sorted(v.window(end=3).out_edges.dst.id) == [3] - assert sorted(v.window(start=2).out_edges.dst.id) == [4] + assert sorted(v.before(end=3).out_edges.dst.id) == [3] + assert sorted(v.after(start=1).out_edges.dst.id) == [4] assert sorted((e.src.id, e.dst.id) for e in v.window(0, 4).edges) == [ (1, 2), (2, 3), (3, 2), (4, 2), ] - assert sorted((e.src.id, e.dst.id) for e in v.window(end=4).edges) == [ + assert sorted((e.src.id, e.dst.id) for e in v.before(end=4).edges) == [ (1, 2), (2, 3), (3, 2), (4, 2), ] - assert sorted((e.src.id, e.dst.id) for e in v.window(start=1).edges) == [ + assert sorted((e.src.id, e.dst.id) for e in v.after(start=0).edges) == [ (1, 2), (2, 3), (2, 4), @@ -1620,9 +1620,9 @@ def test_deletions(): for e in edges: assert g.at(e[0]).has_edge(e[1], e[2]) - assert not g.window(start=11).has_edge(edges[0][1], edges[0][2]) + assert not g.after(start=10).has_edge(edges[0][1], edges[0][2]) for e in edges[1:]: - assert g.window(start=11).has_edge(e[1], e[2]) + assert g.after(start=10).has_edge(e[1], e[2]) assert list(g.edge(edges[0][1], edges[0][2]).explode().latest_time) == [10] diff --git a/raphtory-graphql/src/model/graph/edge.rs b/raphtory-graphql/src/model/graph/edge.rs index ebf942a771..15b90d1ea7 100644 --- a/raphtory-graphql/src/model/graph/edge.rs +++ b/raphtory-graphql/src/model/graph/edge.rs @@ -54,6 +54,18 @@ impl Edge { self.ee.after(time).into() } + async fn shrink_window(&self, start: i64, end: i64) -> Self { + self.ee.shrink_window(start, end).into() + } + + async fn shrink_start(&self, start: i64) -> Self { + self.ee.shrink_start(start).into() + } + + async fn shrink_end(&self, end: i64) -> Self { + self.ee.shrink_end(end).into() + } + async fn earliest_time(&self) -> Option { self.ee.earliest_time() } diff --git a/raphtory-graphql/src/model/graph/graph.rs b/raphtory-graphql/src/model/graph/graph.rs index dd48164469..6441518e74 100644 --- a/raphtory-graphql/src/model/graph/graph.rs +++ b/raphtory-graphql/src/model/graph/graph.rs @@ -95,6 +95,18 @@ impl GqlGraph { GqlGraph::new(self.name.clone(), self.graph.after(time)) } + async fn shrink_window(&self, start: i64, end: i64) -> Self { + GqlGraph::new(self.name.clone(), self.graph.shrink_window(start, end)) + } + + async fn shrink_start(&self, start: i64) -> Self { + GqlGraph::new(self.name.clone(), self.graph.shrink_start(start)) + } + + async fn shrink_end(&self, end: i64) -> Self { + GqlGraph::new(self.name.clone(), self.graph.shrink_end(end)) + } + //////////////////////// //// TIME QUERIES ////// //////////////////////// diff --git a/raphtory-graphql/src/model/graph/node.rs b/raphtory-graphql/src/model/graph/node.rs index 77dce79d9d..74e7d3a239 100644 --- a/raphtory-graphql/src/model/graph/node.rs +++ b/raphtory-graphql/src/model/graph/node.rs @@ -64,6 +64,18 @@ impl Node { self.vv.after(time).into() } + async fn shrink_window(&self, start: i64, end: i64) -> Self { + self.vv.shrink_window(start, end).into() + } + + async fn shrink_start(&self, start: i64) -> Self { + self.vv.shrink_start(start).into() + } + + async fn shrink_end(&self, end: i64) -> Self { + self.vv.shrink_end(end).into() + } + //////////////////////// //// TIME QUERIES ////// //////////////////////// diff --git a/raphtory/src/db/api/view/time.rs b/raphtory/src/db/api/view/time.rs index 1fe31d951f..22dae73bfb 100644 --- a/raphtory/src/db/api/view/time.rs +++ b/raphtory/src/db/api/view/time.rs @@ -21,23 +21,23 @@ pub trait TimeOps<'graph> { fn end(&self) -> Option; /// set the start of the window to the larger of `start` and `self.start()` - fn shrink_start(&self, start: i64) -> Self::WindowedViewType { - let start = max(start, self.start().unwrap_or(i64::MIN)); + fn shrink_start(&self, start: T) -> Self::WindowedViewType { + let start = max(start.into_time(), self.start().unwrap_or(i64::MIN)); let end = self.end().unwrap_or(start); self.window(start, end) } - /// set the end of the winodw to the smaller of `end` and `self.end()` - fn shrink_end(&self, end: i64) -> Self::WindowedViewType { - let end = min(end, self.end().unwrap_or(i64::MAX)); + /// set the end of the window to the smaller of `end` and `self.end()` + fn shrink_end(&self, end: T) -> Self::WindowedViewType { + let end = min(end.into_time(), self.end().unwrap_or(i64::MAX)); let start = self.start().unwrap_or(end); self.window(start, end) } /// shrink both the start and end of the window (same as calling `shrink_start` followed by `shrink_end` but more efficient) - fn shrink_window(&self, start: i64, end: i64) -> Self::WindowedViewType { - let start = max(start, self.start().unwrap_or(i64::MIN)); - let end = min(end, self.end().unwrap_or(i64::MAX)); + fn shrink_window(&self, start: T, end: T) -> Self::WindowedViewType { + let start = max(start.into_time(), self.start().unwrap_or(i64::MIN)); + let end = min(end.into_time(), self.end().unwrap_or(i64::MAX)); self.window(start, end) } diff --git a/raphtory/src/python/types/macros/timeops.rs b/raphtory/src/python/types/macros/timeops.rs index 198424c84f..1809e84d55 100644 --- a/raphtory/src/python/types/macros/timeops.rs +++ b/raphtory/src/python/types/macros/timeops.rs @@ -59,7 +59,7 @@ macro_rules! impl_timeops { /// An expanding window is a window that grows by `step` size at each iteration. /// /// Arguments: - /// step (int): The step size of the window. + /// step (int | str): The step size of the window. /// /// Returns: /// A `WindowSet` object. @@ -72,8 +72,8 @@ macro_rules! impl_timeops { /// A rolling window is a window that moves forward by `step` size at each iteration. /// /// Arguments: - /// window: The size of the window. - /// step: The step size of the window. Defaults to the window size. + /// window (int | str): The size of the window. + /// step (int | str | None): The step size of the window. Defaults to `window`. /// /// Returns: /// A `WindowSet` object. @@ -88,25 +88,24 @@ macro_rules! impl_timeops { #[doc = concat!(r" Create a view of the ", $name, r" including all events between `start` (inclusive) and `end` (exclusive)")] /// /// Arguments: - #[doc = concat!(r" start: The start time of the window. Defaults to the start time of the ", $name, r".")] - #[doc = concat!(r" end: The end time of the window. Defaults to the end time of the ", $name, r".")] + /// start (int | DateTime | str): The start time of the window. + /// end (int | DateTime | str): The end time of the window. /// /// Returns: #[doc = concat!("r A ", $name, " object.")] - #[pyo3(signature = (start = None, end = None))] pub fn window( &self, - start: Option, - end: Option, + start: PyTime, + end: PyTime, ) -> <$base_type as TimeOps<'static>>::WindowedViewType { self.$field - .window(start.unwrap_or(PyTime::MIN), end.unwrap_or(PyTime::MAX)) + .window(start, end) } #[doc = concat!(r" Create a view of the ", $name, r" including all events at `time`.")] /// /// Arguments: - /// time: The time of the window. + /// time (int | DateTime | str): The time of the window. /// /// Returns: #[doc = concat!(r" A ", $name, r" object.")] @@ -117,7 +116,7 @@ macro_rules! impl_timeops { #[doc = concat!(r" Create a view of the ", $name, r" including all events before `end` (exclusive).")] /// /// Arguments: - /// end: The end time of the window. + /// end (int | DateTime | str): The end time of the window. /// /// Returns: #[doc = concat!(r" A ", $name, r" object.")] @@ -128,13 +127,42 @@ macro_rules! impl_timeops { #[doc = concat!(r" Create a view of the ", $name, r" including all events after `start` (exclusive).")] /// /// Arguments: - /// start: The start time of the window. + /// start (int | DateTime | str): The start time of the window. /// /// Returns: #[doc = concat!(r" A ", $name, r" object.")] pub fn after(&self, start: PyTime) -> <$base_type as TimeOps<'static>>::WindowedViewType { self.$field.after(start) } + + /// Set the start of the window to the larger of `start` and `self.start()` + /// + /// Arguments: + /// start (int | DateTime | str): the new start time of the window + /// + /// Returns: + #[doc = concat!(r" A ", $name, r" object.")] + pub fn shrink_start(&self, start: PyTime) -> <$base_type as TimeOps<'static>>::WindowedViewType { + self.$field.shrink_start(start) + } + + /// Set the end of the window to the smaller of `end` and `self.end()` + /// + /// Arguments: + /// end (int | DateTime | str): the new end time of the window + /// Returns: + #[doc = concat!(r" A ", $name, r" object.")] + fn shrink_end(&self, end: PyTime) -> <$base_type as TimeOps<'static>>::WindowedViewType { + self.$field.shrink_end(end) + } + + /// Shrink both the start and end of the window (same as calling `shrink_start` followed by `shrink_end` but more efficient) + /// + /// Arguments: + /// + fn shrink_window(&self, start: PyTime, end: PyTime) -> <$base_type as TimeOps<'static>>::WindowedViewType { + self.$field.shrink_window(start, end) + } } }; }