From 867538a4569c2992e892bda00646d7dfaedcb2db Mon Sep 17 00:00:00 2001 From: shyjsarah <44659226+shyjsarah@users.noreply.github.com> Date: Mon, 27 Jul 2026 04:48:49 -0700 Subject: [PATCH 1/2] fix(datafusion): apply dynamic options to vector search --- .../datafusion/src/sql_context.rs | 15 ++- .../datafusion/src/vector_search.rs | 106 +++++++++++++++++- 2 files changed, 117 insertions(+), 4 deletions(-) diff --git a/crates/integrations/datafusion/src/sql_context.rs b/crates/integrations/datafusion/src/sql_context.rs index 3190487b..eb37966e 100644 --- a/crates/integrations/datafusion/src/sql_context.rs +++ b/crates/integrations/datafusion/src/sql_context.rs @@ -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(), + ); self.catalogs.insert(catalog_name.clone(), catalog); if is_first { self.set_current_catalog(catalog_name).await?; @@ -3257,9 +3262,15 @@ fn register_table_functions( ctx: &SessionContext, catalog: &Arc, 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); diff --git a/crates/integrations/datafusion/src/vector_search.rs b/crates/integrations/datafusion/src/vector_search.rs index cdfed832..ed2bded0 100644 --- a/crates/integrations/datafusion/src/vector_search.rs +++ b/crates/integrations/datafusion/src/vector_search.rs @@ -55,6 +55,7 @@ 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"; @@ -62,16 +63,30 @@ pub fn register_vector_search( ctx: &SessionContext, catalog: Arc, 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, + 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, default_database: String, + dynamic_options: DynamicOptions, } impl Debug for VectorSearchFunction { @@ -84,9 +99,18 @@ impl Debug for VectorSearchFunction { impl VectorSearchFunction { pub fn new(catalog: Arc, default_database: &str) -> Self { + Self::new_with_dynamic_options(catalog, default_database, Default::default()) + } + + pub(crate) fn new_with_dynamic_options( + catalog: Arc, + default_database: &str, + dynamic_options: DynamicOptions, + ) -> Self { Self { catalog, default_database: default_database.to_string(), + dynamic_options, } } } @@ -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", )?; @@ -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, \ + 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::() + .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" + ); + } +} From 9cbfbea7b7a00270285790254c7e8f48f9c21e76 Mon Sep 17 00:00:00 2001 From: shyjsarah <44659226+shyjsarah@users.noreply.github.com> Date: Mon, 27 Jul 2026 05:14:00 -0700 Subject: [PATCH 2/2] fix(datafusion): exclude vector search time travel options --- .../datafusion/src/vector_search.rs | 46 +++++++++++++++---- 1 file changed, 37 insertions(+), 9 deletions(-) diff --git a/crates/integrations/datafusion/src/vector_search.rs b/crates/integrations/datafusion/src/vector_search.rs index ed2bded0..48f9afe6 100644 --- a/crates/integrations/datafusion/src/vector_search.rs +++ b/crates/integrations/datafusion/src/vector_search.rs @@ -45,6 +45,8 @@ use futures::{stream, TryStreamExt}; use paimon::catalog::Catalog; use paimon::spec::{ BigIntType, CoreOptions, DataField, DataType, ROW_ID_FIELD_ID, ROW_ID_FIELD_NAME, + SCAN_SNAPSHOT_ID_OPTION, SCAN_TAG_NAME_OPTION, SCAN_TIMESTAMP_MILLIS_OPTION, + SCAN_VERSION_OPTION, }; use paimon::table::Table; @@ -137,18 +139,16 @@ 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 dynamic_options = vector_search_dynamic_options(&self.dynamic_options); let table = block_on_with_runtime( async move { let table = load_data_table_for_read(&catalog, &identifier, FUNCTION_NAME).await?; - if dynamic_options.is_empty() { - Ok(table) - } else { + let table = if dynamic_options.is_empty() { table - .copy_with_time_travel(dynamic_options) - .await - .map_err(to_datafusion_error) - } + } else { + table.copy_with_options(dynamic_options) + }; + Ok::<_, DataFusionError>(table) }, "vector_search: catalog access thread panicked", )?; @@ -461,6 +461,21 @@ impl ExecutionPlan for VectorSearchExec { } } +/// Vector search currently resolves candidates from the latest snapshot, so forwarding a +/// time-travel selector would search one snapshot and materialize rows from another. +fn vector_search_dynamic_options(dynamic_options: &DynamicOptions) -> HashMap { + let mut options = dynamic_options.read().unwrap().clone(); + for key in [ + SCAN_VERSION_OPTION, + SCAN_TIMESTAMP_MILLIS_OPTION, + SCAN_SNAPSHOT_ID_OPTION, + SCAN_TAG_NAME_OPTION, + ] { + options.remove(key); + } + options +} + /// Projected user columns (+ internal `_ROW_ID`, needed to realign rows to rank). /// Errors if the table has no row tracking, since results then can't be ordered. fn projected_read_fields( @@ -581,7 +596,7 @@ mod tests { use crate::SQLContext; #[tokio::test] - async fn test_vector_search_applies_session_dynamic_options() { + async fn test_vector_search_applies_supported_session_dynamic_options() { let temp_dir = tempfile::tempdir().unwrap(); let mut catalog_options = Options::new(); catalog_options.set( @@ -612,6 +627,10 @@ mod tests { .sql("SET 'paimon.blob-as-descriptor' = 'true'") .await .unwrap(); + sql_context + .sql("SET 'paimon.scan.version' = '1'") + .await + .unwrap(); let state = sql_context.ctx().state(); let table_function = state @@ -635,5 +654,14 @@ mod tests { CoreOptions::new(provider.inner.table().schema().options()).blob_as_descriptor(), "vector_search should apply session dynamic options to the loaded table" ); + assert!( + !provider + .inner + .table() + .schema() + .options() + .contains_key(SCAN_VERSION_OPTION), + "vector_search should not forward unsupported time-travel options" + ); } }