Run SQL on GPU with Polars.

requirements.txt

@@ -1,15 +1,24 @@
+# Dependencies for basic operation.
 fastapi
-matplotlib
-networkx
-numpy
 orjson
-pandas
 pydantic-to-typescript
-scipy
 uvicorn[standard]
 pycrdt
 pycrdt-websocket
-# For llm_ops
+
+# For lynxkite_ops.
+matplotlib
+networkx
+numpy
+pandas
+scipy
+
+# GPU-accelerated graph analytics.
+cugraph-cu12
+cudf-cu12
+nx-cugraph-cu12 
+
+# For llm_ops.
 chromadb
 Jinja2
 openai

run.sh

@@ -1,2 +1,3 @@
 #!/bin/bash -xue
+export NX_CUGRAPH_AUTOCONFIG=True
 uvicorn server.main:app --reload

server/lynxkite_ops.py

@@ -1,5 +1,6 @@
-"""Some operations. To be split into separate files when we have more."""
+"""Graph analytics operations. To be split into separate files when we have more."""
 
+import os
 from . import ops
 from collections import deque
 import dataclasses
@@ -7,6 +8,7 @@ import functools
 import matplotlib
 import networkx as nx
 import pandas as pd
+import polars as pl
 import traceback
 import typing
 
@@ -63,13 +65,27 @@ class Bundle:
             ],
         )
 
+    @classmethod
+    def from_df(cls, df: pd.DataFrame):
+        return cls(dfs={"df": df})
+
     def to_nx(self):
+        # TODO: Use relations.
         graph = nx.from_pandas_edgelist(self.dfs["edges"])
-        nx.set_node_attributes(
-            graph, self.dfs["nodes"].set_index("id").to_dict("index")
-        )
+        if "nodes" in self.dfs:
+            nx.set_node_attributes(
+                graph, self.dfs["nodes"].set_index("id").to_dict("index")
+            )
         return graph
 
+    def copy(self):
+        """Returns a medium depth copy of the bundle. The Bundle is completely new, but the DataFrames and RelationDefinitions are shared."""
+        return Bundle(
+            dfs=dict(self.dfs),
+            relations=list(self.relations),
+            other=dict(self.other) if self.other else None,
+        )
+
 
 def nx_node_attribute_func(name):
     """Decorator for wrapping a function that adds a NetworkX node attribute."""
@@ -99,7 +115,6 @@ def disambiguate_edges(ws):
 @ops.register_executor("LynxKite")
 async def execute(ws):
     catalog = ops.CATALOGS["LynxKite"]
-    # Nodes are responsible for interpreting/executing their child nodes.
     disambiguate_edges(ws)
     outputs = {}
     failed = 0
@@ -115,12 +130,14 @@ async def execute(ws):
                 op = catalog[data.title]
                 params = {**data.params}
                 # Convert inputs.
-                for i, (x, p) in enumerate(zip(inputs, op.inputs.values())):
-                    if p.type == nx.Graph and isinstance(x, Bundle):
-                        inputs[i] = x.to_nx()
-                    elif p.type == Bundle and isinstance(x, nx.Graph):
-                        inputs[i] = Bundle.from_nx(x)
                 try:
+                    for i, (x, p) in enumerate(zip(inputs, op.inputs.values())):
+                        if p.type == nx.Graph and isinstance(x, Bundle):
+                            inputs[i] = x.to_nx()
+                        elif p.type == Bundle and isinstance(x, nx.Graph):
+                            inputs[i] = Bundle.from_nx(x)
+                        elif p.type == Bundle and isinstance(x, pd.DataFrame):
+                            inputs[i] = Bundle.from_df(x)
                     output = op(*inputs, **params)
                 except Exception as e:
                     traceback.print_exc()
@@ -142,10 +159,24 @@ async def execute(ws):
 
 @op("Import Parquet")
 def import_parquet(*, filename: str):
-    """Imports a parquet file."""
+    """Imports a Parquet file."""
     return pd.read_parquet(filename)
 
 
