Speed up connect_edges

datashader/bundling.py

@@ -27,6 +27,7 @@
 import numpy as np
 import pandas as pd
 import param
+from numba import prange
 
 from .utils import ngjit
 
@@ -283,6 +284,27 @@ def accumulate(img, point, accuracy):
         img[int(point[0] * accuracy), int(point[1] * accuracy)] += point[2]
 
 
+@ngjit
+def _edge_val_to_segments(vals):
+    edge_vals = np.full(vals.size * 3, np.nan, np.dtype("float"))
+    for i in prange(len(vals)):
+        pos = int(i * 3)
+        edge_vals[pos:pos+2] = vals[i]
+    return edge_vals
+
+@ngjit
+def _format_graph_segments(nodes, edges):
+    edge_table = np.full((edges.shape[0] * 3, 2), np.nan)
+    for i in prange(edges.shape[0]):
+        pos = int(i * 3)
+        s, t = edges[i]
+        edge_table[pos  , 0] = nodes[s, 0]
+        edge_table[pos  , 1] = nodes[s, 1]
+        edge_table[pos+1, 0] = nodes[t, 0]
+        edge_table[pos+1, 1] = nodes[t, 1]
+    return edge_table
+
+
 def _convert_graph_to_edge_segments(nodes, edges, params):
     """
     Merge graph dataframes into a list of edge segments.
@@ -391,8 +413,17 @@ def __call__(self, nodes, edges, **params):
         a point with NaN as the x or y value.
         """
         p = param.ParamOverrides(self, params)
-        edges, segment_class = _convert_graph_to_edge_segments(nodes, edges, p)
-        return _convert_edge_segments_to_dataframe(edges, segment_class, p)
+        e = pd.DataFrame(
+            _format_graph_segments(
+                nodes[[p.x, p.y]].values,
+                edges[[p.source, p.target]].values),
+            columns=[p.x, p.y]
+        )
+        if p.weight is not None:
+            e[p.weight] = _edge_val_to_segments(edges[p.weight].values)
+        if p.include_edge_id:
+            e["edge_id"] = _edge_val_to_segments(edges["id"].values)
+        return e
 
 directly_connect_edges = connect_edges # For bockwards compatibility; deprecated