ursa_engine/graph/
describe.rs

1//! `SHOW GRAPHS` and `DESCRIBE GRAPH`: what graphs a session has, and what
2//! one of them is.
3
4use std::sync::Arc;
5
6use datafusion::arrow::array::{ArrayRef, RecordBatch, StringArray};
7use datafusion::arrow::datatypes::{DataType, Field, Schema};
8use datafusion::common::TableReference;
9use datafusion::error::{DataFusionError, Result};
10use datafusion::execution::context::SessionState;
11
12use super::define::unknown_graph;
13use super::schema::{edges_plan, graph_table, nodes_plan, EDGES, ID, NODES};
14use super::{registry, GraphRegistry, PropertyGraph};
15
16/// `SHOW GRAPHS`: one row per edge table of each graph (one row for a graph
17/// with no edge tables), sorted by graph name, with its definition.
18///
19/// | column         | value                                                     |
20/// | -------------- | --------------------------------------------------------- |
21/// | `name`         | the graph's name                                          |
22/// | `edge_table`   | the edge table (a file path in single quotes), or NULL    |
23/// | `source`       | its `SOURCE KEY` column                                   |
24/// | `destination`  | its `DESTINATION KEY` column                              |
25/// | `vertex_table` | the vertex table its endpoints refer to, or NULL          |
26/// | `vertex_key`   | that vertex table's key column, or NULL                   |
27///
28/// A graph with one edge table and at most one vertex table, as every graph
29/// before #295, has one row, as before.
30pub fn show_graphs(state: &SessionState) -> Result<RecordBatch> {
31    let graphs = registry(state);
32    let mut rows: Vec<[Option<String>; 6]> = Vec::new();
33    for g in graphs.all() {
34        if g.edge_tables.is_empty() {
35            for v in &g.vertex_tables {
36                rows.push([
37                    Some(g.name.clone()),
38                    None,
39                    None,
40                    None,
41                    Some(table_name(&graphs, &g, &v.table)),
42                    Some(v.key.clone()),
43                ]);
44            }
45            continue;
46        }
47        for e in &g.edge_tables {
48            let vertex = g
49                .endpoint_vertex(e, true)
50                .or_else(|| g.endpoint_vertex(e, false));
51            rows.push([
52                Some(g.name.clone()),
53                Some(table_name(&graphs, &g, &e.table)),
54                Some(e.source.clone()),
55                Some(e.destination.clone()),
56                vertex.map(|v| table_name(&graphs, &g, &v.table)),
57                vertex.map(|v| v.key.clone()),
58            ]);
59        }
60    }
61    let columns = (0..6)
62        .map(|i| Arc::new(rows.iter().map(|r| r[i].clone()).collect::<StringArray>()) as ArrayRef)
63        .collect();
64    let schema = Schema::new(vec![
65        Field::new("name", DataType::Utf8, false),
66        Field::new("edge_table", DataType::Utf8, true),
67        Field::new("source", DataType::Utf8, true),
68        Field::new("destination", DataType::Utf8, true),
69        Field::new("vertex_table", DataType::Utf8, true),
70        Field::new("vertex_key", DataType::Utf8, true),
71    ]);
72    Ok(RecordBatch::try_new(Arc::new(schema), columns)?)
73}
74
75/// `DESCRIBE GRAPH name`: one row per column of `<name>.nodes`, then of
76/// `<name>.edges`, saying what each is and where it comes from.
77///
78/// | column        | value                                                  |
79/// | ------------- | ------------------------------------------------------ |
80/// | `table_name`  | `nodes` or `edges`                                     |
81/// | `column_name` | the column                                             |
82/// | `data_type`   | its type (for `nodes.id`, the graph's node id type)    |
83/// | `role`        | `node id`, `source` or `destination`; NULL for others  |
84/// | `origin`      | the table column it comes from (`ties.src`)            |
85///
86/// `origin` of `nodes.id` is the vertex table's key, or without a vertex
87/// table, the edge endpoints (`ties.src, ties.dst`). So the rows hold the whole
88/// definition, and every row has the same shape.
89///
90/// A graph with several vertex or edge tables has no `nodes` and `edges`
91/// tables: its rows are each element table's columns instead, `table_name`
92/// being the element name and `role` one of `key`, `source` (`→ vertex`) and
93/// `destination` (`→ vertex`).
94pub async fn describe_graph(state: &SessionState, name: &str) -> Result<RecordBatch> {
95    let graphs = registry(state);
96    let graph = graphs
97        .get(name)
98        .ok_or_else(|| unknown_graph(name, &graphs))?;
99    let Some(def) = graph.default_pair() else {
100        return describe_elements(state, &graphs, &graph).await;
101    };
102    let edge_table = table_name(&graphs, &graph, &def.edges.table);
103    let vertex_table = def
104        .vertices
105        .as_ref()
106        .map(|v| (table_name(&graphs, &graph, &v.table), &v.key));
107
108    let unreadable = |e: DataFusionError| {
109        DataFusionError::Plan(format!(
110            "cannot describe graph '{name}': {}",
111            crate::user_message(&e)
112        ))
113    };
114    let mut rows: Vec<[Option<String>; 5]> = Vec::new();
115    let nodes = nodes_plan(state, &def).await.map_err(unreadable)?;
116    for field in nodes.schema().fields() {
117        let column = field.name();
118        let (role, origin) = match (&vertex_table, column == ID) {
119            (Some((table, key)), true) => (Some("node id"), format!("{table}.{key}")),
120            (Some((table, _)), false) => (None, format!("{table}.{column}")),
121            (None, _) => (
122                Some("node id"),
123                format!(
124                    "{edge_table}.{}, {edge_table}.{}",
125                    def.edges.source, def.edges.destination
126                ),
127            ),
128        };
129        rows.push(row(NODES, column, field.data_type(), role, origin));
130    }
131    let edges = edges_plan(state, &def).await.map_err(unreadable)?;
132    for field in edges.schema().fields() {
133        let column = field.name();
134        let role = match (
135            column == &def.edges.source,
136            column == &def.edges.destination,
137        ) {
138            (true, true) => Some("source, destination"),
139            (true, false) => Some("source"),
140            (false, true) => Some("destination"),
141            (false, false) => None,
142        };
143        rows.push(row(
144            EDGES,
145            column,
146            field.data_type(),
147            role,
148            format!("{edge_table}.{column}"),
149        ));
150    }
151
152    describe_batch(rows)
153}
154
155/// [`describe_graph`] for a graph with several element tables.
156async fn describe_elements(
157    state: &SessionState,
158    graphs: &GraphRegistry,
159    graph: &PropertyGraph,
160) -> Result<RecordBatch> {
161    let unreadable = |e: DataFusionError| {
162        DataFusionError::Plan(format!(
163            "cannot describe graph '{}': {}",
164            graph.name,
165            crate::user_message(&e)
166        ))
167    };
168    let mut rows: Vec<[Option<String>; 5]> = Vec::new();
169    for v in &graph.vertex_tables {
170        let schema = graph_table(state, &graph.name, &v.table)
171            .await
172            .map_err(unreadable)?
173            .schema();
174        let table = table_name(graphs, graph, &v.table);
175        for field in schema.fields() {
176            let role = (field.name() == &v.key).then_some("key");
177            rows.push(row(
178                &v.name,
179                field.name(),
180                field.data_type(),
181                role,
182                format!("{table}.{}", field.name()),
183            ));
184        }
185    }
186    for e in &graph.edge_tables {
187        let schema = graph_table(state, &graph.name, &e.table)
188            .await
189            .map_err(unreadable)?
190            .schema();
191        let table = table_name(graphs, graph, &e.table);
192        let target = |source: bool| {
193            graph
194                .endpoint_vertex(e, source)
195                .map(|v| format!(" \u{2192} {}", v.name))
196                .unwrap_or_default()
197        };
198        let (source, destination) = (
199            format!("source{}", target(true)),
200            format!("destination{}", target(false)),
201        );
202        for field in schema.fields() {
203            let column = field.name();
204            let both = format!("{source}, {destination}");
205            let role = match (column == &e.source, column == &e.destination) {
206                (true, true) => Some(both.as_str()),
207                (true, false) => Some(source.as_str()),
208                (false, true) => Some(destination.as_str()),
209                (false, false) => None,
210            };
211            rows.push(row(
212                &e.name,
213                column,
214                field.data_type(),
215                role,
216                format!("{table}.{column}"),
217            ));
218        }
219    }
220    describe_batch(rows)
221}
222
223fn describe_batch(rows: Vec<[Option<String>; 5]>) -> Result<RecordBatch> {
224    let columns = (0..5)
225        .map(|i| Arc::new(rows.iter().map(|r| r[i].clone()).collect::<StringArray>()) as ArrayRef)
226        .collect();
227    let schema = Schema::new(vec![
228        Field::new("table_name", DataType::Utf8, false),
229        Field::new("column_name", DataType::Utf8, false),
230        Field::new("data_type", DataType::Utf8, false),
231        Field::new("role", DataType::Utf8, true),
232        Field::new("origin", DataType::Utf8, false),
233    ]);
234    Ok(RecordBatch::try_new(Arc::new(schema), columns)?)
235}
236
237fn row(
238    table: &str,
239    column: &str,
240    data_type: &DataType,
241    role: Option<&str>,
242    origin: String,
243) -> [Option<String>; 5] {
244    [
245        Some(table.to_string()),
246        Some(column.to_string()),
247        Some(data_type.to_string()),
248        role.map(str::to_string),
249        Some(origin),
250    ]
251}
252
253/// `table` as the definition names it; a file path in single quotes, as it is
254/// written in SQL.
255fn table_name(graphs: &GraphRegistry, graph: &PropertyGraph, table: &TableReference) -> String {
256    if graphs.file_table(&graph.name, table).is_some() {
257        format!("'{}'", table.table().replace('\'', "''"))
258    } else {
259        table.to_string()
260    }
261}
ursa_engine/graph/describe.rs

ursa_engine/graph/
describe.rs
