This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 80a6b65dbb Support consuming Substrait with compound signature
function names (#10653)
80a6b65dbb is described below
commit 80a6b65dbb0d293179a0c32f3d0b6b79cd88d397
Author: Arttu <[email protected]>
AuthorDate: Wed May 29 12:38:49 2024 +0200
Support consuming Substrait with compound signature function names (#10653)
* Support consuming Substrait with compound signature function names
Substrait 0.32.0+ requires functions to be specified using compound names,
which include the function name as well as the arguments it takes.
We don't necessarily need that information while consuming the plans, but
we need to support those compound names.
* Add a test for using "not:bool"
* clippy fixes
* Add license to new file
* Apply suggestions from code review
Co-authored-by: Andrew Lamb <[email protected]>
* remove prost-types dep as it's replaced by pbjson-types
---------
Co-authored-by: Andrew Lamb <[email protected]>
---
datafusion/substrait/Cargo.toml | 5 +-
datafusion/substrait/src/logical_plan/consumer.rs | 15 +++-
datafusion/substrait/src/logical_plan/producer.rs | 10 +--
datafusion/substrait/tests/cases/logical_plans.rs | 63 ++++++++++++++
datafusion/substrait/tests/cases/mod.rs | 1 +
.../tests/testdata/select_not_bool.substrait.json | 98 ++++++++++++++++++++++
6 files changed, 182 insertions(+), 10 deletions(-)
diff --git a/datafusion/substrait/Cargo.toml b/datafusion/substrait/Cargo.toml
index 66600838dd..f614fd6b3f 100644
--- a/datafusion/substrait/Cargo.toml
+++ b/datafusion/substrait/Cargo.toml
@@ -37,11 +37,12 @@ chrono = { workspace = true }
datafusion = { workspace = true, default-features = true }
itertools = { workspace = true }
object_store = { workspace = true }
+pbjson-types = "0.6"
prost = "0.12"
-prost-types = "0.12"
-substrait = "0.34.0"
+substrait = { version = "0.34.0", features = ["serde"] }
[dev-dependencies]
+serde_json = "1.0"
tokio = { workspace = true }
[features]
diff --git a/datafusion/substrait/src/logical_plan/consumer.rs
b/datafusion/substrait/src/logical_plan/consumer.rs
index abebc68123..eb819e2c87 100644
--- a/datafusion/substrait/src/logical_plan/consumer.rs
+++ b/datafusion/substrait/src/logical_plan/consumer.rs
@@ -124,6 +124,15 @@ fn scalar_function_type_from_str(
name: &str,
) -> Result<ScalarFunctionType> {
let s = ctx.state();
+ let name = match name.rsplit_once(':') {
+ // Since 0.32.0, Substrait requires the function names to be in a
compound format
+ // https://substrait.io/extensions/#function-signature-compound-names
+ // for example, `add:i8_i8`.
+ // On the consumer side, we don't really care about the signature
though, just the name.
+ Some((name, _)) => name,
+ None => name,
+ };
+
if let Some(func) = s.scalar_functions().get(name) {
return Ok(ScalarFunctionType::Udf(func.to_owned()));
}
@@ -1525,7 +1534,7 @@ fn from_substrait_literal(
return substrait_err!("Interval year month value is
empty");
};
let value_slice: [u8; 4] =
- raw_val.value.clone().try_into().map_err(|_| {
+ (*raw_val.value).try_into().map_err(|_| {
substrait_datafusion_err!(
"Failed to parse interval year month value"
)
@@ -1537,7 +1546,7 @@ fn from_substrait_literal(
return substrait_err!("Interval day time value is
empty");
};
let value_slice: [u8; 8] =
- raw_val.value.clone().try_into().map_err(|_| {
+ (*raw_val.value).try_into().map_err(|_| {
substrait_datafusion_err!(
"Failed to parse interval day time value"
)
@@ -1549,7 +1558,7 @@ fn from_substrait_literal(
return substrait_err!("Interval month day nano value
is empty");
};
let value_slice: [u8; 16] =
- raw_val.value.clone().try_into().map_err(|_| {
+ (*raw_val.value).try_into().map_err(|_| {
substrait_datafusion_err!(
"Failed to parse interval month day nano value"
)
diff --git a/datafusion/substrait/src/logical_plan/producer.rs
b/datafusion/substrait/src/logical_plan/producer.rs
index 4dd8226366..010386bf97 100644
--- a/datafusion/substrait/src/logical_plan/producer.rs
+++ b/datafusion/substrait/src/logical_plan/producer.rs
@@ -45,7 +45,7 @@ use datafusion::logical_expr::expr::{
};
use datafusion::logical_expr::{expr, Between, JoinConstraint, LogicalPlan,
Operator};
use datafusion::prelude::Expr;
-use prost_types::Any as ProtoAny;
+use pbjson_types::Any as ProtoAny;
use substrait::proto::exchange_rel::{ExchangeKind, RoundRobin, ScatterFields};
use substrait::proto::expression::literal::user_defined::Val;
use substrait::proto::expression::literal::UserDefined;
@@ -547,7 +547,7 @@ pub fn to_substrait_rel(
.serialize_logical_plan(extension_plan.node.as_ref())?;
let detail = ProtoAny {
type_url: extension_plan.node.name().to_string(),
- value: extension_bytes,
+ value: extension_bytes.into(),
};
let mut inputs_rel = extension_plan
.node
@@ -1919,7 +1919,7 @@ fn to_substrait_literal(value: &ScalarValue) ->
Result<Literal> {
}],
val: Some(Val::Value(ProtoAny {
type_url: INTERVAL_YEAR_MONTH_TYPE_URL.to_string(),
- value: bytes.to_vec(),
+ value: bytes.to_vec().into(),
})),
}),
INTERVAL_YEAR_MONTH_TYPE_REF,
@@ -1942,7 +1942,7 @@ fn to_substrait_literal(value: &ScalarValue) ->
Result<Literal> {
type_parameters: vec![i64_param.clone(), i64_param],
val: Some(Val::Value(ProtoAny {
type_url: INTERVAL_MONTH_DAY_NANO_TYPE_URL.to_string(),
- value: bytes.to_vec(),
+ value: bytes.to_vec().into(),
})),
}),
INTERVAL_MONTH_DAY_NANO_TYPE_REF,
@@ -1965,7 +1965,7 @@ fn to_substrait_literal(value: &ScalarValue) ->
Result<Literal> {
}],
val: Some(Val::Value(ProtoAny {
type_url: INTERVAL_DAY_TIME_TYPE_URL.to_string(),
- value: bytes.to_vec(),
+ value: bytes.to_vec().into(),
})),
}),
INTERVAL_DAY_TIME_TYPE_REF,
diff --git a/datafusion/substrait/tests/cases/logical_plans.rs
b/datafusion/substrait/tests/cases/logical_plans.rs
new file mode 100644
index 0000000000..4d485b7f12
--- /dev/null
+++ b/datafusion/substrait/tests/cases/logical_plans.rs
@@ -0,0 +1,63 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Tests for reading substrait plans produced by other systems
+
+#[cfg(test)]
+mod tests {
+ use datafusion::common::Result;
+ use datafusion::prelude::{CsvReadOptions, SessionContext};
+ use datafusion_substrait::logical_plan::consumer::from_substrait_plan;
+ use std::fs::File;
+ use std::io::BufReader;
+ use substrait::proto::Plan;
+
+ #[tokio::test]
+ async fn function_compound_signature() -> Result<()> {
+ // DataFusion currently produces Substrait that refers to functions
only by their name.
+ // However, the Substrait spec requires that functions be identified
by their compound signature.
+ // This test confirms that DataFusion is able to consume plans
following the spec, even though
+ // we don't yet produce such plans.
+ // Once we start producing plans with compound signatures, this test
can be replaced by the roundtrip tests.
+
+ let ctx = create_context().await?;
+
+ // File generated with substrait-java's Isthmus:
+ // ./isthmus-cli/build/graal/isthmus "select not d from data" -c
"create table data (d boolean)"
+ let path = "tests/testdata/select_not_bool.substrait.json";
+ let proto = serde_json::from_reader::<_, Plan>(BufReader::new(
+ File::open(path).expect("file not found"),
+ ))
+ .expect("failed to parse json");
+
+ let plan = from_substrait_plan(&ctx, &proto).await?;
+
+ assert_eq!(
+ format!("{:?}", plan),
+ "Projection: NOT DATA.a\
+ \n TableScan: DATA projection=[a, b, c, d, e, f]"
+ );
+ Ok(())
+ }
+
+ async fn create_context() -> datafusion::common::Result<SessionContext> {
+ let ctx = SessionContext::new();
+ ctx.register_csv("DATA", "tests/testdata/data.csv",
CsvReadOptions::new())
+ .await?;
+ Ok(ctx)
+ }
+}
diff --git a/datafusion/substrait/tests/cases/mod.rs
b/datafusion/substrait/tests/cases/mod.rs
index b17289205f..d049eb2c21 100644
--- a/datafusion/substrait/tests/cases/mod.rs
+++ b/datafusion/substrait/tests/cases/mod.rs
@@ -15,6 +15,7 @@
// specific language governing permissions and limitations
// under the License.
+mod logical_plans;
mod roundtrip_logical_plan;
mod roundtrip_physical_plan;
mod serialize;
diff --git a/datafusion/substrait/tests/testdata/select_not_bool.substrait.json
b/datafusion/substrait/tests/testdata/select_not_bool.substrait.json
new file mode 100644
index 0000000000..e52cf87d50
--- /dev/null
+++ b/datafusion/substrait/tests/testdata/select_not_bool.substrait.json
@@ -0,0 +1,98 @@
+{
+ "extensionUris": [
+ {
+ "extensionUriAnchor": 1,
+ "uri": "/functions_boolean.yaml"
+ }
+ ],
+ "extensions": [
+ {
+ "extensionFunction": {
+ "extensionUriReference": 1,
+ "functionAnchor": 0,
+ "name": "not:bool"
+ }
+ }
+ ],
+ "relations": [
+ {
+ "root": {
+ "input": {
+ "project": {
+ "common": {
+ "emit": {
+ "outputMapping": [
+ 1
+ ]
+ }
+ },
+ "input": {
+ "read": {
+ "common": {
+ "direct": {
+ }
+ },
+ "baseSchema": {
+ "names": [
+ "D"
+ ],
+ "struct": {
+ "types": [
+ {
+ "bool": {
+ "typeVariationReference": 0,
+ "nullability": "NULLABILITY_NULLABLE"
+ }
+ }
+ ],
+ "typeVariationReference": 0,
+ "nullability": "NULLABILITY_REQUIRED"
+ }
+ },
+ "namedTable": {
+ "names": [
+ "DATA"
+ ]
+ }
+ }
+ },
+ "expressions": [
+ {
+ "scalarFunction": {
+ "functionReference": 0,
+ "args": [],
+ "outputType": {
+ "bool": {
+ "typeVariationReference": 0,
+ "nullability": "NULLABILITY_NULLABLE"
+ }
+ },
+ "arguments": [
+ {
+ "value": {
+ "selection": {
+ "directReference": {
+ "structField": {
+ "field": 0
+ }
+ },
+ "rootReference": {
+ }
+ }
+ }
+ }
+ ],
+ "options": []
+ }
+ }
+ ]
+ }
+ },
+ "names": [
+ "EXPR$0"
+ ]
+ }
+ }
+ ],
+ "expectedTypeUrls": []
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
