reorder_sort algo

src/algo/mod.rs

@@ -1,6 +1,7 @@
 use std::collections::BTreeMap;
 
 use anyhow::{anyhow, bail, ensure, Result};
+use either::Either;
 use itertools::Itertools;
 use smartstring::{LazyCompact, SmartString};
 
@@ -14,14 +15,15 @@ use crate::algo::label_propagation::LabelPropagation;
 use crate::algo::louvain::CommunityDetectionLouvain;
 use crate::algo::pagerank::PageRank;
 use crate::algo::prim::MinimumSpanningTreePrim;
+use crate::algo::reorder_sort::ReorderSort;
 use crate::algo::shortest_path_dijkstra::ShortestPathDijkstra;
 use crate::algo::strongly_connected_components::StronglyConnectedComponent;
 use crate::algo::top_sort::TopSort;
 use crate::algo::triangles::ClusteringCoefficients;
 use crate::algo::yen::KShortestPathYen;
-use crate::data::expr::Expr;
+use crate::data::expr::{Expr, OP_LIST};
 use crate::data::id::{EntityId, Validity};
-use crate::data::program::{MagicAlgoRuleArg, MagicSymbol, TripleDir};
+use crate::data::program::{AlgoRuleArg, MagicAlgoRuleArg, MagicSymbol, TripleDir};
 use crate::data::symb::Symbol;
 use crate::data::tuple::{Tuple, TupleIter};
 use crate::data::value::DataValue;
@@ -34,15 +36,16 @@ pub(crate) mod bfs;
 pub(crate) mod degree_centrality;
 pub(crate) mod dfs;
 pub(crate) mod kruskal;
+pub(crate) mod label_propagation;
 pub(crate) mod louvain;
 pub(crate) mod pagerank;
 pub(crate) mod prim;
 pub(crate) mod shortest_path_dijkstra;
+pub(crate) mod reorder_sort;
 pub(crate) mod strongly_connected_components;
 pub(crate) mod top_sort;
 pub(crate) mod triangles;
 pub(crate) mod yen;
