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

ursa_engine/graph/
schema.rs
