indexlake_datafusion/
search.rs

1use std::sync::Arc;
2
3use arrow::datatypes::SchemaRef;
4use datafusion_common::tree_node::TreeNodeRecursion;
5use datafusion_common::{DataFusionError, project_schema};
6use datafusion_execution::{SendableRecordBatchStream, TaskContext};
7use datafusion_physical_expr::{EquivalenceProperties, PhysicalExpr};
8use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType};
9use datafusion_physical_plan::stream::RecordBatchStreamAdapter;
10use datafusion_physical_plan::{
11    DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties,
12};
13use futures::TryStreamExt;
14use indexlake::index::SearchQuery;
15use indexlake::table::TableSearch;
16
17use crate::LazyTable;
18
19#[derive(Debug)]
20pub struct IndexLakeSearchExec {
21    pub lazy_table: LazyTable,
22    pub output_schema: SchemaRef,
23    pub query: Arc<dyn SearchQuery>,
24    pub dynamic_fields: Vec<String>,
25    pub projection: Option<Vec<usize>>,
26    pub limit: Option<usize>,
27    properties: Arc<PlanProperties>,
28}
29
30impl IndexLakeSearchExec {
31    pub fn try_new(
32        lazy_table: LazyTable,
33        output_schema: SchemaRef,
34        query: Arc<dyn SearchQuery>,
35        dynamic_fields: Vec<String>,
36        projection: Option<Vec<usize>>,
37        limit: Option<usize>,
38    ) -> Result<Self, DataFusionError> {
39        let projected_schema = project_schema(&output_schema, projection.as_ref())?;
40        let exec_schema =
41            merge_dynamic_fields(&lazy_table, &query, &projected_schema, &dynamic_fields)?;
42        let properties = Arc::new(PlanProperties::new(
43            EquivalenceProperties::new(exec_schema),
44            Partitioning::UnknownPartitioning(1),
45            EmissionType::Incremental,
46            Boundedness::Bounded,
47        ));
48        Ok(Self {
49            lazy_table,
50            output_schema,
51            query,
52            dynamic_fields,
53            projection,
54            limit,
55            properties,
56        })
57    }
58}
59
60impl ExecutionPlan for IndexLakeSearchExec {
61    fn name(&self) -> &str {
62        "IndexLakeSearchExec"
63    }
64
65    fn properties(&self) -> &Arc<PlanProperties> {
66        &self.properties
67    }
68
69    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
70        vec![]
71    }
72
73    fn with_new_children(
74        self: Arc<Self>,
75        _children: Vec<Arc<dyn ExecutionPlan>>,
76    ) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
77        Ok(self)
78    }
79
80    fn apply_expressions(
81        &self,
82        _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion, DataFusionError>,
83    ) -> Result<TreeNodeRecursion, DataFusionError> {
84        Ok(TreeNodeRecursion::Continue)
85    }
86
87    fn execute(
88        &self,
89        partition: usize,
90        _context: Arc<TaskContext>,
91    ) -> Result<SendableRecordBatchStream, DataFusionError> {
92        if partition != 0 {
93            return Err(DataFusionError::Execution(format!(
94                "partition index out of range: {partition} >= 1"
95            )));
96        }
97
98        let lazy_table = self.lazy_table.clone();
99        let query = self.query.clone();
100        let dynamic_fields = self.dynamic_fields.clone();
101        let projection = self.projection.clone();
102        let limit = self.limit;
103
104        let fut = async move {
105            let table = lazy_table.get_or_load().await?;
106
107            let search = TableSearch {
108                query,
109                projection: projection.clone(),
110                dynamic_fields: dynamic_fields.clone(),
111                limit,
112                concurrency: 8,
113            };
114
115            let stream = table.search(search).await?;
116            let stream = stream.map_err(DataFusionError::from);
117            Ok::<_, DataFusionError>(stream)
118        };
119
120        let stream = futures::stream::once(fut).try_flatten();
121        Ok(Box::pin(RecordBatchStreamAdapter::new(
122            self.schema(),
123            stream,
124        )))
125    }
126
127    fn fetch(&self) -> Option<usize> {
128        self.limit
129    }
130
131    fn with_fetch(&self, limit: Option<usize>) -> Option<Arc<dyn ExecutionPlan>> {
132        match IndexLakeSearchExec::try_new(
133            self.lazy_table.clone(),
134            self.output_schema.clone(),
135            self.query.clone(),
136            self.dynamic_fields.clone(),
137            self.projection.clone(),
138            limit,
139        ) {
140            Ok(exec) => Some(Arc::new(exec)),
141            Err(e) => {
142                log::error!("[indexlake] Failed to create IndexLakeSearchExec with fetch: {e}");
143                None
144            }
145        }
146    }
147}
148
149impl DisplayAs for IndexLakeSearchExec {
150    fn fmt_as(&self, _t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result {
151        write!(
152            f,
153            "IndexLakeSearchExec: table={}.{}, kind={}",
154            self.lazy_table.namespace_name,
155            self.lazy_table.table_name,
156            self.query.index_kind()
157        )?;
158        if !self.dynamic_fields.is_empty() {
159            write!(f, ", dynamic_fields=[{}]", self.dynamic_fields.join(", "))?;
160        }
161        if let Some(ref projection) = self.projection {
162            write!(f, ", projection={projection:?}")?;
163        }
164        if let Some(limit) = self.limit {
165            write!(f, ", limit={limit}")?;
166        }
167        Ok(())
168    }
169}
170
171/// Extend the projected schema with dynamic field columns.
172///
173/// Dynamic field Arrow types are resolved from the [`IndexKind`]
174/// via `lazy_table.client.index_kinds`.
175fn merge_dynamic_fields(
176    lazy_table: &LazyTable,
177    query: &Arc<dyn SearchQuery>,
178    projected_schema: &SchemaRef,
179    dynamic_fields: &[String],
180) -> Result<SchemaRef, DataFusionError> {
181    if dynamic_fields.is_empty() {
182        return Ok(projected_schema.clone());
183    }
184
185    let index_kind = lazy_table
186        .client
187        .index_kinds
188        .get(query.index_kind())
189        .ok_or_else(|| {
190            DataFusionError::Internal(format!("Index kind '{}' not found", query.index_kind()))
191        })?;
192
193    let resolved_fields = index_kind.dynamic_fields();
194
195    let mut fields = projected_schema.fields().to_vec();
196    for name in dynamic_fields {
197        let field = resolved_fields
198            .iter()
199            .find(|f| f.name() == name.as_str())
200            .cloned()
201            .ok_or_else(|| {
202                DataFusionError::Internal(format!(
203                    "Dynamic field '{}' not found in index kind '{}'",
204                    name,
205                    query.index_kind()
206                ))
207            })?;
208        fields.push(field);
209    }
210    Ok(Arc::new(arrow::datatypes::Schema::new_with_metadata(
211        fields,
212        projected_schema.metadata().clone(),
213    )))
214}
indexlake_datafusion/search.rs

indexlake_datafusion/
search.rs