-pub(crate) mod label_propagation;
 
 pub(crate) trait AlgoImpl {
     fn run(
@@ -66,7 +69,11 @@ impl AlgoHandle {
             name: Symbol::from(name),
         }
     }
-    pub(crate) fn arity(&self) -> Result<usize> {
+    pub(crate) fn arity(
+        &self,
+        _args: Either<&[AlgoRuleArg], &[MagicAlgoRuleArg]>,
+        opts: &BTreeMap<SmartString<LazyCompact>, Expr>,
+    ) -> Result<usize> {
         Ok(match &self.name.0 as &str {
             "clustering_coefficients" => 4,
             "degree_centrality" => 4,
@@ -85,6 +92,18 @@ impl AlgoHandle {
             "pagerank" => 2,
             "community_detection_louvain" => 2,
             "label_propagation" => 2,
+            "reorder_sort" => {
+                let out_opts = opts
+                    .get("out")
+                    .ok_or_else(|| anyhow!("'reorder_sort' requires the option 'out'"))?;
+                match out_opts {
+                    Expr::Const(DataValue::List(l)) => l.len() + 1,
+                    Expr::Apply(op, args) if **op == OP_LIST => args.len() + 1,
+                    v => {
+                        bail!("option 'out' of 'reorder_sort' must be a list, got {:?}", v)
+                    }
+                }
+            }
             name => bail!("algorithm '{}' not found", name),
         })
     }
@@ -110,6 +129,7 @@ impl AlgoHandle {
             "pagerank" => Box::new(PageRank),
             "community_detection_louvain" => Box::new(CommunityDetectionLouvain),
             "label_propagation" => Box::new(LabelPropagation),
+            "reorder_sort" => Box::new(ReorderSort),
             name => bail!("algorithm '{}' not found", name),
         })
     }

src/algo/reorder_sort.rs

@@ -0,0 +1,149 @@
+use std::collections::BTreeMap;
+
+use anyhow::{anyhow, bail, ensure, Result};
+use itertools::Itertools;
+use smartstring::{LazyCompact, SmartString};
+
+use crate::algo::AlgoImpl;
+use crate::data::expr::{Expr, OP_LIST};
+use crate::data::program::{MagicAlgoRuleArg, MagicSymbol};
+use crate::data::tuple::Tuple;
+use crate::data::value::DataValue;
+use crate::runtime::derived::DerivedRelStore;
+use crate::runtime::transact::SessionTx;
+
+pub(crate) struct ReorderSort;
+
+impl AlgoImpl for ReorderSort {
+    fn run(
+        &mut self,
+        tx: &SessionTx,
+        rels: &[MagicAlgoRuleArg],
+        opts: &BTreeMap<SmartString<LazyCompact>, Expr>,
+        stores: &BTreeMap<MagicSymbol, DerivedRelStore>,
+        out: &DerivedRelStore,
+    ) -> Result<()> {
+        let in_rel = rels
+            .get(0)
+            .ok_or_else(|| anyhow!("'reorder_sort' requires an input relation"))?;
+
+        let mut out_list = match opts
+            .get("out")
+            .ok_or_else(|| anyhow!("'reorder_sort' requires the option 'out'"))?
+        {
+            Expr::Const(DataValue::List(l)) => {
+                l.iter().map(|d| Expr::Const(d.clone())).collect_vec()
+            }
+            Expr::Apply(op, args) if **op == OP_LIST => args.to_vec(),
+            v => {
+                bail!("option 'out' of 'reorder_sort' must be a list, got {:?}", v)
+            }
+        };
+
+        let mut sort_by = opts
+            .get("sort_by")
+            .cloned()
+            .unwrap_or(Expr::Const(DataValue::Null));
+        let sort_descending = match opts.get("descending") {
+            None => false,
+            Some(Expr::Const(DataValue::Bool(b))) => *b,
+            Some(v) => bail!(
+                "option 'descending' of 'reorder_sort' must be a bool, got {:?}",
+                v
+            ),
+        };
+        let break_ties = match opts.get("break_ties") {
+            None => false,
+            Some(Expr::Const(DataValue::Bool(b))) => *b,
+            Some(v) => bail!(
+                "option 'break_ties' of 'reorder_sort' must be a bool, got {:?}",
+                v
+            ),
+        };
+        let skip = match opts.get("skip") {
+            None => 0,
+            Some(Expr::Const(v)) => v.get_int().ok_or_else(|| {
+                anyhow!(
+                    "option 'skip' of 'reorder_sort' must be an integer, got {:?}",
+                    v
+                )
+            })?,
+            Some(v) => bail!(
+                "option 'skip' of 'reorder_sort' must be an integer, got {:?}",
+                v
+            ),
+        };
+        ensure!(
+            skip >= 0,
+            "option 'skip' of 'reorder_sort' must be non-negative, got {}",
+            skip
+        );
+        let take = match opts.get("take") {
+            None => i64::MAX,
+            Some(Expr::Const(v)) => v.get_int().ok_or_else(|| {
+                anyhow!(
+                    "option 'take' of 'reorder_sort' must be an integer, got {:?}",
+                    v
+                )
+            })?,
+            Some(v) => bail!(
+                "option 'take' of 'reorder_sort' must be an integer, got {:?}",
+                v
+            ),
+        };
+        ensure!(
+            take >= 0,
+            "option 'take' of 'reorder_sort' must be non-negative, got {}",
+            take
+        );
+
+        let binding_map = in_rel.get_binding_map(0);
+        sort_by.fill_binding_indices(&binding_map)?;
+        for out in out_list.iter_mut() {
+            out.fill_binding_indices(&binding_map)?;
+        }
+
+        let mut buffer = vec![];
+        for tuple in in_rel.iter(tx, stores)? {
+            let tuple = tuple?;
+            let sorter = sort_by.eval(&tuple)?;
+            let mut s_tuple: Vec<_> = out_list.iter().map(|ex| ex.eval(&tuple)).try_collect()?;
+            s_tuple.push(sorter);
+            buffer.push(s_tuple);
+        }
+        if sort_descending {
+            buffer.sort_by(|l, r| r.last().cmp(&l.last()));
+        } else {
+            buffer.sort_by(|l, r| l.last().cmp(&r.last()));
+        }
+
+        let mut count = 0usize;
+        let mut rank = 0usize;
+        let mut last = &DataValue::Bottom;
+        let skip = skip as usize;
+        let take_plus_skip = (take as usize).saturating_add(skip);
+        for val in &buffer {
+            let sorter = val.last().unwrap();
+
+            if sorter == last {
+                count += 1;
+            } else {
+                count += 1;
+                rank = count;
+                last = sorter;
+            }
+
+            if count > take_plus_skip {
+                break;
+            }
+
+            if count <= skip {
+                continue;
+            }
+            let mut out_t = vec![DataValue::from(if break_ties { count } else { rank } as i64)];
+            out_t.extend_from_slice(&val[0..val.len() - 1]);
+            out.put(Tuple(out_t), 0);
+        }
+        Ok(())
+    }
+}

src/data/program.rs

@@ -3,6 +3,7 @@ use std::collections::{BTreeMap, BTreeSet};
 use std::fmt::{Debug, Formatter};
 
 use anyhow::{anyhow, bail, ensure, Result};
+use either::{Left, Right};
 use itertools::Itertools;
 use smallvec::SmallVec;
 use smartstring::{LazyCompact, SmartString};
@@ -40,6 +41,12 @@ pub(crate) struct AlgoApply {
     pub(crate) options: BTreeMap<SmartString<LazyCompact>, Expr>,
 }
 
+impl AlgoApply {
+    pub(crate) fn arity(&self) -> Result<usize> {
+        self.algo.arity(Left(&self.rule_args), &self.options)
+    }
+}
+
 impl Debug for AlgoApply {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("AlgoApply")
@@ -57,6 +64,12 @@ pub(crate) struct MagicAlgoApply {
     pub(crate) options: BTreeMap<SmartString<LazyCompact>, Expr>,
 }
 
+impl MagicAlgoApply {
+    pub(crate) fn arity(&self) -> Result<usize> {
+        self.algo.arity(Right(&self.rule_args), &self.options)
+    }
+}
+
 impl Debug for MagicAlgoApply {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("AlgoApply")
@@ -143,7 +156,7 @@ impl InputProgram {
                 .ok_or_else(|| anyhow!("program entry point not found"))?
             {
                 InputRulesOrAlgo::Rules(rules) => rules[0].head.len(),
-                InputRulesOrAlgo::Algo(algo_apply) => algo_apply.algo.arity()?,
+                InputRulesOrAlgo::Algo(algo_apply) => algo_apply.arity()?,
             },
         )
     }
@@ -261,7 +274,7 @@ impl MagicRulesOrAlgo {
     pub(crate) fn arity(&self) -> Result<usize> {
         Ok(match self {
             MagicRulesOrAlgo::Rules(r) => r.first().unwrap().head.len(),
-            MagicRulesOrAlgo::Algo(algo) => algo.algo.arity()?,
+            MagicRulesOrAlgo::Algo(algo) => algo.arity()?,
         })
     }
     pub(crate) fn mut_rules(&mut self) -> Option<&mut Vec<MagicRule>> {