Don't go through JSON

main
Ziyang Hu 2 years ago
parent 8bff4eb92e
commit 525cfeff59

@ -23,7 +23,6 @@ use smartstring::{LazyCompact, SmartString};
use thiserror::Error; use thiserror::Error;
use crate::data::expr::Expr; use crate::data::expr::Expr;
use crate::data::json::JsonValue;
use crate::data::program::{ use crate::data::program::{
FixedRuleOptionNotFoundError, MagicFixedRuleApply, MagicFixedRuleRuleArg, MagicSymbol, FixedRuleOptionNotFoundError, MagicFixedRuleApply, MagicFixedRuleRuleArg, MagicSymbol,
WrongFixedRuleOptionError, WrongFixedRuleOptionError,
@ -567,7 +566,7 @@ pub trait FixedRule: Send + Sync {
/// but implementation is simpler. /// but implementation is simpler.
pub struct SimpleFixedRule { pub struct SimpleFixedRule {
return_arity: usize, return_arity: usize,
rule: Box<dyn Fn(Vec<NamedRows>, JsonValue) -> Result<NamedRows> + Send + Sync + 'static>, rule: Box<dyn Fn(Vec<NamedRows>, BTreeMap<String, DataValue>) -> Result<NamedRows> + Send + Sync + 'static>,
} }
impl SimpleFixedRule { impl SimpleFixedRule {
@ -581,7 +580,7 @@ impl SimpleFixedRule {
// Every row of the returned relation must have length equal to `return_arity`. // Every row of the returned relation must have length equal to `return_arity`.
pub fn new<R>(return_arity: usize, rule: R) -> Self pub fn new<R>(return_arity: usize, rule: R) -> Self
where where
R: Fn(Vec<NamedRows>, JsonValue) -> Result<NamedRows> + Send + Sync + 'static, R: Fn(Vec<NamedRows>, BTreeMap<String, DataValue>) -> Result<NamedRows> + Send + Sync + 'static,
{ {
Self { Self {
return_arity, return_arity,
@ -593,7 +592,7 @@ impl SimpleFixedRule {
return_arity: usize, return_arity: usize,
) -> ( ) -> (
Self, Self,
Receiver<(Vec<NamedRows>, JsonValue, Sender<Result<NamedRows>>)>, Receiver<(Vec<NamedRows>, BTreeMap<String, DataValue>, Sender<Result<NamedRows>>)>,
) { ) {
let (db2app_sender, db2app_receiver) = bounded(0); let (db2app_sender, db2app_receiver) = bounded(0);
( (
@ -628,13 +627,13 @@ impl FixedRule for SimpleFixedRule {
out: &'_ mut RegularTempStore, out: &'_ mut RegularTempStore,
_poison: Poison, _poison: Poison,
) -> Result<()> { ) -> Result<()> {
let options: JsonValue = payload let options: BTreeMap<_, _> = payload
.manifest .manifest
.options .options
.iter() .iter()
.map(|(k, v)| -> Result<_> { .map(|(k, v)| -> Result<_> {
let val = v.clone().eval_to_const()?; let val = v.clone().eval_to_const()?;
Ok((k.to_string(), JsonValue::from(val))) Ok((k.to_string(), val))
}) })
.try_collect()?; .try_collect()?;
let input_arity = payload.manifest.rule_args.len(); let input_arity = payload.manifest.rule_args.len();

@ -69,6 +69,17 @@ fn convert_params(ob: &PyDict) -> PyResult<BTreeMap<String, DataValue>> {
Ok(ret) Ok(ret)
} }
fn options_to_py(opts: BTreeMap<String, DataValue>, py: Python<'_>) -> PyResult<PyObject> {
let ret = PyDict::new(py);
for (k, v) in opts {
let val = value_to_py(v, py);
ret.set_item(k, val)?;
}
Ok(ret.into())
}
fn value_to_py(val: DataValue, py: Python<'_>) -> PyObject { fn value_to_py(val: DataValue, py: Python<'_>) -> PyObject {
match val { match val {
DataValue::Null => py.None(), DataValue::Null => py.None(),
@ -175,23 +186,17 @@ impl CozoDbPy {
) -> PyResult<()> { ) -> PyResult<()> {
if let Some(db) = &self.db { if let Some(db) = &self.db {
let cb: Py<PyAny> = callback.into(); let cb: Py<PyAny> = callback.into();
let (rule_impl, receiver, sender) = SimpleFixedRule::rule_with_channel(arity); let (rule_impl, receiver) = SimpleFixedRule::rule_with_channel(arity);
match db.register_fixed_rule(name, rule_impl) { match db.register_fixed_rule(name, rule_impl) {
Ok(_) => { Ok(_) => {
thread::spawn(move || { thread::spawn(move || {
for (inputs, options) in receiver { for (inputs, options, sender) in receiver {
let res = Python::with_gil(|py| -> Result<NamedRows> { let res = Python::with_gil(|py| -> Result<NamedRows> {
let json_convert =
PyModule::import(py, "json").into_diagnostic()?;
let json_convert: Py<PyAny> =
json_convert.getattr("loads").into_diagnostic()?.into();
let py_inputs = PyList::new( let py_inputs = PyList::new(
py, py,
inputs.into_iter().map(|nr| rows_to_py_rows(nr.rows, py)), inputs.into_iter().map(|nr| rows_to_py_rows(nr.rows, py)),
); );
let opts_str = PyString::new(py, &options.to_string()); let py_opts = options_to_py(options, py).into_diagnostic()?;
let args = PyTuple::new(py, vec![opts_str]);
let py_opts = json_convert.call1(py, args).into_diagnostic()?;
let args = let args =
PyTuple::new(py, vec![PyObject::from(py_inputs), py_opts]); PyTuple::new(py, vec![PyObject::from(py_inputs), py_opts]);
let res = cb.as_ref(py).call1(args).into_diagnostic()?; let res = cb.as_ref(py).call1(args).into_diagnostic()?;

Loading…
Cancel
Save