1use 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#[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 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
142fn 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
168fn 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#[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#[derive(Debug, Clone)]
228struct Edges {
229 graph: String,
230 endpoints: (usize, usize),
232 id_type: DataType,
233 fingerprint: Option<String>,
235}
236
237#[derive(Debug, Clone)]
239struct Shape {
240 edges: Edges,
241 outputs: Vec<Output>,
242 per_node: bool,
243 display: String,
245 schema: SchemaRef,
246}
247
248#[derive(Debug, Clone)]
250struct WeightedInput {
251 plan: Arc<dyn ExecutionPlan>,
252 source: usize,
253 destination: usize,
254 weights: Vec<usize>,
256 fingerprint: Option<String>,
258}
259
260#[derive(Debug)]
263pub struct GraphAlgorithmExec {
264 shape: Shape,
265 input: Arc<dyn ExecutionPlan>,
266 weighted: Option<WeightedInput>,
267 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 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 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#[derive(Debug, Clone)]
418struct TraversalShape {
419 edges: Edges,
420 traversal: Traversal,
421 display: String,
423 schema: SchemaRef,
424}
425
426#[derive(Debug)]
429pub struct GraphTraversalExec {
430 shape: TraversalShape,
431 input: Arc<dyn ExecutionPlan>,
432 weighted: Option<WeightedInput>,
433 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 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
576fn 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 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
629fn 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
637fn 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#[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 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 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
729fn 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
742fn endpoint_keys((source, destination): (usize, usize)) -> Vec<usize> {
744 vec![source, destination]
745}
746
747fn 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
755fn 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
782fn 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
808async 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
820async 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
844async 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 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 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 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 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
931async 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 crate::run_blocking(move || index.stored_weights(&key, n_weights))
943 .await
944 .ok()
945 .flatten()
946}
947
948async 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
961async 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
991struct WeightReader {
996 columns: Vec<usize>,
997 values: Vec<Vec<f64>>,
998 error: Option<DataFusionError>,
999}
1000
1001impl WeightReader {
1002 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 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
1032async 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
1059const INTERN_ROWS: usize = 1 << 17;
1065
1066#[derive(Default)]
1068struct Gather {
1069 src: Vec<ArrayRef>,
1070 dst: Vec<ArrayRef>,
1071 rows: usize,
1072}
1073
1074impl Gather {
1075 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 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
1105fn 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
1117fn 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
1127enum Chunk {
1129 Edges(ArrayRef, ArrayRef),
1130 End,
1132}
1133
1134async 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 #[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 Ok(None)
1169 });
1170 #[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 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 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
1239struct RowCheck<'a> {
1242 index: &'a GraphIndex,
1243 source: Vec<u32>,
1246 destination: Vec<u32>,
1247 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 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 fn complete(&self) -> bool {
1292 self.start == self.source.len()
1293 }
1294}
1295
1296#[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#[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
1319fn 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 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 let swapped = vec![(ints(&[20, 10, 10, 30]), ints(&[30, 20, 30, 10]))];
1379 assert!(!aligned(&built, &swapped).unwrap());
1380 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}