datafusion_opensearch/
table.rs

1//! The [`TableProvider`], the factory that builds one from an index's `_mapping`, and the
2//! [`TableProviderFactory`] behind `CREATE EXTERNAL TABLE … STORED AS OPENSEARCH`.
3
4use std::fmt;
5use std::sync::Arc;
6
7use async_trait::async_trait;
8use datafusion::arrow::datatypes::SchemaRef;
9use datafusion::catalog::{Session, TableProvider, TableProviderFactory};
10use datafusion::error::{DataFusionError, Result as DataFusionResult};
11use datafusion::logical_expr::{CreateExternalTable, Expr, TableProviderFilterPushDown, TableType};
12use datafusion::physical_plan::ExecutionPlan;
13use datafusion::prelude::SessionContext;
14use serde_json::{json, Value};
15
16use crate::client::OpenSearchClient;
17use crate::exec::OpenSearchExec;
18use crate::schema::schema_from_mapping;
19use crate::{pushdown, MalformedMappingSnafu, Result};
20
21/// Hits fetched per `_search` page. Bounded above by the index's `max_result_window` (10 000 by
22/// default); override per table with [`OpenSearchTableProvider::with_page_size`].
23pub const DEFAULT_PAGE_SIZE: usize = 5_000;
24
25/// The `STORED AS` name that selects [`OpenSearchTableProviderFactory`] in `CREATE EXTERNAL TABLE`.
26pub const FILE_TYPE: &str = "OPENSEARCH";
27
28/// A DataFusion table backed by one OpenSearch index.
29pub struct OpenSearchTableProvider {
30    client: OpenSearchClient,
31    index: String,
32    schema: SchemaRef,
33    page_size: usize,
34    /// Cap on rows a scan returns when the query has no LIMIT (`None` = the whole result).
35    max_rows: Option<usize>,
36    /// DataFusion partitions per scan, each a sliced scroll (1 = a plain scroll).
37    partitions: usize,
38    /// Extra query-DSL clauses ANDed into every scan (see [`with_base_filter`]).
39    ///
40    /// [`with_base_filter`]: OpenSearchTableProvider::with_base_filter
41    base_filter: Vec<Value>,
42    /// `sort` pushed into `_search` (field, ascending). Empty → OpenSearch default order.
43    sort: Vec<(String, bool)>,
44    /// `search_after` values for cursor pagination (must align with `sort`).
45    search_after: Option<Vec<Value>>,
46}
47
48impl fmt::Debug for OpenSearchTableProvider {
49    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
50        f.debug_struct("OpenSearchTableProvider")
51            .field("index", &self.index)
52            .field("fields", &self.schema.fields().len())
53            .field("page_size", &self.page_size)
54            .field("max_rows", &self.max_rows)
55            .field("partitions", &self.partitions)
56            .field("base_filter", &self.base_filter.len())
57            .finish()
58    }
59}
60
61impl OpenSearchTableProvider {
62    /// `base_url` e.g. `http://localhost:9200`; `index` e.g. `my-index`; `schema` declares the
63    /// columns (and Arrow types) to read from each hit's `_source`. Supported column types:
64    /// `Utf8`, `Int64`, `Float64`, `Boolean`, `Timestamp(Millisecond, _)` (see
65    /// [`crate::schema::build_batch`]).
66    pub fn new(base_url: impl Into<String>, index: impl Into<String>, schema: SchemaRef) -> Self {
67        Self::with_client_and_schema(OpenSearchClient::new(base_url), index, schema)
68    }
69
70    /// Build over an existing [`OpenSearchClient`] (shared connection pool / auth / TLS).
71    pub fn with_client_and_schema(client: OpenSearchClient, index: impl Into<String>, schema: SchemaRef) -> Self {
72        Self {
73            client,
74            index: index.into(),
75            schema,
76            page_size: DEFAULT_PAGE_SIZE,
77            max_rows: None,
78            partitions: 1,
79            base_filter: Vec::new(),
80            sort: Vec::new(),
81            search_after: None,
82        }
83    }
84
85    /// Push a sort into `_search` — `(field, ascending)` pairs. Required for stable pagination
86    /// (and for a LIMIT to mean "top-N by this order" rather than an arbitrary N).
87    pub fn with_sort(mut self, sort: Vec<(String, bool)>) -> Self {
88        self.sort = sort;
89        self
90    }
91
92    /// Continue after a prior page: OpenSearch `search_after` values, aligned with `with_sort`.
93    /// A scan then reads exactly one page (of the LIMIT, or the page size) — the loop is yours.
94    pub fn with_search_after(mut self, after: Vec<Value>) -> Self {
95        self.search_after = Some(after);
96        self
97    }
98
99    /// Hits per `_search` page while streaming a result (default [`DEFAULT_PAGE_SIZE`]).
100    pub fn with_page_size(mut self, n: usize) -> Self {
101        self.page_size = n.max(1);
102        self
103    }
104
105    /// Cap the rows a LIMIT-less scan returns. By default a scan streams the whole result.
106    pub fn with_max_rows(mut self, n: Option<usize>) -> Self {
107        self.max_rows = n;
108        self
109    }
110
111    /// Read the index through `n` DataFusion partitions in parallel, each an OpenSearch sliced
112    /// scroll (`slice: { id, max }`). Worthwhile for large scans on multi-shard indices; a query
113    /// with a LIMIT always uses one partition so the limit is exact. Default 1.
114    pub fn with_partitions(mut self, n: usize) -> Self {
115        self.partitions = n.max(1);
116        self
117    }
118
119    /// Add query-DSL clauses that are ANDed into the `bool.filter` of EVERY scan, on top of any
120    /// pushed-down WHERE. Use this for constraints the caller must never opt out of — e.g. a
121    /// row-level security predicate, a tenancy filter, or a soft-delete flag. The clauses are
122    /// opaque OpenSearch DSL (`serde_json::Value`), so this crate stays agnostic to what they express.
123    pub fn with_base_filter(mut self, clauses: Vec<Value>) -> Self {
124        self.base_filter = clauses;
125        self
126    }
127
128    /// Supply a preconfigured reqwest client (timeouts, connection pool, TLS roots, headers).
129    pub fn with_client(mut self, http: reqwest::Client) -> Self {
130        self.client = self.client.with_client(http);
131        self
132    }
133
134    /// The index this table reads.
135    pub fn index(&self) -> &str {
136        &self.index
137    }
138
139    /// The `_search` body (without `size`) a scan with these filters and this projection sends.
140    /// Exposed so callers can see exactly what was pushed down.
141    pub fn search_body(&self, projected: &SchemaRef, filters: &[Expr]) -> Value {
142        let mut clauses: Vec<Value> = self.base_filter.clone();
143        clauses.extend(filters.iter().filter_map(pushdown::expr_to_query));
144        let query = if clauses.is_empty() {
145            json!({ "match_all": {} })
146        } else {
147            json!({ "bool": { "filter": clauses } })
148        };
149        let source_fields: Vec<&str> = projected.fields().iter().map(|f| f.name().as_str()).collect();
150        let mut body = json!({ "query": query, "_source": source_fields });
151        if !self.sort.is_empty() {
152            body["sort"] = Value::Array(
153                self.sort
154                    .iter()
155                    .map(|(f, asc)| json!({ f: { "order": if *asc { "asc" } else { "desc" } } }))
156                    .collect(),
157            );
158        }
159        if let Some(after) = &self.search_after {
160            body["search_after"] = Value::Array(after.clone());
161        }
162        body
163    }
164}
165
166#[async_trait]
167impl TableProvider for OpenSearchTableProvider {
168    fn schema(&self) -> SchemaRef {
169        self.schema.clone()
170    }
171
172    fn table_type(&self) -> TableType {
173        TableType::Base
174    }
175
176    /// Report `Exact` for filters [`pushdown::expr_to_query`] can translate (DataFusion drops its
177    /// Filter node), `Unsupported` for the rest (DataFusion keeps filtering them). Same predicate
178    /// as `scan` uses, so the two never disagree.
179    fn supports_filters_pushdown(&self, filters: &[&Expr]) -> DataFusionResult<Vec<TableProviderFilterPushDown>> {
180        Ok(filters
181            .iter()
182            .map(|f| {
183                if pushdown::expr_to_query(f).is_some() {
184                    TableProviderFilterPushDown::Exact
185                } else {
186                    TableProviderFilterPushDown::Unsupported
187                }
188            })
189            .collect())
190    }
191
192    async fn scan(
193        &self,
194        _state: &dyn Session,
195        projection: Option<&Vec<usize>>,
196        filters: &[Expr],
197        limit: Option<usize>,
198    ) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
199        let projected: SchemaRef = match projection {
200            Some(p) => Arc::new(self.schema.project(p)?),
201            None => self.schema.clone(),
202        };
203        let body = self.search_body(&projected, filters);
204        let fetch = match (limit, self.max_rows) {
205            (Some(l), Some(m)) => Some(l.min(m)),
206            (l, m) => l.or(m),
207        };
208        Ok(Arc::new(
209            OpenSearchExec::new(self.client.clone(), self.index.clone(), body, projected)
210                .with_page_size(self.page_size)
211                .with_fetch_limit(fetch)
212                .with_single_page(self.search_after.is_some())
213                .with_partitions(self.partitions),
214        ))
215    }
216}
217
218/// Builds [`OpenSearchTableProvider`]s over one cluster, deriving each table's schema from the
219/// index `_mapping` unless one is declared.
220#[derive(Clone, Debug)]
221pub struct OpenSearchTableFactory {
222    client: OpenSearchClient,
223}
224
225impl OpenSearchTableFactory {
226    pub fn new(client: OpenSearchClient) -> Self {
227        Self { client }
228    }
229
230    /// The client every table built here shares.
231    pub fn client(&self) -> &OpenSearchClient {
232        &self.client
233    }
234
235    /// A table over `index` whose schema is derived from its `_mapping`
236    /// ([`schema_from_mapping`]). Fails if the index does not exist or has no `properties`.
237    pub async fn table_provider(&self, index: impl Into<String>) -> Result<Arc<dyn TableProvider>> {
238        Ok(Arc::new(self.provider(index).await?))
239    }
240
241    /// Like [`table_provider`](Self::table_provider), not type-erased, for the builder methods.
242    pub async fn provider(&self, index: impl Into<String>) -> Result<OpenSearchTableProvider> {
243        let index = index.into();
244        let mapping = self.client.mapping(&index).await?;
245        let schema =
246            schema_from_mapping(&mapping).ok_or_else(|| MalformedMappingSnafu { index: index.clone() }.build())?;
247        Ok(OpenSearchTableProvider::with_client_and_schema(
248            self.client.clone(),
249            index,
250            schema,
251        ))
252    }
253
254    /// A table over `index` with a declared schema (no `_mapping` round-trip; only the listed
255    /// columns are read, and a not-yet-created index reads as empty).
256    pub fn table_provider_with_schema(&self, index: impl Into<String>, schema: SchemaRef) -> Arc<dyn TableProvider> {
257        Arc::new(self.provider_with_schema(index, schema))
258    }
259
260    /// The provider (not type-erased), for callers that want the builder methods.
261    pub fn provider_with_schema(&self, index: impl Into<String>, schema: SchemaRef) -> OpenSearchTableProvider {
262        OpenSearchTableProvider::with_client_and_schema(self.client.clone(), index, schema)
263    }
264}
265
266/// The [`TableProviderFactory`] that makes OpenSearch indices reachable from SQL alone:
267///
268/// ```sql
269/// CREATE EXTERNAL TABLE docs STORED AS OPENSEARCH LOCATION 'http://localhost:9200/my-index';
270/// ```
271///
272/// `LOCATION` is `<base url>/<index>`. Without a column list the schema is derived from the index
273/// `_mapping`; with one, only those columns are read. Recognised `OPTIONS`: `page_size`,
274/// `max_rows`, `partitions`, `sort` (`"field:asc,other:desc"`), `base_filter` (a JSON array of
275/// query-DSL clauses).
276/// Credentials and the HTTP client set on the factory apply to every table it creates.
277#[derive(Clone, Debug, Default)]
278pub struct OpenSearchTableProviderFactory {
279    http: Option<reqwest::Client>,
280    basic: Option<(String, String)>,
281    bearer: Option<String>,
282}
283
284impl OpenSearchTableProviderFactory {
285    pub fn new() -> Self {
286        Self::default()
287    }
288
289    /// Use this [`reqwest::Client`] for every table (timeouts, TLS roots, headers).
290    pub fn with_client(mut self, http: reqwest::Client) -> Self {
291        self.http = Some(http);
292        self
293    }
294
295    /// Send HTTP basic credentials with every request.
296    pub fn with_basic_auth(mut self, user: impl Into<String>, password: impl Into<String>) -> Self {
297        self.basic = Some((user.into(), password.into()));
298        self
299    }
300
301    /// Send a bearer token with every request.
302    pub fn with_bearer_token(mut self, token: impl Into<String>) -> Self {
303        self.bearer = Some(token.into());
304        self
305    }
306
307    /// Register this factory on a [`SessionContext`] as `STORED AS OPENSEARCH`.
308    pub fn register(self, ctx: &SessionContext) {
309        ctx.state_ref()
310            .write()
311            .table_factories_mut()
312            .insert(FILE_TYPE.to_string(), Arc::new(self));
313    }
314
315    fn client_for(&self, base_url: &str) -> OpenSearchClient {
316        let mut c = OpenSearchClient::new(base_url);
317        if let Some(http) = &self.http {
318            c = c.with_client(http.clone());
319        }
320        if let Some((u, p)) = &self.basic {
321            c = c.with_basic_auth(u, p);
322        }
323        if let Some(t) = &self.bearer {
324            c = c.with_bearer_token(t);
325        }
326        c
327    }
328}
329
330/// Split `LOCATION` into `(base_url, index)`: the index is the last path segment.
331fn split_location(location: &str) -> DataFusionResult<(&str, &str)> {
332    let trimmed = location.trim_end_matches('/');
333    let (base, index) = trimmed.rsplit_once('/').ok_or_else(|| {
334        DataFusionError::Plan(format!(
335            "datafusion-opensearch: LOCATION must be '<base url>/<index>', got '{location}'"
336        ))
337    })?;
338    if !base.contains("://") || index.is_empty() || base.ends_with(':') || base.ends_with('/') {
339        return Err(DataFusionError::Plan(format!(
340            "datafusion-opensearch: LOCATION must be '<base url>/<index>', got '{location}'"
341        )));
342    }
343    Ok((base, index))
344}
345
346/// Parse `"field:asc,other:desc"` (direction optional, default ascending).
347fn parse_sort(spec: &str) -> DataFusionResult<Vec<(String, bool)>> {
348    spec.split(',')
349        .map(str::trim)
350        .filter(|s| !s.is_empty())
351        .map(|s| match s.rsplit_once(':') {
352            Some((f, "asc")) => Ok((f.trim().to_string(), true)),
353            Some((f, "desc")) => Ok((f.trim().to_string(), false)),
354            Some((_, other)) => Err(DataFusionError::Plan(format!(
355                "datafusion-opensearch: sort direction must be asc or desc, got '{other}'"
356            ))),
357            None => Ok((s.to_string(), true)),
358        })
359        .collect()
360}
361
362fn parse_usize(options: &std::collections::HashMap<String, String>, key: &str) -> DataFusionResult<Option<usize>> {
363    options
364        .get(key)
365        .map(|v| {
366            v.trim().parse::<usize>().map_err(|_| {
367                DataFusionError::Plan(format!(
368                    "datafusion-opensearch: option '{key}' must be an integer, got '{v}'"
369                ))
370            })
371        })
372        .transpose()
373}
374
375#[async_trait]
376impl TableProviderFactory for OpenSearchTableProviderFactory {
377    async fn create(
378        &self,
379        _state: &dyn Session,
380        cmd: &CreateExternalTable,
381    ) -> DataFusionResult<Arc<dyn TableProvider>> {
382        let (base_url, index) = split_location(&cmd.location)?;
383        // DataFusion qualifies unprefixed OPTIONS keys with `format.`; accept both spellings.
384        let options: std::collections::HashMap<String, String> = cmd
385            .options
386            .iter()
387            .map(|(k, v)| (k.strip_prefix("format.").unwrap_or(k).to_string(), v.clone()))
388            .collect();
389        let factory = OpenSearchTableFactory::new(self.client_for(base_url));
390        let mut provider = if cmd.schema.fields().is_empty() {
391            factory.provider(index).await?
392        } else {
393            factory.provider_with_schema(index, Arc::new(cmd.schema.as_arrow().clone()))
394        };
395        if let Some(n) = parse_usize(&options, "page_size")? {
396            provider = provider.with_page_size(n);
397        }
398        if let Some(n) = parse_usize(&options, "max_rows")? {
399            provider = provider.with_max_rows(Some(n));
400        }
401        if let Some(n) = parse_usize(&options, "partitions")? {
402            provider = provider.with_partitions(n);
403        }
404        if let Some(sort) = options.get("sort") {
405            provider = provider.with_sort(parse_sort(sort)?);
406        }
407        if let Some(raw) = options.get("base_filter") {
408            let clauses: Vec<Value> = serde_json::from_str(raw).map_err(|e| {
409                DataFusionError::Plan(format!(
410                    "datafusion-opensearch: option 'base_filter' must be a JSON array of query clauses: {e}"
411                ))
412            })?;
413            provider = provider.with_base_filter(clauses);
414        }
415        Ok(Arc::new(provider))
416    }
417}
418
419#[cfg(test)]
420mod tests {
421    use super::*;
422    use datafusion::arrow::datatypes::{DataType, Field, Schema};
423    use datafusion::logical_expr::{col, lit};
424
425    fn schema() -> SchemaRef {
426        Arc::new(Schema::new(vec![
427            Field::new("id", DataType::Utf8, true),
428            Field::new("status", DataType::Utf8, true),
429            Field::new("speed", DataType::Float64, true),
430            Field::new("active", DataType::Boolean, true),
431        ]))
432    }
433
434    #[test]
435    fn search_body_composes_base_filter_pushdown_projection_and_sort() {
436        let p = OpenSearchTableProvider::new("http://x", "i", schema())
437            .with_base_filter(vec![json!({ "term": { "tenant": "t1" } })])
438            .with_sort(vec![("speed".into(), false)]);
439        let projected = Arc::new(schema().project(&[0]).unwrap());
440        let body = p.search_body(&projected, &[col("status").eq(lit("OK")), col("a").eq(col("b"))]);
441        assert_eq!(
442            body["query"],
443            json!({ "bool": { "filter": [ { "term": { "tenant": "t1" } }, { "term": { "status": "OK" } } ] } })
444        );
445        assert_eq!(body["_source"], json!(["id"]));
446        assert_eq!(body["sort"], json!([{ "speed": { "order": "desc" } }]));
447        assert!(body.get("size").is_none(), "size belongs to the streaming exec");
448        let plain = OpenSearchTableProvider::new("http://x", "i", schema()).search_body(&schema(), &[]);
449        assert_eq!(plain["query"], json!({ "match_all": {} }));
450        assert_eq!(p.index(), "i");
451    }
452
453    #[test]
454    fn factory_builds_a_declared_schema_provider_without_io() {
455        let f = OpenSearchTableFactory::new(OpenSearchClient::new("http://x:9200"));
456        let p = f
457            .provider_with_schema("idx", schema())
458            .with_page_size(5)
459            .with_max_rows(Some(50));
460        assert_eq!(p.page_size, 5);
461        assert_eq!(p.max_rows, Some(50));
462        assert_eq!(f.table_provider_with_schema("idx", schema()).schema().fields().len(), 4);
463        assert_eq!(f.client().base_url(), "http://x:9200");
464    }
465
466    #[test]
467    fn location_splits_into_base_url_and_index() {
468        assert_eq!(
469            split_location("http://localhost:9200/my-index").unwrap(),
470            ("http://localhost:9200", "my-index")
471        );
472        assert_eq!(
473            split_location("https://os.example.com/logs-2024/").unwrap(),
474            ("https://os.example.com", "logs-2024")
475        );
476        assert!(split_location("my-index").is_err());
477        assert!(split_location("http://localhost:9200/").is_err());
478        assert!(split_location("http://localhost:9200").is_err());
479    }
480
481    #[test]
482    fn sort_option_parses() {
483        assert_eq!(
484            parse_sort("ts:desc, id").unwrap(),
485            vec![("ts".to_string(), false), ("id".to_string(), true)]
486        );
487        assert!(parse_sort("ts:sideways").is_err());
488    }
489
490    #[tokio::test]
491    async fn create_external_table_with_declared_columns_needs_no_io() {
492        let ctx = SessionContext::new();
493        OpenSearchTableProviderFactory::new()
494            .with_basic_auth("u", "p")
495            .register(&ctx);
496        ctx.sql(
497            "CREATE EXTERNAL TABLE docs (id VARCHAR, speed DOUBLE) STORED AS OPENSEARCH \
498             LOCATION 'http://localhost:9200/my-index' OPTIONS ('page_size' '100', 'sort' 'speed:desc')",
499        )
500        .await
501        .unwrap()
502        .collect()
503        .await
504        .unwrap();
505        let t = ctx.table_provider("docs").await.unwrap();
506        assert_eq!(t.schema().fields().len(), 2);
507        let any: &dyn std::any::Any = t.as_ref();
508        let p = any.downcast_ref::<OpenSearchTableProvider>().unwrap();
509        assert_eq!(p.page_size, 100);
510        assert_eq!(p.sort, vec![("speed".to_string(), false)]);
511        assert_eq!(p.index(), "my-index");
512    }
513}
datafusion_opensearch/table.rs

datafusion_opensearch/
table.rs
