shyjsarah commented on code in PR #760:
URL: https://github.com/apache/paimon-rust/pull/760#discussion_r3893611674
##########
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:
Fixed in `14a22f0`. The cache now tracks only partitions that actually
acquire stream leases, keeps the `TaskContext` with each lease, and wraps the
child stream so cancellation/end/error drops the inner stream before releasing
the cache lease. Eviction occurs when no streams remain and either the
execution context is otherwise gone or all declared partitions completed. Added
regressions for subset execution cleanup and cancellation drop ordering; the
sequential-partition snapshot reuse integration test remains green.
##########
crates/paimon/src/table/vector_search_builder.rs:
##########
@@ -1194,14 +1216,32 @@ impl<'a> BatchVectorSearchBuilder<'a> {
self
}
- /// Attach a residual scalar predicate applied *after* vector recall on the
- /// primary-key vector path, shared across every query in the batch.
Mirrors
- /// the single [`VectorSearchBuilder::with_filter`]: only the primary-key
- /// vector path (via [`execute_read`](Self::execute_read)) consumes it,
and only
- /// when the table exposes physical rows directly (deletion vectors without
- /// merge-on-read); otherwise the query fails loud.
+ /// Attach one scalar predicate shared by every query in the batch and
applied
+ /// before vector Top-K. See [`VectorSearchBuilder::with_filter`] for the
+ /// primary-key and data-evolution execution semantics.
pub fn with_filter(&mut self, filter: Predicate) -> &mut Self {
self.filter = Some(filter);
+ self.include_row_ids = None;
+ self
+ }
+
+ /// Reuse row IDs from a previously prepared scalar pre-filter.
+ ///
+ /// The builder's table must be [`PreparedVectorSearchFilter::table`] (or
an
+ /// equivalent copy pinned to the same snapshot).
+ pub fn with_include_row_ids(&mut self, include_row_ids: RoaringTreemap) ->
&mut Self {
+ self.include_row_ids = Some(Arc::new(include_row_ids));
+ self.filter = None;
+ self
+ }
+
+ /// Reuse a shared row-ID allow-list without copying its bitmap.
+ pub fn with_shared_include_row_ids(
Review Comment:
Fixed in `14a22f0`. `BatchVectorSearchBuilder::with_prepared_filter` now
accepts the complete `PreparedVectorSearchFilter` and executes against its
pinned table/snapshot atomically. DataFusion callers were migrated to this API;
the raw bitmap setter remains documented as a lower-level, snapshot-unbound API.
##########
crates/paimon/src/table/vector_search_builder.rs:
##########
@@ -3046,7 +3313,7 @@ impl RawScoringPlan {
.collect();
for (query_index, vector_search) in vector_searches.iter().enumerate()
{
Review Comment:
Fixed in `14a22f0`. The raw scoring plan preserves a shared filter as one
`Arc<RoaringTreemap>` plus one query-index vector, avoiding `O(Q×B)` row/query
associations. Before planning the raw read, filtered queries now intersect
unindexed/detail ranges with the shared filter (or the union for distinct
filters). Added regressions for bounded planning state and sparse range pruning.
##########
crates/paimon/src/vindex/reader.rs:
##########
@@ -472,7 +472,7 @@ fn prepare_search(
),
};
- let filter_bytes = if let Some(include_ids) =
&vector_search.include_row_ids {
+ let filter_bytes = if let Some(include_ids) =
vector_search.effective_include_row_ids() {
Review Comment:
Fixed in `14a22f0`. Batch preparation detects a shared bitmap by `Arc`
identity, serializes it once into one `Arc<[u8]>`, reuses that allocation for
every prepared query, and groups native batch searches by params plus
serialized-filter identity rather than byte-vector comparison. Added a
128-query/100k-row regression asserting all prepared searches share the same
serialized allocation.
--
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]
