use std::sync::Arc;
use anyhow::{Context, Result, bail};
use serde_json::Value as Json;
use crate::dsl::{
Value,
arrow_value::ipc_list_literal_to_value,
expr::{ColumnRef, Expr},
ops::{
list::{ListOp, SetOp},
scalar::{CastType, ScalarOp, parse_binary_op},
string::StringOp,
struct_::StructOp,
},
};
pub(crate) fn parse_polars_json(json: &str) -> Result<Expr<ColumnRef>> {
parse_expr(&serde_json::from_str(json)?)
}
fn parse_expr(value: &Json) -> Result<Expr<ColumnRef>> {
if value.as_str() == Some("Element") {
return Ok(Expr::Element);
}
let object = value.as_object().context("expected expression object")?;
if let Some(column) = object.get("Column") {
return parse_column(column.as_str().context("Column must be a string")?);
}
if let Some(literal) = object.get("Literal") {
return Ok(Expr::Literal(parse_literal(literal)?));
}
if let Some(binary) = object.get("BinaryExpr").and_then(Json::as_object) {
return Ok(Expr::Scalar(
parse_binary_op(
binary
.get("op")
.and_then(Json::as_str)
.context("binary op must be a string")?,
)?,
vec![
parse_expr(binary.get("left").context("BinaryExpr missing left")?)?,
parse_expr(binary.get("right").context("BinaryExpr missing right")?)?,
],
));
}
if let Some(cast) = object.get("Cast").and_then(Json::as_object) {
return Ok(Expr::Scalar(
ScalarOp::Cast(parse_cast_type(
cast.get("dtype").context("Cast missing dtype")?,
)?),
vec![parse_expr(cast.get("expr").context("Cast missing expr")?)?],
));
}
if let Some(alias) = object.get("Alias").and_then(Json::as_array) {
ensure_arity(alias, 2)?;
let name = alias[1].as_str().context("Alias name must be a string")?;
return Ok(Expr::Alias(
Box::new(parse_expr(&alias[0])?),
name.to_owned(),
));
}
if let Some(ternary) = object.get("Ternary") {
return parse_ternary(ternary);
}
if let Some(rx) = object.get("Rx") {
return parse_rx_expr(rx);
}
if let Some(function) = object.get("Function") {
return parse_function(function);
}
if let Some(eval) = object.get("Eval") {
return parse_eval(eval);
}
if let Some(eval) = object.get("StructEval") {
return parse_struct_eval(eval);
}
if let Some(filter) = object.get("Filter") {
return parse_filter(filter);
}
if let Some(explode) = object.get("Explode") {
return parse_explode(explode);
}
bail!("unsupported Polars expression JSON: {value}")
}
fn parse_column(name: &str) -> Result<Expr<ColumnRef>> {
if name.is_empty() {
return Ok(Expr::Element);
}
Ok(Expr::Column(match name {
"src.id" => ColumnRef::SrcId,
"dest.id" => ColumnRef::DestId,
"edge.id" => ColumnRef::EdgeId,
_ => {
let (scope, field) = name
.split_once('.')
.ok_or_else(|| anyhow::anyhow!("column {name:?} must use scope.field"))?;
match scope {
"src" => ColumnRef::SrcField(field.to_owned()),
"dest" => ColumnRef::DestField(field.to_owned()),
"edge" => ColumnRef::EdgeField(field.to_owned()),
"state" => ColumnRef::State(field.to_owned()),
_ => bail!("unknown column scope {scope:?} in {name:?}"),
}
}
}))
}
fn parse_rx_expr(value: &Json) -> Result<Expr<ColumnRef>> {
let object = value
.as_object()
.context("Rx expression must be an object")?;
let value = object
.get("MaskIfAny")
.context("unsupported Rx expression")?
.as_object()
.context("MaskIfAny expression must be an object")?;
Ok(Expr::Scalar(
ScalarOp::MaskIfAny,
vec![
parse_expr(value.get("value").context("MaskIfAny missing value")?)?,
parse_expr(value.get("mask").context("MaskIfAny missing mask")?)?,
parse_expr(
value
.get("then")
.or_else(|| value.get("then_value"))
.context("MaskIfAny missing then")?,
)?,
],
))
}
fn parse_literal(value: &Json) -> Result<Value> {
parse_scalar_literal(
value
.pointer("/Value")
.or_else(|| value.pointer("/Scalar"))
.or_else(|| value.pointer("/Dyn"))
.context("unsupported literal")?,
)
}
fn parse_scalar_literal(value: &Json) -> Result<Value> {
let object = value.as_object().context("literal must be object")?;
if let Some(value) = object.get("String") {
return Ok(Value::Str(Arc::from(
value.as_str().context("String literal must be string")?,
)));
}
if let Some(value) = object.get("Boolean") {
return Ok(Value::Bool(
value.as_bool().context("Boolean literal must be bool")?,
));
}
if object.contains_key("Null") {
return Ok(Value::Null);
}
if let Some(value) = object.get("Int") {
return Ok(Value::I64(
value.as_i64().context("Int literal must be i64")?,
));
}
if let Some(value) = object.get("UInt") {
return Ok(Value::U64(
value.as_u64().context("UInt literal must be u64")?,
));
}
if let Some(value) = object.get("Float") {
return Ok(Value::F64(
value.as_f64().context("Float literal must be f64")?,
));
}
if let Some(value) = object.get("List") {
return ipc_list_literal_to_value(&parse_byte_array(value)?);
}
bail!("unsupported literal {value}")
}
fn parse_byte_array(value: &Json) -> Result<Vec<u8>> {
value
.as_array()
.context("List literal must be an Arrow IPC byte array")?
.iter()
.map(|value| {
let byte = value
.as_u64()
.context("Arrow IPC byte array entries must be unsigned integers")?;
u8::try_from(byte).context("Arrow IPC byte array entries must fit in u8")
})
.collect()
}
fn parse_function(value: &Json) -> Result<Expr<ColumnRef>> {
let object = value.as_object().context("Function must be an object")?;
let inputs = object
.get("input")
.and_then(Json::as_array)
.context("Function missing input array")?;
let function = object
.get("function")
.context("Function missing function")?;
if function.pointer("/Boolean").and_then(Json::as_str) == Some("Not") {
return unary(inputs, |input| Expr::Scalar(ScalarOp::Not, vec![input]));
}
if function.pointer("/Boolean").and_then(Json::as_str) == Some("IsNull") {
return unary(inputs, |input| Expr::Scalar(ScalarOp::IsNull, vec![input]));
}
if function.pointer("/Boolean").and_then(Json::as_str) == Some("IsNotNull") {
return unary(inputs, |input| {
Expr::Scalar(ScalarOp::IsNotNull, vec![input])
});
}
if function.pointer("/StringExpr/Contains").is_some() {
return op(inputs, |args| Expr::String(StringOp::Contains, args));
}
if function.pointer("/StringExpr/StartsWith").is_some() {
return op(inputs, |args| Expr::String(StringOp::StartsWith, args));
}
if function.pointer("/StringExpr/EndsWith").is_some() {
return op(inputs, |args| Expr::String(StringOp::EndsWith, args));
}
if let Some(list) = function.get("ListExpr") {
return parse_list_function(list, inputs);
}
if let Some(struct_) = function.get("StructExpr") {
return parse_struct_function(struct_, inputs);
}
bail!("unsupported Polars function {function}")
}
fn parse_list_function(value: &Json, inputs: &[Json]) -> Result<Expr<ColumnRef>> {
let list_op = match value.as_str() {
Some("Concat") => ListOp::Concat,
Some("Length") => ListOp::Len,
Some("Sum") => ListOp::Sum,
Some("Min") => ListOp::Min,
Some("Max") => ListOp::Max,
Some("ArgMin") => ListOp::ArgMin,
Some("ArgMax") => ListOp::ArgMax,
Some("Mean") => ListOp::Mean,
Some("Median") => ListOp::Median,
Some("Reverse") => ListOp::Reverse,
Some("DropNulls") => ListOp::DropNulls,
Some("CountMatches") => ListOp::CountMatches,
Some("NUnique") => ListOp::NUnique,
Some("Shift") => ListOp::Shift,
Some("GatherEvery") => ListOp::GatherEvery,
Some("Slice") => ListOp::Slice,
Some(other) => bail!("unsupported Polars list function {other:?}"),
None => parse_list_object(value)?,
};
op(inputs, |args| Expr::List(list_op, args))
}
fn parse_list_object(value: &Json) -> Result<ListOp> {
let object = value
.as_object()
.context("ListExpr must be string or object")?;
if let Some(value) = object.get("Contains") {
return Ok(ListOp::Contains {
nulls_equal: value
.get("nulls_equal")
.and_then(Json::as_bool)
.unwrap_or(true),
});
}
if let Some(value) = object.get("Get") {
return Ok(ListOp::Get {
null_on_oob: value.as_bool().unwrap_or(false),
});
}
if object.contains_key("Unique") {
return Ok(ListOp::Unique);
}
if let Some(sort) = object.get("Sort").and_then(Json::as_object) {
return Ok(ListOp::Sort {
descending: sort
.get("descending")
.and_then(Json::as_bool)
.unwrap_or(false),
nulls_last: sort
.get("nulls_last")
.and_then(Json::as_bool)
.unwrap_or(false),
});
}
if let Some(value) = object.get("SetOperation").and_then(Json::as_str) {
return Ok(ListOp::Set(match value {
"Union" => SetOp::Union,
"Intersection" => SetOp::Intersection,
"Difference" => SetOp::Difference,
"SymmetricDifference" => SetOp::SymmetricDifference,
other => bail!("unsupported list set operation {other:?}"),
}));
}
if let Some(value) = object.get("Join") {
return Ok(ListOp::Join {
ignore_nulls: value.as_bool().unwrap_or(true),
});
}
if let Some(fields) = object.get("ToStruct").and_then(Json::as_array) {
return Ok(ListOp::ToStruct(
fields
.iter()
.map(|field| {
field
.as_str()
.map(str::to_owned)
.context("ToStruct field names must be strings")
})
.collect::<Result<_>>()?,
));
}
bail!("unsupported Polars list function {value}")
}
fn parse_struct_function(value: &Json, inputs: &[Json]) -> Result<Expr<ColumnRef>> {
let struct_op = if value.as_str() == Some("JsonEncode") {
StructOp::JsonEncode
} else if let Some(name) = value.pointer("/FieldByName").and_then(Json::as_str) {
StructOp::FieldByName(name.to_owned())
} else if let Some(names) = value.pointer("/RenameFields").and_then(Json::as_array) {
StructOp::RenameFields(
names
.iter()
.map(|name| {
name.as_str()
.map(str::to_owned)
.context("RenameFields names must be strings")
})
.collect::<Result<_>>()?,
)
} else {
bail!("unsupported Polars struct function {value}")
};
op(inputs, |args| Expr::Struct(struct_op, args))
}
fn parse_eval(value: &Json) -> Result<Expr<ColumnRef>> {
let object = value.as_object().context("Eval must be an object")?;
let list = parse_expr(object.get("expr").context("Eval missing expr")?)?;
let evaluation = object
.get("evaluation")
.context("Eval missing evaluation")?;
match object
.get("variant")
.and_then(Json::as_str)
.context("Eval missing variant")?
{
"List" => {
if let Some(filter) = evaluation.get("Filter") {
return parse_filter_with_list(list, filter);
}
Ok(Expr::List(
ListOp::Eval,
vec![list, parse_expr(evaluation)?],
))
}
"ListAgg" => {
let op = list_agg_op(evaluation)?;
Ok(Expr::List(op, vec![list]))
}
variant => bail!("unsupported Eval variant {variant:?}"),
}
}
fn parse_ternary(value: &Json) -> Result<Expr<ColumnRef>> {
let object = value.as_object().context("Ternary must be an object")?;
Ok(Expr::Ternary {
predicate: Box::new(parse_expr(
object
.get("predicate")
.context("Ternary missing predicate")?,
)?),
truthy: Box::new(parse_expr(
object.get("truthy").context("Ternary missing truthy")?,
)?),
falsy: Box::new(parse_expr(
object.get("falsy").context("Ternary missing falsy")?,
)?),
})
}
fn parse_explode(value: &Json) -> Result<Expr<ColumnRef>> {
let object = value.as_object().context("Explode must be an object")?;
let options = object.get("options").and_then(Json::as_object);
Ok(Expr::List(
ListOp::Explode {
empty_as_null: options
.and_then(|options| options.get("empty_as_null"))
.and_then(Json::as_bool)
.unwrap_or(true),
keep_nulls: options
.and_then(|options| options.get("keep_nulls"))
.and_then(Json::as_bool)
.unwrap_or(true),
},
vec![parse_expr(
object.get("input").context("Explode missing input")?,
)?],
))
}
fn parse_filter(value: &Json) -> Result<Expr<ColumnRef>> {
parse_filter_with_list(Expr::Element, value)
}
fn parse_filter_with_list(list: Expr<ColumnRef>, value: &Json) -> Result<Expr<ColumnRef>> {
let object = value.as_object().context("Filter must be an object")?;
Ok(Expr::List(
ListOp::Filter,
vec![
list,
parse_expr(object.get("by").context("Filter missing by")?)?,
],
))
}
fn list_agg_op(value: &Json) -> Result<ListOp> {
let function = value
.get("Function")
.and_then(Json::as_object)
.and_then(|object| object.get("function"))
.context("ListAgg must contain a Function")?;
if let Some(all) = function.pointer("/Boolean/All").and_then(Json::as_object) {
return Ok(ListOp::All {
ignore_nulls: all
.get("ignore_nulls")
.and_then(Json::as_bool)
.unwrap_or(true),
});
}
if let Some(any) = function.pointer("/Boolean/Any").and_then(Json::as_object) {
return Ok(ListOp::Any {
ignore_nulls: any
.get("ignore_nulls")
.and_then(Json::as_bool)
.unwrap_or(true),
});
}
bail!("unsupported ListAgg expression {value}")
}
fn parse_struct_eval(value: &Json) -> Result<Expr<ColumnRef>> {
let object = value.as_object().context("StructEval must be an object")?;
let mut args = vec![parse_expr(
object.get("expr").context("StructEval missing expr")?,
)?];
let mut names = Vec::new();
for expr in object
.get("evaluation")
.and_then(Json::as_array)
.context("StructEval missing evaluation")?
{
let Expr::Alias(expr, name) = parse_expr(expr)? else {
bail!("StructEval fields must be aliases")
};
names.push(name);
args.push(*expr);
}
Ok(Expr::Struct(StructOp::WithFields(names), args))
}
fn parse_cast_type(value: &Json) -> Result<CastType> {
let dtype = value
.get("Literal")
.and_then(Json::as_str)
.context("only literal cast dtypes are supported")?;
Ok(match dtype {
"Boolean" => CastType::Bool,
"Int64" | "Int32" | "Int16" | "Int8" => CastType::Int64,
"UInt64" | "UInt32" | "UInt16" | "UInt8" => CastType::UInt64,
"Float64" | "Float32" => CastType::Float64,
"String" | "Utf8" => CastType::String,
other => bail!("unsupported cast dtype {other:?}"),
})
}
fn unary(
inputs: &[Json],
wrap: impl FnOnce(Expr<ColumnRef>) -> Expr<ColumnRef>,
) -> Result<Expr<ColumnRef>> {
ensure_arity(inputs, 1)?;
Ok(wrap(parse_expr(&inputs[0])?))
}
fn op(
inputs: &[Json],
wrap: impl FnOnce(Vec<Expr<ColumnRef>>) -> Expr<ColumnRef>,
) -> Result<Expr<ColumnRef>> {
inputs
.iter()
.map(parse_expr)
.collect::<Result<Vec<_>>>()
.map(wrap)
}
fn ensure_arity(inputs: &[Json], len: usize) -> Result<()> {
if inputs.len() == len {
Ok(())
} else {
bail!("expected {len} function inputs, got {}", inputs.len())
}
}