paimon_datafusion/
vector_search.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use std::fmt::Debug;
19use std::sync::Arc;
20
21use async_trait::async_trait;
22use datafusion::arrow::datatypes::SchemaRef as ArrowSchemaRef;
23use datafusion::catalog::Session;
24use datafusion::catalog::TableFunctionImpl;
25use datafusion::common::project_schema;
26use datafusion::datasource::{TableProvider, TableType};
27use datafusion::error::{DataFusionError, Result as DFResult};
28use datafusion::logical_expr::{Expr, TableProviderFilterPushDown};
29use datafusion::physical_plan::empty::EmptyExec;
30use datafusion::physical_plan::ExecutionPlan;
31use datafusion::prelude::SessionContext;
32use paimon::catalog::Catalog;
33
34use crate::error::to_datafusion_error;
35use crate::runtime::{await_with_runtime, block_on_with_runtime};
36use crate::table::{PaimonScanBuilder, PaimonTableProvider};
37use crate::table_function_args::{
38    extract_int_literal, extract_string_literal, parse_table_identifier,
39};
40use crate::table_loader::load_data_table_for_read;
41
42const FUNCTION_NAME: &str = "vector_search";
43
44pub fn register_vector_search(
45    ctx: &SessionContext,
46    catalog: Arc<dyn Catalog>,
47    default_database: &str,
48) {
49    ctx.register_udtf(
50        "vector_search",
51        Arc::new(VectorSearchFunction::new(catalog, default_database)),
52    );
53}
54
55pub struct VectorSearchFunction {
56    catalog: Arc<dyn Catalog>,
57    default_database: String,
58}
59
60impl Debug for VectorSearchFunction {
61    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
62        f.debug_struct("VectorSearchFunction")
63            .field("default_database", &self.default_database)
64            .finish()
65    }
66}
67
68impl VectorSearchFunction {
69    pub fn new(catalog: Arc<dyn Catalog>, default_database: &str) -> Self {
70        Self {
71            catalog,
72            default_database: default_database.to_string(),
73        }
74    }
75}
76
77impl TableFunctionImpl for VectorSearchFunction {
78    fn call(&self, args: &[Expr]) -> DFResult<Arc<dyn TableProvider>> {
79        if args.len() != 4 {
80            return Err(datafusion::error::DataFusionError::Plan(
81                "vector_search requires 4 arguments: (table_name, column_name, query_vector_json, limit)".to_string(),
82            ));
83        }
84
85        let table_name = extract_string_literal(FUNCTION_NAME, &args[0], "table_name")?;
86        let column_name = extract_string_literal(FUNCTION_NAME, &args[1], "column_name")?;
87        let limit = extract_int_literal(FUNCTION_NAME, &args[3], "limit")?;
88
89        if limit <= 0 {
90            return Err(DataFusionError::Plan(
91                "vector_search: limit must be positive".to_string(),
92            ));
93        }
94
95        let identifier =
96            parse_table_identifier(FUNCTION_NAME, &table_name, &self.default_database)?;
97
98        let catalog = Arc::clone(&self.catalog);
99        let table = block_on_with_runtime(
100            async move { load_data_table_for_read(&catalog, &identifier, FUNCTION_NAME).await },
101            "vector_search: catalog access thread panicked",
102        )?;
103
104        let inner = PaimonTableProvider::try_new(table)?;
105        let query_vector_json =
106            match extract_string_literal(FUNCTION_NAME, &args[2], "query_vector_json") {
107                Ok(value) => value,
108                Err(_) if matches!(args[2], Expr::Column(_)) => {
109                    return Ok(Arc::new(LateralVectorSearchTableProvider {
110                        inner,
111                        column_name,
112                        query_vector_expr: args[2].clone(),
113                        limit: limit as usize,
114                    }));
115                }
116                Err(err) => return Err(err),
117            };
118
119        let query_vector: Vec<f32> = serde_json::from_str(&query_vector_json).map_err(|e| {
120            DataFusionError::Plan(format!(
121                "vector_search: query_vector_json must be a JSON array of floats, got '{}': {}",
122                query_vector_json, e
123            ))
124        })?;
125
126        if query_vector.is_empty() {
127            return Err(DataFusionError::Plan(
128                "vector_search: query vector cannot be empty".to_string(),
129            ));
130        }
131
132        Ok(Arc::new(VectorSearchTableProvider {
133            inner,
134            column_name,
135            query_vector,
136            limit: limit as usize,
137        }))
138    }
139}
140
141#[derive(Debug)]
142pub(crate) struct LateralVectorSearchTableProvider {
143    inner: PaimonTableProvider,
144    column_name: String,
145    query_vector_expr: Expr,
146    limit: usize,
147}
148
149impl LateralVectorSearchTableProvider {
150    pub(crate) fn inner(&self) -> &PaimonTableProvider {
151        &self.inner
152    }
153
154    pub(crate) fn column_name(&self) -> &str {
155        &self.column_name
156    }
157
158    pub(crate) fn query_vector_expr(&self) -> &Expr {
159        &self.query_vector_expr
160    }
161
162    pub(crate) fn limit(&self) -> usize {
163        self.limit
164    }
165}
166
167#[async_trait]
168impl TableProvider for LateralVectorSearchTableProvider {
169    fn schema(&self) -> ArrowSchemaRef {
170        self.inner.schema()
171    }
172
173    fn table_type(&self) -> TableType {
174        TableType::Base
175    }
176
177    async fn scan(
178        &self,
179        _state: &dyn Session,
180        _projection: Option<&Vec<usize>>,
181        _filters: &[Expr],
182        _limit: Option<usize>,
183    ) -> DFResult<Arc<dyn ExecutionPlan>> {
184        Err(DataFusionError::Plan(
185            "lateral vector_search must be planned through a lateral join".to_string(),
186        ))
187    }
188}
189
190#[derive(Debug)]
191struct VectorSearchTableProvider {
192    inner: PaimonTableProvider,
193    column_name: String,
194    query_vector: Vec<f32>,
195    limit: usize,
196}
197
198#[async_trait]
199impl TableProvider for VectorSearchTableProvider {
200    fn schema(&self) -> ArrowSchemaRef {
201        self.inner.schema()
202    }
203
204    fn table_type(&self) -> TableType {
205        TableType::Base
206    }
207
208    async fn scan(
209        &self,
210        state: &dyn Session,
211        projection: Option<&Vec<usize>>,
212        _filters: &[Expr],
213        limit: Option<usize>,
214    ) -> DFResult<Arc<dyn ExecutionPlan>> {
215        let table = self.inner.table();
216
217        let row_ranges = await_with_runtime(async {
218            let mut builder = table.new_vector_search_builder();
219            builder
220                .with_vector_column(&self.column_name)
221                .with_query_vector(self.query_vector.clone())
222                .with_limit(self.limit);
223            builder.execute().await.map_err(to_datafusion_error)
224        })
225        .await?;
226
227        if row_ranges.is_empty() {
228            let schema = project_schema(&self.schema(), projection)?;
229            return Ok(Arc::new(EmptyExec::new(schema)));
230        }
231
232        let mut read_builder = table.new_read_builder();
233        if let Some(limit) = limit {
234            read_builder.with_limit(limit);
235        }
236        let scan = read_builder.new_scan().with_row_ranges(row_ranges);
237        let plan = await_with_runtime(scan.plan())
238            .await
239            .map_err(to_datafusion_error)?;
240
241        let target = state.config_options().execution.target_partitions;
242        PaimonScanBuilder {
243            table,
244            schema: &self.schema(),
245            plan: &plan,
246            scan_trace: None,
247            projection,
248            pushed_predicate: None,
249            limit,
250            target_partitions: target,
251            filter_exact: false,
252            case_sensitive: true,
253        }
254        .build()
255    }
256
257    fn supports_filters_pushdown(
258        &self,
259        filters: &[&Expr],
260    ) -> DFResult<Vec<TableProviderFilterPushDown>> {
261        Ok(vec![
262            TableProviderFilterPushDown::Unsupported;
263            filters.len()
264        ])
265    }
266}
paimon_datafusion/vector_search.rs

paimon_datafusion/
vector_search.rs
