(&self, args: &[Expr])
| 76 | |
| 77 | impl 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 query_vector_json = |
| 88 | extract_string_literal(FUNCTION_NAME, &args[2], "query_vector_json")?; |
| 89 | let limit = extract_int_literal(FUNCTION_NAME, &args[3], "limit")?; |
| 90 | |
| 91 | if limit <= 0 { |
| 92 | return Err(datafusion::error::DataFusionError::Plan( |
| 93 | "vector_search: limit must be positive".to_string(), |
| 94 | )); |
| 95 | } |
| 96 | |
| 97 | let query_vector: Vec<f32> = serde_json::from_str(&query_vector_json).map_err(|e| { |
| 98 | datafusion::error::DataFusionError::Plan(format!( |
| 99 | "vector_search: query_vector_json must be a JSON array of floats, got '{}': {}", |
| 100 | query_vector_json, e |
| 101 | )) |
| 102 | })?; |
| 103 | |
| 104 | if query_vector.is_empty() { |
| 105 | return Err(datafusion::error::DataFusionError::Plan( |
| 106 | "vector_search: query vector cannot be empty".to_string(), |
| 107 | )); |
| 108 | } |
| 109 | |
| 110 | let identifier = |
| 111 | parse_table_identifier(FUNCTION_NAME, &table_name, &self.default_database)?; |
| 112 | |
| 113 | let catalog = Arc::clone(&self.catalog); |
| 114 | let table = block_on_with_runtime( |
| 115 | async move { catalog.get_table(&identifier).await }, |
| 116 | "vector_search: catalog access thread panicked", |
| 117 | ) |
| 118 | .map_err(to_datafusion_error)?; |
| 119 | |
| 120 | let inner = PaimonTableProvider::try_new(table)?; |
| 121 | |
| 122 | Ok(Arc::new(VectorSearchTableProvider { |
| 123 | inner, |
| 124 | column_name, |
| 125 | query_vector, |
| 126 | limit: limit as usize, |
| 127 | })) |
| 128 | } |
| 129 | } |
| 130 | |
| 131 | #[derive(Debug)] |
nothing calls this directly
no test coverage detected