ursa_engine/
lib.rs

1//! Ursa's SQL engine session.
2//!
3//! [`new_session`] builds a DataFusion [`SessionContext`] on which a file is a
4//! table: `SELECT * FROM 'edges.parquet'`, `'edges.csv.gz'`, `'edges/'`,
5//! `'edges/*.parquet'`, `'events.jsonl'`, and the same over `s3://`, `gs://`,
6//! `az://` and `https://`, with object stores registered the first time a URL
7//! is read. `read_parquet`, `read_csv` and `read_json` read the same with
8//! options ([`read_functions`]).
9//!
10//! Graphs are defined per session ([`graph`]). A session may have a `.ursa`
11//! database ([`database`]), which keeps its catalog between processes.
12
13use std::sync::Arc;
14
15use datafusion::catalog::DynamicFileCatalog;
16use datafusion::prelude::{SessionConfig, SessionContext};
17
18mod blocking;
19pub mod database;
20pub mod error;
21mod file_table;
22pub mod graph;
23pub mod options;
24pub mod paths;
25pub mod read_functions;
26pub mod serial;
27pub mod stores;
28
29pub use blocking::{run_blocking, BlockingError};
30pub use error::user_message;
31pub use file_table::{open_file_table, open_files, FileFormat, FileOptions};
32pub use stores::{ensure_object_store, ensure_object_store_for};
33
34/// A session with Ursa's defaults.
35pub fn new_session() -> SessionContext {
36    new_session_with_config(SessionConfig::new())
37}
38
39/// A session with Ursa's file tables over `config`. `information_schema` is
40/// turned on, so `SHOW TABLES` and `SHOW COLUMNS` work.
41///
42/// In the browser (wasm32) the session plans for one partition. DataFusion runs
43/// each extra partition as a spawned Tokio task, and the browser build has no
44/// runtime to spawn on; its futures are driven from JavaScript's event loop, as
45/// in DataFusion's own `wasmtest` crate (#278).
46pub fn new_session_with_config(config: SessionConfig) -> SessionContext {
47    #[cfg(target_arch = "wasm32")]
48    let config = config.with_target_partitions(1);
49    let config = graph::with_algorithms(graph::with_registry(config.with_information_schema(true)));
50    let config = database::with_database_slot(options::with_options(config));
51    let config = read_functions::with_read_tables(config);
52    let ctx = SessionContext::new_with_config(config);
53    // Graph schemas are added to and dropped from the unwrapped catalogs.
54    graph::registry_of(&ctx).set_catalogs(Arc::clone(ctx.state().catalog_list()));
55    // As `SessionContext::enable_url_table` does, with Ursa's factory in place
56    // of DataFusion's.
57    let factory = Arc::new(file_table::FileTableFactory::default());
58    let catalogs = Arc::new(DynamicFileCatalog::new(
59        Arc::clone(ctx.state().catalog_list()),
60        Arc::clone(&factory) as _,
61    ));
62    let session_id = ctx.session_id();
63    let ctx = ctx
64        .into_state_builder()
65        .with_session_id(session_id)
66        .with_catalog_list(catalogs)
67        .with_query_planner(graph::query_planner())
68        .with_physical_optimizer_rule(graph::edge_order_rule());
69    // In the browser, last: every plan runs on the task that polls it.
70    #[cfg(target_arch = "wasm32")]
71    let ctx = ctx.with_physical_optimizer_rule(Arc::new(serial::SingleTask));
72    let mut state = ctx.build();
73    for ext in ["jsonl", "ndjson"] {
74        state
75            .register_file_format(Arc::new(file_table::JsonLines(ext)), false)
76            .expect("registering a file format cannot fail");
77    }
78    let ctx: SessionContext = state.into();
79    graph::register_functions(&ctx).expect("registering a relation planner cannot fail");
80    let reads = read_functions::ReadFunctions::new(ctx.state().config())
81        .expect("the session config has a place for read_* tables");
82    ctx.register_relation_planner(Arc::new(reads))
83        .expect("registering a relation planner cannot fail");
84    factory.attach(ctx.state_weak_ref());
85    ctx
86}
ursa_engine/lib.rs

ursa_engine/
lib.rs
