jerry-024 commented on code in PR #760:
URL: https://github.com/apache/paimon-rust/pull/760#discussion_r3893241665
##########
crates/integrations/datafusion/src/lateral_vector_search.rs:
##########
@@ -352,9 +450,141 @@ 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,
+ unfinished_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>,
+ prepared_filter: Arc<OnceCell<PreparedVectorSearchFilter>>,
+ partition: usize,
+ ) -> Self {
+ Self {
+ prepared_filter: Arc::clone(&prepared_filter),
+ _completion: Arc::new(ExecutionPartitionCompletion {
+ cache: Arc::downgrade(cache),
+ prepared_filter,
+ partition,
+ }),
+ }
+ }
+
+ fn prepared_filter(&self) -> &OnceCell<PreparedVectorSearchFilter> {
+ &self.prepared_filter
+ }
+}
+
+struct ExecutionPartitionCompletion {
+ cache: Weak<ExecutionPreparedFilterCache>,
+ 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.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.unfinished_partitions.insert(partition);
+ *entry.active_partition_leases.entry(partition).or_default()
+= 1;
+ return ExecutionPreparedFilterLease::new(
+ self,
+ 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,
+ unfinished_partitions: (0..partition_count).collect(),
Review Comment:
<!-- dlf-review -->
**[MAJOR]** Do not require every declared partition to start before cleanup
This pre-populates completion state for every declared partition, but an
entry is cleared only when each partition has acquired and dropped a lease. A
consumer may execute only a subset of partitions, or stop early, so unstarted
indexes can never call finish_partition; the prepared bitmap and pinned table
then remain retained for the physical plan's lifetime. Please tie eviction to
execution-level completion or cancellation instead of requiring every declared
partition to start.
--
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]
