use std::path::Path;
use std::time::Duration;
use datafusion::arrow::util::display::{ArrayFormatter, FormatOptions};
use datafusion::prelude::SessionContext;
use ursa_engine::database::OpenOptions;
use ursa_engine::graph::{index_builds, index_loads};
use ursa_sql::{execute_with_text, open_session, parse_sql_with_text_for};
async fn run(ctx: &SessionContext, sql: &str) -> datafusion::error::Result<Vec<String>> {
let mut lines = Vec::new();
for (statement, text) in parse_sql_with_text_for(ctx, sql)? {
let batches = execute_with_text(ctx, statement, &text)
.await?
.collect()
.await?;
lines.clear();
let options = FormatOptions::default().with_null("NULL");
for batch in &batches {
let formatters: Vec<_> = batch
.columns()
.iter()
.map(|c| ArrayFormatter::try_new(c.as_ref(), &options).unwrap())
.collect();
for row in 0..batch.num_rows() {
let cells: Vec<_> = formatters
.iter()
.map(|f| f.value(row).to_string())
.collect();
lines.push(cells.join(","));
}
}
}
Ok(lines)
}
async fn ok(ctx: &SessionContext, sql: &str) -> Vec<String> {
run(ctx, sql)
.await
.unwrap_or_else(|e| panic!("{sql}: {}", ursa_engine::user_message(&e)))
}
const PEOPLE: &str = "
CREATE VIEW people AS SELECT * FROM (VALUES
(1, 'ann'), (2, 'bob'), (3, 'cy'), (4, 'dee'), (5, 'eve')) t(id, name);
CREATE VIEW follows AS SELECT * FROM (VALUES
(1, 2, 1.0), (9, 1, 5.0), (2, 3, 2.0), (3, 1, 3.0), (3, 9, 4.0)) t(src, dst, w);
CREATE PROPERTY GRAPH g
VERTEX TABLES (people KEY (id))
EDGE TABLES (follows SOURCE KEY (src) REFERENCES people (id)
DESTINATION KEY (dst) REFERENCES people (id));";
async fn people() -> SessionContext {
let ctx = ursa_engine::new_session();
ok(&ctx, PEOPLE).await;
ctx
}
#[tokio::test]
async fn every_algorithm_sees_the_vertex_tables_nodes() {
let ctx = people().await;
let nodes = ["1", "2", "3", "4", "5"];
assert_eq!(ok(&ctx, "SELECT id FROM g.nodes ORDER BY id").await, nodes);
assert_eq!(
ok(
&ctx,
"SELECT * FROM degree(g, direction => 'both') ORDER BY id"
)
.await,
["1,2", "2,2", "3,2", "4,0", "5,0"]
);
for f in [
"pagerank(g)",
"connected_components(g)",
"triangle_count(g)",
"closeness(g)",
"betweenness(g)",
"label_propagation(g)",
"louvain(g, seed => 1)",
"clustering_coefficient(g)",
"algorithms(g, degree(), pagerank())",
"neighbors_agg(g, attribute => length(name), agg => 'sum')",
] {
assert_eq!(
ok(&ctx, &format!("SELECT id FROM {f} ORDER BY id")).await,
nodes,
"{f}"
);
}
assert_eq!(
ok(
&ctx,
"SELECT count(DISTINCT connected_components) FROM connected_components(g)"
)
.await,
["3"]
);
assert_eq!(
ok(
&ctx,
"SELECT id, neighbors_agg FROM neighbors_agg(g, attribute => id, agg => 'sum', \
direction => 'both') WHERE id >= 4 ORDER BY id"
)
.await,
["4,0.0", "5,0.0"]
);
}
#[tokio::test]
async fn edges_to_or_from_no_vertex_are_dropped_and_counted() {
let ctx = people().await;
assert_eq!(ok(&ctx, "SELECT count(*) FROM g.edges").await, ["5"]);
assert_eq!(
ok(&ctx, "SELECT * FROM describe(g, full => true)").await,
["5,3,0.15,0.6,3,2"]
);
assert_eq!(
ok(
&ctx,
"SELECT n_nodes, n_edges, n_components, n_dropped_edges \
FROM describe(g, full => true, nodes => 'touched')"
)
.await,
["3,3,1,2"]
);
assert_eq!(
ok(
&ctx,
"SELECT src, dst, cost FROM shortest_path(g, source => 1, target => 3, weight => w)"
)
.await,
["1,2,1.0", "2,3,3.0"]
);
assert_eq!(
ok(
&ctx,
"SELECT a, b, edges FROM GRAPH_TABLE (g MATCH p = (x)-[e]->{2}(y) \
COLUMNS (x.id AS a, y.id AS b, edges(p) AS edges)) ORDER BY a"
)
.await,
["1,3,[0, 2]", "2,1,[2, 3]", "3,2,[3, 0]"]
);
ok(
&ctx,
"CREATE PROPERTY GRAPH bare EDGE TABLES (follows SOURCE KEY (src) DESTINATION KEY (dst))",
)
.await;
assert_eq!(
ok(
&ctx,
"SELECT n_nodes, n_edges, n_dropped_edges FROM describe(bare)"
)
.await,
["4,5,0"]
);
}
#[tokio::test]
async fn a_subgraph_keeps_every_vertex_unless_touched() {
let ctx = people().await;
ok(
&ctx,
"CREATE VIEW first AS SELECT * FROM follows WHERE src = 1;
CREATE PROPERTY GRAPH sub
VERTEX TABLES (people KEY (id))
EDGE TABLES (first SOURCE KEY (src) REFERENCES people (id)
DESTINATION KEY (dst) REFERENCES people (id));",
)
.await;
assert_eq!(
ok(&ctx, "SELECT n_nodes, n_edges FROM describe(sub)").await,
["5,1"]
);
assert_eq!(
ok(
&ctx,
"SELECT n_nodes, n_edges FROM describe(sub, nodes => 'touched')"
)
.await,
["2,1"]
);
assert_eq!(ok(&ctx, "SELECT count(*) FROM degree(sub)").await, ["5"]);
}
#[tokio::test]
async fn a_node_on_no_edge_reaches_nothing() {
let ctx = people().await;
assert_eq!(
ok(
&ctx,
"SELECT * FROM hop(g, seeds => [4], direction => 'both')"
)
.await,
Vec::<String>::new()
);
assert_eq!(
ok(
&ctx,
"SELECT * FROM shortest_path(g, source => 4, target => 1)"
)
.await,
Vec::<String>::new()
);
assert_eq!(
ok(
&ctx,
"SELECT * FROM random_walk(g, start => [5], steps => 3, seed => 1)"
)
.await,
["0,0,5"]
);
assert_eq!(
ok(
&ctx,
"SELECT * FROM hop(g, seeds => [9], direction => 'both')"
)
.await,
Vec::<String>::new()
);
}
#[tokio::test]
async fn graph_table_matches_the_same_nodes() {
let ctx = people().await;
assert_eq!(
ok(
&ctx,
"SELECT id FROM GRAPH_TABLE (g MATCH (a) COLUMNS (a.id AS id)) ORDER BY id"
)
.await,
ok(&ctx, "SELECT id FROM g.nodes ORDER BY id").await
);
assert_eq!(
ok(
&ctx,
"SELECT a, b FROM GRAPH_TABLE (g MATCH (x)-[e]->(y) COLUMNS (x.id AS a, y.id AS b)) \
ORDER BY a"
)
.await,
["1,2", "2,3", "3,1"]
);
assert_eq!(
ok(
&ctx,
"SELECT a, b FROM GRAPH_TABLE (g MATCH (x)-[e]->{0,1}(y) \
COLUMNS (x.id AS a, y.id AS b)) ORDER BY a, b"
)
.await,
["1,1", "1,2", "2,2", "2,3", "3,1", "3,3", "4,4", "5,5"]
);
}
#[tokio::test]
async fn a_multi_table_graph_runs_on_the_named_vertex_table() {
let ctx = people().await;
ok(
&ctx,
"CREATE VIEW teams AS SELECT * FROM (VALUES (10), (11)) t(id);
CREATE VIEW plays AS SELECT * FROM (VALUES (1, 10), (2, 10), (2, 12)) t(person, team);
CREATE PROPERTY GRAPH club
VERTEX TABLES (people KEY (id), teams KEY (id))
EDGE TABLES (
follows SOURCE KEY (src) REFERENCES people (id)
DESTINATION KEY (dst) REFERENCES people (id) LABEL knows,
plays SOURCE KEY (person) REFERENCES people (id)
DESTINATION KEY (team) REFERENCES teams (id));",
)
.await;
assert_eq!(
ok(
&ctx,
"SELECT id, degree FROM degree(club, people, knows, direction => 'both') ORDER BY id"
)
.await,
["1,2", "2,2", "3,2", "4,0", "5,0"]
);
assert_eq!(
ok(
&ctx,
"SELECT n_nodes, n_edges, n_dropped_edges FROM describe(club, people, knows)"
)
.await,
["5,3,2"]
);
assert_eq!(
ok(
&ctx,
"SELECT a, b, edges FROM GRAPH_TABLE (club MATCH p = (x:people)-[e:plays]->(y:teams) \
COLUMNS (x.id AS a, y.id AS b, edges(p) AS edges)) ORDER BY a, b"
)
.await,
["1,10,[0]", "2,10,[1]"]
);
assert_eq!(
ok(
&ctx,
"SELECT count(*) FROM GRAPH_TABLE (club MATCH (x:people)-[e:knows]->{0,0}(y:people) \
COLUMNS (x.id AS a))"
)
.await,
["5"]
);
}
#[tokio::test]
async fn a_changed_vertex_file_rebuilds_the_stored_index() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("club.ursa");
let people = dir.path().join("people.csv");
let edges = dir.path().join("edges.csv");
std::fs::write(&edges, "src,dst\n1,2\n2,3\n3,1\n3,9\n").unwrap();
std::fs::write(&people, "id\n1\n2\n3\n").unwrap();
let options = OpenOptions {
lock_timeout: Duration::from_millis(500),
..OpenOptions::default()
};
let define = format!(
"CREATE VIEW people AS SELECT * FROM '{}';
CREATE VIEW ties AS SELECT * FROM '{}';
CREATE PROPERTY GRAPH g VERTEX TABLES (people KEY (id))
EDGE TABLES (ties SOURCE KEY (src) REFERENCES people (id)
DESTINATION KEY (dst) REFERENCES people (id));",
people.display(),
edges.display()
);
let query = "SELECT n_nodes, n_edges, n_dropped_edges FROM describe(g)";
let open = |path: &Path| {
let path = path.to_path_buf();
let options = options.clone();
async move { open_session(&path, &options).await.unwrap() }
};
let ctx = open(&path).await;
ok(&ctx, &define).await;
assert_eq!(ok(&ctx, query).await, ["3,3,1"]);
assert_eq!((index_builds(&ctx), index_loads(&ctx)), (1, 0));
drop(ctx);
let ctx = open(&path).await;
assert_eq!(ok(&ctx, query).await, ["3,3,1"]);
assert_eq!((index_builds(&ctx), index_loads(&ctx)), (0, 1));
std::fs::write(&people, "id\n1\n2\n3\n9\n4\n").unwrap();
assert_eq!(ok(&ctx, query).await, ["5,4,0"]);
assert_eq!((index_builds(&ctx), index_loads(&ctx)), (1, 1));
drop(ctx);
let ctx = open(&path).await;
assert_eq!(ok(&ctx, query).await, ["5,4,0"]);
assert_eq!(
ok(&ctx, "SELECT id, degree FROM degree(g) ORDER BY id").await,
["1,1", "2,1", "3,2", "4,0", "9,0"]
);
assert_eq!((index_builds(&ctx), index_loads(&ctx)), (0, 1));
}