Skip to content

Commit 3982bab

Browse files
authored
Add value filtering of nodes and edges (#571)
1 parent ba7dc58 commit 3982bab

File tree

14 files changed

+544
-229
lines changed

14 files changed

+544
-229
lines changed

partiql-ast-passes/src/name_resolver.rs

Lines changed: 50 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ use crate::error::{AstTransformError, AstTransformationError};
22
use fnv::FnvBuildHasher;
33
use indexmap::{IndexMap, IndexSet};
44
use partiql_ast::ast;
5-
use partiql_ast::ast::{GroupByExpr, GroupKey};
5+
use partiql_ast::ast::{GraphPattern, GroupByExpr, GroupKey};
66
use partiql_ast::visit::{Traverse, Visit, Visitor};
77
use partiql_catalog::catalog::Catalog;
88
use partiql_common::node::NodeId;
@@ -507,6 +507,55 @@ impl<'ast> Visitor<'ast> for NameResolver<'_> {
507507
self.schema.insert(id, KeySchema { consume, produce });
508508
Traverse::Continue
509509
}
510+
511+
fn enter_graph_pattern(&mut self, _: &'ast GraphPattern) -> Traverse {
512+
self.enter_keyref();
513+
Traverse::Continue
514+
}
515+
516+
fn exit_graph_pattern(&mut self, _: &'ast GraphPattern) -> Traverse {
517+
let id = *self.current_node();
518+
let KeyRefs {
519+
consume,
520+
produce_required,
521+
..
522+
} = match self.exit_keyref() {
523+
Ok(kr) => kr,
524+
Err(e) => {
525+
self.errors.push(e);
526+
return Traverse::Stop;
527+
}
528+
};
529+
let produce: Names = produce_required;
530+
self.schema.insert(id, KeySchema { consume, produce });
531+
532+
self.in_scope.entry(id).or_default().push(id);
533+
Traverse::Continue
534+
}
535+
536+
fn exit_graph_match_node(&mut self, node: &'ast ast::GraphMatchNode) -> Traverse {
537+
if let Some(sym) = &node.variable {
538+
let var = Symbol::Known(sym.clone());
539+
self.keyref_stack
540+
.last_mut()
541+
.unwrap()
542+
.produce_required
543+
.insert(var);
544+
}
545+
Traverse::Continue
546+
}
547+
548+
fn exit_graph_match_edge(&mut self, edge: &'ast ast::GraphMatchEdge) -> Traverse {
549+
if let Some(sym) = &edge.variable {
550+
let var = Symbol::Known(sym.clone());
551+
self.keyref_stack
552+
.last_mut()
553+
.unwrap()
554+
.produce_required
555+
.insert(var);
556+
}
557+
Traverse::Continue
558+
}
510559
}
511560

512561
/// Attempt to infer an alias for a simple variable reference expression.

partiql-eval/src/eval/eval_expr_wrapper.rs

Lines changed: 84 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -471,6 +471,26 @@ impl<E: 'static, F: 'static> EvalExprWrapper<E, F> {
471471
let expr = ArgCheckEvalExpr::<STRICT, N, _, ArgC>::new(types, args, expr);
472472
Ok(Box::new(expr))
473473
}
474+
475+
#[inline]
476+
pub(crate) fn create_checked_with_ctx<
477+
const STRICT: bool,
478+
const N: usize,
479+
ArgC: 'static + ArgChecker,
480+
>(
481+
ident: E,
482+
types: [PartiqlShape; N],
483+
args: Vec<Box<dyn EvalExpr>>,
484+
f: F,
485+
) -> Result<Box<dyn EvalExpr>, BindError>
486+
where
487+
EvalExprWrapper<E, F>: ExecuteEvalExpr<N>,
488+
{
489+
let args = unwrap_args(args)?;
490+
let expr = Self { ident, f };
491+
let expr = ArgCheckEvalExpr::<STRICT, N, _, ArgC>::new(types, args, expr);
492+
Ok(Box::new(expr))
493+
}
474494
}
475495

476496
/// An [`ExecuteEvalExpr`] over a single [`Value`] argument
@@ -479,19 +499,19 @@ pub(crate) struct UnaryValueExpr {}
479499

480500
impl<F> ExecuteEvalExpr<1> for EvalExprWrapper<UnaryValueExpr, F>
481501
where
482-
F: Fn(&Value) -> Value,
502+
F: Fn(&Value, &dyn EvalContext) -> Value,
483503
{
484504
#[inline]
485505
fn evaluate<'a, 'c, 'o>(
486506
&'a self,
487507
args: [Cow<'a, Value>; 1],
488-
_ctx: &'c dyn EvalContext,
508+
ctx: &'c dyn EvalContext,
489509
) -> Cow<'a, Value>
490510
where
491511
'c: 'a,
492512
{
493513
let [arg] = args;
494-
Cow::Owned((self.f)(arg.borrow()))
514+
Cow::Owned((self.f)(arg.borrow(), ctx))
495515
}
496516
}
497517

