From 4d9ae54d84bd5184ce59f8afee3ce9199092cd51 Mon Sep 17 00:00:00 2001 From: Josh Pschorr Date: Thu, 7 Mar 2024 10:51:18 -0800 Subject: [PATCH] Add session context to evaluation (#446) --- .github/workflows/ci_build_test.yml | 4 +- CHANGELOG.md | 2 + Cargo.toml | 2 +- .../src/lib.rs | 17 +- partiql-catalog/src/context.rs | 27 +++ partiql-catalog/src/lib.rs | 9 +- partiql-conformance-tests/tests/mod.rs | 19 +- partiql-eval/Cargo.toml | 1 + partiql-eval/benches/bench_eval.rs | 11 +- partiql-eval/src/env.rs | 11 +- partiql-eval/src/error.rs | 11 +- partiql-eval/src/eval/eval_expr_wrapper.rs | 62 +++-- partiql-eval/src/eval/evaluable.rs | 75 +++--- partiql-eval/src/eval/expr/base_table.rs | 11 +- partiql-eval/src/eval/expr/control_flow.rs | 9 +- partiql-eval/src/eval/expr/data_types.rs | 36 ++- partiql-eval/src/eval/expr/mod.rs | 8 +- partiql-eval/src/eval/expr/operators.rs | 18 +- partiql-eval/src/eval/expr/path.rs | 50 +++- partiql-eval/src/eval/expr/strings.rs | 9 +- partiql-eval/src/eval/mod.rs | 92 ++++++-- partiql-eval/src/lib.rs | 46 +++- partiql-logical-planner/src/lib.rs | 11 +- partiql-logical-planner/src/lower.rs | 2 +- partiql-value/src/datetime.rs | 4 + partiql/benches/bench_agg.rs | 11 +- partiql/benches/bench_eval_multi_like.rs | 11 +- partiql/src/lib.rs | 27 ++- partiql/tests/user_context.rs | 215 ++++++++++++++++++ 29 files changed, 660 insertions(+), 151 deletions(-) create mode 100644 partiql-catalog/src/context.rs create mode 100644 partiql/tests/user_context.rs diff --git a/.github/workflows/ci_build_test.yml b/.github/workflows/ci_build_test.yml index e82b3945..94fe6e8b 100644 --- a/.github/workflows/ci_build_test.yml +++ b/.github/workflows/ci_build_test.yml @@ -92,7 +92,7 @@ jobs: - name: Rust Toolchain uses: dtolnay/rust-toolchain@master with: - toolchain: nightly-2023-03-09 + toolchain: nightly-2023-06-09 - uses: actions/cache@v3 id: restore-build with: @@ -132,7 +132,7 @@ jobs: - name: Rust Toolchain uses: dtolnay/rust-toolchain@master with: - toolchain: nightly-2023-03-09 + toolchain: nightly-2023-06-09 - uses: actions/cache@v3 id: restore-build-and-conformance with: diff --git a/CHANGELOG.md b/CHANGELOG.md index 8787f29a..7c38c1bb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,9 +8,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Changed - Adds quotes to the attributes of PartiQL tuple's debug output so it can be read and transformed using Kotlin `partiql-cli` +- [breaking] Changes the interface to `EvalPlan` to accept an `EvalContext` ### Added - Add `partiql-extension-visualize` for visualizing AST and logical plan +- Add a `SessionContext` containing both a system-level and a user-level context object usable by expression evaluation ### Fixed - Fixed `ORDER BY`'s ability to see into projection aliases diff --git a/Cargo.toml b/Cargo.toml index 8e33e3a1..93f94f14 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,9 +4,9 @@ homepage = "https://github.com/partiql/partiql-lang-rust" repository = "https://github.com/partiql/partiql-lang-rust" version = "0.6.0" edition = "2021" -resolver = 2 [workspace] +resolver = "2" members = [ "partiql", diff --git a/extension/partiql-extension-ion-functions/src/lib.rs b/extension/partiql-extension-ion-functions/src/lib.rs index abc92dfa..034e190d 100644 --- a/extension/partiql-extension-ion-functions/src/lib.rs +++ b/extension/partiql-extension-ion-functions/src/lib.rs @@ -11,6 +11,7 @@ use partiql_logical as logical; use partiql_value::Value; use std::borrow::Cow; +use partiql_catalog::context::SessionContext; use std::error::Error; use std::fmt::Debug; use std::fs::File; @@ -98,7 +99,11 @@ impl BaseTableFunctionInfo for ReadIonFunction { pub(crate) struct EvalFnReadIon {} impl BaseTableExpr for EvalFnReadIon { - fn evaluate(&self, args: &[Cow]) -> BaseTableExprResult { + fn evaluate<'c>( + &self, + args: &[Cow], + _ctx: &'c dyn SessionContext<'c>, + ) -> BaseTableExprResult<'c> { if let Some(arg1) = args.first() { match arg1.as_ref() { Value::String(path) => parse_ion_file(path), @@ -155,11 +160,13 @@ fn parse_ion_buff<'a, I: 'a + ToIonDataSource>(input: I) -> BaseTableExprResult< mod tests { use super::*; + use partiql_catalog::context::SystemContext; use partiql_catalog::{Catalog, Extension, PartiqlCatalog}; use partiql_eval::env::basic::MapBindings; + use partiql_eval::eval::BasicContext; use partiql_eval::plan::EvaluationMode; use partiql_parser::{Parsed, ParserResult}; - use partiql_value::{bag, tuple, Value}; + use partiql_value::{bag, tuple, DateTime, Value}; #[track_caller] #[inline] @@ -189,7 +196,11 @@ mod tests { let mut plan = planner.compile(&logical).expect("Expect no plan error"); - if let Ok(out) = plan.execute_mut(bindings) { + let sys = SystemContext { + now: DateTime::from_system_now_utc(), + }; + let ctx = BasicContext::new(bindings, sys); + if let Ok(out) = plan.execute_mut(&ctx) { out.result } else { Value::Missing diff --git a/partiql-catalog/src/context.rs b/partiql-catalog/src/context.rs new file mode 100644 index 00000000..2094f51d --- /dev/null +++ b/partiql-catalog/src/context.rs @@ -0,0 +1,27 @@ +use partiql_value::{BindingsName, DateTime, Tuple, Value}; +use std::any::Any; +use std::fmt::Debug; + +pub trait Bindings: Debug { + fn get(&self, name: &BindingsName) -> Option<&T>; +} + +impl Bindings for Tuple { + fn get(&self, name: &BindingsName) -> Option<&Value> { + self.get(name) + } +} + +#[derive(Debug)] +pub struct SystemContext { + pub now: DateTime, +} + +/// Represents a session context that is used during evaluation of a plan. +pub trait SessionContext<'a>: Debug { + fn bindings(&self) -> &dyn Bindings; + + fn system_context(&self) -> &SystemContext; + + fn user_context(&self, name: &str) -> Option<&(dyn Any)>; +} diff --git a/partiql-catalog/src/lib.rs b/partiql-catalog/src/lib.rs index e0ffa04b..62f92a21 100644 --- a/partiql-catalog/src/lib.rs +++ b/partiql-catalog/src/lib.rs @@ -4,6 +4,7 @@ use partiql_types::PartiqlType; use partiql_value::Value; use std::borrow::Cow; +use crate::context::SessionContext; use std::collections::HashMap; use std::error::Error; use std::fmt::Debug; @@ -13,6 +14,8 @@ use unicase::UniCase; pub mod call_defs; +pub mod context; + pub trait Extension: Debug { fn name(&self) -> String; fn load(&self, catalog: &mut dyn Catalog) -> Result<(), Box>; @@ -49,7 +52,11 @@ pub type BaseTableExprResult<'a> = Result, BaseTableExprResultError>; pub trait BaseTableExpr: Debug { - fn evaluate(&self, args: &[Cow]) -> BaseTableExprResult; + fn evaluate<'c>( + &self, + args: &[Cow], + ctx: &'c dyn SessionContext<'c>, + ) -> BaseTableExprResult<'c>; } pub trait BaseTableFunctionInfo: Debug { diff --git a/partiql-conformance-tests/tests/mod.rs b/partiql-conformance-tests/tests/mod.rs index 70d38d16..46510126 100644 --- a/partiql-conformance-tests/tests/mod.rs +++ b/partiql-conformance-tests/tests/mod.rs @@ -1,13 +1,14 @@ use partiql_ast_passes::error::AstTransformationError; use partiql_catalog::{Catalog, PartiqlCatalog}; use partiql_eval as eval; -use partiql_eval::env::basic::MapBindings; + use partiql_eval::error::{EvalErr, PlanErr}; -use partiql_eval::eval::{EvalPlan, EvalResult, Evaluated}; +use partiql_eval::eval::{BasicContext, EvalContext, EvalPlan, EvalResult, Evaluated}; use partiql_logical as logical; use partiql_parser::{Parsed, ParserError, ParserResult}; -use partiql_value::Value; +use partiql_value::DateTime; +use partiql_catalog::context::SystemContext; use thiserror::Error; mod test_value; @@ -58,8 +59,8 @@ pub(crate) fn compile( #[track_caller] #[inline] -pub(crate) fn evaluate(mut plan: EvalPlan, bindings: MapBindings) -> EvalResult { - plan.execute_mut(bindings) +pub(crate) fn evaluate<'c>(mut plan: EvalPlan, ctx: &'c dyn EvalContext<'c>) -> EvalResult { + plan.execute_mut(ctx) } #[track_caller] @@ -161,9 +162,15 @@ pub(crate) fn eval<'a>( let parsed = parse(statement)?; let lowered = lower(&catalog, &parsed)?; + let bindings = env.as_ref().map(|e| (&e.value).into()).unwrap_or_default(); + let sys = SystemContext { + now: DateTime::from_system_now_utc(), + }; + let ctx = BasicContext::new(bindings, sys); let plan = compile(mode, &catalog, lowered)?; - Ok(evaluate(plan, bindings)?) + + Ok(evaluate(plan, &ctx)?) } #[track_caller] diff --git a/partiql-eval/Cargo.toml b/partiql-eval/Cargo.toml index c9ce4c0e..46135f5c 100644 --- a/partiql-eval/Cargo.toml +++ b/partiql-eval/Cargo.toml @@ -36,6 +36,7 @@ assert_matches = "1.5.*" regex = "1.7" regex-syntax = "0.6" rustc-hash = "1" +delegate = "0.12" [dev-dependencies] criterion = "0.4" diff --git a/partiql-eval/benches/bench_eval.rs b/partiql-eval/benches/bench_eval.rs index 89d514ce..f1d7ecf5 100644 --- a/partiql-eval/benches/bench_eval.rs +++ b/partiql-eval/benches/bench_eval.rs @@ -2,10 +2,11 @@ use std::borrow::Cow; use std::time::Duration; use criterion::{black_box, criterion_group, criterion_main, Criterion}; +use partiql_catalog::context::SystemContext; use partiql_catalog::PartiqlCatalog; use partiql_eval::env::basic::MapBindings; -use partiql_eval::eval::EvalPlan; +use partiql_eval::eval::{BasicContext, EvalPlan}; use partiql_eval::plan; use partiql_eval::plan::EvaluationMode; use partiql_logical as logical; @@ -13,7 +14,7 @@ use partiql_logical::BindingsOp::{Project, ProjectAll}; use partiql_logical::{ BinaryOp, BindingsOp, JoinKind, LogicalPlan, PathComponent, ValueExpr, VarRefType, }; -use partiql_value::{bag, list, tuple, BindingsName, Value}; +use partiql_value::{bag, list, tuple, BindingsName, DateTime, Value}; fn data() -> MapBindings { let hr = tuple![( @@ -134,7 +135,11 @@ fn eval_plan(logical: &LogicalPlan) -> EvalPlan { } fn evaluate(mut plan: EvalPlan, bindings: MapBindings) -> Value { - if let Ok(out) = plan.execute_mut(bindings) { + let sys = SystemContext { + now: DateTime::from_system_now_utc(), + }; + let ctx = BasicContext::new(bindings, sys); + if let Ok(out) = plan.execute_mut(&ctx) { out.result } else { Value::Missing diff --git a/partiql-eval/src/env.rs b/partiql-eval/src/env.rs index 102529ca..7d1f2cd6 100644 --- a/partiql-eval/src/env.rs +++ b/partiql-eval/src/env.rs @@ -1,17 +1,8 @@ +use partiql_catalog::context::Bindings; use partiql_value::{BindingsName, Tuple, Value}; use std::fmt::Debug; use unicase::UniCase; -pub trait Bindings: Debug { - fn get(&self, name: &BindingsName) -> Option<&T>; -} - -impl Bindings for Tuple { - fn get(&self, name: &BindingsName) -> Option<&Value> { - self.get(name) - } -} - pub mod basic { use super::*; use std::collections::HashMap; diff --git a/partiql-eval/src/error.rs b/partiql-eval/src/error.rs index fde7d210..2bdde323 100644 --- a/partiql-eval/src/error.rs +++ b/partiql-eval/src/error.rs @@ -56,7 +56,7 @@ impl ErrorNode { } impl Evaluable for ErrorNode { - fn evaluate(&mut self, _ctx: &dyn EvalContext) -> Value { + fn evaluate<'a, 'c>(&mut self, _ctx: &'c dyn EvalContext<'c>) -> Value { panic!("ErrorNode will not be evaluated") } @@ -66,7 +66,14 @@ impl Evaluable for ErrorNode { } impl EvalExpr for ErrorNode { - fn evaluate<'a>(&'a self, _bindings: &'a Tuple, _ctx: &'a dyn EvalContext) -> Cow<'a, Value> { + fn evaluate<'a, 'c>( + &'a self, + _bindings: &'a Tuple, + _ctx: &'c dyn EvalContext<'c>, + ) -> Cow<'a, Value> + where + 'c: 'a, + { panic!("ErrorNode will not be evaluated") } } diff --git a/partiql-eval/src/eval/eval_expr_wrapper.rs b/partiql-eval/src/eval/eval_expr_wrapper.rs index 0ca68faf..ada92854 100644 --- a/partiql-eval/src/eval/eval_expr_wrapper.rs +++ b/partiql-eval/src/eval/eval_expr_wrapper.rs @@ -66,11 +66,13 @@ pub(crate) fn unwrap_args( /// An expression that is evaluated over `N` input arguments pub(crate) trait ExecuteEvalExpr: Debug { /// Evaluate the expression - fn evaluate<'a>( + fn evaluate<'a, 'c>( &'a self, args: [Cow<'a, Value>; N], - ctx: &'a dyn EvalContext, - ) -> Cow<'a, Value>; + ctx: &'c dyn EvalContext<'c>, + ) -> Cow<'a, Value> + where + 'c: 'a; } /// Used to tell argument checking whether it should exit early or go on as usual. @@ -249,11 +251,14 @@ impl, ArgC: ArgChecker /// /// If type-checking fails, the appropriate failure case of [`ArgCheckControlFlow`] is returned, /// else [`ArgCheckControlFlow::Continue`] is returned containing the `N` values. - pub fn evaluate_args<'a>( + pub fn evaluate_args<'a, 'c>( &'a self, bindings: &'a Tuple, - ctx: &'a dyn EvalContext, - ) -> ControlFlow; N]> { + ctx: &'c dyn EvalContext<'c>, + ) -> ControlFlow; N]> + where + 'c: 'a, + { let err_arg_count_mismatch = |args: Vec<_>| { if STRICT { ctx.add_error(EvaluationError::IllegalState(format!( @@ -326,7 +331,14 @@ impl, ArgC: ArgChecker impl, ArgC: ArgChecker> EvalExpr for ArgCheckEvalExpr { - fn evaluate<'a>(&'a self, bindings: &'a Tuple, ctx: &'a dyn EvalContext) -> Cow<'a, Value> { + fn evaluate<'a, 'c>( + &'a self, + bindings: &'a Tuple, + ctx: &'c dyn EvalContext<'c>, + ) -> Cow<'a, Value> + where + 'c: 'a, + { if STRICT && ctx.has_errors() { return Cow::Owned(Missing); } @@ -379,11 +391,14 @@ where F: Fn(&Value) -> Value, { #[inline] - fn evaluate<'a>( + fn evaluate<'a, 'c>( &'a self, args: [Cow<'a, Value>; 1], - _ctx: &'a dyn EvalContext, - ) -> Cow<'a, Value> { + _ctx: &'c dyn EvalContext<'c>, + ) -> Cow<'a, Value> + where + 'c: 'a, + { let [arg] = args; Cow::Owned((self.f)(arg.borrow())) } @@ -440,11 +455,14 @@ where F: Fn(&Value, &Value) -> Value, { #[inline] - fn evaluate<'a>( + fn evaluate<'a, 'c>( &'a self, args: [Cow<'a, Value>; 2], - _ctx: &'a dyn EvalContext, - ) -> Cow<'a, Value> { + _ctx: &'c dyn EvalContext<'c>, + ) -> Cow<'a, Value> + where + 'c: 'a, + { let [arg1, arg2] = args; Cow::Owned((self.f)(arg1.borrow(), arg2.borrow())) } @@ -501,11 +519,14 @@ where F: Fn(&Value, &Value, &Value) -> Value, { #[inline] - fn evaluate<'a>( + fn evaluate<'a, 'c>( &'a self, args: [Cow<'a, Value>; 3], - _ctx: &'a dyn EvalContext, - ) -> Cow<'a, Value> { + _ctx: &'c dyn EvalContext<'c>, + ) -> Cow<'a, Value> + where + 'c: 'a, + { let [arg1, arg2, arg3] = args; Cow::Owned((self.f)(arg1.borrow(), arg2.borrow(), arg3.borrow())) } @@ -562,11 +583,14 @@ where F: Fn(&Value, &Value, &Value, &Value) -> Value, { #[inline] - fn evaluate<'a>( + fn evaluate<'a, 'c>( &'a self, args: [Cow<'a, Value>; 4], - _ctx: &'a dyn EvalContext, - ) -> Cow<'a, Value> { + _ctx: &'c dyn EvalContext<'c>, + ) -> Cow<'a, Value> + where + 'c: 'a, + { let [arg1, arg2, arg3, arg4] = args; Cow::Owned((self.f)( arg1.borrow(), diff --git a/partiql-eval/src/eval/evaluable.rs b/partiql-eval/src/eval/evaluable.rs index 9ac9f7e0..95c5db2b 100644 --- a/partiql-eval/src/eval/evaluable.rs +++ b/partiql-eval/src/eval/evaluable.rs @@ -1,7 +1,7 @@ use crate::env::basic::MapBindings; use crate::error::EvaluationError; use crate::eval::expr::EvalExpr; -use crate::eval::{EvalContext, EvalPlan}; +use crate::eval::{EvalContext, EvalPlan, NestedContext}; use itertools::Itertools; use partiql_value::Value::{Boolean, Missing, Null}; use partiql_value::{ @@ -40,7 +40,7 @@ pub enum EvalType { /// `Evaluable` represents each evaluation operator in the evaluation plan as an evaluable entity. pub trait Evaluable: Debug { - fn evaluate(&mut self, ctx: &dyn EvalContext) -> Value; + fn evaluate<'c>(&mut self, ctx: &'c dyn EvalContext<'c>) -> Value; fn update_input(&mut self, input: Value, branch_num: u8, ctx: &dyn EvalContext); fn get_vars(&self) -> Option<&[String]> { None @@ -102,7 +102,7 @@ impl EvalScan { } impl Evaluable for EvalScan { - fn evaluate(&mut self, ctx: &dyn EvalContext) -> Value { + fn evaluate<'a, 'c>(&mut self, ctx: &'c dyn EvalContext<'c>) -> Value { let input_value = self.input.take().unwrap_or(Missing); let bindings = match input_value { @@ -197,7 +197,7 @@ impl EvalJoin { } impl Evaluable for EvalJoin { - fn evaluate(&mut self, ctx: &dyn EvalContext) -> Value { + fn evaluate<'a, 'c>(&mut self, ctx: &'c dyn EvalContext<'c>) -> Value { /// Creates a `Tuple` with attributes `attrs`, each with value `Null` #[inline] fn tuple_with_null_vals(attrs: I) -> Tuple @@ -640,7 +640,7 @@ impl EvalGroupBy { } #[inline] - fn group_key(&self, bindings: &Tuple, ctx: &dyn EvalContext) -> GroupKey { + fn group_key<'a, 'c>(&'a self, bindings: &'a Tuple, ctx: &'c dyn EvalContext<'c>) -> GroupKey { self.group .iter() .map(|expr| match expr.evaluate(bindings, ctx).as_ref() { @@ -652,7 +652,7 @@ impl EvalGroupBy { } impl Evaluable for EvalGroupBy { - fn evaluate(&mut self, ctx: &dyn EvalContext) -> Value { + fn evaluate<'a, 'c>(&mut self, ctx: &'c dyn EvalContext<'c>) -> Value { let group_as_alias = &self.group_as_alias; let input_value = take_input!(self.input.take(), ctx); @@ -770,7 +770,7 @@ impl EvalPivot { } impl Evaluable for EvalPivot { - fn evaluate(&mut self, ctx: &dyn EvalContext) -> Value { + fn evaluate<'a, 'c>(&mut self, ctx: &'c dyn EvalContext<'c>) -> Value { let input_value = take_input!(self.input.take(), ctx); let tuple: Tuple = input_value @@ -827,7 +827,7 @@ impl EvalUnpivot { } impl Evaluable for EvalUnpivot { - fn evaluate(&mut self, ctx: &dyn EvalContext) -> Value { + fn evaluate<'a, 'c>(&mut self, ctx: &'c dyn EvalContext<'c>) -> Value { let tuple = match self.expr.evaluate(&Tuple::new(), ctx).into_owned() { Value::Tuple(tuple) => *tuple, other => other.coerce_into_tuple(), @@ -871,7 +871,7 @@ impl EvalFilter { } #[inline] - fn eval_filter(&self, bindings: &Tuple, ctx: &dyn EvalContext) -> bool { + fn eval_filter<'a, 'c>(&'a self, bindings: &'a Tuple, ctx: &'c dyn EvalContext<'c>) -> bool { let result = self.expr.evaluate(bindings, ctx); match result.as_ref() { Boolean(bool_val) => *bool_val, @@ -884,7 +884,7 @@ impl EvalFilter { } impl Evaluable for EvalFilter { - fn evaluate(&mut self, ctx: &dyn EvalContext) -> Value { + fn evaluate<'a, 'c>(&mut self, ctx: &'c dyn EvalContext<'c>) -> Value { let input_value = take_input!(self.input.take(), ctx); let filtered = input_value @@ -914,7 +914,7 @@ impl EvalHaving { } #[inline] - fn eval_having(&self, bindings: &Tuple, ctx: &dyn EvalContext) -> bool { + fn eval_having<'a, 'c>(&'a self, bindings: &'a Tuple, ctx: &'c dyn EvalContext<'c>) -> bool { let result = self.expr.evaluate(bindings, ctx); match result.as_ref() { Boolean(bool_val) => *bool_val, @@ -929,7 +929,7 @@ impl EvalHaving { } impl Evaluable for EvalHaving { - fn evaluate(&mut self, ctx: &dyn EvalContext) -> Value { + fn evaluate<'a, 'c>(&mut self, ctx: &'c dyn EvalContext<'c>) -> Value { let input_value = take_input!(self.input.take(), ctx); let filtered = input_value @@ -967,7 +967,7 @@ pub(crate) struct EvalOrderBy { impl EvalOrderBy { #[inline] - fn compare(&self, l: &Value, r: &Value, ctx: &dyn EvalContext) -> Ordering { + fn compare<'c>(&self, l: &Value, r: &Value, ctx: &'c dyn EvalContext<'c>) -> Ordering { let l = l.as_tuple_ref(); let r = r.as_tuple_ref(); self.cmp @@ -1005,7 +1005,7 @@ impl EvalOrderBy { } impl Evaluable for EvalOrderBy { - fn evaluate(&mut self, ctx: &dyn EvalContext) -> Value { + fn evaluate<'a, 'c>(&mut self, ctx: &'c dyn EvalContext<'c>) -> Value { let input_value = take_input!(self.input.take(), ctx); let mut values = input_value.into_iter().collect_vec(); @@ -1027,7 +1027,7 @@ pub(crate) struct EvalLimitOffset { } impl Evaluable for EvalLimitOffset { - fn evaluate(&mut self, ctx: &dyn EvalContext) -> Value { + fn evaluate<'a, 'c>(&mut self, ctx: &'c dyn EvalContext<'c>) -> Value { let input_value = take_input!(self.input.take(), ctx); let empty_bindings = Tuple::new(); @@ -1096,7 +1096,7 @@ impl EvalSelectValue { } impl Evaluable for EvalSelectValue { - fn evaluate(&mut self, ctx: &dyn EvalContext) -> Value { + fn evaluate<'a, 'c>(&mut self, ctx: &'c dyn EvalContext<'c>) -> Value { let input_value = take_input!(self.input.take(), ctx); let ordered = input_value.is_ordered(); @@ -1148,7 +1148,7 @@ impl Debug for EvalSelect { } impl Evaluable for EvalSelect { - fn evaluate(&mut self, ctx: &dyn EvalContext) -> Value { + fn evaluate<'a, 'c>(&mut self, ctx: &'c dyn EvalContext<'c>) -> Value { let input_value = take_input!(self.input.take(), ctx); let ordered = input_value.is_ordered(); @@ -1192,7 +1192,7 @@ impl EvalSelectAll { } impl Evaluable for EvalSelectAll { - fn evaluate(&mut self, ctx: &dyn EvalContext) -> Value { + fn evaluate<'a, 'c>(&mut self, ctx: &'c dyn EvalContext<'c>) -> Value { let input_value = take_input!(self.input.take(), ctx); let ordered = input_value.is_ordered(); @@ -1231,7 +1231,7 @@ impl EvalExprQuery { } impl Evaluable for EvalExprQuery { - fn evaluate(&mut self, ctx: &dyn EvalContext) -> Value { + fn evaluate<'a, 'c>(&mut self, ctx: &'c dyn EvalContext<'c>) -> Value { let input_value = self.input.take().unwrap_or(Value::Null).coerce_into_tuple(); self.expr.evaluate(&input_value, ctx).into_owned() @@ -1255,7 +1255,7 @@ impl EvalDistinct { } impl Evaluable for EvalDistinct { - fn evaluate(&mut self, ctx: &dyn EvalContext) -> Value { + fn evaluate<'a, 'c>(&mut self, ctx: &'c dyn EvalContext<'c>) -> Value { let input_value = take_input!(self.input.take(), ctx); let ordered = input_value.is_ordered(); @@ -1277,7 +1277,7 @@ pub(crate) struct EvalSink { } impl Evaluable for EvalSink { - fn evaluate(&mut self, _ctx: &dyn EvalContext) -> Value { + fn evaluate<'a, 'c>(&mut self, _ctx: &'c dyn EvalContext<'c>) -> Value { self.input.take().unwrap_or_else(|| Missing) } @@ -1308,15 +1308,24 @@ impl EvalSubQueryExpr { } impl EvalExpr for EvalSubQueryExpr { - fn evaluate<'a>(&'a self, bindings: &'a Tuple, _ctx: &'a dyn EvalContext) -> Cow<'a, Value> { - let value = if let Ok(evaluated) = self - .plan - .borrow_mut() - .execute_mut(MapBindings::from(bindings)) - { - evaluated.result - } else { - Missing + fn evaluate<'a, 'c>( + &'a self, + bindings: &'a Tuple, + ctx: &'c dyn EvalContext<'c>, + ) -> Cow<'a, Value> + where + 'c: 'a, + { + let bindings = MapBindings::from(bindings); + let value = { + let nested_ctx: NestedContext = NestedContext::new(bindings, ctx); + + let mut plan = self.plan.borrow_mut(); + if let Ok(evaluated) = plan.execute_mut(&nested_ctx) { + evaluated.result + } else { + Missing + } }; Cow::Owned(value) } @@ -1357,7 +1366,7 @@ impl EvalOuterUnion { } impl Evaluable for EvalOuterUnion { - fn evaluate(&mut self, _ctx: &dyn EvalContext) -> Value { + fn evaluate<'a, 'c>(&mut self, _ctx: &'c dyn EvalContext<'c>) -> Value { let lhs = bagop_iter(self.l_input.take().unwrap_or(Missing)); let rhs = bagop_iter(self.r_input.take().unwrap_or(Missing)); let chained = lhs.chain(rhs); @@ -1398,7 +1407,7 @@ impl EvalOuterIntersect { } impl Evaluable for EvalOuterIntersect { - fn evaluate(&mut self, _ctx: &dyn EvalContext) -> Value { + fn evaluate<'a, 'c>(&mut self, _ctx: &'c dyn EvalContext<'c>) -> Value { let lhs = bagop_iter(self.l_input.take().unwrap_or(Missing)); let rhs = bagop_iter(self.r_input.take().unwrap_or(Missing)); @@ -1454,7 +1463,7 @@ impl EvalOuterExcept { } impl Evaluable for EvalOuterExcept { - fn evaluate(&mut self, _ctx: &dyn EvalContext) -> Value { + fn evaluate<'a, 'c>(&mut self, _ctx: &'c dyn EvalContext<'c>) -> Value { let lhs = bagop_iter(self.l_input.take().unwrap_or(Missing)); let rhs = bagop_iter(self.r_input.take().unwrap_or(Missing)); diff --git a/partiql-eval/src/eval/expr/base_table.rs b/partiql-eval/src/eval/expr/base_table.rs index 8f3c26fb..7a9e955b 100644 --- a/partiql-eval/src/eval/expr/base_table.rs +++ b/partiql-eval/src/eval/expr/base_table.rs @@ -18,13 +18,20 @@ pub(crate) struct EvalFnBaseTableExpr { impl EvalExpr for EvalFnBaseTableExpr { #[inline] - fn evaluate<'a>(&'a self, bindings: &'a Tuple, ctx: &'a dyn EvalContext) -> Cow<'a, Value> { + fn evaluate<'a, 'c>( + &'a self, + bindings: &'a Tuple, + ctx: &'c dyn EvalContext<'c>, + ) -> Cow<'a, Value> + where + 'c: 'a, + { let args = self .args .iter() .map(|arg| arg.evaluate(bindings, ctx)) .collect_vec(); - let results = self.expr.evaluate(&args); + let results = self.expr.evaluate(&args, ctx.as_session()); let result = match results { Ok(it) => { let bag: Result = it.collect(); diff --git a/partiql-eval/src/eval/expr/control_flow.rs b/partiql-eval/src/eval/expr/control_flow.rs index 92b111a3..9f628d1d 100644 --- a/partiql-eval/src/eval/expr/control_flow.rs +++ b/partiql-eval/src/eval/expr/control_flow.rs @@ -13,7 +13,14 @@ pub(crate) struct EvalSearchedCaseExpr { } impl EvalExpr for EvalSearchedCaseExpr { - fn evaluate<'a>(&'a self, bindings: &'a Tuple, ctx: &'a dyn EvalContext) -> Cow<'a, Value> { + fn evaluate<'a, 'c>( + &'a self, + bindings: &'a Tuple, + ctx: &'c dyn EvalContext<'c>, + ) -> Cow<'a, Value> + where + 'c: 'a, + { for (when_expr, then_expr) in &self.cases { let when_expr_evaluated = when_expr.evaluate(bindings, ctx); if when_expr_evaluated.as_ref() == &Value::Boolean(true) { diff --git a/partiql-eval/src/eval/expr/data_types.rs b/partiql-eval/src/eval/expr/data_types.rs index c4d9c193..aa9c83f5 100644 --- a/partiql-eval/src/eval/expr/data_types.rs +++ b/partiql-eval/src/eval/expr/data_types.rs @@ -20,7 +20,14 @@ pub(crate) struct EvalTupleExpr { } impl EvalExpr for EvalTupleExpr { - fn evaluate<'a>(&'a self, bindings: &'a Tuple, ctx: &'a dyn EvalContext) -> Cow<'a, Value> { + fn evaluate<'a, 'c>( + &'a self, + bindings: &'a Tuple, + ctx: &'c dyn EvalContext<'c>, + ) -> Cow<'a, Value> + where + 'c: 'a, + { let tuple = self .attrs .iter() @@ -52,7 +59,14 @@ pub(crate) struct EvalListExpr { } impl EvalExpr for EvalListExpr { - fn evaluate<'a>(&'a self, bindings: &'a Tuple, ctx: &'a dyn EvalContext) -> Cow<'a, Value> { + fn evaluate<'a, 'c>( + &'a self, + bindings: &'a Tuple, + ctx: &'c dyn EvalContext<'c>, + ) -> Cow<'a, Value> + where + 'c: 'a, + { let values = self .elements .iter() @@ -70,7 +84,14 @@ pub(crate) struct EvalBagExpr { } impl EvalExpr for EvalBagExpr { - fn evaluate<'a>(&'a self, bindings: &'a Tuple, ctx: &'a dyn EvalContext) -> Cow<'a, Value> { + fn evaluate<'a, 'c>( + &'a self, + bindings: &'a Tuple, + ctx: &'c dyn EvalContext<'c>, + ) -> Cow<'a, Value> + where + 'c: 'a, + { let values = self .elements .iter() @@ -89,7 +110,14 @@ pub(crate) struct EvalIsTypeExpr { } impl EvalExpr for EvalIsTypeExpr { - fn evaluate<'a>(&'a self, bindings: &'a Tuple, ctx: &'a dyn EvalContext) -> Cow<'a, Value> { + fn evaluate<'a, 'c>( + &'a self, + bindings: &'a Tuple, + ctx: &'c dyn EvalContext<'c>, + ) -> Cow<'a, Value> + where + 'c: 'a, + { let expr = self.expr.evaluate(bindings, ctx); let expr = expr.as_ref(); let result = match self.is_type { diff --git a/partiql-eval/src/eval/expr/mod.rs b/partiql-eval/src/eval/expr/mod.rs index a79852f5..8e107e19 100644 --- a/partiql-eval/src/eval/expr/mod.rs +++ b/partiql-eval/src/eval/expr/mod.rs @@ -26,7 +26,13 @@ use thiserror::Error; /// A trait for expressions that require evaluation, e.g. `a + b` or `c > 2`. pub trait EvalExpr: Debug { - fn evaluate<'a>(&'a self, bindings: &'a Tuple, ctx: &'a dyn EvalContext) -> Cow<'a, Value>; + fn evaluate<'a, 'c>( + &'a self, + bindings: &'a Tuple, + ctx: &'c dyn EvalContext<'c>, + ) -> Cow<'a, Value> + where + 'c: 'a; } #[derive(Error, Debug, Clone, PartialEq)] diff --git a/partiql-eval/src/eval/expr/operators.rs b/partiql-eval/src/eval/expr/operators.rs index b2ef7430..a5ddb018 100644 --- a/partiql-eval/src/eval/expr/operators.rs +++ b/partiql-eval/src/eval/expr/operators.rs @@ -42,17 +42,27 @@ impl BindEvalExpr for EvalLitExpr { } impl EvalExpr for EvalLitExpr { - fn evaluate<'a>(&'a self, _: &'a Tuple, _: &'a dyn EvalContext) -> Cow<'a, Value> { + fn evaluate<'a, 'c>( + &'a self, + _bindings: &'a Tuple, + _ctx: &'c dyn EvalContext<'c>, + ) -> Cow<'a, Value> + where + 'c: 'a, + { Cow::Borrowed(&self.lit) } } impl ExecuteEvalExpr<0> for Value { - fn evaluate<'a>( + fn evaluate<'a, 'c>( &'a self, _args: [Cow<'a, Value>; 0], - _ctx: &'a dyn EvalContext, - ) -> Cow<'a, Value> { + _ctx: &'c dyn EvalContext<'c>, + ) -> Cow<'a, Value> + where + 'c: 'a, + { Cow::Borrowed(self) } } diff --git a/partiql-eval/src/eval/expr/path.rs b/partiql-eval/src/eval/expr/path.rs index ed88ecd5..c0bf8801 100644 --- a/partiql-eval/src/eval/expr/path.rs +++ b/partiql-eval/src/eval/expr/path.rs @@ -1,5 +1,3 @@ -use crate::env::Bindings; - pub use core::borrow::Borrow; use crate::eval::expr::{BindError, BindEvalExpr, EvalExpr}; @@ -8,6 +6,7 @@ use crate::eval::EvalContext; use partiql_value::Value::Missing; use partiql_value::{BindingsName, Tuple, Value}; +use partiql_catalog::context::Bindings; use std::borrow::Cow; use std::fmt::{Debug, Formatter}; @@ -80,11 +79,11 @@ fn as_int(v: &Value) -> Option { impl EvalPathComponent { #[inline] - fn get_val<'a>( + fn get_val<'a, 'c>( &self, value: &'a Value, bindings: &'a Tuple, - ctx: &dyn EvalContext, + ctx: &'c dyn EvalContext<'c>, ) -> Option<&'a Value> { match (self, value) { (EvalPathComponent::Key(k), Value::Tuple(tuple)) => tuple.get(k), @@ -100,7 +99,12 @@ impl EvalPathComponent { } #[inline] - fn take_val(&self, value: Value, bindings: &Tuple, ctx: &dyn EvalContext) -> Option { + fn take_val<'c>( + &self, + value: Value, + bindings: &Tuple, + ctx: &'c dyn EvalContext<'c>, + ) -> Option { match (self, value) { (EvalPathComponent::Key(k), Value::Tuple(tuple)) => tuple.take_val(k), (EvalPathComponent::Index(idx), Value::List(list)) => list.take_val(*idx), @@ -116,7 +120,14 @@ impl EvalPathComponent { } impl EvalExpr for EvalPath { - fn evaluate<'a>(&'a self, bindings: &'a Tuple, ctx: &'a dyn EvalContext) -> Cow<'a, Value> { + fn evaluate<'a, 'c>( + &'a self, + bindings: &'a Tuple, + ctx: &'c dyn EvalContext<'c>, + ) -> Cow<'a, Value> + where + 'c: 'a, + { let value = self.expr.evaluate(bindings, ctx); let mut path_componenents = self.components.iter(); match value { @@ -138,7 +149,14 @@ pub(crate) struct EvalDynamicLookup { } impl EvalExpr for EvalDynamicLookup { - fn evaluate<'a>(&'a self, bindings: &'a Tuple, ctx: &'a dyn EvalContext) -> Cow<'a, Value> { + fn evaluate<'a, 'c>( + &'a self, + bindings: &'a Tuple, + ctx: &'c dyn EvalContext<'c>, + ) -> Cow<'a, Value> + where + 'c: 'a, + { let mut lookups = self.lookups.iter().filter_map(|lookup| { let val = lookup.evaluate(bindings, ctx); match val.as_ref() { @@ -182,7 +200,14 @@ pub(crate) struct EvalLocalVarRef { } impl EvalExpr for EvalLocalVarRef { - fn evaluate<'a>(&'a self, bindings: &'a Tuple, _: &'a dyn EvalContext) -> Cow<'a, Value> { + fn evaluate<'a, 'c>( + &'a self, + bindings: &'a Tuple, + _ctx: &'c dyn EvalContext<'c>, + ) -> Cow<'a, Value> + where + 'c: 'a, + { borrow_or_missing(Bindings::get(bindings, &self.name)) } } @@ -212,7 +237,14 @@ impl Debug for EvalGlobalVarRef { } impl EvalExpr for EvalGlobalVarRef { - fn evaluate<'a>(&'a self, _: &'a Tuple, ctx: &'a dyn EvalContext) -> Cow<'a, Value> { + fn evaluate<'a, 'c>( + &'a self, + _bindings: &'a Tuple, + ctx: &'c dyn EvalContext<'c>, + ) -> Cow<'a, Value> + where + 'c: 'a, + { borrow_or_missing(ctx.bindings().get(&self.name)) } } diff --git a/partiql-eval/src/eval/expr/strings.rs b/partiql-eval/src/eval/expr/strings.rs index b3cbdb41..1fcffbb2 100644 --- a/partiql-eval/src/eval/expr/strings.rs +++ b/partiql-eval/src/eval/expr/strings.rs @@ -66,11 +66,14 @@ where R: Into, { #[inline] - fn evaluate<'a>( + fn evaluate<'a, 'c>( &'a self, args: [Cow<'a, Value>; 1], - _ctx: &'a dyn EvalContext, - ) -> Cow<'a, Value> { + _ctx: &'c dyn EvalContext<'c>, + ) -> Cow<'a, Value> + where + 'c: 'a, + { let [value] = args; Cow::Owned(match value.borrow() { Value::String(s) => ((self.f)(s)).into(), diff --git a/partiql-eval/src/eval/mod.rs b/partiql-eval/src/eval/mod.rs index a6fafe6d..c12794c8 100644 --- a/partiql-eval/src/eval/mod.rs +++ b/partiql-eval/src/eval/mod.rs @@ -1,6 +1,9 @@ use itertools::Itertools; +use std::any::Any; use std::cell::RefCell; +use std::collections::HashMap; +use delegate::delegate; use std::fmt::Debug; use petgraph::algo::toposort; @@ -11,12 +14,13 @@ use petgraph::{Directed, Outgoing}; use partiql_value::Value; use crate::env::basic::MapBindings; -use crate::env::Bindings; use petgraph::graph::NodeIndex; use crate::error::{EvalErr, EvaluationError}; +use partiql_catalog::context::{Bindings, SessionContext, SystemContext}; use petgraph::visit::EdgeRef; +use unicase::UniCase; use crate::eval::evaluable::{EvalType, Evaluable}; @@ -62,8 +66,7 @@ impl EvalPlan { /// Executes the plan while mutating its state by changing the inputs and outputs of plan /// operators. - pub fn execute_mut(&mut self, bindings: MapBindings) -> Result { - let ctx: Box = Box::new(BasicContext::new(bindings)); + pub fn execute_mut<'c>(&mut self, ctx: &'c dyn EvalContext<'c>) -> Result { // We are only interested in DAGs that can be used as execution plans, which leads to the // following definition. // A DAG is a directed, cycle-free graph G = (V, E) with a denoted root node v0 ∈ V such @@ -95,7 +98,7 @@ impl EvalPlan { }); if graph_managed { let src = self.get_node(idx)?; - result = Some(src.evaluate(&*ctx)); + result = Some(src.evaluate(ctx)); // return on first evaluation error if ctx.has_errors() { @@ -114,7 +117,7 @@ impl EvalPlan { let res = res.ok_or_else(|| err_illegal_state("Error in retrieving source value"))?; - self.get_node(dst_id)?.update_input(res, branch_num, &*ctx); + self.get_node(dst_id)?.update_input(res, branch_num, ctx); } } } @@ -140,33 +143,53 @@ pub struct Evaluated { } /// Represents an evaluation context that is used during evaluation of a plan. -pub trait EvalContext { - fn bindings(&self) -> &dyn Bindings; +pub trait EvalContext<'a>: SessionContext<'a> + Debug { + fn as_session(&'a self) -> &'a dyn SessionContext<'a>; + fn add_error(&self, error: EvaluationError); fn has_errors(&self) -> bool; fn errors(&self) -> Vec; } -#[derive(Default, Debug)] -pub struct BasicContext { - bindings: MapBindings, - errors: RefCell>, +#[derive(Debug)] +pub struct BasicContext<'a> { + pub bindings: MapBindings, + + pub sys: SystemContext, + pub user: HashMap, &'a (dyn Any)>, + + pub errors: RefCell>, } -impl BasicContext { - pub fn new(bindings: MapBindings) -> Self { +impl<'a> BasicContext<'a> { + pub fn new(bindings: MapBindings, sys: SystemContext) -> Self { BasicContext { bindings, + sys, + user: Default::default(), errors: RefCell::new(vec![]), } } } - -impl EvalContext for BasicContext { +impl<'a> SessionContext<'a> for BasicContext<'a> { fn bindings(&self) -> &dyn Bindings { &self.bindings } + fn system_context(&self) -> &SystemContext { + &self.sys + } + + fn user_context(&self, name: &str) -> Option<&(dyn Any)> { + let key = name.into(); + self.user.get(&key).copied() + } +} +impl<'a> EvalContext<'a> for BasicContext<'a> { + fn as_session(&'a self) -> &'a dyn SessionContext<'a> { + self + } + fn add_error(&self, error: EvaluationError) { self.errors.borrow_mut().push(error) } @@ -179,3 +202,42 @@ impl EvalContext for BasicContext { self.errors.take() } } + +#[derive(Debug)] +pub struct NestedContext<'a, 'c> { + pub bindings: MapBindings, + pub parent: &'a dyn EvalContext<'c>, +} + +impl<'a, 'c> NestedContext<'a, 'c> { + pub fn new(bindings: MapBindings, parent: &'a dyn EvalContext<'c>) -> Self { + NestedContext { bindings, parent } + } +} + +impl<'a, 'c> SessionContext<'a> for NestedContext<'a, 'c> { + fn bindings(&self) -> &dyn Bindings { + &self.bindings + } + + delegate! { + to self.parent { + fn system_context(&self) -> &SystemContext; + fn user_context(&self, name: &str) -> Option<& (dyn Any )>; + } + } +} + +impl<'a, 'c> EvalContext<'a> for NestedContext<'a, 'c> { + fn as_session(&'a self) -> &'a dyn SessionContext<'a> { + self + } + + delegate! { + to self.parent { + fn add_error(&self, error: EvaluationError); + fn has_errors(&self) -> bool; + fn errors(&self) -> Vec; + } + } +} diff --git a/partiql-eval/src/lib.rs b/partiql-eval/src/lib.rs index 96435f36..824d4079 100644 --- a/partiql-eval/src/lib.rs +++ b/partiql-eval/src/lib.rs @@ -9,12 +9,14 @@ mod tests { use crate::env::basic::MapBindings; use crate::plan; + use partiql_catalog::context::SystemContext; use partiql_catalog::PartiqlCatalog; use rust_decimal_macros::dec; use partiql_logical as logical; use partiql_logical::BindingsOp::{Distinct, Project, ProjectAll, ProjectValue}; + use crate::eval::BasicContext; use crate::plan::EvaluationMode; use partiql_logical::{ BagExpr, BetweenExpr, BinaryOp, BindingsOp, CoalesceExpr, ExprQuery, IsTypeExpr, JoinKind, @@ -22,14 +24,17 @@ mod tests { }; use partiql_value as value; use partiql_value::Value::{Missing, Null}; - use partiql_value::{bag, list, tuple, Bag, BindingsName, List, Tuple, Value}; + use partiql_value::{bag, list, tuple, Bag, BindingsName, DateTime, List, Tuple, Value}; fn evaluate(logical: LogicalPlan, bindings: MapBindings) -> Value { let catalog = PartiqlCatalog::default(); let mut planner = plan::EvaluatorPlanner::new(EvaluationMode::Permissive, &catalog); let mut plan = planner.compile(&logical).expect("Expect no plan error"); - - if let Ok(out) = plan.execute_mut(bindings) { + let sys = SystemContext { + now: DateTime::from_system_now_utc(), + }; + let ctx = BasicContext::new(bindings, sys); + if let Ok(out) = plan.execute_mut(&ctx) { out.result } else { Missing @@ -2212,7 +2217,8 @@ mod tests { use crate::eval::evaluable::{EvalScan, Evaluable}; use crate::eval::expr::{EvalGlobalVarRef, EvalPath, EvalPathComponent}; use crate::eval::BasicContext; - use partiql_value::{bag, list, BindingsName}; + use partiql_catalog::context::SystemContext; + use partiql_value::{bag, list, BindingsName, DateTime}; use super::*; @@ -2230,7 +2236,10 @@ mod tests { let mut p0: MapBindings = MapBindings::default(); p0.insert("someOrderedTable", some_ordered_table().into()); - let ctx = BasicContext::new(p0); + let sys = SystemContext { + now: DateTime::from_system_now_utc(), + }; + let ctx = BasicContext::new(p0, sys); let mut scan = EvalScan::new_with_at_key( Box::new(EvalGlobalVarRef { @@ -2256,7 +2265,10 @@ mod tests { let mut p0: MapBindings = MapBindings::default(); p0.insert("someUnorderedTable", some_unordered_table().into()); - let ctx = BasicContext::new(p0); + let sys = SystemContext { + now: DateTime::from_system_now_utc(), + }; + let ctx = BasicContext::new(p0, sys); let mut scan = EvalScan::new_with_at_key( Box::new(EvalGlobalVarRef { @@ -2300,7 +2312,10 @@ mod tests { }; let mut scan = EvalScan::new(Box::new(path_to_scalar), "x"); - let ctx = BasicContext::new(p0); + let sys = SystemContext { + now: DateTime::from_system_now_utc(), + }; + let ctx = BasicContext::new(p0, sys); let scan_res = scan.evaluate(&ctx); let expected = bag![tuple![("x", 0)]]; @@ -2325,7 +2340,10 @@ mod tests { }; let mut scan = EvalScan::new(Box::new(path_to_scalar), "x"); - let ctx = BasicContext::new(p0); + let sys = SystemContext { + now: DateTime::from_system_now_utc(), + }; + let ctx = BasicContext::new(p0, sys); let res = scan.evaluate(&ctx); let expected = bag![tuple![("x", value::Value::Missing)]]; @@ -2334,7 +2352,7 @@ mod tests { } mod clause_unpivot { - use partiql_value::{bag, BindingsName, Tuple}; + use partiql_value::{bag, BindingsName, DateTime, Tuple}; use crate::eval::evaluable::{EvalUnpivot, Evaluable}; use crate::eval::expr::EvalGlobalVarRef; @@ -2360,7 +2378,10 @@ mod tests { Some("symbol".into()), ); - let ctx = BasicContext::new(p0); + let sys = SystemContext { + now: DateTime::from_system_now_utc(), + }; + let ctx = BasicContext::new(p0, sys); let res = unpivot.evaluate(&ctx); let expected = bag![ @@ -2384,7 +2405,10 @@ mod tests { Some("y".into()), ); - let ctx = BasicContext::new(p0); + let sys = SystemContext { + now: DateTime::from_system_now_utc(), + }; + let ctx = BasicContext::new(p0, sys); let res = unpivot.evaluate(&ctx); let expected = bag![tuple![("x", 1), ("y", "_1")]]; diff --git a/partiql-logical-planner/src/lib.rs b/partiql-logical-planner/src/lib.rs index 322664f1..fc13857a 100644 --- a/partiql-logical-planner/src/lib.rs +++ b/partiql-logical-planner/src/lib.rs @@ -38,9 +38,11 @@ impl<'c> LogicalPlanner<'c> { mod tests { use assert_matches::assert_matches; use partiql_ast_passes::error::AstTransformationError; + use partiql_catalog::context::SystemContext; use partiql_catalog::PartiqlCatalog; use partiql_eval::env::basic::MapBindings; + use partiql_eval::eval::BasicContext; use partiql_eval::plan; use partiql_eval::plan::EvaluationMode; @@ -49,7 +51,7 @@ mod tests { use partiql_logical as logical; use partiql_logical::{BindingsOp, LogicalPlan}; use partiql_parser::{Parsed, Parser}; - use partiql_value::{bag, tuple, Value}; + use partiql_value::{bag, tuple, DateTime, Value}; #[track_caller] fn parse(text: &str) -> Parsed { @@ -71,8 +73,11 @@ mod tests { let mut planner = plan::EvaluatorPlanner::new(EvaluationMode::Permissive, &catalog); let mut plan = planner.compile(&logical).expect("Expect no plan error"); println!("{}", plan.to_dot_graph()); - - if let Ok(out) = plan.execute_mut(bindings) { + let sys = SystemContext { + now: DateTime::from_system_now_utc(), + }; + let ctx = BasicContext::new(bindings, sys); + if let Ok(out) = plan.execute_mut(&ctx) { out.result } else { Value::Missing diff --git a/partiql-logical-planner/src/lower.rs b/partiql-logical-planner/src/lower.rs index 86feb936..be953db3 100644 --- a/partiql-logical-planner/src/lower.rs +++ b/partiql-logical-planner/src/lower.rs @@ -1160,7 +1160,7 @@ impl<'a, 'ast> Visitor<'ast> for AstToLogical<'a> { not_yet_implemented_fault!(self, "PositionalType call argument".to_string()); } CallArg::NamedType(_) => { - not_yet_implemented_fault!(self, "PositionalType call argument".to_string()); + not_yet_implemented_fault!(self, "NamedType call argument".to_string()); } } Traverse::Continue diff --git a/partiql-value/src/datetime.rs b/partiql-value/src/datetime.rs index b69c53ef..7cd65ef5 100644 --- a/partiql-value/src/datetime.rs +++ b/partiql-value/src/datetime.rs @@ -17,6 +17,10 @@ pub enum DateTime { } impl DateTime { + pub fn from_system_now_utc() -> Self { + DateTime::TimestampWithTz(time::OffsetDateTime::now_utc()) + } + pub fn from_hms(hour: u8, minute: u8, second: u8) -> Self { DateTime::Time(time::Time::from_hms(hour, minute, second).expect("valid time value")) } diff --git a/partiql/benches/bench_agg.rs b/partiql/benches/bench_agg.rs index b4bfd6cb..d28b4c62 100644 --- a/partiql/benches/bench_agg.rs +++ b/partiql/benches/bench_agg.rs @@ -2,16 +2,17 @@ use std::time::Duration; use criterion::{black_box, criterion_group, criterion_main, Criterion}; use itertools::Itertools; +use partiql_catalog::context::SystemContext; use partiql_catalog::{Catalog, PartiqlCatalog}; use partiql_eval::env::basic::MapBindings; -use partiql_eval::eval::EvalPlan; +use partiql_eval::eval::{BasicContext, EvalPlan}; use partiql_eval::plan::{EvaluationMode, EvaluatorPlanner}; use partiql_logical::{BindingsOp, LogicalPlan}; use partiql_logical_planner::LogicalPlanner; use partiql_parser::{Parser, ParserResult}; -use partiql_value::{tuple, Bag, Value}; +use partiql_value::{tuple, Bag, DateTime, Value}; fn numbers() -> impl Iterator { (0..1000i64).map(Value::from) @@ -49,7 +50,11 @@ fn plan(catalog: &dyn Catalog, logical: &LogicalPlan) -> EvalPlan { } #[inline] pub(crate) fn evaluate(mut eval: EvalPlan, bindings: MapBindings) -> Value { - if let Ok(out) = eval.execute_mut(bindings) { + let sys = SystemContext { + now: DateTime::from_system_now_utc(), + }; + let ctx = BasicContext::new(bindings, sys); + if let Ok(out) = eval.execute_mut(&ctx) { out.result } else { Value::Missing diff --git a/partiql/benches/bench_eval_multi_like.rs b/partiql/benches/bench_eval_multi_like.rs index 6348808f..8de96226 100644 --- a/partiql/benches/bench_eval_multi_like.rs +++ b/partiql/benches/bench_eval_multi_like.rs @@ -2,17 +2,18 @@ use std::time::Duration; use criterion::{black_box, criterion_group, criterion_main, Criterion}; use itertools::Itertools; +use partiql_catalog::context::SystemContext; use partiql_catalog::{Catalog, PartiqlCatalog}; use rand::{Rng, SeedableRng}; use partiql_eval::env::basic::MapBindings; -use partiql_eval::eval::EvalPlan; +use partiql_eval::eval::{BasicContext, EvalPlan}; use partiql_eval::plan::{EvaluationMode, EvaluatorPlanner}; use partiql_logical::{BindingsOp, LogicalPlan}; use partiql_logical_planner::LogicalPlanner; use partiql_parser::{Parser, ParserResult}; -use partiql_value::{tuple, Bag, Value}; +use partiql_value::{tuple, Bag, DateTime, Value}; // Benchmarks: // - parsing, @@ -347,7 +348,11 @@ fn plan(catalog: &dyn Catalog, logical: &LogicalPlan) -> EvalPlan { } #[inline] pub(crate) fn evaluate(mut eval: EvalPlan, bindings: MapBindings) -> Value { - if let Ok(out) = eval.execute_mut(bindings) { + let sys = SystemContext { + now: DateTime::from_system_now_utc(), + }; + let ctx = BasicContext::new(bindings, sys); + if let Ok(out) = eval.execute_mut(&ctx) { out.result } else { Value::Missing diff --git a/partiql/src/lib.rs b/partiql/src/lib.rs index ac5da82a..602fd95c 100644 --- a/partiql/src/lib.rs +++ b/partiql/src/lib.rs @@ -1,15 +1,16 @@ #[cfg(test)] mod tests { use partiql_ast_passes::error::AstTransformationError; + use partiql_catalog::context::SystemContext; use partiql_catalog::{Catalog, PartiqlCatalog}; use partiql_eval as eval; use partiql_eval::env::basic::MapBindings; use partiql_eval::error::{EvalErr, PlanErr}; - use partiql_eval::eval::{EvalPlan, EvalResult, Evaluated}; + use partiql_eval::eval::{BasicContext, EvalPlan, EvalResult, Evaluated}; use partiql_eval::plan::EvaluationMode; use partiql_logical as logical; use partiql_parser::{Parsed, ParserError, ParserResult}; - use partiql_value::Value; + use partiql_value::{DateTime, Value}; use thiserror::Error; #[derive(Error, Debug)] @@ -78,7 +79,11 @@ mod tests { #[track_caller] #[inline] fn evaluate(mut plan: EvalPlan, bindings: MapBindings) -> EvalResult { - plan.execute_mut(bindings) + let sys = SystemContext { + now: DateTime::from_system_now_utc(), + }; + let ctx = BasicContext::new(bindings, sys); + plan.execute_mut(&ctx) } #[track_caller] @@ -97,14 +102,14 @@ mod tests { fn order_by_count() { let query = "select foo, count(1) as n from << - { 'foo': 'foo' }, - { 'foo': 'bar' }, - { 'foo': 'qux' }, - { 'foo': 'bar' }, - { 'foo': 'baz' }, - { 'foo': 'bar' }, - { 'foo': 'baz' } - >> group by foo order by n desc"; + { 'foo': 'foo' }, + { 'foo': 'bar' }, + { 'foo': 'qux' }, + { 'foo': 'bar' }, + { 'foo': 'baz' }, + { 'foo': 'bar' }, + { 'foo': 'baz' } + >> group by foo order by n desc"; let res = eval(query, EvaluationMode::Permissive); assert!(res.is_ok()); diff --git a/partiql/tests/user_context.rs b/partiql/tests/user_context.rs new file mode 100644 index 00000000..bd35c776 --- /dev/null +++ b/partiql/tests/user_context.rs @@ -0,0 +1,215 @@ +use std::any::Any; +use std::borrow::Cow; +use std::cell::RefCell; + +use std::error::Error; + +use thiserror::Error; + +use partiql_catalog::call_defs::{CallDef, CallSpec, CallSpecArg}; +use partiql_catalog::context::{SessionContext, SystemContext}; +use partiql_catalog::{ + BaseTableExpr, BaseTableExprResult, BaseTableExprResultError, BaseTableFunctionInfo, Catalog, + Extension, PartiqlCatalog, TableFunction, +}; +use partiql_eval::env::basic::MapBindings; +use partiql_eval::eval::BasicContext; +use partiql_eval::plan::EvaluationMode; +use partiql_parser::{Parsed, ParserResult}; +use partiql_value::{bag, tuple, DateTime, Value}; + +use partiql_logical as logical; + +#[derive(Debug)] +pub struct UserCtxTestExtension {} + +impl partiql_catalog::Extension for UserCtxTestExtension { + fn name(&self) -> String { + "ion".into() + } + + fn load(&self, catalog: &mut dyn Catalog) -> Result<(), Box> { + match catalog + .add_table_function(TableFunction::new(Box::new(TestUserContextFunction::new()))) + { + Ok(_) => Ok(()), + Err(e) => Err(Box::new(e) as Box), + } + } +} + +#[derive(Debug)] +pub(crate) struct TestUserContextFunction { + call_def: CallDef, +} + +impl TestUserContextFunction { + pub fn new() -> Self { + TestUserContextFunction { + call_def: CallDef { + names: vec!["test_user_context"], + overloads: vec![CallSpec { + input: vec![CallSpecArg::Positional], + output: Box::new(|args| { + logical::ValueExpr::Call(logical::CallExpr { + name: logical::CallName::ByName("test_user_context".to_string()), + arguments: args, + }) + }), + }], + }, + } + } +} + +impl BaseTableFunctionInfo for TestUserContextFunction { + fn call_def(&self) -> &CallDef { + &self.call_def + } + + fn plan_eval(&self) -> Box { + Box::new(EvalTestCtxTable {}) + } +} + +#[derive(Error, Debug)] +#[non_exhaustive] +pub enum UserCtxError { + #[error("unknown error")] + Unknown, +} + +#[derive(Debug)] +pub(crate) struct EvalTestCtxTable {} + +impl BaseTableExpr for EvalTestCtxTable { + fn evaluate<'c>( + &self, + args: &[Cow], + ctx: &'c dyn SessionContext<'c>, + ) -> BaseTableExprResult<'c> { + if let Some(arg1) = args.first() { + match arg1.as_ref() { + Value::String(name) => generated_data(name.to_string(), ctx), + _ => { + let error = UserCtxError::Unknown; + Err(Box::new(error) as BaseTableExprResultError) + } + } + } else { + let error = UserCtxError::Unknown; + Err(Box::new(error) as BaseTableExprResultError) + } + } +} + +struct TestDataGen<'a> { + ctx: &'a dyn SessionContext<'a>, + name: String, +} + +impl<'a> Iterator for TestDataGen<'a> { + type Item = Result; + + fn next(&mut self) -> Option { + if let Some(cv) = self.ctx.user_context(&self.name) { + if let Some(counter) = cv.downcast_ref::() { + let mut n = counter.data.borrow_mut(); + + if *n > 0 { + *n -= 1; + + let idx: u8 = (5 - *n) as u8; + let id = format!("id_{idx}"); + let m = idx % 2; + + return Some(Ok(tuple![("foo", m), ("bar", id)].into())); + } + } + } + None + } +} + +fn generated_data<'a>(name: String, ctx: &'a dyn SessionContext<'a>) -> BaseTableExprResult<'a> { + Ok(Box::new(TestDataGen { ctx, name })) +} + +#[derive(Debug)] +pub struct Counter { + data: RefCell, +} + +#[track_caller] +#[inline] +pub(crate) fn parse(statement: &str) -> ParserResult { + partiql_parser::Parser::default().parse(statement) +} + +#[track_caller] +#[inline] +pub(crate) fn lower( + catalog: &dyn Catalog, + parsed: &Parsed, +) -> partiql_logical::LogicalPlan { + let planner = partiql_logical_planner::LogicalPlanner::new(catalog); + planner.lower(parsed).expect("lower") +} + +#[track_caller] +#[inline] +pub(crate) fn evaluate( + catalog: &dyn Catalog, + logical: partiql_logical::LogicalPlan, + bindings: MapBindings, + ctx_vals: &[(String, &(dyn Any))], +) -> Value { + let mut planner = + partiql_eval::plan::EvaluatorPlanner::new(EvaluationMode::Permissive, catalog); + + let mut plan = planner.compile(&logical).expect("Expect no plan error"); + + let sys = SystemContext { + now: DateTime::from_system_now_utc(), + }; + let mut ctx = BasicContext::new(bindings, sys); + for (k, v) in ctx_vals { + ctx.user.insert(k.as_str().into(), *v); + } + if let Ok(out) = plan.execute_mut(&ctx) { + out.result + } else { + Value::Missing + } +} +#[test] +fn test_context() { + let expected: Value = bag![ + tuple![("foo", 1), ("bar", "id_1")], + tuple![("foo", 0), ("bar", "id_2")], + tuple![("foo", 1), ("bar", "id_3")], + tuple![("foo", 0), ("bar", "id_4")], + tuple![("foo", 1), ("bar", "id_5")], + ] + .into(); + + let query = "SELECT foo, bar from test_user_context('counter') as data"; + + let mut catalog = PartiqlCatalog::default(); + let ext = UserCtxTestExtension {}; + ext.load(&mut catalog).expect("extension load to succeed"); + + let parsed = parse(query); + let lowered = lower(&catalog, &parsed.expect("parse")); + let bindings = Default::default(); + + let counter = Counter { + data: RefCell::new(5), + }; + let ctx: [(String, &dyn Any); 1] = [("counter".to_string(), &counter)]; + let out = evaluate(&catalog, lowered, bindings, &ctx); + + assert!(out.is_bag()); + assert_eq!(&out, &expected); + assert_eq!(*counter.data.borrow(), 0); +}