+@op("Import CSV")
+def import_csv(
+    *, filename: str, columns: str = "<from file>", separator: str = "<auto>"
+):
+    """Imports a CSV file."""
+    return pd.read_csv(
+        filename,
+        names=pd.api.extensions.no_default
+        if columns == "<from file>"
+        else columns.split(","),
+        sep=pd.api.extensions.no_default if separator == "<auto>" else separator,
+    )
+
+
 @op("Create scale-free graph")
 def create_scale_free_graph(*, nodes: int = 10):
     """Creates a scale-free graph with the given number of nodes."""
@@ -158,6 +189,12 @@ def compute_pagerank(graph: nx.Graph, *, damping=0.85, iterations=100):
     return nx.pagerank(graph, alpha=damping, max_iter=iterations)
 
 
+@op("Compute betweenness centrality")
+@nx_node_attribute_func("betweenness_centrality")
+def compute_betweenness_centrality(graph: nx.Graph, *, k=10):
+    return nx.betweenness_centrality(graph, k=k, backend="cugraph")
+
+
 @op("Discard loop edges")
 def discard_loop_edges(graph: nx.Graph):
     graph = graph.copy()
@@ -165,6 +202,35 @@ def discard_loop_edges(graph: nx.Graph):
     return graph
 
 
+@op("SQL")
+def sql(bundle: Bundle, *, query: ops.LongStr, save_as: str = "result"):
+    """Run a SQL query on the DataFrames in the bundle. Save the results as a new DataFrame."""
+    bundle = bundle.copy()
+    if os.environ.get("NX_CUGRAPH_AUTOCONFIG", "").strip().lower() == "true":
+        with pl.Config() as cfg:
+            cfg.set_verbose(True)
+            res = (
+                pl.SQLContext(bundle.dfs)
+                .execute(query)
+                .collect(engine="gpu")
+                .to_pandas()
+            )
+            # TODO: Currently `collect()` moves the data from cuDF to Polars. Then we convert it to Pandas,
+            # which (hopefully) puts it back into cuDF. Hopefully we will be able to keep it in cuDF.
+    else:
+        res = pl.SQLContext(bundle.dfs).execute(query)
+    bundle.dfs[save_as] = res
+    return bundle
+
+
+@op("Organize bundle")
+def organize_bundle(bundle: Bundle, *, code: ops.LongStr):
+    """Lets you rename/copy/delete DataFrames, and modify relations. TODO: Use a declarative solution instead of Python code. Add UI."""
+    bundle = bundle.copy()
+    exec(code, globals(), {"bundle": bundle})
+    return bundle
+
+
 @op("Sample graph")
 def sample_graph(graph: nx.Graph, *, nodes: int = 100):
     """Takes a (preferably connected) subgraph."""
@@ -237,13 +303,21 @@ def visualize_graph(graph: Bundle, *, color_nodes_by: ops.NodeAttribute = None):
     return v
 
 
+def collect(df: pd.DataFrame):
+    if isinstance(df, pl.LazyFrame):
+        df = df.collect()
+    if isinstance(df, pl.DataFrame):
+        return [[d[c] for c in df.columns] for d in df.to_dicts()]
+    return df.values.tolist()
+
+
 @op("View tables", view="table_view")
 def view_tables(bundle: Bundle):
     v = {
         "dataframes": {
             name: {
                 "columns": [str(c) for c in df.columns],
-                "data": df.values.tolist(),
+                "data": collect(df),
             }
             for name, df in bundle.dfs.items()
         },

web/package-lock.json

@@ -8,9 +8,11 @@
       "name": "lynxkite",
       "version": "0.0.0",
       "dependencies": {
+        "@esbuild/linux-x64": "^0.24.0",
         "@iconify-json/tabler": "^1.2.10",
         "@svgr/core": "^8.1.0",
         "@svgr/plugin-jsx": "^8.1.0",
+        "@swc/core": "^1.10.1",
         "@syncedstore/core": "^0.6.0",
         "@syncedstore/react": "^0.6.0",
         "@types/node": "^22.10.1",
@@ -42,6 +44,9 @@
         "typescript": "~5.6.2",
         "typescript-eslint": "^8.15.0",
         "vite": "^6.0.1"
+      },
+      "optionalDependencies": {
+        "@rollup/rollup-linux-x64-gnu": "^4.28.1"
       }
     },
     "node_modules/@alloc/quick-lru": {
@@ -351,6 +356,21 @@
         "node": ">=18"
       }
     },
