JingsongLi commented on code in PR #760:
URL: https://github.com/apache/paimon-rust/pull/760#discussion_r3893836306
##########
crates/integrations/datafusion/src/lateral_vector_search.rs:
##########
@@ -352,9 +452,198 @@ 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: Arc<OnceCell<PreparedVectorSearchFilter>>,
+ partition_count: usize,
+ completed_partitions: HashSet<usize>,
+ active_partition_leases: HashMap<usize, usize>,
+}
+
+#[derive(Debug, Default)]
+struct ExecutionPreparedFilterCache {
+ // DataFusion passes the same TaskContext Arc to every partition of one
+ // execution. Keep the prepared filter alive for that TaskContext so
+ // sequential partitions resolve the same target snapshot. Completion
+ // leases remove the exact entry as soon as every partition finishes.
+ entries: Mutex<Vec<ExecutionPreparedFilterEntry>>,
+}
+
+#[derive(Clone)]
+struct ExecutionPreparedFilterLease {
+ prepared_filter: Arc<OnceCell<PreparedVectorSearchFilter>>,
+ _completion: Arc<ExecutionPartitionCompletion>,
+}
+
+impl ExecutionPreparedFilterLease {
+ fn new(
+ cache: &Arc<ExecutionPreparedFilterCache>,
+ context: Arc<TaskContext>,
+ prepared_filter: Arc<OnceCell<PreparedVectorSearchFilter>>,
+ partition: usize,
+ ) -> Self {
+ Self {
+ prepared_filter: Arc::clone(&prepared_filter),
+ _completion: Arc::new(ExecutionPartitionCompletion {
+ cache: Arc::downgrade(cache),
+ context,
+ prepared_filter,
+ partition,
+ }),
+ }
+ }
+}
+
+struct ExecutionPartitionCompletion {
+ cache: Weak<ExecutionPreparedFilterCache>,
+ context: Arc<TaskContext>,
+ prepared_filter: Arc<OnceCell<PreparedVectorSearchFilter>>,
+ partition: usize,
+}
+
+impl Drop for ExecutionPartitionCompletion {
+ fn drop(&mut self) {
+ if let Some(cache) = self.cache.upgrade() {
+ cache.finish_partition(&self.context, &self.prepared_filter,
self.partition);
+ }
+ }
+}
+
+impl ExecutionPreparedFilterCache {
+ fn for_execution(
+ self: &Arc<Self>,
+ context: &Arc<TaskContext>,
+ partition: usize,
+ partition_count: usize,
+ ) -> ExecutionPreparedFilterLease {
+ debug_assert!(partition < partition_count);
+ let mut entries = self
+ .entries
+ .lock()
+ .unwrap_or_else(std::sync::PoisonError::into_inner);
+ entries.retain(|entry| entry.context.strong_count() > 0);
+ for entry in entries.iter_mut() {
+ let Some(entry_context) = entry.context.upgrade() else {
+ continue;
+ };
+ if Arc::ptr_eq(&entry_context, context) && entry.partition_count
== partition_count {
+ entry.completed_partitions.remove(&partition);
+ *entry.active_partition_leases.entry(partition).or_default()
+= 1;
+ return ExecutionPreparedFilterLease::new(
+ self,
+ Arc::clone(context),
+ Arc::clone(&entry.prepared_filter),
+ partition,
+ );
+ }
+ }
+
+ let prepared_filter = Arc::new(OnceCell::new());
+ let mut active_partition_leases = HashMap::new();
+ active_partition_leases.insert(partition, 1);
+ entries.push(ExecutionPreparedFilterEntry {
+ context: Arc::downgrade(context),
+ prepared_filter: Arc::clone(&prepared_filter),
+ partition_count,
+ completed_partitions: HashSet::new(),
+ active_partition_leases,
+ });
+ ExecutionPreparedFilterLease::new(self, Arc::clone(context),
prepared_filter, partition)
+ }
+
+ fn finish_partition(
+ &self,
+ context: &Arc<TaskContext>,
+ prepared_filter: &Arc<OnceCell<PreparedVectorSearchFilter>>,
+ partition: usize,
+ ) {
+ let mut entries = self
+ .entries
+ .lock()
+ .unwrap_or_else(std::sync::PoisonError::into_inner);
+ let Some(entry_index) = entries
+ .iter()
+ .position(|entry| Arc::ptr_eq(&entry.prepared_filter,
prepared_filter))
+ else {
+ return;
+ };
+ let entry = &mut entries[entry_index];
+ let Some(active_leases) =
entry.active_partition_leases.get_mut(&partition) else {
+ return;
+ };
+ *active_leases -= 1;
+ if *active_leases == 0 {
+ entry.active_partition_leases.remove(&partition);
+ entry.completed_partitions.insert(partition);
+ }
+ if entry.active_partition_leases.is_empty()
+ && (Arc::strong_count(context) == 1
Review Comment:
[P1] Do not make cleanup depend on the context already being dropped
This still retains a subset execution when the caller keeps its normal
`Arc<TaskContext>` until after the returned stream finishes. For example:
create a 4-partition execution, execute/consume only partition 0, drop its
stream/lease while the caller still owns the context, then drop the caller’s
context. At `finish_partition` time the strong count is greater than 1 and not
all declared partitions completed, so the entry stays; the later context drop
has no callback to re-run this condition. The weak entry is only pruned by a
future `for_execution`, leaving the pinned table and potentially large bitmap
alive for the physical plan’s idle lifetime. The new subset test avoids this
order by dropping `context` before `lease`, so it does not cover the common
retained-context case. Please tie the cached value to an
execution/context-owned lifetime (or another actual completion hook) and add
the stream-first/context-second regression.
--
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]
