datafusion_opensearch/
exec.rs

1//! The physical operator: a `_search` that streams its result page by page.
2//!
3//! One `_search` can return at most `index.max_result_window` hits (10 000 by default), so a scan
4//! that simply set `size` would silently truncate any larger result. [`OpenSearchExec`] instead
5//! opens a scroll cursor and emits one `RecordBatch` per page until the result is exhausted or the
6//! plan's `fetch` (a pushed-down `LIMIT`) is satisfied, so memory stays bounded by the page size
7//! and the result is complete. With more than one partition the scan is a **sliced scroll**: each
8//! DataFusion partition reads one slice (`slice: { id, max }`) in parallel, which is how providers
9//! over sharded stores fan out. A `LIMIT` or `search_after` read always uses one partition, so a
10//! limit is exact.
11
12use std::fmt;
13use std::sync::Arc;
14
15use datafusion::arrow::datatypes::SchemaRef;
16use datafusion::arrow::record_batch::RecordBatch;
17use datafusion::error::Result;
18use datafusion::execution::TaskContext;
19use datafusion::physical_expr::EquivalenceProperties;
20use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
21use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet};
22use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
23use datafusion::physical_plan::{
24    DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, SendableRecordBatchStream,
25};
26use futures::StreamExt;
27use serde_json::Value;
28
29use crate::client::OpenSearchClient;
30use crate::schema::build_batch;
31
32/// How long the server keeps a scroll cursor alive between pages.
33const KEEP_ALIVE: &str = "1m";
34
35/// A scan of one index, streamed as pages of `_source` documents turned into record batches.
36#[derive(Debug)]
37pub struct OpenSearchExec {
38    inner: Arc<Inner>,
39    properties: Arc<PlanProperties>,
40    metrics: ExecutionPlanMetricsSet,
41}
42
43#[derive(Clone, Debug)]
44struct Inner {
45    client: OpenSearchClient,
46    index: String,
47    /// The `_search` body without `size`: query, `_source`, sort, search_after.
48    body: Value,
49    schema: SchemaRef,
50    page_size: usize,
51    /// Stop after this many rows (a pushed-down LIMIT or the table's row cap).
52    fetch: Option<usize>,
53    /// `search_after` pagination is the caller's loop: read exactly one page.
54    single_page: bool,
55    /// Sliced-scroll partitions (1 = a plain scroll).
56    partitions: usize,
57}
58
59impl OpenSearchExec {
60    /// A scan of `index` sending `body` (the `_search` request without `size`), producing
61    /// `schema`. Tune it with the `with_*` builders; by default it streams the whole result
62    /// through one scroll in pages of 5 000.
63    pub fn new(client: OpenSearchClient, index: String, body: Value, schema: SchemaRef) -> Self {
64        let inner = Inner {
65            client,
66            index,
67            body,
68            schema,
69            page_size: crate::table::DEFAULT_PAGE_SIZE,
70            fetch: None,
71            single_page: false,
72            partitions: 1,
73        };
74        Self::from_inner(inner)
75    }
76
77    fn from_inner(inner: Inner) -> Self {
78        // A fetch or a single-page read needs one stream so the row cap stays exact.
79        let partitions = if inner.fetch.is_some() || inner.single_page {
80            1
81        } else {
82            inner.partitions.max(1)
83        };
84        let inner = Inner { partitions, ..inner };
85        let properties = Arc::new(PlanProperties::new(
86            EquivalenceProperties::new(inner.schema.clone()),
87            Partitioning::UnknownPartitioning(partitions),
88            EmissionType::Incremental,
89            Boundedness::Bounded,
90        ));
91        Self {
92            inner: Arc::new(inner),
93            properties,
94            metrics: ExecutionPlanMetricsSet::new(),
95        }
96    }
97
98    fn rebuild(&self, f: impl FnOnce(&mut Inner)) -> Self {
99        let mut inner = self.inner.as_ref().clone();
100        f(&mut inner);
101        Self::from_inner(inner)
102    }
103
104    /// Hits per page (at least 1).
105    pub fn with_page_size(&self, n: usize) -> Self {
106        self.rebuild(|i| i.page_size = n.max(1))
107    }
108
109    /// Stop after this many rows (a pushed-down LIMIT or the table's row cap). Forces one partition.
110    pub fn with_fetch_limit(&self, fetch: Option<usize>) -> Self {
111        self.rebuild(|i| i.fetch = fetch)
112    }
113
114    /// Read exactly one page (caller-driven `search_after` pagination). Forces one partition.
115    pub fn with_single_page(&self, single: bool) -> Self {
116        self.rebuild(|i| i.single_page = single)
117    }
118
119    /// Fan the scan out over `n` sliced scrolls, one per DataFusion partition.
120    pub fn with_partitions(&self, n: usize) -> Self {
121        self.rebuild(|i| i.partitions = n.max(1))
122    }
123
124    /// The index being scanned.
125    pub fn index(&self) -> &str {
126        &self.inner.index
127    }
128
129    /// The `_search` body (without `size`) this scan sends — useful to inspect what was pushed down.
130    pub fn body(&self) -> &Value {
131        &self.inner.body
132    }
133}
134
135impl DisplayAs for OpenSearchExec {
136    fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter<'_>) -> fmt::Result {
137        let i = &self.inner;
138        match t {
139            DisplayFormatType::Default | DisplayFormatType::TreeRender => {
140                write!(
141                    f,
142                    "OpenSearchExec: index={}, partitions={}, fetch={:?}",
143                    i.index, i.partitions, i.fetch
144                )
145            }
146            DisplayFormatType::Verbose => write!(
147                f,
148                "OpenSearchExec: index={}, partitions={}, fetch={:?}, page_size={}, body={}",
149                i.index, i.partitions, i.fetch, i.page_size, i.body
150            ),
151        }
152    }
153}
154
155struct PageState {
156    inner: Arc<Inner>,
157    partition: usize,
158    scroll_id: Option<String>,
159    remaining: Option<usize>,
160    started: bool,
161    done: bool,
162}
163
164async fn next_page(mut st: PageState) -> Result<Option<(RecordBatch, PageState)>> {
165    if st.done || st.remaining == Some(0) {
166        if let Some(id) = st.scroll_id.take() {
167            st.inner.client.clear_scroll(&id).await;
168        }
169        return Ok(None);
170    }
171    let size = st.remaining.map_or(st.inner.page_size, |r| r.min(st.inner.page_size));
172    let page = if !st.started {
173        st.started = true;
174        let mut body = st.inner.body.clone();
175        body["size"] = Value::from(size);
176        if st.inner.partitions > 1 {
177            body["slice"] = serde_json::json!({ "id": st.partition, "max": st.inner.partitions });
178        }
179        let keep_alive = if st.inner.single_page { None } else { Some(KEEP_ALIVE) };
180        st.inner.client.search_page(&st.inner.index, &body, keep_alive).await?
181    } else {
182        match &st.scroll_id {
183            Some(id) => st.inner.client.scroll(id, KEEP_ALIVE).await?,
184            None => return Ok(None),
185        }
186    };
187    if page.scroll_id.is_some() {
188        st.scroll_id = page.scroll_id;
189    }
190    if page.hits.is_empty() {
191        st.done = true;
192        if let Some(id) = st.scroll_id.take() {
193            st.inner.client.clear_scroll(&id).await;
194        }
195        return Ok(None);
196    }
197    // A scroll keeps the first request's page size, so the last page can overshoot the fetch:
198    // keep only what is still wanted.
199    let hits = truncate_to_fetch(page.hits, st.remaining);
200    if let Some(r) = st.remaining.as_mut() {
201        *r = r.saturating_sub(hits.len());
202    }
203    // A short page, a single-page read, or an exhausted fetch ends the scan after this batch.
204    if st.inner.single_page || hits.len() < size || st.remaining == Some(0) {
205        st.done = true;
206    }
207    let batch = build_batch(&st.inner.schema, &hits)?;
208    Ok(Some((batch, st)))
209}
210
211/// Drop hits beyond the rows still wanted (`None` = unlimited).
212fn truncate_to_fetch(mut hits: Vec<Value>, remaining: Option<usize>) -> Vec<Value> {
213    if let Some(r) = remaining {
214        hits.truncate(r);
215    }
216    hits
217}
218
219impl ExecutionPlan for OpenSearchExec {
220    fn name(&self) -> &str {
221        "OpenSearchExec"
222    }
223
224    fn properties(&self) -> &Arc<PlanProperties> {
225        &self.properties
226    }
227
228    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
229        vec![]
230    }
231
232    fn with_new_children(self: Arc<Self>, _children: Vec<Arc<dyn ExecutionPlan>>) -> Result<Arc<dyn ExecutionPlan>> {
233        Ok(self)
234    }
235
236    fn execute(&self, partition: usize, _context: Arc<TaskContext>) -> Result<SendableRecordBatchStream> {
237        if partition >= self.inner.partitions {
238            return Err(datafusion::error::DataFusionError::Internal(format!(
239                "OpenSearchExec: partition {partition} requested of {}",
240                self.inner.partitions
241            )));
242        }
243        let state = PageState {
244            inner: self.inner.clone(),
245            partition,
246            scroll_id: None,
247            remaining: self.inner.fetch,
248            started: false,
249            done: false,
250        };
251        let baseline = BaselineMetrics::new(&self.metrics, partition);
252        let stream = futures::stream::try_unfold(state, next_page)
253            .map(move |r| {
254                if let Ok(batch) = &r {
255                    baseline.record_output(batch.num_rows());
256                }
257                r
258            })
259            .boxed();
260        Ok(Box::pin(RecordBatchStreamAdapter::new(
261            self.inner.schema.clone(),
262            stream,
263        )))
264    }
265
266    fn metrics(&self) -> Option<MetricsSet> {
267        Some(self.metrics.clone_inner())
268    }
269
270    fn supports_limit_pushdown(&self) -> bool {
271        true
272    }
273
274    fn with_fetch(&self, limit: Option<usize>) -> Option<Arc<dyn ExecutionPlan>> {
275        let fetch = match (self.inner.fetch, limit) {
276            (Some(a), Some(b)) => Some(a.min(b)),
277            (a, b) => a.or(b),
278        };
279        Some(Arc::new(self.with_fetch_limit(fetch)))
280    }
281
282    fn fetch(&self) -> Option<usize> {
283        self.inner.fetch
284    }
285}
286
287#[cfg(test)]
288mod tests {
289    use super::*;
290    use datafusion::arrow::datatypes::{DataType, Field, Schema};
291    use serde_json::json;
292
293    fn exec(fetch: Option<usize>) -> OpenSearchExec {
294        partitioned(fetch, 1)
295    }
296
297    fn partitioned(fetch: Option<usize>, partitions: usize) -> OpenSearchExec {
298        let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Utf8, true)]));
299        OpenSearchExec::new(
300            OpenSearchClient::new("http://x:9200"),
301            "idx".into(),
302            json!({ "query": { "match_all": {} } }),
303            schema,
304        )
305        .with_page_size(500)
306        .with_fetch_limit(fetch)
307        .with_partitions(partitions)
308    }
309
310    #[test]
311    fn properties_single_bounded_incremental_partition() {
312        let e = exec(None);
313        assert_eq!(e.properties().partitioning.partition_count(), 1);
314        assert_eq!(e.properties().emission_type, EmissionType::Incremental);
315        assert_eq!(e.properties().boundedness, Boundedness::Bounded);
316        assert_eq!(e.schema().fields().len(), 1);
317        assert!(e.children().is_empty());
318    }
319
320    #[test]
321    fn fetch_narrows_but_never_widens() {
322        let e = exec(Some(100));
323        assert_eq!(e.fetch(), Some(100));
324        assert!(e.supports_limit_pushdown());
325        let tighter = e.with_fetch(Some(10)).unwrap();
326        assert_eq!(tighter.fetch(), Some(10));
327        let looser = e.with_fetch(Some(1000)).unwrap();
328        assert_eq!(looser.fetch(), Some(100));
329        assert_eq!(exec(None).with_fetch(Some(7)).unwrap().fetch(), Some(7));
330    }
331
332    #[test]
333    fn the_last_page_is_truncated_to_the_fetch() {
334        let hits: Vec<Value> = (0..5).map(|i| json!({ "n": i })).collect();
335        assert_eq!(truncate_to_fetch(hits.clone(), Some(2)).len(), 2);
336        assert_eq!(truncate_to_fetch(hits.clone(), Some(9)).len(), 5);
337        assert_eq!(truncate_to_fetch(hits.clone(), None).len(), 5);
338        assert!(truncate_to_fetch(hits, Some(0)).is_empty());
339    }
340
341    #[test]
342    fn display_names_the_index() {
343        let e = exec(Some(3));
344        let s = datafusion::physical_plan::displayable(&e).one_line().to_string();
345        assert!(
346            s.contains("OpenSearchExec: index=idx, partitions=1, fetch=Some(3)"),
347            "{s}"
348        );
349        assert_eq!(e.index(), "idx");
350        assert_eq!(e.body()["query"], json!({ "match_all": {} }));
351    }
352}
datafusion_opensearch/exec.rs

datafusion_opensearch/
exec.rs
