shyjsarah commented on code in PR #760:
URL: https://github.com/apache/paimon-rust/pull/760#discussion_r3900326138
##########
crates/paimon/src/table/vector_search_builder.rs:
##########
@@ -1221,7 +1268,12 @@ impl<'a> BatchVectorSearchBuilder<'a> {
// returns data-derived row ids/scores outside `TableScan`/`TableRead`,
// so it must refuse a `query-auth.enabled` table before any fast path
// (an empty snapshot would otherwise return empty results and bypass
it).
- let core = CoreOptions::new(self.table.schema().options());
+ let execution_table = self
+ .prepared_filter
+ .as_ref()
+ .map(PreparedVectorSearchFilter::table)
Review Comment:
Fixed in `d11eaaf`. `BatchVectorSearchBuilder::execute` now validates the
prepared filter against the builder target before selecting the pinned table:
normalized table location (ignoring trailing `/`) and branch must match,
otherwise it fails closed with `DataInvalid`. Added an A/B table regression
proving a prepared filter cannot retarget another builder.
##########
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:
Fixed in `d11eaaf`. Because DataFusion 54 exposes no `TaskContext` drop
hook, the execution cache now runs a single cache-level reaper over weak
contexts while entries exist. It removes subset executions after the caller
drops its context even when the stream/lease was dropped first; all-partition
completion and context-already-gone cleanup remain immediate. The reaper is
started only after an entry is installed to avoid a concurrent empty-cache
restart race. Added the requested stream-first/context-second regression.
##########
crates/paimon/src/vindex/reader.rs:
##########
@@ -421,16 +421,37 @@ fn search_vindex(
Ok(Some(id_to_scores))
}
-#[derive(Clone, PartialEq, Eq)]
+#[derive(Clone)]
struct PreparedSearch {
params: VectorSearchParams,
- filter_bytes: Option<Vec<u8>>,
+ filter_bytes: Option<Arc<[u8]>>,
}
+impl PreparedSearch {
+ fn same_batch_group(&self, other: &Self) -> bool {
+ self.params == other.params
+ && match (&self.filter_bytes, &other.filter_bytes) {
+ (Some(left), Some(right)) => Arc::ptr_eq(left, right),
Review Comment:
Fixed in `d11eaaf`. Batch grouping keeps the `Arc::ptr_eq` fast path for the
crate-internal shared allocation and falls back to serialized-filter content
equality for equal filters supplied through the public `with_include_row_ids`
API. The regression now again builds two separate public filters and verifies
filtered native batching/probed-list I/O reuse against scalar execution.
--
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]