+    "node_modules/@esbuild/linux-x64": {
+      "version": "0.24.0",
+      "resolved": "https://registry.npmjs.org/@esbuild/linux-x64/-/linux-x64-0.24.0.tgz",
+      "integrity": "sha512-vbutsFqQ+foy3wSSbmjBXXIJ6PL3scghJoM8zCL142cGaZKAdCZHyf+Bpu/MmX9zT9Q0zFBVKb36Ma5Fzfa8xA==",
+      "cpu": [
+        "x64"
+      ],
+      "license": "MIT",
+      "os": [
+        "linux"
+      ],
+      "engines": {
+        "node": ">=18"
+      }
+    },
     "node_modules/@eslint-community/eslint-utils": {
       "version": "4.4.1",
       "resolved": "https://registry.npmjs.org/@eslint-community/eslint-utils/-/eslint-utils-4.4.1.tgz",
@@ -748,6 +768,19 @@
         "darwin"
       ]
     },
+    "node_modules/@rollup/rollup-linux-x64-gnu": {
+      "version": "4.28.1",
+      "resolved": "https://registry.npmjs.org/@rollup/rollup-linux-x64-gnu/-/rollup-linux-x64-gnu-4.28.1.tgz",
+      "integrity": "sha512-fzgeABz7rrAlKYB0y2kSEiURrI0691CSL0+KXwKwhxvj92VULEDQLpBYLHpF49MSiPG4sq5CK3qHMnb9tlCjBw==",
+      "cpu": [
+        "x64"
+      ],
+      "license": "MIT",
+      "optional": true,
+      "os": [
+        "linux"
+      ]
+    },
     "node_modules/@svgr/babel-plugin-add-jsx-attribute": {
       "version": "8.0.0",
       "resolved": "https://registry.npmjs.org/@svgr/babel-plugin-add-jsx-attribute/-/babel-plugin-add-jsx-attribute-8.0.0.tgz",
@@ -965,7 +998,6 @@
       "version": "1.10.1",
       "resolved": "https://registry.npmjs.org/@swc/core/-/core-1.10.1.tgz",
       "integrity": "sha512-rQ4dS6GAdmtzKiCRt3LFVxl37FaY1cgL9kSUTnhQ2xc3fmHOd7jdJK/V4pSZMG1ruGTd0bsi34O2R0Olg9Zo/w==",
-      "dev": true,
       "hasInstallScript": true,
       "license": "Apache-2.0",
       "dependencies": {
@@ -1007,7 +1039,6 @@
       "cpu": [
         "arm64"
       ],
-      "dev": true,
       "license": "Apache-2.0 AND MIT",
       "optional": true,
       "os": [
@@ -1017,18 +1048,160 @@
         "node": ">=10"
       }
     },
+    "node_modules/@swc/core-darwin-x64": {
+      "version": "1.10.1",
+      "resolved": "https://registry.npmjs.org/@swc/core-darwin-x64/-/core-darwin-x64-1.10.1.tgz",
+      "integrity": "sha512-L4BNt1fdQ5ZZhAk5qoDfUnXRabDOXKnXBxMDJ+PWLSxOGBbWE6aJTnu4zbGjJvtot0KM46m2LPAPY8ttknqaZA==",
+      "cpu": [
+        "x64"
+      ],
+      "license": "Apache-2.0 AND MIT",
+      "optional": true,
+      "os": [
+        "darwin"
+      ],
+      "engines": {
+        "node": ">=10"
+      }
+    },
+    "node_modules/@swc/core-linux-arm-gnueabihf": {
+      "version": "1.10.1",
+      "resolved": "https://registry.npmjs.org/@swc/core-linux-arm-gnueabihf/-/core-linux-arm-gnueabihf-1.10.1.tgz",
+      "integrity": "sha512-Y1u9OqCHgvVp2tYQAJ7hcU9qO5brDMIrA5R31rwWQIAKDkJKtv3IlTHF0hrbWk1wPR0ZdngkQSJZple7G+Grvw==",
+      "cpu": [
+        "arm"
+      ],
+      "license": "Apache-2.0",
+      "optional": true,
+      "os": [
+        "linux"
+      ],
+      "engines": {
+        "node": ">=10"
+      }
+    },
+    "node_modules/@swc/core-linux-arm64-gnu": {
+      "version": "1.10.1",
+      "resolved": "https://registry.npmjs.org/@swc/core-linux-arm64-gnu/-/core-linux-arm64-gnu-1.10.1.tgz",
+      "integrity": "sha512-tNQHO/UKdtnqjc7o04iRXng1wTUXPgVd8Y6LI4qIbHVoVPwksZydISjMcilKNLKIwOoUQAkxyJ16SlOAeADzhQ==",
+      "cpu": [
+        "arm64"
+      ],
+      "license": "Apache-2.0 AND MIT",
+      "optional": true,
+      "os": [
+        "linux"
+      ],
+      "engines": {
+        "node": ">=10"
+      }
+    },
+    "node_modules/@swc/core-linux-arm64-musl": {
+      "version": "1.10.1",
+      "resolved": "https://registry.npmjs.org/@swc/core-linux-arm64-musl/-/core-linux-arm64-musl-1.10.1.tgz",
+      "integrity": "sha512-x0L2Pd9weQ6n8dI1z1Isq00VHFvpBClwQJvrt3NHzmR+1wCT/gcYl1tp9P5xHh3ldM8Cn4UjWCw+7PaUgg8FcQ==",
+      "cpu": [
+        "arm64"
+      ],
+      "license": "Apache-2.0 AND MIT",
+      "optional": true,
+      "os": [
+        "linux"
+      ],
+      "engines": {
+        "node": ">=10"
+      }
+    },
+    "node_modules/@swc/core-linux-x64-gnu": {
+      "version": "1.10.1",
+      "resolved": "https://registry.npmjs.org/@swc/core-linux-x64-gnu/-/core-linux-x64-gnu-1.10.1.tgz",
+      "integrity": "sha512-yyYEwQcObV3AUsC79rSzN9z6kiWxKAVJ6Ntwq2N9YoZqSPYph+4/Am5fM1xEQYf/kb99csj0FgOelomJSobxQA==",
+      "cpu": [
+        "x64"
+      ],
+      "license": "Apache-2.0 AND MIT",
+      "optional": true,
+      "os": [
+        "linux"
+      ],
+      "engines": {
+        "node": ">=10"
+      }
+    },
+    "node_modules/@swc/core-linux-x64-musl": {
+      "version": "1.10.1",
+      "resolved": "https://registry.npmjs.org/@swc/core-linux-x64-musl/-/core-linux-x64-musl-1.10.1.tgz",
+      "integrity": "sha512-tcaS43Ydd7Fk7sW5ROpaf2Kq1zR+sI5K0RM+0qYLYYurvsJruj3GhBCaiN3gkzd8m/8wkqNqtVklWaQYSDsyqA==",
+      "cpu": [
+        "x64"
+      ],
+      "license": "Apache-2.0 AND MIT",
+      "optional": true,
+      "os": [
+        "linux"
+      ],
+      "engines": {
+        "node": ">=10"
+      }
+    },
+    "node_modules/@swc/core-win32-arm64-msvc": {
+      "version": "1.10.1",
+      "resolved": "https://registry.npmjs.org/@swc/core-win32-arm64-msvc/-/core-win32-arm64-msvc-1.10.1.tgz",
+      "integrity": "sha512-D3Qo1voA7AkbOzQ2UGuKNHfYGKL6eejN8VWOoQYtGHHQi1p5KK/Q7V1ku55oxXBsj79Ny5FRMqiRJpVGad7bjQ==",
+      "cpu": [
+        "arm64"
+      ],
+      "license": "Apache-2.0 AND MIT",
+      "optional": true,
+      "os": [
+        "win32"
+      ],
+      "engines": {
+        "node": ">=10"
+      }
+    },
+    "node_modules/@swc/core-win32-ia32-msvc": {
+      "version": "1.10.1",
+      "resolved": "https://registry.npmjs.org/@swc/core-win32-ia32-msvc/-/core-win32-ia32-msvc-1.10.1.tgz",
+      "integrity": "sha512-WalYdFoU3454Og+sDKHM1MrjvxUGwA2oralknXkXL8S0I/8RkWZOB++p3pLaGbTvOO++T+6znFbQdR8KRaa7DA==",
+      "cpu": [
+        "ia32"
+      ],
+      "license": "Apache-2.0 AND MIT",
+      "optional": true,
+      "os": [
+        "win32"
+      ],
+      "engines": {
+        "node": ">=10"
+      }
+    },
+    "node_modules/@swc/core-win32-x64-msvc": {
+      "version": "1.10.1",
+      "resolved": "https://registry.npmjs.org/@swc/core-win32-x64-msvc/-/core-win32-x64-msvc-1.10.1.tgz",
+      "integrity": "sha512-JWobfQDbTnoqaIwPKQ3DVSywihVXlQMbDuwik/dDWlj33A8oEHcjPOGs4OqcA3RHv24i+lfCQpM3Mn4FAMfacA==",
+      "cpu": [
+        "x64"
+      ],
+      "license": "Apache-2.0 AND MIT",
+      "optional": true,
+      "os": [
+        "win32"
+      ],
+      "engines": {
+        "node": ">=10"
+      }
+    },
     "node_modules/@swc/counter": {
       "version": "0.1.3",
       "resolved": "https://registry.npmjs.org/@swc/counter/-/counter-0.1.3.tgz",
       "integrity": "sha512-e2BR4lsJkkRlKZ/qCHPw9ZaSxc0MVUd7gtbtaB7aMvHeJVYe8sOB8DBZkP2DtISHGSku9sCK6T6cnY0CtXrOCQ==",
-      "dev": true,
       "license": "Apache-2.0"
     },
     "node_modules/@swc/types": {
       "version": "0.1.17",
       "resolved": "https://registry.npmjs.org/@swc/types/-/types-0.1.17.tgz",
       "integrity": "sha512-V5gRru+aD8YVyCOMAjMpWR1Ui577DD5KSJsHP8RAxopAH22jFz6GZd/qxqjO6MJHQhcsjvjOFXyDhyLQUnMveQ==",
-      "dev": true,
       "license": "Apache-2.0",
       "dependencies": {
         "@swc/counter": "^0.1.3"

web/package.json

@@ -10,9 +10,11 @@
     "preview": "vite preview"
   },
   "dependencies": {
+    "@esbuild/linux-x64": "^0.24.0",
     "@iconify-json/tabler": "^1.2.10",
     "@svgr/core": "^8.1.0",
     "@svgr/plugin-jsx": "^8.1.0",
+    "@swc/core": "^1.10.1",
     "@syncedstore/core": "^0.6.0",
     "@syncedstore/react": "^0.6.0",
     "@types/node": "^22.10.1",
@@ -44,5 +46,8 @@
     "typescript": "~5.6.2",
     "typescript-eslint": "^8.15.0",
     "vite": "^6.0.1"
+  },
+  "optionalDependencies": {
+    "@rollup/rollup-linux-x64-gnu": "^4.28.1"
   }
 }

web/src/workspace/Workspace.tsx

@@ -57,7 +57,7 @@ function LynxKiteFlow() {
     const state = syncedStore({ workspace: {} as Workspace });
     setState(state);
     const doc = getYjsDoc(state);
-    const wsProvider = new WebsocketProvider("ws://localhost:8000/ws/crdt", path!, doc);
+    const wsProvider = new WebsocketProvider("ws://localhost:5173/ws/crdt", path!, doc);
     const onChange = (_update: any, origin: any, _doc: any, _tr: any) => {
       if (origin === wsProvider) {
         // An update from the CRDT. Apply it to the local state.
@@ -125,6 +125,7 @@ function LynxKiteFlow() {
       const edgeIndex = wedges.findIndex((e) => e.id === ch.id);
       if (ch.type === 'remove') {
         wedges.splice(edgeIndex, 1);
+      } else if (ch.type === 'select') {
       } else {
         console.log('Unknown edge change', ch);
       }

web/src/workspace/nodes/LynxKiteNode.tsx

@@ -45,9 +45,6 @@ export default function LynxKiteNode(props: LynxKiteNodeProps) {
   const data = props.data;
   const expanded = !data.collapsed;
   const handles = getHandles(data.meta?.inputs || {}, data.meta?.outputs || {});
-  function asPx(n: number | undefined) {
-    return (n ? n + 'px' : undefined) || '200px';
-  }
   function titleClicked() {
     reactFlow.updateNodeData(props.id, { collapsed: expanded });
   }
@@ -55,7 +52,7 @@ export default function LynxKiteNode(props: LynxKiteNodeProps) {
 
   return (
     <div className={'node-container ' + (expanded ? 'expanded' : 'collapsed')}
-      style={{ width: asPx(props.width), height: asPx(expanded ? props.height : undefined) }}>
+      style={{ width: props.width || 200, height: expanded ? props.height || 200 : undefined }}>
       <div className="lynxkite-node" style={props.nodeStyle}>
         <div className="title bg-primary" onClick={titleClicked}>
           {data.title}
@@ -67,14 +64,6 @@ export default function LynxKiteNode(props: LynxKiteNodeProps) {
             <div className="error">{data.error}</div>
           }
           {props.children}
-          {handles.map(handle => (
-            <Handle
-              key={handle.name}
-              id={handle.name} type={handle.type} position={handle.position as Position}
-              style={{ [handleOffsetDirection[handle.position]]: handle.offsetPercentage + '%' }}>
-              {handle.showLabel && <span className="handle-name">{handle.name.replace(/_/g, " ")}</span>}
-            </Handle >
-          ))}
           <NodeResizeControl
             minWidth={100}
             minHeight={50}
@@ -83,6 +72,14 @@ export default function LynxKiteNode(props: LynxKiteNodeProps) {
             <ChevronDownRight className="node-resizer" />
           </NodeResizeControl>
         </>}
+        {handles.map(handle => (
+          <Handle
+            key={handle.name}
+            id={handle.name} type={handle.type} position={handle.position as Position}
+            style={{ [handleOffsetDirection[handle.position]]: handle.offsetPercentage + '%' }}>
+            {handle.showLabel && <span className="handle-name">{handle.name.replace(/_/g, " ")}</span>}
+          </Handle >
+        ))}
       </div>
     </div>
   );

web/src/workspace/nodes/NodeWithTableView.tsx

@@ -22,12 +22,12 @@ export default function NodeWithTableView(props: any) {
     <LynxKiteNode {...props}>
       {display && [
         Object.entries(display.dataframes || {}).map(([name, df]: [string, any]) => <React.Fragment key={name}>
-          {!single && <div key={name} className="df-head" onClick={() => setOpen({ ...open, [name]: !open[name] })}>{name}</div>}
+          {!single && <div key={name + '-header'} className="df-head" onClick={() => setOpen({ ...open, [name]: !open[name] })}>{name}</div>}
           {(single || open[name]) &&
             (df.data.length > 1 ?
-              <Table key={name} columns={df.columns} data={df.data} />
+              <Table key={name + '-table'} columns={df.columns} data={df.data} />
               :
-              <dl key={name}>
+              <dl key={name + '-dl'}>
                 {df.columns.map((c: string, i: number) =>
                   <React.Fragment key={name + '-' + c}>
                     <dt>{c}</dt>