diff --git a/crates/lance-csr/src/csr_expand.rs b/crates/lance-csr/src/csr_expand.rs index 49a8f90..da229fe 100644 --- a/crates/lance-csr/src/csr_expand.rs +++ b/crates/lance-csr/src/csr_expand.rs @@ -903,6 +903,10 @@ impl EdgeRows { } impl EdgeRows { + fn in_memory(&self) -> bool { + matches!(self.data, TableData::Memory(_)) + } + async fn take(&self, rows: &[u64], columns: &[usize]) -> Result> { self.data.take(&self.schema, rows, columns, true).await } @@ -1012,10 +1016,11 @@ impl fmt::Debug for EdgeFilterState { } } -/// Switch from fetching filter columns per expanded edge to one read of the whole -/// column once the expanded edges pass 1/`EDGE_MASK_FRACTION` of the relationship's -/// edges. A batch of scattered rows from a large Lance table touches most of the -/// column's pages, so each batch fetched before the switch costs about one full read. +/// For a Lance relationship table, switch from fetching filter columns per expanded +/// edge to one read of the whole column once the expanded edges pass +/// 1/`EDGE_MASK_FRACTION` of the relationship's edges. A batch of scattered rows from +/// a large Lance table touches most of the column's pages, so each batch fetched before +/// the switch costs about one full read. In-memory tables never switch. // shortcut: fixed share, measured on IC5 (half of a 1.6M-edge table) and IC11; tune if hops in between regress. const EDGE_MASK_FRACTION: u64 = 256; @@ -2197,7 +2202,12 @@ impl ExecutionPlan for CsrExpandExec { let target_filter = self.target_filter.clone().map(Arc::new); let edge_filter = self.edge_filter.clone().map(Arc::new); let edge_state = self.edge_state.clone(); - let mask_after = (self.adjacency.num_edges() / EDGE_MASK_FRACTION).max(1); + // In memory, each partition gathers its chunk's filter columns from loaded arrays, + // in parallel; the shared mask only beats scattered Lance reads. + let mask_after = match &self.edge { + Some(edge) if edge.rows.in_memory() => u64::MAX, + _ => (self.adjacency.num_edges() / EDGE_MASK_FRACTION).max(1), + }; let target_state = self.target_state.clone(); // Every non-id column the target fetches: its output and its filter's columns. let scan_columns: Arc> = Arc::new( diff --git a/docs/csr.md b/docs/csr.md index c502b15..e24d8be 100644 --- a/docs/csr.md +++ b/docs/csr.md @@ -179,9 +179,10 @@ nodes are separate rows, each with its own properties, as on the join path. A `WHERE` condition that reads only `k`'s own columns, such as `k.since > 2020`, moves into the expansion (`edge_filter=` in `explain()`), so failing edges are dropped before anything else is fetched. A small hop -fetches just the filter's columns for each batch of expanded edges. Once a hop has expanded more than 1/256 -of the relationship's edges, it reads the filter's columns once, in order, and skips failing edges from then -on. Conditions that mix `k` with other variables stay after the expansion. +fetches just the filter's columns for each batch of expanded edges. On a Lance dataset, once a hop has +expanded more than 1/256 of the relationship's edges, it reads the filter's columns once, in order, and skips +failing edges from then on. In-memory PyArrow tables always filter each batch, since fetching their columns +is an in-memory gather. Conditions that mix `k` with other variables stay after the expansion. !!! note "Large hops on lazily read tables" diff --git a/python/python/tests/test_engine_csr.py b/python/python/tests/test_engine_csr.py index 248ce60..28aa28e 100644 --- a/python/python/tests/test_engine_csr.py +++ b/python/python/tests/test_engine_csr.py @@ -176,8 +176,9 @@ def test_csr_operator_used_only_for_supported_shapes(datasets): @pytest.mark.parametrize("lazy", [False, True]) def test_relationship_filter_over_many_batches(tmp_path, lazy): - # Many small input batches, so the expansion passes the point where it switches - # from per-batch filter fetches to one read of the filter column. + # Many small input batches. From Lance, the expansion passes the point where it + # switches from per-batch filter fetches to one read of the filter column; in + # memory it fetches per batch throughout. n = 300 people = pa.table({"id": list(range(n)), "age": [i % 90 for i in range(n)]}) knows = pa.table(