jerry-024 commented on code in PR #760:
URL: https://github.com/apache/paimon-rust/pull/760#discussion_r3892666239
##########
crates/integrations/datafusion/src/lateral_vector_search.rs:
##########
@@ -352,9 +450,57 @@ struct LateralVectorSearchExec {
query_vector_expr: Arc<dyn PhysicalExpr>,
limit: usize,
output_schema: ArrowSchemaRef,
+ filter: Option<Predicate>,
+ prepared_filter_cache: Arc<ExecutionPreparedFilterCache>,
plan_properties: Arc<PlanProperties>,
}
+#[derive(Debug)]
+struct ExecutionPreparedFilterEntry {
+ context: Weak<TaskContext>,
+ prepared_filter: Weak<OnceCell<PreparedVectorSearchFilter>>,
+}
+
+#[derive(Debug, Default)]
+struct ExecutionPreparedFilterCache {
+ // DataFusion passes the same TaskContext Arc to every partition of one
+ // execution. Weak references let those partitions share one prepared
filter
+ // while ensuring a later execution of the reusable plan starts fresh.
+ entries: Mutex<Vec<ExecutionPreparedFilterEntry>>,
+}
+
+impl ExecutionPreparedFilterCache {
+ fn for_execution(
+ &self,
+ context: &Arc<TaskContext>,
+ ) -> Arc<OnceCell<PreparedVectorSearchFilter>> {
+ let mut entries = self
+ .entries
+ .lock()
+ .unwrap_or_else(std::sync::PoisonError::into_inner);
+ entries.retain(|entry| {
+ entry.context.strong_count() > 0 &&
entry.prepared_filter.strong_count() > 0
+ });
+ for entry in entries.iter() {
+ let Some(entry_context) = entry.context.upgrade() else {
+ continue;
+ };
+ if Arc::ptr_eq(&entry_context, context) {
+ if let Some(prepared_filter) = entry.prepared_filter.upgrade()
{
+ return prepared_filter;
+ }
+ }
+ }
+
+ let prepared_filter = Arc::new(OnceCell::new());
+ entries.push(ExecutionPreparedFilterEntry {
+ context: Arc::downgrade(context),
+ prepared_filter: Arc::downgrade(&prepared_filter),
+ });
Review Comment:
<!-- dlf-review -->
**[MAJOR]** The execution-scoped prefilter cache only retains a weak
reference to the prepared filter, so one logical execution can resolve
different target snapshots across partitions. Once an earlier partition stream
is exhausted and drops the last strong reference, a later partition using the
same live `TaskContext` creates a new `OnceCell` and resolves the table
snapshot again.
Keep the prepared filter strongly owned for the complete execution across
all partitions, releasing it only when that execution finishes. Please add a
sequential-partition test that updates the target between partitions.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
