ursa_engine/
file_table.rs

1//! Files as tables: the `UrlTableFactory` behind `SELECT * FROM 'file.parquet'`.
2//!
3//! This does what DataFusion's `DynamicListTableFactory` (installed by
4//! `SessionContext::enable_url_table`) does, plus three things it does not:
5//!
6//! - registers the object store for a URL the first time it is read (as
7//!   datafusion-cli's `DynamicObjectStoreCatalog` does);
8//! - expands globs in URLs, not just in local paths (see [`crate::paths`]);
9//! - reports why a location could not be opened. DataFusion's factory turns any
10//!   failure into "table not found", which hides a missing file or an unknown
11//!   extension.
12
13use std::sync::{Arc, OnceLock, Weak};
14
15use async_trait::async_trait;
16use datafusion::catalog::{TableProvider, UrlTableFactory};
17use datafusion::datasource::listing::{ListingTable, ListingTableConfig, ListingTableConfigExt};
18use datafusion::error::{DataFusionError, Result};
19use datafusion::execution::context::SessionState;
20use parking_lot::RwLock;
21
22use crate::error::{matched, user_message};
23use crate::paths::{looks_like_location, table_url};
24use crate::stores::ensure_object_store;
25
26/// Opens a location named in `FROM` as a listing table.
27#[derive(Debug, Default)]
28pub(crate) struct FileTableFactory {
29    /// The session this factory belongs to, set once the session is built.
30    session: OnceLock<Weak<RwLock<SessionState>>>,
31}
32
33impl FileTableFactory {
34    pub(crate) fn attach(&self, session: Weak<RwLock<SessionState>>) {
35        let _ = self.session.set(session);
36    }
37}
38
39#[async_trait]
40impl UrlTableFactory for FileTableFactory {
41    async fn try_new(&self, name: &str) -> Result<Option<Arc<dyn TableProvider>>> {
42        if !looks_like_location(name) {
43            return Ok(None);
44        }
45        let state = self
46            .session
47            .get()
48            .and_then(Weak::upgrade)
49            .map(|state| state.read().clone())
50            .ok_or_else(|| DataFusionError::Internal("the session has been dropped".into()))?;
51        open_file_table(&state, name).await.map(Some)
52    }
53}
54
55/// Open `location` (a path, glob, directory or URL) as a listing table. The
56/// format and compression come from the file extension; Hive-style partition
57/// directories (`year=2024/`) become columns. A relative path is relative to
58/// the current directory, in a session with a database too: a stored
59/// definition's paths are made relative to the database's directory when it is
60/// stored, and resolved against it on replay (see [`crate::database`]).
61pub async fn open_file_table(
62    state: &SessionState,
63    location: &str,
64) -> Result<Arc<dyn TableProvider>> {
65    let url = table_url(location)?;
66    ensure_object_store(state.runtime_env(), url.get_url()).await?;
67    let config = ListingTableConfig::new(url)
68        .infer_options(state)
69        .await
70        .map_err(|e| explain(location, &e))?
71        .infer_partitions_from_path(state)
72        .await
73        .map_err(|e| explain(location, &e))?
74        .infer_schema(state)
75        .await
76        .map_err(|e| explain(location, &e))?;
77    Ok(Arc::new(ListingTable::try_new(config)?))
78}
79
80/// Say why `location` could not be opened.
81fn explain(location: &str, err: &DataFusionError) -> DataFusionError {
82    let message = user_message(err);
83    if message.contains(matched::NO_FILES_FOR_TABLE) {
84        DataFusionError::Plan(format!("no file found at '{location}'"))
85    } else if message.contains(matched::NO_FILE_FORMAT_FOR_EXTENSION) {
86        DataFusionError::Plan(format!(
87            "cannot tell the format of '{location}' from its extension; Ursa reads .parquet, \
88             .csv, .json and .arrow files, and compressed CSV and JSON (.gz, .bz2, .xz, .zst)"
89        ))
90    } else {
91        DataFusionError::Plan(format!("cannot read '{location}': {message}"))
92    }
93}
ursa_engine/file_table.rs

ursa_engine/
file_table.rs
