Skip to content
Open
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 13 additions & 2 deletions crates/integrations/datafusion/src/sql_context.rs
Original file line number Diff line number Diff line change
Expand Up @@ -198,7 +198,12 @@ impl SQLContext {
Some(session_state),
)),
);
register_table_functions(&self.ctx, &catalog, default_db.unwrap_or("default"));
register_table_functions(
&self.ctx,
&catalog,
default_db.unwrap_or("default"),
self.dynamic_options.clone(),
Comment thread
shyjsarah marked this conversation as resolved.
);
self.catalogs.insert(catalog_name.clone(), catalog);
if is_first {
self.set_current_catalog(catalog_name).await?;
Expand Down Expand Up @@ -3257,9 +3262,15 @@ fn register_table_functions(
ctx: &SessionContext,
catalog: &Arc<dyn Catalog>,
default_database: &str,
dynamic_options: DynamicOptions,
) {
crate::blob_view::register_blob_view(ctx, Arc::clone(catalog), default_database);
crate::vector_search::register_vector_search(ctx, Arc::clone(catalog), default_database);
crate::vector_search::register_vector_search_with_dynamic_options(
ctx,
Arc::clone(catalog),
default_database,
dynamic_options,
);
#[cfg(feature = "fulltext")]
crate::full_text_search::register_full_text_search(ctx, Arc::clone(catalog), default_database);
crate::hybrid_search::register_hybrid_search(ctx, Arc::clone(catalog), default_database);
Expand Down
106 changes: 104 additions & 2 deletions crates/integrations/datafusion/src/vector_search.rs
Original file line number Diff line number Diff line change
Expand Up @@ -55,23 +55,38 @@ use crate::table_function_args::{
extract_int_literal, extract_string_literal, parse_table_identifier,
};
use crate::table_loader::load_data_table_for_read;
use crate::DynamicOptions;

const FUNCTION_NAME: &str = "vector_search";

pub fn register_vector_search(
ctx: &SessionContext,
catalog: Arc<dyn Catalog>,
default_database: &str,
) {
register_vector_search_with_dynamic_options(ctx, catalog, default_database, Default::default());
}

pub(crate) fn register_vector_search_with_dynamic_options(
ctx: &SessionContext,
catalog: Arc<dyn Catalog>,
default_database: &str,
dynamic_options: DynamicOptions,
) {
ctx.register_udtf(
"vector_search",
Arc::new(VectorSearchFunction::new(catalog, default_database)),
Arc::new(VectorSearchFunction::new_with_dynamic_options(
catalog,
default_database,
dynamic_options,
)),
);
}

pub struct VectorSearchFunction {
catalog: Arc<dyn Catalog>,
default_database: String,
dynamic_options: DynamicOptions,
}

impl Debug for VectorSearchFunction {
Expand All @@ -84,9 +99,18 @@ impl Debug for VectorSearchFunction {

impl VectorSearchFunction {
pub fn new(catalog: Arc<dyn Catalog>, default_database: &str) -> Self {
Self::new_with_dynamic_options(catalog, default_database, Default::default())
}

pub(crate) fn new_with_dynamic_options(
catalog: Arc<dyn Catalog>,
default_database: &str,
dynamic_options: DynamicOptions,
) -> Self {
Self {
catalog,
default_database: default_database.to_string(),
dynamic_options,
}
}
}
Expand All @@ -113,8 +137,19 @@ impl TableFunctionImpl for VectorSearchFunction {
parse_table_identifier(FUNCTION_NAME, &table_name, &self.default_database)?;

let catalog = Arc::clone(&self.catalog);
let dynamic_options = self.dynamic_options.read().unwrap().clone();
let table = block_on_with_runtime(
async move { load_data_table_for_read(&catalog, &identifier, FUNCTION_NAME).await },
async move {
let table = load_data_table_for_read(&catalog, &identifier, FUNCTION_NAME).await?;
if dynamic_options.is_empty() {
Ok(table)
} else {
table
.copy_with_time_travel(dynamic_options)
.await
.map_err(to_datafusion_error)
}
},
"vector_search: catalog access thread panicked",
)?;

Expand Down Expand Up @@ -535,3 +570,70 @@ fn gather_rows_by_rank(
RecordBatch::try_new_with_options(Arc::clone(output_schema), columns, &options)
.map_err(DataFusionError::from)
}

#[cfg(test)]
mod tests {
use datafusion::catalog::TableFunctionArgs;
use datafusion::logical_expr::lit;
use paimon::{CatalogOptions, FileSystemCatalog, Options};

use super::*;
use crate::SQLContext;

#[tokio::test]
async fn test_vector_search_applies_session_dynamic_options() {
let temp_dir = tempfile::tempdir().unwrap();
let mut catalog_options = Options::new();
catalog_options.set(
CatalogOptions::WAREHOUSE,
format!("file://{}", temp_dir.path().display()),
);
let catalog = Arc::new(FileSystemCatalog::new(catalog_options).unwrap());

let mut sql_context = SQLContext::new();
sql_context
.register_catalog("paimon", catalog)
.await
.unwrap();
sql_context
.sql(
"CREATE TABLE paimon.default.vector_blob (\
id INT, \
embedding ARRAY<FLOAT>, \
picture BLOB\
) WITH (\
'data-evolution.enabled' = 'true', \
'row-tracking.enabled' = 'true'\
)",
)
.await
.unwrap();
sql_context
.sql("SET 'paimon.blob-as-descriptor' = 'true'")
.await
.unwrap();

let state = sql_context.ctx().state();
let table_function = state
.table_functions()
.get(FUNCTION_NAME)
.expect("vector_search should be registered");
let args = [
lit("paimon.default.vector_blob"),
lit("embedding"),
lit("[1.0]"),
lit(1_i64),
];
let provider = table_function
.create_table_provider_with_args(TableFunctionArgs::new(&args, &state))
.unwrap();
let provider = provider
.downcast_ref::<VectorSearchTableProvider>()
.expect("vector_search should return its table provider");

assert!(
CoreOptions::new(provider.inner.table().schema().options()).blob_as_descriptor(),
"vector_search should apply session dynamic options to the loaded table"
);
}
}
Loading