Skip to main content

ursa_engine/graph/
exec.rs

1//! Executing a graph node.
2//!
3//! [`GraphPlanner`] maps a [`GraphAlgorithmNode`] to a [`GraphAlgorithmExec`]
4//! whose children are the planned edge plans. When it executes, it looks up
5//! the index in the session's [`IndexCache`]; on a miss it reads a child,
6//! partition by partition, and builds the index as the rows arrive. Then it runs the calls and
7//! emits one batch.
8//!
9//! Row order matters to the index: dense node ids are given in first-seen
10//! order, and they decide component labels and the order floating-point sums
11//! run in. The child's partitions are read in order, the node asks for no
12//! repartitioning of its input, and the session turns off file-scan work
13//! stealing, so a file's edges arrive in the file's own order on any thread
14//! count, as `ursa-io` reads them for the Python library.
15//!
16//! ## Weights
17//!
18//! Weights are looked up by edge row, so they must be in the order of the
19//! rows the index was built from. A weighted node reads only its weighted
20//! child (the endpoints and every weight expression, in one pass, evaluating
21//! the weights a batch at a time as it goes), and:
22//!
23//! - on an index-cache miss, builds the index from those same rows, so the
24//!   weights line up by construction;
25//! - on a hit, checks that the rows' endpoints are the index's, row for row
26//!   (an edge plan over a join need not produce its rows in the same order
27//!   twice). If they are not, it reads the child again and builds a private
28//!   index from those rows rather than mis-assign a weight.
29//!
30//! Each distinct weight expression is evaluated once per node, cast to
31//! `Float64`, and checked as `ursa-plan` checks weights (no nulls, no negative
32//! values, with the same messages). Weights that line up with a cached index
33//! are stored on it, keyed by the weighted plan's fingerprint and file
34//! identities, so the next query with the same weights reads no edges at all.
35//! In a database, they are stored with the index on disk too, so a weighted
36//! query over a stored index in a later process reads no edges either.
37
38use std::fmt;
39use std::sync::{Arc, Mutex};
40
41use async_trait::async_trait;
42use datafusion::arrow::array::{new_empty_array, Array, ArrayRef, Float64Array};
43use datafusion::arrow::compute::SortOptions;
44use datafusion::arrow::compute::{cast, cast_with_options, concat, CastOptions};
45use datafusion::arrow::datatypes::{DataType, SchemaRef};
46use datafusion::arrow::record_batch::RecordBatch;
47use datafusion::catalog::Session;
48use datafusion::common::config::ConfigOptions;
49use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion};
50use datafusion::common::ScalarValue;
51use datafusion::error::{DataFusionError, Result};
52use datafusion::execution::context::QueryPlanner;
53use datafusion::execution::TaskContext;
54use datafusion::logical_expr::physical_planning_context::PhysicalPlanningContext;
55use datafusion::logical_expr::{LogicalPlan, UserDefinedLogicalNode};
56use datafusion::physical_expr::expressions::Column;
57use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr};
58use datafusion::physical_expr::{LexOrdering, PhysicalSortExpr};
59use datafusion::physical_optimizer::PhysicalOptimizerRule;
60use datafusion::physical_plan::aggregates::AggregateExec;
61use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
62use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
63use datafusion::physical_plan::repartition::RepartitionExec;
64use datafusion::physical_plan::replace_children_if_necessary;
65use datafusion::physical_plan::sorts::sort::SortExec;
66use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
67use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
68use datafusion::physical_plan::union::{InterleaveExec, UnionExec};
69use datafusion::physical_plan::{
70    DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, Partitioning,
71    PlanProperties, SendableRecordBatchStream, StatisticsArgs, StatisticsContext,
72};
73use datafusion::physical_planner::{DefaultPhysicalPlanner, ExtensionPlanner, PhysicalPlanner};
74use futures::{FutureExt, StreamExt};
75use ursa_core::{EdgeInterner, IdError, Topology};
76
77use super::index::{file_identities, fingerprint, FileIdentity, GraphIndex, IndexCache, Weights};
78use super::kernels;
79use super::node::{
80    weight_column, GraphAlgorithmNode, GraphTraversalNode, Output, WEIGHTED_DESTINATION,
81    WEIGHTED_SOURCE,
82};
83use super::traversal::{self, Traversal, TraversalParam, TraversalParamType};
84
85/// Plans Ursa's graph nodes.
86#[derive(Debug)]
87pub struct GraphPlanner;
88
89#[async_trait]
90impl ExtensionPlanner for GraphPlanner {
91    async fn plan_extension(
92        &self,
93        _planner: &dyn PhysicalPlanner,
94        node: &dyn UserDefinedLogicalNode,
95        logical_inputs: &[&LogicalPlan],
96        physical_inputs: &[Arc<dyn ExecutionPlan>],
97        _session_state: &dyn Session,
98        _planning_ctx: &PhysicalPlanningContext,
99    ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
100        if let Some(node) = node.as_any().downcast_ref::<GraphTraversalNode>() {
101            return plan_traversal(node, logical_inputs, physical_inputs).map(Some);
102        }
103        let Some(node) = node.as_any().downcast_ref::<GraphAlgorithmNode>() else {
104            return Ok(None);
105        };
106        let (source, destination) = node.endpoint_indices()?;
107        let weighted = match (
108            &node.weighted,
109            logical_inputs.get(1),
110            physical_inputs.get(1),
111        ) {
112            (Some(_), Some(logical), Some(physical)) => {
113                Some(weighted_input(logical, physical, node.weights.len())?)
114            }
115            _ => None,
116        };
117        // A neighbour aggregation's node values are the last input.
118        let values = node
119            .values
120            .as_ref()
121            .and_then(|_| physical_inputs.last().cloned());
122        Ok(Some(Arc::new(GraphAlgorithmExec::new(
123            Shape {
124                edges: Edges {
125                    graph: node.graph.clone(),
126                    endpoints: (source, destination),
127                    id_type: node.id_type.clone(),
128                    fingerprint: fingerprint(logical_inputs[0]),
129                },
130                outputs: node.outputs.clone(),
131                per_node: node.per_node,
132                display: node.display_outputs(),
133                schema: Arc::new(node.schema().as_arrow().clone()),
134            },
135            Arc::clone(&physical_inputs[0]),
136            weighted,
137            values,
138        ))))
139    }
140}
141
142/// The weighted edge plan `physical` (planned from `logical`), with
143/// `n_weights` weight columns.
144fn weighted_input(
145    logical: &LogicalPlan,
146    physical: &Arc<dyn ExecutionPlan>,
147    n_weights: usize,
148) -> Result<WeightedInput> {
149    let schema = logical.schema();
150    let column = |name: &str| schema.index_of_column_by_name(None, name);
151    let missing =
152        |name: &str| DataFusionError::Internal(format!("the weighted edge plan has no {name}"));
153    let weights = (0..n_weights)
154        .map(|i| {
155            let name = weight_column(i);
156            column(&name).ok_or_else(|| missing(&name))
157        })
158        .collect::<Result<Vec<_>>>()?;
159    Ok(WeightedInput {
160        plan: Arc::clone(physical),
161        source: column(WEIGHTED_SOURCE).ok_or_else(|| missing(WEIGHTED_SOURCE))?,
162        destination: column(WEIGHTED_DESTINATION).ok_or_else(|| missing(WEIGHTED_DESTINATION))?,
163        weights,
164        fingerprint: fingerprint(logical),
165    })
166}
167
168/// The exec of a traversal node.
169fn plan_traversal(
170    node: &GraphTraversalNode,
171    logical_inputs: &[&LogicalPlan],
172    physical_inputs: &[Arc<dyn ExecutionPlan>],
173) -> Result<Arc<dyn ExecutionPlan>> {
174    let expected = 1 + usize::from(node.weighted.is_some()) + node.nodes.len();
175    if logical_inputs.len() != expected || physical_inputs.len() != expected {
176        return Err(DataFusionError::Internal(format!(
177            "GraphTraversal has {expected} inputs; planned {}",
178            physical_inputs.len()
179        )));
180    }
181    let (source, destination) = node.endpoint_indices()?;
182    let mut next = 1;
183    let weighted = match node.weighted {
184        Some(_) => {
185            next += 1;
186            Some(weighted_input(logical_inputs[1], &physical_inputs[1], 1)?)
187        }
188        None => None,
189    };
190    Ok(Arc::new(GraphTraversalExec::new(
191        TraversalShape {
192            edges: Edges {
193                graph: node.graph.clone(),
194                endpoints: (source, destination),
195                id_type: node.id_type.clone(),
196                fingerprint: fingerprint(logical_inputs[0]),
197            },
198            traversal: node.traversal,
199            display: node.display_call(),
200            schema: Arc::new(node.schema().as_arrow().clone()),
201        },
202        Arc::clone(&physical_inputs[0]),
203        weighted,
204        physical_inputs[next..].to_vec(),
205    )))
206}
207
208/// A query planner whose physical planner knows Ursa's graph nodes.
209#[derive(Debug)]
210pub struct UrsaQueryPlanner;
211
212#[async_trait]
213impl QueryPlanner for UrsaQueryPlanner {
214    async fn create_physical_plan(
215        &self,
216        logical_plan: &LogicalPlan,
217        session_state: &dyn Session,
218    ) -> Result<Arc<dyn ExecutionPlan>> {
219        DefaultPhysicalPlanner::with_extension_planners(vec![Arc::new(GraphPlanner)])
220            .create_physical_plan(logical_plan, session_state)
221            .await
222    }
223}
224
225/// A graph node's edge input: where its endpoints are, and how its index is
226/// cached.
227#[derive(Debug, Clone)]
228struct Edges {
229    graph: String,
230    /// The source and destination columns of the first input.
231    endpoints: (usize, usize),
232    id_type: DataType,
233    /// The edge plan's fingerprint, or `None` if it cannot be cached.
234    fingerprint: Option<String>,
235}
236
237/// What a [`GraphAlgorithmExec`] computes, besides its children.
238#[derive(Debug, Clone)]
239struct Shape {
240    edges: Edges,
241    outputs: Vec<Output>,
242    per_node: bool,
243    /// The outputs as `EXPLAIN` shows them.
244    display: String,
245    schema: SchemaRef,
246}
247
248/// The weighted edge plan, and where its columns are.
249#[derive(Debug, Clone)]
250struct WeightedInput {
251    plan: Arc<dyn ExecutionPlan>,
252    source: usize,
253    destination: usize,
254    /// One column per weight expression, in the node's order.
255    weights: Vec<usize>,
256    /// The weighted plan's fingerprint, or `None` if it cannot be cached.
257    fingerprint: Option<String>,
258}
259
260/// Builds or fetches a graph index from its child's edge rows, and runs
261/// graph functions over it.
262#[derive(Debug)]
263pub struct GraphAlgorithmExec {
264    shape: Shape,
265    input: Arc<dyn ExecutionPlan>,
266    weighted: Option<WeightedInput>,
267    /// A neighbour aggregation's node values: node ids, then values.
268    values: Option<Arc<dyn ExecutionPlan>>,
269    properties: Arc<PlanProperties>,
270}
271
272impl GraphAlgorithmExec {
273    fn new(
274        shape: Shape,
275        input: Arc<dyn ExecutionPlan>,
276        weighted: Option<WeightedInput>,
277        values: Option<Arc<dyn ExecutionPlan>>,
278    ) -> Self {
279        let properties = Arc::new(PlanProperties::new(
280            EquivalenceProperties::new(Arc::clone(&shape.schema)),
281            Partitioning::UnknownPartitioning(1),
282            EmissionType::Final,
283            Boundedness::Bounded,
284        ));
285        Self {
286            shape,
287            input,
288            weighted,
289            values,
290            properties,
291        }
292    }
293}
294
295impl DisplayAs for GraphAlgorithmExec {
296    fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result {
297        write!(
298            f,
299            "GraphAlgorithmExec: graph={}, outputs=[{}]",
300            self.shape.edges.graph, self.shape.display
301        )
302    }
303}
304
305impl ExecutionPlan for GraphAlgorithmExec {
306    fn name(&self) -> &str {
307        "GraphAlgorithmExec"
308    }
309
310    fn properties(&self) -> &Arc<PlanProperties> {
311        &self.properties
312    }
313
314    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
315        let mut children = vec![&self.input];
316        children.extend(self.weighted.as_ref().map(|w| &w.plan));
317        children.extend(self.values.as_ref());
318        children
319    }
320
321    /// The edges are read partition by partition in order; splitting them
322    /// further would only scramble that order.
323    fn benefits_from_input_partitioning(&self) -> Vec<bool> {
324        vec![false; self.children().len()]
325    }
326
327    fn apply_expressions(
328        &self,
329        _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
330    ) -> Result<TreeNodeRecursion> {
331        Ok(TreeNodeRecursion::Continue)
332    }
333
334    fn with_new_children(
335        self: Arc<Self>,
336        children: Vec<Arc<dyn ExecutionPlan>>,
337    ) -> Result<Arc<dyn ExecutionPlan>> {
338        let expected = self.children().len();
339        if children.len() != expected {
340            return Err(DataFusionError::Internal(format!(
341                "GraphAlgorithmExec has {expected} children; got {}",
342                children.len()
343            )));
344        }
345        let mut children = children.into_iter();
346        let input = children.next().expect("checked above");
347        let weighted = self.weighted.clone().map(|w| WeightedInput {
348            plan: children.next().expect("checked above"),
349            ..w
350        });
351        let values = self
352            .values
353            .as_ref()
354            .map(|_| children.next().expect("checked above"));
355        Ok(Arc::new(Self::new(
356            self.shape.clone(),
357            input,
358            weighted,
359            values,
360        )))
361    }
362
363    fn execute(
364        &self,
365        partition: usize,
366        context: Arc<TaskContext>,
367    ) -> Result<SendableRecordBatchStream> {
368        if partition != 0 {
369            return Err(DataFusionError::Internal(format!(
370                "GraphAlgorithmExec has one partition; asked for {partition}"
371            )));
372        }
373        let shape = self.shape.clone();
374        let input = Arc::clone(&self.input);
375        let weighted = self.weighted.clone();
376        let values_plan = self.values.clone();
377        let schema = Arc::clone(&shape.schema);
378        let run = async move {
379            let (index, weights) = match weighted {
380                None => (unweighted_index(&shape.edges, input, &context).await?, None),
381                Some(w) => {
382                    let (index, weights) = weighted_index(&shape.edges, input, w, &context).await?;
383                    (index, Some(weights))
384                }
385            };
386            let values = match values_plan {
387                Some(plan) => {
388                    let batches = drain(&plan, &[0, 1], &context).await?;
389                    Some(kernels::node_values(&index, &shape.edges.graph, batches)?)
390                }
391                None => None,
392            };
393            // The kernels are CPU-bound and parallel (rayon): keep them off
394            // the async workers (inline in the browser, which has none).
395            crate::run_blocking(move || {
396                let weights = weights.as_deref().map_or(&[][..], Vec::as_slice);
397                kernels::batch(
398                    &index,
399                    &shape.outputs,
400                    weights,
401                    values.as_deref(),
402                    shape.per_node,
403                    shape.schema,
404                )
405            })
406            .await
407            .map_err(|e| DataFusionError::Execution(format!("a graph kernel failed: {e}")))?
408        };
409        Ok(Box::pin(RecordBatchStreamAdapter::new(
410            schema,
411            futures::stream::once(run),
412        )))
413    }
414}
415
416/// What a [`GraphTraversalExec`] computes, besides its children.
417#[derive(Debug, Clone)]
418struct TraversalShape {
419    edges: Edges,
420    traversal: Traversal,
421    /// The call as `EXPLAIN` shows it.
422    display: String,
423    schema: SchemaRef,
424}
425
426/// Builds or fetches a graph index from its first child's edge rows, reads
427/// the node ids its other children produce, and runs a traversal.
428#[derive(Debug)]
429pub struct GraphTraversalExec {
430    shape: TraversalShape,
431    input: Arc<dyn ExecutionPlan>,
432    weighted: Option<WeightedInput>,
433    /// One plan per node-id parameter, in the spec's order.
434    nodes: Vec<Arc<dyn ExecutionPlan>>,
435    properties: Arc<PlanProperties>,
436}
437
438impl GraphTraversalExec {
439    fn new(
440        shape: TraversalShape,
441        input: Arc<dyn ExecutionPlan>,
442        weighted: Option<WeightedInput>,
443        nodes: Vec<Arc<dyn ExecutionPlan>>,
444    ) -> Self {
445        let properties = Arc::new(PlanProperties::new(
446            EquivalenceProperties::new(Arc::clone(&shape.schema)),
447            Partitioning::UnknownPartitioning(1),
448            EmissionType::Final,
449            Boundedness::Bounded,
450        ));
451        Self {
452            shape,
453            input,
454            weighted,
455            nodes,
456            properties,
457        }
458    }
459}
460
461impl DisplayAs for GraphTraversalExec {
462    fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result {
463        write!(
464            f,
465            "GraphTraversalExec: graph={}, call={}",
466            self.shape.edges.graph, self.shape.display
467        )
468    }
469}
470
471impl ExecutionPlan for GraphTraversalExec {
472    fn name(&self) -> &str {
473        "GraphTraversalExec"
474    }
475
476    fn properties(&self) -> &Arc<PlanProperties> {
477        &self.properties
478    }
479
480    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
481        let mut children = vec![&self.input];
482        children.extend(self.weighted.as_ref().map(|w| &w.plan));
483        children.extend(self.nodes.iter());
484        children
485    }
486
487    /// The edges are read in order (see [`GraphAlgorithmExec`]), and the node
488    /// ids in the order their query gives them, which sets the order of the
489    /// rows.
490    fn benefits_from_input_partitioning(&self) -> Vec<bool> {
491        vec![false; self.children().len()]
492    }
493
494    fn apply_expressions(
495        &self,
496        _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
497    ) -> Result<TreeNodeRecursion> {
498        Ok(TreeNodeRecursion::Continue)
499    }
500
501    fn with_new_children(
502        self: Arc<Self>,
503        children: Vec<Arc<dyn ExecutionPlan>>,
504    ) -> Result<Arc<dyn ExecutionPlan>> {
505        let expected = self.children().len();
506        if children.len() != expected {
507            return Err(DataFusionError::Internal(format!(
508                "GraphTraversalExec has {expected} children; got {}",
509                children.len()
510            )));
511        }
512        let mut children = children.into_iter();
513        let input = children.next().expect("checked above");
514        let weighted = self.weighted.clone().map(|w| WeightedInput {
515            plan: children.next().expect("checked above"),
516            ..w
517        });
518        Ok(Arc::new(Self::new(
519            self.shape.clone(),
520            input,
521            weighted,
522            children.collect(),
523        )))
524    }
525
526    fn execute(
527        &self,
528        partition: usize,
529        context: Arc<TaskContext>,
530    ) -> Result<SendableRecordBatchStream> {
531        if partition != 0 {
532            return Err(DataFusionError::Internal(format!(
533                "GraphTraversalExec has one partition; asked for {partition}"
534            )));
535        }
536        let shape = self.shape.clone();
537        let input = Arc::clone(&self.input);
538        let weighted = self.weighted.clone();
539        let node_plans = self.nodes.clone();
540        let schema = Arc::clone(&shape.schema);
541        let run = async move {
542            let (index, weights) = match weighted {
543                None => (unweighted_index(&shape.edges, input, &context).await?, None),
544                Some(w) => {
545                    let (index, weights) = weighted_index(&shape.edges, input, w, &context).await?;
546                    (index, Some(weights))
547                }
548            };
549            let spec = shape.traversal.spec();
550            let mut nodes = Vec::with_capacity(node_plans.len());
551            for (param, plan) in spec.node_params().zip(&node_plans) {
552                let batches = drain(plan, &[0], &context).await?;
553                nodes.push(node_ids(&shape, param, &index, batches)?);
554            }
555            crate::run_blocking(move || {
556                let weights = weights.as_ref().map(|w| w[0].as_slice());
557                traversal::batch(
558                    &index.topology,
559                    &index.ids,
560                    &shape.traversal,
561                    &nodes,
562                    weights,
563                    shape.schema,
564                )
565            })
566            .await
567            .map_err(|e| DataFusionError::Execution(format!("a graph kernel failed: {e}")))?
568        };
569        Ok(Box::pin(RecordBatchStreamAdapter::new(
570            schema,
571            futures::stream::once(run),
572        )))
573    }
574}
575
576/// The ids a node-id input produced, as dense indexes of `index` (`None` for a
577/// null or an id not in the graph), cast to the graph's id type first.
578fn node_ids(
579    shape: &TraversalShape,
580    param: &TraversalParam,
581    index: &GraphIndex,
582    batches: Vec<Vec<ArrayRef>>,
583) -> Result<Vec<Option<u32>>> {
584    let (function, name, graph) = (shape.traversal.name(), param.name, &shape.edges.graph);
585    let id_type = index.ids.user_type();
586    let options = CastOptions {
587        safe: false,
588        ..CastOptions::default()
589    };
590    let arrays = batches
591        .iter()
592        .map(|b| {
593            // No float, boolean or date, and a string must be exactly an
594            // integer on a graph with integer ids (#271).
595            if let Some(row) = traversal::first_refused_node_id(&b[0], &id_type) {
596                let got = sql_text(&b[0], row)
597                    .unwrap_or_else(|| format!("a value of type {}", b[0].data_type()));
598                return Err(DataFusionError::Execution(traversal::not_node_ids(
599                    function, param, graph, &id_type, &got,
600                )));
601            }
602            cast_with_options(&b[0], &id_type, &options).map_err(|e| {
603                let got = first_uncast(&b[0], &id_type).unwrap_or_else(|| {
604                    crate::user_message(&DataFusionError::ArrowError(Box::new(e), None))
605                });
606                DataFusionError::Execution(traversal::not_node_ids(
607                    function, param, graph, &id_type, &got,
608                ))
609            })
610        })
611        .collect::<Result<Vec<_>>>()?;
612    let refs: Vec<&dyn Array> = arrays.iter().map(|a| a.as_ref()).collect();
613    let ids = match refs.as_slice() {
614        [] => new_empty_array(&id_type),
615        _ => concat(&refs)?,
616    };
617    if param.ty == TraversalParamType::Node && ids.len() > 1 {
618        return Err(DataFusionError::Execution(format!(
619            "{function}({name} =>) must be one node id; its query returned {} rows",
620            ids.len()
621        )));
622    }
623    index
624        .ids
625        .dense_from_array(ids.as_ref())
626        .map_err(|e| DataFusionError::Execution(format!("{function}({name} =>): {e}")))
627}
628
629/// The first value of `values` that does not cast to `id_type`, as SQL writes
630/// it (`'Valjean'`, `2.5`), for a message.
631fn first_uncast(values: &ArrayRef, id_type: &DataType) -> Option<String> {
632    let lenient = cast(values, id_type).ok()?;
633    let row = (0..values.len()).find(|&i| values.is_valid(i) && lenient.is_null(i))?;
634    sql_text(values, row)
635}
636
637/// Row `row` of `values` as SQL writes it (`'Valjean'`, `2.5`).
638fn sql_text(values: &ArrayRef, row: usize) -> Option<String> {
639    let mut value = ScalarValue::try_from_array(values, row).ok()?;
640    while let ScalarValue::Dictionary(_, inner) = value {
641        value = *inner;
642    }
643    Some(match value.try_as_str() {
644        Some(Some(s)) => format!("'{}'", s.replace('\'', "''")),
645        _ => value.to_string(),
646    })
647}
648
649/// Keeps a graph node's edge rows in the order the plan produces them, and
650/// gives them a fixed order where the plan has none.
651///
652/// DataFusion's `EnforceDistribution` puts a round-robin repartition under an
653/// operator that benefits from parallel input, such as the `FilterExec` of a
654/// view with a `WHERE` (a subgraph) or the `ProjectionExec` that computes a
655/// weight expression. Draining its partitions in order would then read the
656/// batches out of file order, changing dense ids (and so labels and float
657/// summation order) from what `ursa-io` gives the Python library. This rule
658/// runs after the default rules and removes every round-robin repartition
659/// below a graph node: a round-robin repartition only adds parallelism, so
660/// removing it never changes a result.
661///
662/// What remains may still not preserve order: a join, an aggregate, a hash
663/// repartition, a merge of partitions or a `UNION` produces its rows in an
664/// order that can change between runs and thread counts (#263). Over such a
665/// plan, the rule sorts the edge rows by `(source, destination)`, and a
666/// weighted node's rows by their weights after that, so that rows that tie
667/// are identical in every column the node reads. A traversal's node-id
668/// inputs over such a plan are sorted by value. Order-preserving plans
669/// (scans, views over files, filters, projections) are left as they are,
670/// in file order. The sort is a `SortExec` under the node, so `EXPLAIN`
671/// shows it.
672#[derive(Debug)]
673pub struct KeepEdgeOrder;
674
675impl PhysicalOptimizerRule for KeepEdgeOrder {
676    fn optimize(
677        &self,
678        plan: Arc<dyn ExecutionPlan>,
679        _config: &ConfigOptions,
680    ) -> Result<Arc<dyn ExecutionPlan>> {
681        plan.transform_up(|node| {
682            // The columns each child is sorted by, if its order is not fixed.
683            let keys: Vec<Vec<usize>> =
684                if let Some(exec) = node.downcast_ref::<GraphAlgorithmExec>() {
685                    let mut keys = vec![endpoint_keys(exec.shape.edges.endpoints)];
686                    keys.extend(exec.weighted.as_ref().map(weighted_keys));
687                    // A neighbour aggregation's values give one value per
688                    // node whatever their order: no sort.
689                    keys.extend(exec.values.as_ref().map(|_| vec![]));
690                    keys
691                } else if let Some(exec) = node.downcast_ref::<GraphTraversalExec>() {
692                    let mut keys = vec![endpoint_keys(exec.shape.edges.endpoints)];
693                    keys.extend(exec.weighted.as_ref().map(weighted_keys));
694                    keys.extend(exec.nodes.iter().map(|_| vec![0]));
695                    keys
696                } else {
697                    return Ok(Transformed::no(node));
698                };
699            let children = node
700                .children()
701                .into_iter()
702                .zip(keys)
703                .map(|(child, keys)| {
704                    let child = without_round_robin(Arc::clone(child))?;
705                    if preserves_order(&child) {
706                        Ok(child)
707                    } else {
708                        sorted(child, &keys)
709                    }
710                })
711                .collect::<Result<Vec<_>>>()?;
712            Ok(Transformed::yes(replace_children_if_necessary(
713                Arc::clone(&node),
714                children,
715            )?))
716        })
717        .map(|t| t.data)
718    }
719
720    fn name(&self) -> &str {
721        "keep_edge_order"
722    }
723
724    fn schema_check(&self) -> bool {
725        true
726    }
727}
728
729/// `plan` with its round-robin repartitions removed.
730fn without_round_robin(plan: Arc<dyn ExecutionPlan>) -> Result<Arc<dyn ExecutionPlan>> {
731    plan.transform_up(|node| {
732        if let Some(repartition) = node.downcast_ref::<RepartitionExec>() {
733            if matches!(repartition.partitioning(), Partitioning::RoundRobinBatch(_)) {
734                return Ok(Transformed::yes(Arc::clone(repartition.input())));
735            }
736        }
737        Ok(Transformed::no(node))
738    })
739    .map(|t| t.data)
740}
741
742/// The sort key of an edge input: its source, then its destination.
743fn endpoint_keys((source, destination): (usize, usize)) -> Vec<usize> {
744    vec![source, destination]
745}
746
747/// The sort key of a weighted edge input: its endpoints, then every weight,
748/// so rows that tie are the same in every column the node reads.
749fn weighted_keys(weighted: &WeightedInput) -> Vec<usize> {
750    let mut keys = endpoint_keys((weighted.source, weighted.destination));
751    keys.extend(&weighted.weights);
752    keys
753}
754
755/// Whether reading `plan`'s partitions in order gives its rows in an order
756/// that does not depend on the run or the thread count: true unless it has
757/// a join, an aggregate, a repartition, a merge of partitions or a union. A
758/// nested graph node emits its rows in a fixed order, whatever its inputs.
759fn preserves_order(plan: &Arc<dyn ExecutionPlan>) -> bool {
760    let mut preserves = true;
761    let _ = plan.apply(|node| {
762        if node.downcast_ref::<GraphAlgorithmExec>().is_some()
763            || node.downcast_ref::<GraphTraversalExec>().is_some()
764        {
765            return Ok(TreeNodeRecursion::Jump);
766        }
767        let reorders = node.downcast_ref::<RepartitionExec>().is_some()
768            || node.downcast_ref::<CoalescePartitionsExec>().is_some()
769            || node.downcast_ref::<AggregateExec>().is_some()
770            || node.downcast_ref::<UnionExec>().is_some()
771            || node.downcast_ref::<InterleaveExec>().is_some()
772            || node.name().contains("Join");
773        if reorders {
774            preserves = false;
775            return Ok(TreeNodeRecursion::Stop);
776        }
777        Ok(TreeNodeRecursion::Continue)
778    });
779    preserves
780}
781
782/// `plan`'s rows sorted by the columns `keys`, ascending with nulls last,
783/// as one partition: each partition is sorted, then they are merged.
784fn sorted(plan: Arc<dyn ExecutionPlan>, keys: &[usize]) -> Result<Arc<dyn ExecutionPlan>> {
785    let schema = plan.schema();
786    let exprs = keys.iter().map(|&i| {
787        PhysicalSortExpr::new(
788            Arc::new(Column::new(schema.field(i).name(), i)) as Arc<dyn PhysicalExpr>,
789            SortOptions {
790                descending: false,
791                nulls_first: false,
792            },
793        )
794    });
795    let Some(ordering) = LexOrdering::new(exprs) else {
796        return Ok(plan);
797    };
798    let partitions = plan.output_partitioning().partition_count();
799    let sort =
800        Arc::new(SortExec::new(ordering.clone(), plan).with_preserve_partitioning(partitions > 1));
801    Ok(if partitions > 1 {
802        Arc::new(SortPreservingMergeExec::new(ordering, sort))
803    } else {
804        sort
805    })
806}
807
808/// The index cache key for the edges `input` reads, if they can be cached.
809async fn index_key(
810    shape: &Edges,
811    input: &Arc<dyn ExecutionPlan>,
812    context: &TaskContext,
813) -> Result<Option<(String, Vec<FileIdentity>)>> {
814    Ok(match &shape.fingerprint {
815        Some(f) => Some((f.clone(), file_identities(input, context).await?)),
816        None => None,
817    })
818}
819
820/// The index, cached or built from the endpoint-only input.
821async fn unweighted_index(
822    shape: &Edges,
823    input: Arc<dyn ExecutionPlan>,
824    context: &Arc<TaskContext>,
825) -> Result<Arc<GraphIndex>> {
826    let build = {
827        let (shape, input, context) = (shape.clone(), Arc::clone(&input), Arc::clone(context));
828        async move {
829            let mut no_weights = WeightReader::new(&[], 0);
830            build_streamed(&shape, &input, shape.endpoints, &context, &mut no_weights).await
831        }
832    };
833    match context.session_config().get_extension::<IndexCache>() {
834        Some(cache) => {
835            let key = index_key(shape, &input, context).await?;
836            cache
837                .get_or_build_stored(context.session_config(), key, build.boxed())
838                .await
839        }
840        None => build.await.map(Arc::new),
841    }
842}
843
844/// The index and the weight columns, for a node that reads weights (see the
845/// module docs).
846async fn weighted_index(
847    shape: &Edges,
848    input: Arc<dyn ExecutionPlan>,
849    weighted: WeightedInput,
850    context: &Arc<TaskContext>,
851) -> Result<(Arc<GraphIndex>, Weights)> {
852    let cache = context.session_config().get_extension::<IndexCache>();
853    let key = match &cache {
854        Some(_) => index_key(shape, &input, context).await?,
855        None => None,
856    };
857    let weights_key = match &weighted.fingerprint {
858        Some(f) => Some(format!(
859            "{f}\n{:?}",
860            file_identities(&weighted.plan, context).await?
861        )),
862        None => None,
863    };
864    let n_weights = weighted.weights.len();
865    // Weights already evaluated over the cached index's rows: nothing to read.
866    if let (Some(cache), Some(key), Some(weights_key)) = (&cache, &key, &weights_key) {
867        if let Some(index) = cache.get(key) {
868            if let Some(weights) = known_weights(&index, weights_key, n_weights).await {
869                return Ok((index, weights));
870            }
871        }
872    }
873
874    let (index, weights, shared) = match &cache {
875        Some(cache) => {
876            // Weights read by this node's own build, if the cache ran it.
877            let slot = Arc::new(Mutex::new(None));
878            let build = {
879                let (shape, weighted) = (shape.clone(), weighted.clone());
880                let (context, slot) = (Arc::clone(context), Arc::clone(&slot));
881                async move {
882                    let (index, weights) = build_weighted(&shape, &weighted, &context).await?;
883                    *slot.lock().unwrap_or_else(|e| e.into_inner()) = Some(weights);
884                    Ok(index)
885                }
886            };
887            let index = cache
888                .get_or_build_stored(context.session_config(), key, build.boxed())
889                .await?;
890            let built_here = slot.lock().unwrap_or_else(|e| e.into_inner()).take();
891            // A stored index may have these weights stored with it.
892            if built_here.is_none() {
893                if let Some(weights_key) = &weights_key {
894                    if let Some(weights) = known_weights(&index, weights_key, n_weights).await {
895                        return Ok((index, weights));
896                    }
897                }
898            }
899            match built_here {
900                Some(weights) => (index, weights, true),
901                None => {
902                    let mut weights =
903                        WeightReader::new(&weighted.weights, index.topology.n_edges());
904                    if read_aligned(&index, &weighted, &shape.id_type, context, &mut weights)
905                        .await?
906                    {
907                        (index, weights, true)
908                    } else {
909                        // The cached index's rows are in another order: these
910                        // weights belong to an index of these rows.
911                        cache.record_build();
912                        let (private, weights) = build_weighted(shape, &weighted, context).await?;
913                        (Arc::new(private), weights, false)
914                    }
915                }
916            }
917        }
918        None => {
919            let (index, weights) = build_weighted(shape, &weighted, context).await?;
920            (Arc::new(index), weights, false)
921        }
922    };
923
924    let weights: Weights = Arc::new(weights.finish()?.into_iter().map(Arc::new).collect());
925    if let (true, Some(weights_key)) = (shared, weights_key) {
926        index.store_weights(weights_key, Arc::clone(&weights));
927    }
928    Ok((index, weights))
929}
930
931/// The weights `index` holds under `key` (with `n_weights` columns), in
932/// memory or stored with it on disk.
933async fn known_weights(index: &Arc<GraphIndex>, key: &str, n_weights: usize) -> Option<Weights> {
934    if let Some(weights) = index.weights(key) {
935        return Some(weights);
936    }
937    if !index.has_stored_weights(key) {
938        return None;
939    }
940    let (index, key) = (Arc::clone(index), key.to_string());
941    // Reading a file is blocking work.
942    crate::run_blocking(move || index.stored_weights(&key, n_weights))
943        .await
944        .ok()
945        .flatten()
946}
947
948/// Build the index of the weighted rows, reading their weights in the same
949/// pass, so they line up by construction.
950async fn build_weighted(
951    shape: &Edges,
952    weighted: &WeightedInput,
953    context: &Arc<TaskContext>,
954) -> Result<(GraphIndex, WeightReader)> {
955    let mut weights = WeightReader::new(&weighted.weights, row_hint(&weighted.plan));
956    let endpoints = (weighted.source, weighted.destination);
957    let index = build_streamed(shape, &weighted.plan, endpoints, context, &mut weights).await?;
958    Ok((index, weights))
959}
960
961/// Read the weighted rows once, reading their weights, and say whether
962/// their endpoints are the rows `index` was built from, row for row. Stops
963/// at the first row that is not.
964async fn read_aligned(
965    index: &GraphIndex,
966    weighted: &WeightedInput,
967    id_type: &DataType,
968    context: &Arc<TaskContext>,
969    weights: &mut WeightReader,
970) -> Result<bool> {
971    let mut rows = RowCheck::new(index);
972    let plan = &weighted.plan;
973    for partition in 0..plan.output_partitioning().partition_count() {
974        let mut stream = plan.execute(partition, Arc::clone(context))?;
975        while let Some(batch) = stream.next().await {
976            let batch = batch?;
977            if batch.num_rows() == 0 {
978                continue;
979            }
980            let src = cast(batch.column(weighted.source), id_type)?;
981            let dst = cast(batch.column(weighted.destination), id_type)?;
982            if !rows.check(&src, &dst)? {
983                return Ok(false);
984            }
985            weights.push(&batch);
986        }
987    }
988    Ok(rows.complete())
989}
990
991/// The weight columns of a weighted node's rows, evaluated a batch at a
992/// time as the rows are read. The first bad weight is reported once every
993/// row has been read, after any error in building the index from those
994/// rows, as when the weights were evaluated after the build.
995struct WeightReader {
996    columns: Vec<usize>,
997    values: Vec<Vec<f64>>,
998    error: Option<DataFusionError>,
999}
1000
1001impl WeightReader {
1002    /// A reader of the columns `columns`, sized for `rows` rows.
1003    fn new(columns: &[usize], rows: usize) -> Self {
1004        Self {
1005            columns: columns.to_vec(),
1006            values: columns.iter().map(|_| Vec::with_capacity(rows)).collect(),
1007            error: None,
1008        }
1009    }
1010
1011    fn push(&mut self, batch: &RecordBatch) {
1012        if self.error.is_some() {
1013            return;
1014        }
1015        for (&c, values) in self.columns.iter().zip(&mut self.values) {
1016            if let Err(e) = extend_weights(batch.column(c), values) {
1017                self.error = Some(e);
1018                return;
1019            }
1020        }
1021    }
1022
1023    /// One `f64` per row read, per column, or the first bad weight.
1024    fn finish(self) -> Result<Vec<Vec<f64>>> {
1025        match self.error {
1026            Some(e) => Err(e),
1027            None => Ok(self.values),
1028        }
1029    }
1030}
1031
1032/// Every batch of `input`, partition by partition in order, as the columns
1033/// `columns` of each non-empty batch.
1034async fn drain(
1035    input: &Arc<dyn ExecutionPlan>,
1036    columns: &[usize],
1037    context: &Arc<TaskContext>,
1038) -> Result<Vec<Vec<ArrayRef>>> {
1039    let mut out = Vec::new();
1040    let partitions = input.output_partitioning().partition_count();
1041    for partition in 0..partitions {
1042        let mut stream = input.execute(partition, Arc::clone(context))?;
1043        while let Some(batch) = stream.next().await {
1044            let batch = batch?;
1045            if batch.num_rows() == 0 {
1046                continue;
1047            }
1048            out.push(
1049                columns
1050                    .iter()
1051                    .map(|&c| Arc::clone(batch.column(c)))
1052                    .collect(),
1053            );
1054        }
1055    }
1056    Ok(out)
1057}
1058
1059/// Rows the index build interns at a time. A scan's batches are 8,192 rows,
1060/// and `EdgeInterner` resolves known ids in parallel only in larger chunks,
1061/// so the build gathers batches up to this many rows first: what `ursa-io`'s
1062/// scan does for the Python library (its `SCAN_BATCH_ROWS`). A chunk is
1063/// 2 MiB of `Int64` endpoints.
1064const INTERN_ROWS: usize = 1 << 17;
1065
1066/// Endpoint chunks gathered up to [`INTERN_ROWS`] rows, in row order.
1067#[derive(Default)]
1068struct Gather {
1069    src: Vec<ArrayRef>,
1070    dst: Vec<ArrayRef>,
1071    rows: usize,
1072}
1073
1074impl Gather {
1075    /// Add a chunk; once [`INTERN_ROWS`] rows are gathered, take them.
1076    fn push(&mut self, src: ArrayRef, dst: ArrayRef) -> Result<Option<(ArrayRef, ArrayRef)>> {
1077        self.rows += src.len();
1078        self.src.push(src);
1079        self.dst.push(dst);
1080        if self.rows >= INTERN_ROWS {
1081            self.take()
1082        } else {
1083            Ok(None)
1084        }
1085    }
1086
1087    /// The rows gathered so far, as one chunk, if there are any.
1088    fn take(&mut self) -> Result<Option<(ArrayRef, ArrayRef)>> {
1089        self.rows = 0;
1090        let (src, dst) = (std::mem::take(&mut self.src), std::mem::take(&mut self.dst));
1091        Ok(match src.len() {
1092            0 => None,
1093            1 => src.into_iter().zip(dst).next(),
1094            _ => {
1095                let join = |arrays: &[ArrayRef]| {
1096                    let refs: Vec<&dyn Array> = arrays.iter().map(|a| a.as_ref()).collect();
1097                    concat(&refs)
1098                };
1099                Some((join(&src)?, join(&dst)?))
1100            }
1101        })
1102    }
1103}
1104
1105/// The index of everything `interner` was given; an empty graph of
1106/// `id_type` ids if it was given no edges.
1107fn finish_index(interner: EdgeInterner, id_type: &DataType) -> GraphIndex {
1108    if interner.is_empty() {
1109        let (topology, ids) = ursa_io::topology::empty_topology(id_type);
1110        return GraphIndex::new(topology, ids);
1111    }
1112    let (ids, source, destination) = interner.finish();
1113    let topology = Topology::build(ids.len(), source, destination);
1114    GraphIndex::new(Arc::new(topology), Arc::new(ids))
1115}
1116
1117/// The number of rows `plan` will produce, if its statistics say: sizes the
1118/// interner, as `ursa-io`'s scan sizes it from the Parquet metadata.
1119fn row_hint(plan: &Arc<dyn ExecutionPlan>) -> usize {
1120    StatisticsContext::new()
1121        .compute(plan.as_ref(), &StatisticsArgs::new())
1122        .ok()
1123        .and_then(|s| s.num_rows.get_value().copied())
1124        .unwrap_or(0)
1125}
1126
1127/// What the reader hands the interning thread.
1128enum Chunk {
1129    Edges(ArrayRef, ArrayRef),
1130    /// Every row has been read: build the index.
1131    End,
1132}
1133
1134/// Build the index of the edge rows `input` produces, interning each chunk
1135/// as it is read rather than collecting every endpoint array first, so the
1136/// Arrow input never accumulates: the build holds the dense endpoints and a
1137/// chunk or two in flight, as `ursa-io`'s scan does. Reading (here) and
1138/// interning (on a blocking thread) overlap. The rows are interned in the
1139/// order [`drain`] would collect them, so the index is the one
1140/// `ursa_io::build_topology_batches` gives over the collected rows. A
1141/// weighted node's weights are read from the same rows (`weights`).
1142async fn build_streamed(
1143    shape: &Edges,
1144    input: &Arc<dyn ExecutionPlan>,
1145    (source, destination): (usize, usize),
1146    context: &Arc<TaskContext>,
1147    weights: &mut WeightReader,
1148) -> Result<GraphIndex> {
1149    let (graph, id_type) = (shape.graph.clone(), shape.id_type.clone());
1150    let hint = row_hint(input);
1151    let (tx, mut rx) = tokio::sync::mpsc::channel::<Chunk>(2);
1152
1153    // Native: the interner runs on a blocking thread, so reading and interning
1154    // overlap on two threads. It is spawned here, before the reader starts
1155    // (`run_blocking` would start it only when awaited, after the reader).
1156    #[cfg(not(target_arch = "wasm32"))]
1157    let interning = tokio::task::spawn_blocking(move || -> Result<Option<GraphIndex>> {
1158        let mut interner = EdgeInterner::with_capacity(hint);
1159        while let Some(chunk) = rx.blocking_recv() {
1160            match chunk {
1161                Chunk::Edges(src, dst) => interner
1162                    .push(src.as_ref(), dst.as_ref())
1163                    .map_err(|e| build_error(&graph, e))?,
1164                Chunk::End => return Ok(Some(finish_index(interner, &id_type))),
1165            }
1166        }
1167        // The reader stopped before the end (its error is the one reported).
1168        Ok(None)
1169    });
1170    // The browser (wasm32) has one thread: the interner is a future polled
1171    // alongside the reader, taking each chunk as it is sent. The channel's
1172    // bound still keeps at most two chunks in flight, and the chunks arrive in
1173    // the same order, so the index is the same (#278).
1174    #[cfg(target_arch = "wasm32")]
1175    let interning = async move {
1176        let mut interner = EdgeInterner::with_capacity(hint);
1177        while let Some(chunk) = rx.recv().await {
1178            match chunk {
1179                Chunk::Edges(src, dst) => interner
1180                    .push(src.as_ref(), dst.as_ref())
1181                    .map_err(|e| build_error(&graph, e))?,
1182                Chunk::End => return Ok(Some(finish_index(interner, &id_type))),
1183            }
1184        }
1185        Ok::<_, DataFusionError>(None)
1186    };
1187
1188    let read = async {
1189        // The reader owns the sender, so the interner sees the channel close
1190        // when the reader stops, on an error as at the end.
1191        let tx = tx;
1192        let mut gather = Gather::default();
1193        for partition in 0..input.output_partitioning().partition_count() {
1194            let mut stream = input.execute(partition, Arc::clone(context))?;
1195            while let Some(batch) = stream.next().await {
1196                let batch = batch?;
1197                if batch.num_rows() == 0 {
1198                    continue;
1199                }
1200                let src = cast(batch.column(source), &shape.id_type)?;
1201                let dst = cast(batch.column(destination), &shape.id_type)?;
1202                weights.push(&batch);
1203                drop(batch);
1204                if let Some((src, dst)) = gather.push(src, dst)? {
1205                    if tx.send(Chunk::Edges(src, dst)).await.is_err() {
1206                        // The interner stopped on an error: it is reported below.
1207                        return Ok(());
1208                    }
1209                }
1210            }
1211        }
1212        if let Some((src, dst)) = gather.take()? {
1213            if tx.send(Chunk::Edges(src, dst)).await.is_err() {
1214                return Ok(());
1215            }
1216        }
1217        let _ = tx.send(Chunk::End).await;
1218        Ok::<_, DataFusionError>(())
1219    };
1220    #[cfg(not(target_arch = "wasm32"))]
1221    let (read, built) = {
1222        let read = read.await;
1223        let built = interning.await.map_err(|e| {
1224            DataFusionError::Execution(format!("building graph '{}' failed: {e}", shape.graph))
1225        })?;
1226        (read, built)
1227    };
1228    #[cfg(target_arch = "wasm32")]
1229    let (read, built) = futures::join!(read, interning);
1230    read?;
1231    built?.ok_or_else(|| {
1232        DataFusionError::Internal(format!(
1233            "building graph '{}' stopped before the last edge",
1234            shape.graph
1235        ))
1236    })
1237}
1238
1239/// Checks, a chunk at a time, that rows are the edge rows an index was
1240/// built from, in order.
1241struct RowCheck<'a> {
1242    index: &'a GraphIndex,
1243    /// Each row's endpoints as dense ids, from the CSR and its row
1244    /// permutation.
1245    source: Vec<u32>,
1246    destination: Vec<u32>,
1247    /// Rows checked so far.
1248    start: usize,
1249}
1250
1251impl<'a> RowCheck<'a> {
1252    fn new(index: &'a GraphIndex) -> Self {
1253        let topology = &index.topology;
1254        let n = topology.n_edges();
1255        let out = topology.out();
1256        let (mut source, mut destination) = (vec![0u32; n], vec![0u32; n]);
1257        for u in 0..topology.n_nodes() {
1258            let slots = out.offsets[u] as usize..out.offsets[u + 1] as usize;
1259            for slot in slots {
1260                let row = out.edge_ids[slot] as usize;
1261                source[row] = u as u32;
1262                destination[row] = out.targets[slot];
1263            }
1264        }
1265        Self {
1266            index,
1267            source,
1268            destination,
1269            start: 0,
1270        }
1271    }
1272
1273    /// Whether the next rows are `src` → `dst`. Each chunk's ids are
1274    /// compared on their own, so no column of every row's ids is made.
1275    fn check(&mut self, src: &ArrayRef, dst: &ArrayRef) -> Result<bool> {
1276        let rows = self.start..self.start + src.len();
1277        if rows.end > self.source.len() {
1278            return Ok(false);
1279        }
1280        self.start = rows.end;
1281        let ids = &self.index.ids;
1282        let user_type = ids.user_type();
1283        Ok(
1284            ids.gather_user(&self.source[rows.clone()]).as_ref() == cast(src, &user_type)?.as_ref()
1285                && ids.gather_user(&self.destination[rows]).as_ref()
1286                    == cast(dst, &user_type)?.as_ref(),
1287        )
1288    }
1289
1290    /// Whether every row of the index has been checked.
1291    fn complete(&self) -> bool {
1292        self.start == self.source.len()
1293    }
1294}
1295
1296/// Whether `endpoints` are the edge rows `index` was built from, row for row.
1297#[cfg(test)]
1298fn aligned(index: &GraphIndex, endpoints: &[(ArrayRef, ArrayRef)]) -> Result<bool> {
1299    let mut rows = RowCheck::new(index);
1300    for (src, dst) in endpoints {
1301        if !rows.check(src, dst)? {
1302            return Ok(false);
1303        }
1304    }
1305    Ok(rows.complete())
1306}
1307
1308/// A weight expression's values, one `f64` per edge row, checked as
1309/// `ursa-plan`'s `evaluate_weight` checks them, with its messages.
1310#[cfg(test)]
1311fn edge_weights(columns: &[ArrayRef]) -> Result<Vec<f64>> {
1312    let mut out = Vec::with_capacity(columns.iter().map(|c| c.len()).sum());
1313    for column in columns {
1314        extend_weights(column, &mut out)?;
1315    }
1316    Ok(out)
1317}
1318
1319/// Append one chunk of a weight expression's values to `out`, checked as
1320/// `ursa-plan`'s `evaluate_weight` checks them, with its messages.
1321fn extend_weights(column: &ArrayRef, out: &mut Vec<f64>) -> Result<()> {
1322    let column = cast(column, &DataType::Float64)?;
1323    let column = column
1324        .as_any()
1325        .downcast_ref::<Float64Array>()
1326        .expect("cast to Float64 yields Float64Array");
1327    if column.null_count() > 0 {
1328        return Err(DataFusionError::Execution(
1329            "weight expression produced a null value; weights must be non-null".into(),
1330        ));
1331    }
1332    let values = column.values();
1333    if let Some(&w) = values.iter().find(|&&w| w < 0.0) {
1334        return Err(DataFusionError::Execution(format!(
1335            "weight expression produced a negative value ({w}); weights must be \
1336             non-negative"
1337        )));
1338    }
1339    out.extend_from_slice(values);
1340    Ok(())
1341}
1342
1343fn build_error(graph: &str, e: IdError) -> DataFusionError {
1344    DataFusionError::Execution(match e {
1345        IdError::Null => format!(
1346            "graph '{graph}': an edge has a null source or destination; filter those rows out \
1347             of the edge table (for example, with a view that keeps rows WHERE both are NOT NULL)"
1348        ),
1349        other => format!("graph '{graph}': {other}"),
1350    })
1351}
1352
1353#[cfg(test)]
1354mod tests {
1355    use super::*;
1356    use datafusion::arrow::array::{Int64Array, StringArray};
1357
1358    fn index(src: &dyn Array, dst: &dyn Array) -> GraphIndex {
1359        let (topology, ids) = ursa_io::build_topology_batches(&[(src, dst)]).unwrap();
1360        GraphIndex::new(topology, ids)
1361    }
1362
1363    fn ints(v: &[i64]) -> ArrayRef {
1364        Arc::new(Int64Array::from(v.to_vec()))
1365    }
1366
1367    #[test]
1368    fn weights_line_up_only_with_the_rows_the_index_was_built_from() {
1369        let (src, dst) = (ints(&[10, 20, 10, 30]), ints(&[20, 30, 30, 10]));
1370        let built = index(src.as_ref(), dst.as_ref());
1371        // The same rows, in two batches: aligned.
1372        let split = vec![
1373            (src.slice(0, 1), dst.slice(0, 1)),
1374            (src.slice(1, 3), dst.slice(1, 3)),
1375        ];
1376        assert!(aligned(&built, &split).unwrap());
1377        // The same edge set in another order: not.
1378        let swapped = vec![(ints(&[20, 10, 10, 30]), ints(&[30, 20, 30, 10]))];
1379        assert!(!aligned(&built, &swapped).unwrap());
1380        // Fewer rows: not.
1381        assert!(!aligned(&built, &split[1..]).unwrap());
1382
1383        let names = |v: &[&str]| Arc::new(StringArray::from(v.to_vec())) as ArrayRef;
1384        let (s, d) = (names(&["a", "b"]), names(&["b", "c"]));
1385        let built = index(s.as_ref(), d.as_ref());
1386        assert!(aligned(&built, &[(Arc::clone(&s), Arc::clone(&d))]).unwrap());
1387        assert!(!aligned(&built, &[(d, s)]).unwrap());
1388    }
1389
1390    #[test]
1391    fn weights_are_checked_as_ursa_plan_checks_them() {
1392        let ok = edge_weights(&[ints(&[1, 2]), ints(&[0])]).unwrap();
1393        assert_eq!(ok, vec![1.0, 2.0, 0.0]);
1394        let negative = edge_weights(&[ints(&[1, -1])]).unwrap_err();
1395        assert_eq!(
1396            crate::user_message(&negative),
1397            "weight expression produced a negative value (-1); weights must be non-negative"
1398        );
1399        let null: ArrayRef = Arc::new(Int64Array::from(vec![Some(1), None]));
1400        assert_eq!(
1401            crate::user_message(&edge_weights(&[null]).unwrap_err()),
1402            "weight expression produced a null value; weights must be non-null"
1403        );
1404    }
1405}