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 four 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//! - reads `.jsonl` and `.ndjson` files (and `.jsonl.gz`, …) as
10//!   newline-delimited JSON, as DataFusion reads `.json`;
11//! - reports why a location could not be opened. DataFusion's factory turns any
12//!   failure into "table not found", which hides a missing file or an unknown
13//!   extension.
14//!
15//! [`open_files`] is the same with options: what `read_csv(...)` and its
16//! siblings plan (see [`crate::read_functions`]).
17
18use std::collections::HashMap;
19use std::sync::{Arc, OnceLock, Weak};
20
21use async_trait::async_trait;
22use datafusion::catalog::Session;
23use datafusion::catalog::{TableProvider, UrlTableFactory};
24use datafusion::common::parsers::CompressionTypeVariant;
25use datafusion::common::GetExt;
26use datafusion::datasource::file_format::json::JsonFormatFactory;
27use datafusion::datasource::file_format::{FileFormat as DataFusionFileFormat, FileFormatFactory};
28use datafusion::datasource::listing::{ListingOptions, ListingTable, ListingTableConfig};
29use datafusion::error::{DataFusionError, Result};
30use datafusion::execution::context::SessionState;
31use futures::StreamExt;
32use parking_lot::RwLock;
33
34use crate::error::user_message;
35use crate::paths::{looks_like_location, table_url};
36use crate::stores::ensure_object_store;
37
38/// Opens a location named in `FROM` as a listing table.
39#[derive(Debug, Default)]
40pub(crate) struct FileTableFactory {
41    /// The session this factory belongs to, set once the session is built.
42    session: OnceLock<Weak<RwLock<SessionState>>>,
43}
44
45impl FileTableFactory {
46    pub(crate) fn attach(&self, session: Weak<RwLock<SessionState>>) {
47        let _ = self.session.set(session);
48    }
49}
50
51#[async_trait]
52impl UrlTableFactory for FileTableFactory {
53    async fn try_new(&self, name: &str) -> Result<Option<Arc<dyn TableProvider>>> {
54        if !looks_like_location(name) {
55            return Ok(None);
56        }
57        let state = self
58            .session
59            .get()
60            .and_then(Weak::upgrade)
61            .map(|state| state.read().clone())
62            .ok_or_else(|| DataFusionError::Internal("the session has been dropped".into()))?;
63        open_file_table(&state, name).await.map(Some)
64    }
65}
66
67/// Open `location` (a path, glob, directory or URL) as a listing table. The
68/// format and compression come from the file extension; Hive-style partition
69/// directories (`year=2024/`) become columns. A relative path is relative to
70/// the current directory, in a session with a database too: a stored
71/// definition's paths are made relative to the database's directory when it is
72/// stored, and resolved against it on replay (see [`crate::database`]).
73pub async fn open_file_table(
74    state: &SessionState,
75    location: &str,
76) -> Result<Arc<dyn TableProvider>> {
77    open_files(state, &[location.to_string()], &FileOptions::default()).await
78}
79
80/// A file format Ursa reads.
81#[derive(Debug, Clone, Copy, PartialEq, Eq)]
82pub enum FileFormat {
83    Parquet,
84    Csv,
85    /// Newline-delimited JSON, or a JSON array of objects with
86    /// [`FileOptions::newline_delimited`] `false`.
87    Json,
88    Arrow,
89}
90
91impl FileFormat {
92    /// The format a file extension names (without its dot or compression
93    /// suffix). `json`, `jsonl` and `ndjson` are all newline-delimited JSON.
94    pub fn from_extension(ext: &str) -> Option<Self> {
95        match ext.to_ascii_lowercase().as_str() {
96            "parquet" => Some(Self::Parquet),
97            "csv" => Some(Self::Csv),
98            "json" | "jsonl" | "ndjson" => Some(Self::Json),
99            "arrow" => Some(Self::Arrow),
100            _ => None,
101        }
102    }
103
104    /// The extension DataFusion registers the format's factory under.
105    fn factory(self) -> &'static str {
106        match self {
107            Self::Parquet => "parquet",
108            Self::Csv => "csv",
109            Self::Json => "json",
110            Self::Arrow => "arrow",
111        }
112    }
113
114    /// Whether the format's files are compressed whole (`.csv.gz`).
115    fn compressible(self) -> bool {
116        matches!(self, Self::Csv | Self::Json)
117    }
118}
119
120/// How to read files, beyond what their extension says: the options of
121/// `read_csv(...)` and its siblings. `None` leaves a choice to the extension
122/// or to the session's defaults.
123#[derive(Debug, Clone, Default, PartialEq, Eq)]
124pub struct FileOptions {
125    /// The format, whatever the extension.
126    pub format: Option<FileFormat>,
127    /// The compression of CSV or JSON files, whatever the extension.
128    pub compression: Option<CompressionTypeVariant>,
129    /// CSV: whether the first line is a header.
130    pub has_header: Option<bool>,
131    /// CSV: the field delimiter.
132    pub delimiter: Option<u8>,
133    /// CSV: the quote character.
134    pub quote: Option<u8>,
135    /// JSON: one object per line (`true`), or one array of objects (`false`).
136    pub newline_delimited: Option<bool>,
137}
138
139/// Open `locations` (paths, globs, directories or URLs) as one listing table,
140/// read as `options` say. What `options` leaves open comes from the extension
141/// of the first file found, as DataFusion infers it; a directory or glob
142/// lists only the files with that extension. Each location must name a file.
143pub async fn open_files(
144    state: &SessionState,
145    locations: &[String],
146    options: &FileOptions,
147) -> Result<Arc<dyn TableProvider>> {
148    let shown = locations.join("', '");
149    let mut urls = Vec::with_capacity(locations.len());
150    let mut first_file = None;
151    // Each location's first file's extensions (`.jsonl.gz`), to list by.
152    let mut extensions = Vec::with_capacity(locations.len());
153    for location in locations {
154        let url = table_url(location)?;
155        ensure_object_store(state.runtime_env(), url.get_url()).await?;
156        let store = state.runtime_env().object_store(&url)?;
157        let file = url
158            .list_all_files(state, store.as_ref(), "")
159            .await
160            .map_err(|e| explain(location, &e))?
161            .next()
162            .await
163            .transpose()
164            .map_err(|e| explain(location, &e))?
165            .ok_or_else(|| DataFusionError::Plan(format!("no file found at '{location}'")))?;
166        let (ext, suffix) = extensions_of(file.location.as_ref())?;
167        // A directory or glob lists the files with the first file's extension
168        // and compression suffix, as DataFusion's `infer_options` does; a file
169        // whose extension names no format (`read_csv('data.txt')`) by its own.
170        extensions.push(match (&ext, &suffix) {
171            (Some(ext), Some(suffix)) => format!(".{ext}.{suffix}"),
172            (Some(ext), None) => format!(".{ext}"),
173            (None, _) => String::new(),
174        });
175        first_file.get_or_insert((ext, suffix));
176        urls.push(url);
177    }
178    let (ext, suffix) =
179        first_file.ok_or_else(|| DataFusionError::Plan("no file to read".into()))?;
180    let named = ext.as_deref().and_then(FileFormat::from_extension);
181    let format = options.format.or(named).ok_or_else(|| {
182        DataFusionError::Plan(format!(
183            "cannot tell the format of '{shown}' from its extension; Ursa reads .parquet, \
184             .csv, .json, .jsonl, .ndjson and .arrow files, and compressed CSV and JSON (.gz, \
185             .bz2, .xz, .zst). read_parquet(), read_csv() and read_json() read a file \
186             whatever its extension"
187        ))
188    })?;
189    let mut format_options = HashMap::new();
190    let mut set = |key: &str, value: String| {
191        format_options.insert(format!("format.{key}"), value);
192    };
193    if format.compressible() {
194        let compression = match (options.compression, &suffix) {
195            (Some(c), _) => Some(c),
196            (None, Some(s)) => Some(s.parse::<CompressionTypeVariant>()?),
197            (None, None) => None,
198        };
199        if let Some(c) = compression {
200            set("compression", compression_name(c));
201        }
202    }
203    if format == FileFormat::Csv {
204        if let Some(header) = options.has_header {
205            set("has_header", header.to_string());
206        }
207        if let Some(delimiter) = options.delimiter {
208            set("delimiter", char::from(delimiter).to_string());
209        }
210        if let Some(quote) = options.quote {
211            set("quote", char::from(quote).to_string());
212        }
213    }
214    if format == FileFormat::Json {
215        if let Some(newline_delimited) = options.newline_delimited {
216            set("newline_delimited", newline_delimited.to_string());
217        }
218    }
219    let file_format = state
220        .get_file_format_factory(format.factory())
221        .ok_or_else(|| {
222            DataFusionError::Internal(format!("no {} file format factory", format.factory()))
223        })?
224        .create(state, &format_options)?;
225    // One listing extension for all the locations: theirs if they agree
226    // (`['a.jsonl', 'b.ndjson']` do not), else none.
227    let extension = if extensions.windows(2).all(|w| w[0] == w[1]) {
228        extensions.swap_remove(0)
229    } else {
230        String::new()
231    };
232    let config = ListingTableConfig::new_with_multi_paths(urls)
233        .with_listing_options(ListingOptions::new(file_format).with_file_extension(extension))
234        .infer_partitions_from_path(state)
235        .await
236        .map_err(|e| explain(&shown, &e))?
237        .infer_schema(state)
238        .await
239        .map_err(|e| explain(&shown, &e))?;
240    Ok(Arc::new(ListingTable::try_new(config)?))
241}
242
243/// JSON under another extension, for DataFusion's own lookups of a format by
244/// extension: `COPY … TO 'out.jsonl'` writes newline-delimited JSON, as `COPY
245/// … TO 'out.json'` does.
246#[derive(Debug)]
247pub(crate) struct JsonLines(pub(crate) &'static str);
248
249impl FileFormatFactory for JsonLines {
250    fn create(
251        &self,
252        state: &dyn Session,
253        format_options: &HashMap<String, String>,
254    ) -> Result<Arc<dyn DataFusionFileFormat>> {
255        JsonFormatFactory::new().create(state, format_options)
256    }
257
258    fn default(&self) -> Arc<dyn DataFusionFileFormat> {
259        JsonFormatFactory::new().default()
260    }
261}
262
263impl GetExt for JsonLines {
264    fn get_ext(&self) -> String {
265        self.0.to_string()
266    }
267}
268
269/// A file's extension and compression suffix, as DataFusion infers them
270/// (`x.jsonl.gz` has `jsonl` and `gz`); no extension for a file without one.
271fn extensions_of(path: &str) -> Result<(Option<String>, Option<String>)> {
272    let (ext, suffix) = ListingTableConfig::infer_file_extension_and_compression_type(path)?;
273    // For a file without an extension, DataFusion gives the whole path.
274    let ext = (!ext.contains('/') && ext != path).then_some(ext);
275    Ok((ext, suffix))
276}
277
278/// The value DataFusion's `format.compression` option takes for `c`.
279fn compression_name(c: CompressionTypeVariant) -> String {
280    match c {
281        CompressionTypeVariant::UNCOMPRESSED => "uncompressed".to_string(),
282        other => other.to_string().to_lowercase(),
283    }
284}
285
286/// Say why `location` could not be opened.
287fn explain(location: &str, err: &DataFusionError) -> DataFusionError {
288    DataFusionError::Plan(format!("cannot read '{location}': {}", user_message(err)))
289}
ursa_engine/file_table.rs

ursa_engine/
file_table.rs
