ursa_engine/graph/
schema.rs

1//! A graph as a schema: `<graph>.edges` and `<graph>.nodes`.
2//!
3//! Each graph registers a [`GraphSchema`] in the default catalog under its
4//! name. Its two tables are views ([`ViewTable`]s) whose plans are built from
5//! the graph's implied pair ([`PropertyGraph::default_pair`]: its edge table
6//! and the vertex table that edge table refers to) every time a query names
7//! them, so they always read the current contents of the tables (and files)
8//! the graph is defined over. A graph with several vertex or edge tables has no
9//! implied pair, and so no `nodes` or `edges` table: its element tables are
10//! session tables already, and `MATCH` reads across them.
11//!
12//! - `edges` is the edge table: every row and column, unchanged.
13//! - `nodes` is one row per node, the node id in a column named `id`:
14//!   - without a vertex table, the distinct non-null endpoints (`source` and
15//!     `destination` values) of the edges;
16//!   - with one, every row of the vertex table (its key renamed to `id`,
17//!     followed by its other columns), and nothing else: the vertex table
18//!     defines the node set, as in SQL/PGQ and DuckPGQ, and an edge whose
19//!     source or destination is not one of its keys is dropped by the
20//!     algorithms (#311). When the key and the endpoints have different types
21//!     (`Int32` and `Int64`, `Utf8` and a dictionary), `id` has their common
22//!     type, a dictionary as its values.
23//!
24//! The plan builders ([`edges_plan`], [`nodes_plan`] and their `_from`
25//! variants, which take table sources already resolved) are public so the graph
26//! algorithms can reuse them.
27//!
28//! When a table under the graph cannot be read (it was deregistered through
29//! the DataFusion API, say), `edges` and `nodes` are still listed, with the
30//! columns they last had, and fail with the reason when a query reads them.
31//! Session listings (`SHOW TABLES`, `information_schema`) look every table up,
32//! so an error here would break them for the whole session.
33
34use std::collections::HashMap;
35use std::sync::{Arc, Mutex, Weak};
36
37use async_trait::async_trait;
38use datafusion::arrow::datatypes::{DataType, Schema, SchemaRef};
39use datafusion::catalog::{SchemaProvider, Session, TableProvider};
40use datafusion::common::{Column, TableReference};
41use datafusion::datasource::{provider_as_source, TableType, ViewTable};
42use datafusion::error::{DataFusionError, Result};
43use datafusion::execution::context::SessionState;
44use datafusion::logical_expr::expr_fn::cast;
45use datafusion::logical_expr::type_coercion::binary::comparison_coercion;
46use datafusion::logical_expr::{col, Expr, LogicalPlan, LogicalPlanBuilder, TableSource};
47use datafusion::physical_plan::ExecutionPlan;
48use parking_lot::RwLock;
49
50use super::{registry, GraphDef, PropertyGraph};
51
52/// The name of a graph's edge table in its schema.
53pub const EDGES: &str = "edges";
54/// The name of a graph's node table in its schema.
55pub const NODES: &str = "nodes";
56/// The node id column of `<graph>.nodes`.
57pub const ID: &str = "id";
58
59/// The schema of one graph: its `edges` and `nodes` tables.
60#[derive(Debug)]
61pub struct GraphSchema {
62    graph: String,
63    session: Weak<RwLock<SessionState>>,
64    /// The columns `edges` and `nodes` had when last planned.
65    last_columns: Mutex<HashMap<&'static str, SchemaRef>>,
66}
67
68impl GraphSchema {
69    pub(crate) fn new(graph: String, session: Weak<RwLock<SessionState>>) -> Self {
70        Self {
71            graph,
72            session,
73            last_columns: Mutex::default(),
74        }
75    }
76
77    /// The graph, if the session and the graph still exist.
78    fn graph(&self) -> Option<(SessionState, std::sync::Arc<PropertyGraph>)> {
79        let state = self.session.upgrade()?.read().clone();
80        let graph = registry(&state).get(&self.graph)?;
81        Some((state, graph))
82    }
83
84    /// Whether the graph has `nodes` and `edges` tables (an implied pair).
85    fn has_tables(&self) -> bool {
86        self.graph()
87            .is_some_and(|(_, graph)| graph.default_pair().is_some())
88    }
89}
90
91#[async_trait]
92impl SchemaProvider for GraphSchema {
93    fn table_names(&self) -> Vec<String> {
94        if self.has_tables() {
95            vec![EDGES.to_string(), NODES.to_string()]
96        } else {
97            Vec::new()
98        }
99    }
100
101    async fn table(&self, name: &str) -> Result<Option<Arc<dyn TableProvider>>> {
102        let name = match name {
103            EDGES => EDGES,
104            NODES => NODES,
105            _ => return Ok(None),
106        };
107        let state = self
108            .session
109            .upgrade()
110            .map(|state| state.read().clone())
111            .ok_or_else(|| DataFusionError::Internal("the session has been dropped".into()))?;
112        let Some(def) = registry(&state)
113            .get(&self.graph)
114            .and_then(|graph| graph.default_pair())
115        else {
116            return Ok(None);
117        };
118        let plan = if name == EDGES {
119            edges_plan(&state, &def).await
120        } else {
121            nodes_plan(&state, &def).await
122        };
123        let mut last_columns = self.last_columns.lock().unwrap_or_else(|e| e.into_inner());
124        match plan {
125            Ok(plan) => {
126                let view = ViewTable::new(plan, None);
127                last_columns.insert(name, view.schema());
128                Ok(Some(Arc::new(view)))
129            }
130            Err(e) => Ok(Some(Arc::new(UnreadableTable {
131                schema: last_columns
132                    .get(name)
133                    .cloned()
134                    .unwrap_or_else(|| Arc::new(Schema::empty())),
135                message: format!(
136                    "cannot read {}.{name}: {}",
137                    def.name,
138                    crate::user_message(&e)
139                ),
140            }))),
141        }
142    }
143
144    fn table_exist(&self, name: &str) -> bool {
145        (name == EDGES || name == NODES) && self.has_tables()
146    }
147}
148
149/// A graph table whose plan cannot be built: listed with the columns it last
150/// had, and failing with `message` when a query reads it.
151#[derive(Debug)]
152struct UnreadableTable {
153    schema: SchemaRef,
154    message: String,
155}
156
157#[async_trait]
158impl TableProvider for UnreadableTable {
159    fn schema(&self) -> SchemaRef {
160        Arc::clone(&self.schema)
161    }
162
163    fn table_type(&self) -> TableType {
164        TableType::View
165    }
166
167    async fn scan(
168        &self,
169        _state: &dyn Session,
170        _projection: Option<&Vec<usize>>,
171        _filters: &[Expr],
172        _limit: Option<usize>,
173    ) -> Result<Arc<dyn ExecutionPlan>> {
174        Err(DataFusionError::Plan(self.message.clone()))
175    }
176}
177
178/// The table `table` names for the graph called `graph`: one the graph opened
179/// from a file path when it was created, or else the table or view in `state`.
180pub async fn graph_table(
181    state: &SessionState,
182    graph: &str,
183    table: &TableReference,
184) -> Result<Arc<dyn TableProvider>> {
185    match registry(state).file_table(graph, table) {
186        Some(provider) => Ok(provider),
187        None => resolve_table(state, table).await,
188    }
189}
190
191/// The table provider `table` names in `state`, or an error saying it does not
192/// exist.
193pub async fn resolve_table(
194    state: &SessionState,
195    table: &TableReference,
196) -> Result<Arc<dyn TableProvider>> {
197    lookup_table(state, table)
198        .await?
199        .ok_or_else(|| DataFusionError::Plan(format!("table '{table}' not found")))
200}
201
202/// The table provider `table` names in `state`, if there is one.
203pub async fn lookup_table(
204    state: &SessionState,
205    table: &TableReference,
206) -> Result<Option<Arc<dyn TableProvider>>> {
207    match state.schema_for_ref(table.clone()) {
208        Ok(schema) => schema.table(table.table()).await,
209        Err(_) => Ok(None),
210    }
211}
212
213/// The plan of `<graph>.edges`: a scan of the edge table, as it is now.
214pub async fn edges_plan(state: &SessionState, def: &GraphDef) -> Result<LogicalPlan> {
215    let edges = graph_table(state, &def.name, &def.edges.table).await?;
216    edges_plan_from(def, provider_as_source(edges))
217}
218
219/// [`edges_plan`], with the edge table's source already resolved.
220pub fn edges_plan_from(def: &GraphDef, edges: Arc<dyn TableSource>) -> Result<LogicalPlan> {
221    LogicalPlanBuilder::scan(def.edges.table.clone(), edges, None)?.build()
222}
223
224/// The plan of `<graph>.nodes` (see the module docs).
225pub async fn nodes_plan(state: &SessionState, def: &GraphDef) -> Result<LogicalPlan> {
226    let edges = provider_as_source(graph_table(state, &def.name, &def.edges.table).await?);
227    let vertices = match &def.vertices {
228        Some(v) => Some(provider_as_source(
229            graph_table(state, &def.name, &v.table).await?,
230        )),
231        None => None,
232    };
233    nodes_plan_from(def, edges, vertices)
234}
235
236/// [`nodes_plan`], with the tables' sources already resolved. `vertices` must
237/// be given when the graph has a vertex table.
238pub fn nodes_plan_from(
239    def: &GraphDef,
240    edges: Arc<dyn TableSource>,
241    vertices: Option<Arc<dyn TableSource>>,
242) -> Result<LogicalPlan> {
243    let endpoint_type = column_type(&edges.schema(), &def.edges.table, &def.edges.source)?;
244    let (vertex_def, vertices) = match (&def.vertices, vertices) {
245        (None, _) => {
246            return endpoints(def, &edges)?.build();
247        }
248        (Some(v), Some(source)) => (v, source),
249        (Some(v), None) => {
250            return Err(DataFusionError::Internal(format!(
251                "nodes_plan_from: graph {} has vertex table {} but no source was given for it",
252                def.name, v.table
253            )))
254        }
255    };
256
257    let vertex_schema = vertices.schema();
258    let key_type = column_type(&vertex_schema, &vertex_def.table, &vertex_def.key)?;
259    let id_type = comparison_coercion(&key_type, &endpoint_type).ok_or_else(|| {
260        DataFusionError::Plan(format!(
261            "the key of vertex table '{}' is {key_type} and the edge endpoints are \
262             {endpoint_type}; they cannot be compared",
263            vertex_def.table
264        ))
265    })?;
266    // The key and the endpoints differ (say Int32 and Int64): `id` takes their
267    // common type, a dictionary as its values.
268    let id_type = match (key_type == endpoint_type, id_type) {
269        (true, _) => key_type.clone(),
270        (false, DataType::Dictionary(_, values)) => *values,
271        (false, other) => other,
272    };
273
274    // The vertex table: `key AS id`, then its other columns.
275    let others: Vec<_> = vertex_schema
276        .fields()
277        .iter()
278        .filter(|f| f.name() != &vertex_def.key)
279        .cloned()
280        .collect();
281    let mut vertex_exprs = vec![cast_to(
282        Expr::Column(Column::new(Some(vertex_def.table.clone()), &vertex_def.key)),
283        &key_type,
284        &id_type,
285    )
286    .alias(ID)];
287    vertex_exprs.extend(
288        others
289            .iter()
290            .map(|f| Expr::Column(Column::new(Some(vertex_def.table.clone()), f.name()))),
291    );
292    // Only the vertex table's rows: an endpoint that is not one of its keys
293    // is no node, and its edges are dropped (#311).
294    LogicalPlanBuilder::scan(vertex_def.table.clone(), vertices, None)?
295        .project(vertex_exprs)?
296        .build()
297}
298
299/// The distinct non-null endpoints of the edges, as one column `id`.
300fn endpoints(def: &GraphDef, edges: &Arc<dyn TableSource>) -> Result<LogicalPlanBuilder> {
301    let side = |column: &str| -> Result<LogicalPlan> {
302        LogicalPlanBuilder::scan(def.edges.table.clone(), Arc::clone(edges), None)?
303            .project(vec![Expr::Column(Column::new(
304                Some(def.edges.table.clone()),
305                column,
306            ))
307            .alias(ID)])?
308            .build()
309    };
310    LogicalPlanBuilder::from(side(&def.edges.source)?)
311        .union_distinct(side(&def.edges.destination)?)?
312        .filter(col(ID).is_not_null())
313}
314
315fn cast_to(expr: Expr, from: &DataType, to: &DataType) -> Expr {
316    if from == to {
317        expr
318    } else {
319        cast(expr, to.clone())
320    }
321}
322
323fn column_type(schema: &Schema, table: &TableReference, column: &str) -> Result<DataType> {
324    schema
325        .field_with_name(column)
326        .map(|f| f.data_type().clone())
327        .map_err(|_| {
328            DataFusionError::Plan(format!("column '{column}' not found in table '{table}'"))
329        })
330}
331
332#[cfg(test)]
333mod tests {
334    use datafusion::arrow::array::{DictionaryArray, Int64Array, RecordBatch, StringArray};
335    use datafusion::arrow::datatypes::Int32Type;
336    use datafusion::datasource::MemTable;
337
338    use super::*;
339    use crate::graph::{create_graph, drop_graph, EdgeTable, VertexTable};
340
341    fn table(batch: RecordBatch) -> Arc<MemTable> {
342        Arc::new(MemTable::try_new(batch.schema(), vec![vec![batch]]).unwrap())
343    }
344
345    #[tokio::test]
346    async fn dictionary_ids_and_a_vertex_table_without_every_node() {
347        let ctx = crate::new_session();
348        let ids = |values: Vec<&str>| -> Arc<dyn datafusion::arrow::array::Array> {
349            Arc::new(values.into_iter().collect::<DictionaryArray<Int32Type>>())
350        };
351        let edges = RecordBatch::try_from_iter(vec![
352            ("a", ids(vec!["x", "y", "z", "z"])),
353            ("b", ids(vec!["y", "z", "x", "w"])),
354        ])
355        .unwrap();
356        let people = RecordBatch::try_from_iter(vec![
357            (
358                "name",
359                Arc::new(StringArray::from(vec!["x", "y", "z"])) as _,
360            ),
361            ("age", Arc::new(Int64Array::from(vec![30, 40, 50])) as _),
362        ])
363        .unwrap();
364        ctx.register_table("follows", table(edges)).unwrap();
365        ctx.register_table("people", table(people)).unwrap();
366
367        let def = GraphDef {
368            name: "social".into(),
369            edges: EdgeTable::new(TableReference::bare("follows"), "a", "b"),
370            vertices: Some(VertexTable::new(TableReference::bare("people"), "name")),
371        };
372        create_graph(&ctx, def.clone(), false, false).await.unwrap();
373
374        // The plans the algorithms can reuse.
375        let state = ctx.state();
376        let edges = edges_plan(&state, &def).await.unwrap();
377        assert_eq!(edges.schema().fields().len(), 2);
378        let nodes = nodes_plan(&state, &def).await.unwrap();
379        let columns: Vec<_> = nodes
380            .schema()
381            .fields()
382            .iter()
383            .map(|f| (f.name().clone(), f.data_type().clone()))
384            .collect();
385        assert_eq!(
386            columns,
387            [
388                (ID.to_string(), DataType::Utf8),
389                ("age".to_string(), DataType::Int64)
390            ]
391        );
392
393        // The vertex table's rows only: 'w', an endpoint that is not a key,
394        // is no node (#311).
395        let batches = ctx
396            .sql("SELECT id, age FROM social.nodes ORDER BY id")
397            .await
398            .unwrap()
399            .collect()
400            .await
401            .unwrap();
402        let formatted = datafusion::arrow::util::pretty::pretty_format_batches(&batches)
403            .unwrap()
404            .to_string();
405        assert_eq!(
406            formatted,
407            "\
408+----+-----+
409| id | age |
410+----+-----+
411| x  | 30  |
412| y  | 40  |
413| z  | 50  |
414+----+-----+"
415        );
416
417        drop_graph(&ctx, "social", false).unwrap();
418        assert!(ctx.sql("SELECT * FROM social.nodes").await.is_err());
419    }
420}
ursa_engine/graph/schema.rs

ursa_engine/graph/
schema.rs
