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 notices;
24pub mod options;
25pub mod paths;
26pub mod read_functions;
27pub mod serial;
28pub mod stores;
29
30pub use blocking::{run_blocking, BlockingError};
31pub use error::user_message;
32pub use file_table::{open_file_table, open_files, FileFormat, FileOptions};
33pub use stores::{ensure_object_store, ensure_object_store_for};
34
35/// A session with Ursa's defaults.
36pub fn new_session() -> SessionContext {
37    new_session_with_config(SessionConfig::new())
38}
39
40/// A session with Ursa's file tables over `config`. `information_schema` is
41/// turned on, so `SHOW TABLES` and `SHOW COLUMNS` work.
42///
43/// In the browser (wasm32) the session plans for one partition. DataFusion runs
44/// each extra partition as a spawned Tokio task, and the browser build has no
45/// runtime to spawn on; its futures are driven from JavaScript's event loop, as
46/// in DataFusion's own `wasmtest` crate (#278).
47pub fn new_session_with_config(config: SessionConfig) -> SessionContext {
48    #[cfg(target_arch = "wasm32")]
49    let config = config.with_target_partitions(1);
50    let config = graph::with_algorithms(graph::with_registry(config.with_information_schema(true)));
51    let config = database::with_database_slot(options::with_options(config));
52    let config = notices::with_notice_slot(config);
53    let config = read_functions::with_read_tables(config);
54    let ctx = SessionContext::new_with_config(config);
55    // Graph schemas are added to and dropped from the unwrapped catalogs.
56    graph::registry_of(&ctx).set_catalogs(Arc::clone(ctx.state().catalog_list()));
57    // As `SessionContext::enable_url_table` does, with Ursa's factory in place
58    // of DataFusion's.
59    let factory = Arc::new(file_table::FileTableFactory::default());
60    let catalogs = Arc::new(DynamicFileCatalog::new(
61        Arc::clone(ctx.state().catalog_list()),
62        Arc::clone(&factory) as _,
63    ));
64    let session_id = ctx.session_id();
65    let ctx = ctx
66        .into_state_builder()
67        .with_session_id(session_id)
68        .with_catalog_list(catalogs)
69        .with_query_planner(graph::query_planner())
70        .with_physical_optimizer_rule(graph::edge_order_rule());
71    // In the browser, last: every plan runs on the task that polls it.
72    #[cfg(target_arch = "wasm32")]
73    let ctx = ctx.with_physical_optimizer_rule(Arc::new(serial::SingleTask));
74    let mut state = ctx.build();
75    for ext in ["jsonl", "ndjson"] {
76        state
77            .register_file_format(Arc::new(file_table::JsonLines(ext)), false)
78            .expect("registering a file format cannot fail");
79    }
80    let ctx: SessionContext = state.into();
81    graph::register_functions(&ctx).expect("registering a relation planner cannot fail");
82    let reads = read_functions::ReadFunctions::new(ctx.state().config())
83        .expect("the session config has a place for read_* tables");
84    ctx.register_relation_planner(Arc::new(reads))
85        .expect("registering a relation planner cannot fail");
86    factory.attach(ctx.state_weak_ref());
87    ctx
88}
ursa_engine/lib.rs

ursa_engine/
lib.rs
