Does DataFusion Support JSON Path Filtering Like jsonb_path_exists in PostgreSQL?

[dadepo]

PostgreSQL has a handy function called jsonb_path_exists that makes it easy to filter JSON data. For example:

SELECT * FROM table WHERE jsonb_path_exists(properties, '$.person.age ? (@ == 18)')

I’m wondering if there’s a way to do something similar in DataFusion using its Expression API. From what I’ve seen in the docs, it doesn’t look like this is currently supported, but I just wanted to check:

• Is there a way to achieve this with existing DataFusion expressions?
• If not, directly via the Expression function, is there a way to achieve it, maybe by using more low level primitives?

Would love to hear any thoughts on this!

[alamb]

You could create a function called jsonb_path_exists that takes a binary column and a json path string perhaps?

You could take a look at how it is done https://github.com/datafusion-contrib/datafusion-functions-json

There is a lower level API that lets you customize how expressions are planned, here: https://docs.rs/datafusion/latest/datafusion/logical_expr/planner/trait.ExprPlanner.html

You can register them via
https://docs.rs/datafusion/latest/datafusion/execution/trait.FunctionRegistry.html#method.register_expr_planner

Though I don't think we have a good example of using them yet,..

[alamb]

Actually, here is a good example :

datafusion/datafusion/core/tests/user_defined/expr_planner.rs

Lines 35 to 124 in 87eec43

	impl ExprPlanner for MyCustomPlanner {
	fn plan_binary_op(
	&self,
	expr: RawBinaryExpr,
	_schema: &DFSchema,
	) -> Result<PlannerResult<RawBinaryExpr>> {
	match &expr.op {
	BinaryOperator::Arrow => {
	Ok(PlannerResult::Planned(Expr::BinaryExpr(BinaryExpr {
	left: Box::new(expr.left.clone()),
	right: Box::new(expr.right.clone()),
	op: Operator::StringConcat,
	})))
	}
	BinaryOperator::LongArrow => {
	Ok(PlannerResult::Planned(Expr::BinaryExpr(BinaryExpr {
	left: Box::new(expr.left.clone()),
	right: Box::new(expr.right.clone()),
	op: Operator::Plus,
	})))
	}
	BinaryOperator::Question => {
	Ok(PlannerResult::Planned(Expr::Alias(Alias::new(
	Expr::Literal(ScalarValue::Boolean(Some(true))),
	None::<&str>,
	format!("{} ? {}", expr.left, expr.right),
	))))
	}
	_ => Ok(PlannerResult::Original(expr)),
	}
	}
	}
	async fn plan_and_collect(sql: &str) -> Result<Vec<RecordBatch>> {
	let config =
	SessionConfig::new().set_str("datafusion.sql_parser.dialect", "postgres");
	let mut ctx = SessionContext::new_with_config(config);
	ctx.register_expr_planner(Arc::new(MyCustomPlanner))?;
	ctx.sql(sql).await?.collect().await
	}
	#[tokio::test]
	async fn test_custom_operators_arrow() {
	let actual = plan_and_collect("select 'foo'->'bar';").await.unwrap();
	let expected = [
	"+----------------------------+",
	"| Utf8(\"foo\") || Utf8(\"bar\") |",
	"+----------------------------+",
	"| foobar |",
	"+----------------------------+",
	];
	assert_batches_eq!(&expected, &actual);
	}
	#[tokio::test]
	async fn test_custom_operators_long_arrow() {
	let actual = plan_and_collect("select 1->>2;").await.unwrap();
	let expected = [
	"+---------------------+",
	"| Int64(1) + Int64(2) |",
	"+---------------------+",
	"| 3 |",
	"+---------------------+",
	];
	assert_batches_eq!(&expected, &actual);
	}
	#[tokio::test]
	async fn test_question_select() {
	let actual = plan_and_collect("select a ? 2 from (select 1 as a);")
	.await
	.unwrap();
	let expected = [
	"+--------------+",
	"| a ? Int64(2) |",
	"+--------------+",
	"| true |",
	"+--------------+",
	];
	assert_batches_eq!(&expected, &actual);
	}
	#[tokio::test]
	async fn test_question_filter() {
	let actual = plan_and_collect("select a from (select 1 as a) where a ? 2;")
	.await
	.unwrap();
	let expected = ["+---+", "| a |", "+---+", "| 1 |", "+---+"];
	assert_batches_eq!(&expected, &actual);
	}

[alamb]

I also filed #15267 to track adding some better documentation here

[dadepo]

You could create a function called jsonb_path_exists that takes a binary column and a json path string perhaps?

I think what I am missing is how this can be used via the expression API. Using sql, I know I can do something like:

SessionContext.sql("select jsonb_path_exists(d_col, path)")

But via the expression API, how will that look? ie

df.filter(col('d_col')...)

How will that look?

[alamb]

I think you can use ScalarUDF::call

let json_path_exist = ScalarUDF::from(Arc::new(impl));
json_path_exist.call(d_col, path)