@@ -517,9 +537,22 @@ impl UnaryValueExpr {
517537
) -> Result<Box<dyn EvalExpr>, BindError>
518538
where
519539
F: 'static + Fn(&Value) -> Value,
540+
{
541+
Self::create_typed_with_ctx::<{ STRICT }, _>(types, args, move |val, _| f(val))
542+
}
543+
544+
#[allow(dead_code)]
545+
#[inline]
546+
pub(crate) fn create_typed_with_ctx<const STRICT: bool, F>(
547+
types: [PartiqlShape; 1],
548+
args: Vec<Box<dyn EvalExpr>>,
549+
f: F,
550+
) -> Result<Box<dyn EvalExpr>, BindError>
551+
where
552+
F: 'static + Fn(&Value, &dyn EvalContext) -> Value,
520553
{
521554
type Check<const STRICT: bool> = DefaultArgChecker<STRICT, PropagateMissing<true>>;
522-
Self::create_checked::<{ STRICT }, Check<STRICT>, F>(types, args, f)
555+
Self::create_checked_with_ctx::<{ STRICT }, Check<STRICT>, F>(types, args, f)
523556
}
524557

525558
#[allow(dead_code)]
@@ -533,7 +566,26 @@ impl UnaryValueExpr {
533566
F: 'static + Fn(&Value) -> Value,
534567
ArgC: 'static + ArgChecker,
535568
{
536-
EvalExprWrapper::create_checked::<{ STRICT }, 1, ArgC>(Self::default(), types, args, f)
569+
Self::create_checked_with_ctx::<{ STRICT }, ArgC, _>(types, args, move |val, _| f(val))
570+
}
571+
572+
#[allow(dead_code)]
573+
#[inline]
574+
pub(crate) fn create_checked_with_ctx<const STRICT: bool, ArgC, F>(
575+
types: [PartiqlShape; 1],
576+
args: Vec<Box<dyn EvalExpr>>,
577+
f: F,
578+
) -> Result<Box<dyn EvalExpr>, BindError>
579+
where
580+
F: 'static + Fn(&Value, &dyn EvalContext) -> Value,
581+
ArgC: 'static + ArgChecker,
582+
{
583+
EvalExprWrapper::create_checked_with_ctx::<{ STRICT }, 1, ArgC>(
584+
Self::default(),
585+
types,
586+
args,
587+
f,
588+
)
537589
}
538590
}
539591

@@ -543,19 +595,19 @@ pub(crate) struct BinaryValueExpr {}
543595

544596
impl<F> ExecuteEvalExpr<2> for EvalExprWrapper<BinaryValueExpr, F>
545597
where
546-
F: Fn(&Value, &Value) -> Value,
598+
F: Fn(&Value, &Value, &dyn EvalContext) -> Value,
547599
{
548600
#[inline]
549601
fn evaluate<'a, 'c, 'o>(
550602
&'a self,
551603
args: [Cow<'a, Value>; 2],
552-
_ctx: &'c dyn EvalContext,
604+
ctx: &'c dyn EvalContext,
553605
) -> Cow<'a, Value>
554606
where
555607
'c: 'a,
556608
{
557609
let [arg1, arg2] = args;
558-
Cow::Owned((self.f)(arg1.borrow(), arg2.borrow()))
610+
Cow::Owned((self.f)(arg1.borrow(), arg2.borrow(), ctx))
559611
}
560612
}
561613

@@ -597,7 +649,12 @@ impl BinaryValueExpr {
597649
F: 'static + Fn(&Value, &Value) -> Value,
598650
ArgC: 'static + ArgChecker,
599651
{
600-
EvalExprWrapper::create_checked::<{ STRICT }, 2, ArgC>(Self::default(), types, args, f)
652+
EvalExprWrapper::create_checked::<{ STRICT }, 2, ArgC>(
653+
Self::default(),
654+
types,
655+
args,
656+
move |v1, v2, _| f(v1, v2),
657+
)
601658
}
602659
}
603660

@@ -607,19 +664,19 @@ pub(crate) struct TernaryValueExpr {}
607664

608665
impl<F> ExecuteEvalExpr<3> for EvalExprWrapper<TernaryValueExpr, F>
609666
where
610-
F: Fn(&Value, &Value, &Value) -> Value,
667+
F: Fn(&Value, &Value, &Value, &dyn EvalContext) -> Value,
611668
{
612669
#[inline]
613670
fn evaluate<'a, 'c, 'o>(
614671
&'a self,
615672
args: [Cow<'a, Value>; 3],
616-
_ctx: &'c dyn EvalContext,
673+
ctx: &'c dyn EvalContext,
617674
) -> Cow<'a, Value>
618675
where
619676
'c: 'a,
620677
{
621678
let [arg1, arg2, arg3] = args;
622-
Cow::Owned((self.f)(arg1.borrow(), arg2.borrow(), arg3.borrow()))
679+
Cow::Owned((self.f)(arg1.borrow(), arg2.borrow(), arg3.borrow(), ctx))
623680
}
624681
}
625682

@@ -661,7 +718,12 @@ impl TernaryValueExpr {
661718
F: 'static + Fn(&Value, &Value, &Value) -> Value,
662719
ArgC: 'static + ArgChecker,
663720
{
664-
EvalExprWrapper::create_checked::<{ STRICT }, 3, ArgC>(Self::default(), types, args, f)
721+
EvalExprWrapper::create_checked::<{ STRICT }, 3, ArgC>(
722+
Self::default(),
723+
types,
724+
args,
725+
move |v1, v2, v3, _| f(v1, v2, v3),
726+
)
665727
}
666728
}
667729

@@ -671,13 +733,13 @@ pub(crate) struct QuaternaryValueExpr {}
671733

672734
impl<F> ExecuteEvalExpr<4> for EvalExprWrapper<QuaternaryValueExpr, F>
673735
where
674-
F: Fn(&Value, &Value, &Value, &Value) -> Value,
736+
F: Fn(&Value, &Value, &Value, &Value, &dyn EvalContext) -> Value,
675737
{
676738
#[inline]
677739
fn evaluate<'a, 'c, 'o>(
678740
&'a self,
679741
args: [Cow<'a, Value>; 4],
680-
_ctx: &'c dyn EvalContext,
742+
ctx: &'c dyn EvalContext,
681743
) -> Cow<'a, Value>
682744
where
683745
'c: 'a,
@@ -688,6 +750,7 @@ where
688750
arg2.borrow(),
689751
arg3.borrow(),
690752
arg4.borrow(),
753+
ctx,
691754
))
692755
}
693756
}
@@ -730,6 +793,11 @@ impl QuaternaryValueExpr {
730793
F: 'static + Fn(&Value, &Value, &Value, &Value) -> Value,
731794
ArgC: 'static + ArgChecker,
732795
{
733-
EvalExprWrapper::create_checked::<{ STRICT }, 4, ArgC>(Self::default(), types, args, f)
796+
EvalExprWrapper::create_checked::<{ STRICT }, 4, ArgC>(
797+
Self::default(),
798+
types,
799+
args,
800+
move |v1, v2, v3, v4, _| f(v1, v2, v3, v4),
801+
)
734802
}
735803
}

partiql-eval/src/eval/evaluable.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,7 @@ pub trait Evaluable: Debug {
5353
}
5454

5555
/// Represents an evaluation `Scan` operator; `Scan` operator scans the given bindings from its
56-
/// input and and the environment and outputs a bag of binding tuples for tuples/values matching the
56+
/// input and the environment and outputs a bag of binding tuples for tuples/values matching the
5757
/// scan `expr`, e.g. an SQL expression `table1` in SQL expression `FROM table1`.
5858
pub(crate) struct EvalScan {
5959
pub(crate) expr: Box<dyn EvalExpr>,

partiql-eval/src/eval/expr/graph_match.rs

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -28,18 +28,20 @@ impl BindEvalExpr for EvalGraphMatch {
2828
) -> Result<Box<dyn EvalExpr>, BindError> {
2929
// use DummyShapeBuilder, as we don't care about shape Ids for evaluation dispatch
3030
let mut bld = PartiqlNoIdShapeBuilder::default();
31-
UnaryValueExpr::create_typed::<{ STRICT }, _>([type_graph!(bld)], args, move |value| {
32-
match value {
31+
UnaryValueExpr::create_typed_with_ctx::<{ STRICT }, _>(
32+
[type_graph!(bld)],
33+
args,
34+
move |value, ctx| match value {
3335
Value::Graph(graph) => match graph.as_ref() {
3436
Graph::Simple(g) => {
3537
let engine = SimpleGraphEngine::new(g.clone());
3638
let ge = GraphEvaluator::new(engine);
37-
ge.eval(&self.pattern)
39+
ge.eval(&self.pattern, ctx)
3840
}
3941
},
4042
_ => Missing,
41-
}
42-
})
43+
},
44+
)
4345
}
4446
}
4547

0 commit comments

Comments
 (0)