This is an automated email from the ASF dual-hosted git repository.
jonah 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 cd69e37332 support recursive CTEs logical plans in datafusion-proto
(#13314)
cd69e37332 is described below
commit cd69e373320a1682dfa7c32d14d5294add8b5fee
Author: Leonardo Yvens <[email protected]>
AuthorDate: Tue Nov 12 04:42:56 2024 +0000
support recursive CTEs logical plans in datafusion-proto (#13314)
* support LogicaPlan::RecursiveQuery in datafusion-proto
* fixed and failing test roundtrip_recursive_query
* fix rebase artifact
* add node for CteWorkTableScan in datafusion-proto
* Use Arc::clone
---------
Co-authored-by: jonahgao <[email protected]>
---
datafusion/core/src/datasource/cte_worktable.rs | 12 +-
datafusion/proto/proto/datafusion.proto | 14 +
datafusion/proto/src/generated/pbjson.rs | 281 +++++++++++++++++++++
datafusion/proto/src/generated/prost.rs | 26 +-
datafusion/proto/src/logical_plan/mod.rs | 82 +++++-
.../proto/tests/cases/roundtrip_logical_plan.rs | 28 ++
6 files changed, 434 insertions(+), 9 deletions(-)
diff --git a/datafusion/core/src/datasource/cte_worktable.rs
b/datafusion/core/src/datasource/cte_worktable.rs
index 23f57b12ae..9504721cec 100644
--- a/datafusion/core/src/datasource/cte_worktable.rs
+++ b/datafusion/core/src/datasource/cte_worktable.rs
@@ -39,8 +39,6 @@ use crate::datasource::{TableProvider, TableType};
#[derive(Debug)]
pub struct CteWorkTable {
/// The name of the CTE work table
- // WIP, see https://github.com/apache/datafusion/issues/462
- #[allow(dead_code)]
name: String,
/// This schema must be shared across both the static and recursive terms
of a recursive query
table_schema: SchemaRef,
@@ -56,6 +54,16 @@ impl CteWorkTable {
table_schema,
}
}
+
+ /// The user-provided name of the CTE
+ pub fn name(&self) -> &str {
+ &self.name
+ }
+
+ /// The schema of the recursive term of the query
+ pub fn schema(&self) -> SchemaRef {
+ Arc::clone(&self.table_schema)
+ }
}
#[async_trait]
diff --git a/datafusion/proto/proto/datafusion.proto
b/datafusion/proto/proto/datafusion.proto
index d6fa129edc..b9a1cff94d 100644
--- a/datafusion/proto/proto/datafusion.proto
+++ b/datafusion/proto/proto/datafusion.proto
@@ -59,6 +59,8 @@ message LogicalPlanNode {
DistinctOnNode distinct_on = 28;
CopyToNode copy_to = 29;
UnnestNode unnest = 30;
+ RecursiveQueryNode recursive_query = 31;
+ CteWorkTableScanNode cte_work_table_scan = 32;
}
}
@@ -1249,3 +1251,15 @@ message PartitionStats {
int64 num_bytes = 3;
repeated datafusion_common.ColumnStats column_stats = 4;
}
+
+message RecursiveQueryNode {
+ string name = 1;
+ LogicalPlanNode static_term = 2;
+ LogicalPlanNode recursive_term = 3;
+ bool is_distinct = 4;
+}
+
+message CteWorkTableScanNode {
+ string name = 1;
+ datafusion_common.Schema schema = 2;
+}
diff --git a/datafusion/proto/src/generated/pbjson.rs
b/datafusion/proto/src/generated/pbjson.rs
index 16f14d9ddf..52ba1ea8aa 100644
--- a/datafusion/proto/src/generated/pbjson.rs
+++ b/datafusion/proto/src/generated/pbjson.rs
@@ -4153,6 +4153,114 @@ impl<'de> serde::Deserialize<'de> for CsvSinkExecNode {
deserializer.deserialize_struct("datafusion.CsvSinkExecNode", FIELDS,
GeneratedVisitor)
}
}
+impl serde::Serialize for CteWorkTableScanNode {
+ #[allow(deprecated)]
+ fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok,
S::Error>
+ where
+ S: serde::Serializer,
+ {
+ use serde::ser::SerializeStruct;
+ let mut len = 0;
+ if !self.name.is_empty() {
+ len += 1;
+ }
+ if self.schema.is_some() {
+ len += 1;
+ }
+ let mut struct_ser =
serializer.serialize_struct("datafusion.CteWorkTableScanNode", len)?;
+ if !self.name.is_empty() {
+ struct_ser.serialize_field("name", &self.name)?;
+ }
+ if let Some(v) = self.schema.as_ref() {
+ struct_ser.serialize_field("schema", v)?;
+ }
+ struct_ser.end()
+ }
+}
+impl<'de> serde::Deserialize<'de> for CteWorkTableScanNode {
+ #[allow(deprecated)]
+ fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
+ where
+ D: serde::Deserializer<'de>,
+ {
+ const FIELDS: &[&str] = &[
+ "name",
+ "schema",
+ ];
+
+ #[allow(clippy::enum_variant_names)]
+ enum GeneratedField {
+ Name,
+ Schema,
+ }
+ impl<'de> serde::Deserialize<'de> for GeneratedField {
+ fn deserialize<D>(deserializer: D) ->
std::result::Result<GeneratedField, D::Error>
+ where
+ D: serde::Deserializer<'de>,
+ {
+ struct GeneratedVisitor;
+
+ impl<'de> serde::de::Visitor<'de> for GeneratedVisitor {
+ type Value = GeneratedField;
+
+ fn expecting(&self, formatter: &mut
std::fmt::Formatter<'_>) -> std::fmt::Result {
+ write!(formatter, "expected one of: {:?}", &FIELDS)
+ }
+
+ #[allow(unused_variables)]
+ fn visit_str<E>(self, value: &str) ->
std::result::Result<GeneratedField, E>
+ where
+ E: serde::de::Error,
+ {
+ match value {
+ "name" => Ok(GeneratedField::Name),
+ "schema" => Ok(GeneratedField::Schema),
+ _ => Err(serde::de::Error::unknown_field(value,
FIELDS)),
+ }
+ }
+ }
+ deserializer.deserialize_identifier(GeneratedVisitor)
+ }
+ }
+ struct GeneratedVisitor;
+ impl<'de> serde::de::Visitor<'de> for GeneratedVisitor {
+ type Value = CteWorkTableScanNode;
+
+ fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) ->
std::fmt::Result {
+ formatter.write_str("struct datafusion.CteWorkTableScanNode")
+ }
+
+ fn visit_map<V>(self, mut map_: V) ->
std::result::Result<CteWorkTableScanNode, V::Error>
+ where
+ V: serde::de::MapAccess<'de>,
+ {
+ let mut name__ = None;
+ let mut schema__ = None;
+ while let Some(k) = map_.next_key()? {
+ match k {
+ GeneratedField::Name => {
+ if name__.is_some() {
+ return
Err(serde::de::Error::duplicate_field("name"));
+ }
+ name__ = Some(map_.next_value()?);
+ }
+ GeneratedField::Schema => {
+ if schema__.is_some() {
+ return
Err(serde::de::Error::duplicate_field("schema"));
+ }
+ schema__ = map_.next_value()?;
+ }
+ }
+ }
+ Ok(CteWorkTableScanNode {
+ name: name__.unwrap_or_default(),
+ schema: schema__,
+ })
+ }
+ }
+ deserializer.deserialize_struct("datafusion.CteWorkTableScanNode",
FIELDS, GeneratedVisitor)
+ }
+}
impl serde::Serialize for CubeNode {
#[allow(deprecated)]
fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok,
S::Error>
@@ -10602,6 +10710,12 @@ impl serde::Serialize for LogicalPlanNode {
logical_plan_node::LogicalPlanType::Unnest(v) => {
struct_ser.serialize_field("unnest", v)?;
}
+ logical_plan_node::LogicalPlanType::RecursiveQuery(v) => {
+ struct_ser.serialize_field("recursiveQuery", v)?;
+ }
+ logical_plan_node::LogicalPlanType::CteWorkTableScan(v) => {
+ struct_ser.serialize_field("cteWorkTableScan", v)?;
+ }
}
}
struct_ser.end()
@@ -10656,6 +10770,10 @@ impl<'de> serde::Deserialize<'de> for LogicalPlanNode {
"copy_to",
"copyTo",
"unnest",
+ "recursive_query",
+ "recursiveQuery",
+ "cte_work_table_scan",
+ "cteWorkTableScan",
];
#[allow(clippy::enum_variant_names)]
@@ -10689,6 +10807,8 @@ impl<'de> serde::Deserialize<'de> for LogicalPlanNode {
DistinctOn,
CopyTo,
Unnest,
+ RecursiveQuery,
+ CteWorkTableScan,
}
impl<'de> serde::Deserialize<'de> for GeneratedField {
fn deserialize<D>(deserializer: D) ->
std::result::Result<GeneratedField, D::Error>
@@ -10739,6 +10859,8 @@ impl<'de> serde::Deserialize<'de> for LogicalPlanNode {
"distinctOn" | "distinct_on" =>
Ok(GeneratedField::DistinctOn),
"copyTo" | "copy_to" => Ok(GeneratedField::CopyTo),
"unnest" => Ok(GeneratedField::Unnest),
+ "recursiveQuery" | "recursive_query" =>
Ok(GeneratedField::RecursiveQuery),
+ "cteWorkTableScan" | "cte_work_table_scan" =>
Ok(GeneratedField::CteWorkTableScan),
_ => Err(serde::de::Error::unknown_field(value,
FIELDS)),
}
}
@@ -10962,6 +11084,20 @@ impl<'de> serde::Deserialize<'de> for LogicalPlanNode {
return
Err(serde::de::Error::duplicate_field("unnest"));
}
logical_plan_type__ =
map_.next_value::<::std::option::Option<_>>()?.map(logical_plan_node::LogicalPlanType::Unnest)
+;
+ }
+ GeneratedField::RecursiveQuery => {
+ if logical_plan_type__.is_some() {
+ return
Err(serde::de::Error::duplicate_field("recursiveQuery"));
+ }
+ logical_plan_type__ =
map_.next_value::<::std::option::Option<_>>()?.map(logical_plan_node::LogicalPlanType::RecursiveQuery)
+;
+ }
+ GeneratedField::CteWorkTableScan => {
+ if logical_plan_type__.is_some() {
+ return
Err(serde::de::Error::duplicate_field("cteWorkTableScan"));
+ }
+ logical_plan_type__ =
map_.next_value::<::std::option::Option<_>>()?.map(logical_plan_node::LogicalPlanType::CteWorkTableScan)
;
}
}
@@ -17486,6 +17622,151 @@ impl<'de> serde::Deserialize<'de> for
RecursionUnnestOption {
deserializer.deserialize_struct("datafusion.RecursionUnnestOption",
FIELDS, GeneratedVisitor)
}
}
+impl serde::Serialize for RecursiveQueryNode {
+ #[allow(deprecated)]
+ fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok,
S::Error>
+ where
+ S: serde::Serializer,
+ {
+ use serde::ser::SerializeStruct;
+ let mut len = 0;
+ if !self.name.is_empty() {
+ len += 1;
+ }
+ if self.static_term.is_some() {
+ len += 1;
+ }
+ if self.recursive_term.is_some() {
+ len += 1;
+ }
+ if self.is_distinct {
+ len += 1;
+ }
+ let mut struct_ser =
serializer.serialize_struct("datafusion.RecursiveQueryNode", len)?;
+ if !self.name.is_empty() {
+ struct_ser.serialize_field("name", &self.name)?;
+ }
+ if let Some(v) = self.static_term.as_ref() {
+ struct_ser.serialize_field("staticTerm", v)?;
+ }
+ if let Some(v) = self.recursive_term.as_ref() {
+ struct_ser.serialize_field("recursiveTerm", v)?;
+ }
+ if self.is_distinct {
+ struct_ser.serialize_field("isDistinct", &self.is_distinct)?;
+ }
+ struct_ser.end()
+ }
+}
+impl<'de> serde::Deserialize<'de> for RecursiveQueryNode {
+ #[allow(deprecated)]
+ fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
+ where
+ D: serde::Deserializer<'de>,
+ {
+ const FIELDS: &[&str] = &[
+ "name",
+ "static_term",
+ "staticTerm",
+ "recursive_term",
+ "recursiveTerm",
+ "is_distinct",
+ "isDistinct",
+ ];
+
+ #[allow(clippy::enum_variant_names)]
+ enum GeneratedField {
+ Name,
+ StaticTerm,
+ RecursiveTerm,
+ IsDistinct,
+ }
+ impl<'de> serde::Deserialize<'de> for GeneratedField {
+ fn deserialize<D>(deserializer: D) ->
std::result::Result<GeneratedField, D::Error>
+ where
+ D: serde::Deserializer<'de>,
+ {
+ struct GeneratedVisitor;
+
+ impl<'de> serde::de::Visitor<'de> for GeneratedVisitor {
+ type Value = GeneratedField;
+
+ fn expecting(&self, formatter: &mut
std::fmt::Formatter<'_>) -> std::fmt::Result {
+ write!(formatter, "expected one of: {:?}", &FIELDS)
+ }
+
+ #[allow(unused_variables)]
+ fn visit_str<E>(self, value: &str) ->
std::result::Result<GeneratedField, E>
+ where
+ E: serde::de::Error,
+ {
+ match value {
+ "name" => Ok(GeneratedField::Name),
+ "staticTerm" | "static_term" =>
Ok(GeneratedField::StaticTerm),
+ "recursiveTerm" | "recursive_term" =>
Ok(GeneratedField::RecursiveTerm),
+ "isDistinct" | "is_distinct" =>
Ok(GeneratedField::IsDistinct),
+ _ => Err(serde::de::Error::unknown_field(value,
FIELDS)),
+ }
+ }
+ }
+ deserializer.deserialize_identifier(GeneratedVisitor)
+ }
+ }
+ struct GeneratedVisitor;
+ impl<'de> serde::de::Visitor<'de> for GeneratedVisitor {
+ type Value = RecursiveQueryNode;
+
+ fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) ->
std::fmt::Result {
+ formatter.write_str("struct datafusion.RecursiveQueryNode")
+ }
+
+ fn visit_map<V>(self, mut map_: V) ->
std::result::Result<RecursiveQueryNode, V::Error>
+ where
+ V: serde::de::MapAccess<'de>,
+ {
+ let mut name__ = None;
+ let mut static_term__ = None;
+ let mut recursive_term__ = None;
+ let mut is_distinct__ = None;
+ while let Some(k) = map_.next_key()? {
+ match k {
+ GeneratedField::Name => {
+ if name__.is_some() {
+ return
Err(serde::de::Error::duplicate_field("name"));
+ }
+ name__ = Some(map_.next_value()?);
+ }
+ GeneratedField::StaticTerm => {
+ if static_term__.is_some() {
+ return
Err(serde::de::Error::duplicate_field("staticTerm"));
+ }
+ static_term__ = map_.next_value()?;
+ }
+ GeneratedField::RecursiveTerm => {
+ if recursive_term__.is_some() {
+ return
Err(serde::de::Error::duplicate_field("recursiveTerm"));
+ }
+ recursive_term__ = map_.next_value()?;
+ }
+ GeneratedField::IsDistinct => {
+ if is_distinct__.is_some() {
+ return
Err(serde::de::Error::duplicate_field("isDistinct"));
+ }
+ is_distinct__ = Some(map_.next_value()?);
+ }
+ }
+ }
+ Ok(RecursiveQueryNode {
+ name: name__.unwrap_or_default(),
+ static_term: static_term__,
+ recursive_term: recursive_term__,
+ is_distinct: is_distinct__.unwrap_or_default(),
+ })
+ }
+ }
+ deserializer.deserialize_struct("datafusion.RecursiveQueryNode",
FIELDS, GeneratedVisitor)
+ }
+}
impl serde::Serialize for RepartitionExecNode {
#[allow(deprecated)]
fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok,
S::Error>
diff --git a/datafusion/proto/src/generated/prost.rs
b/datafusion/proto/src/generated/prost.rs
index 59a90eb31a..c7f5606049 100644
--- a/datafusion/proto/src/generated/prost.rs
+++ b/datafusion/proto/src/generated/prost.rs
@@ -5,7 +5,7 @@
pub struct LogicalPlanNode {
#[prost(
oneof = "logical_plan_node::LogicalPlanType",
- tags = "1, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18,
19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30"
+ tags = "1, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18,
19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32"
)]
pub logical_plan_type:
::core::option::Option<logical_plan_node::LogicalPlanType>,
}
@@ -71,6 +71,10 @@ pub mod logical_plan_node {
CopyTo(::prost::alloc::boxed::Box<super::CopyToNode>),
#[prost(message, tag = "30")]
Unnest(::prost::alloc::boxed::Box<super::UnnestNode>),
+ #[prost(message, tag = "31")]
+ RecursiveQuery(::prost::alloc::boxed::Box<super::RecursiveQueryNode>),
+ #[prost(message, tag = "32")]
+ CteWorkTableScan(super::CteWorkTableScanNode),
}
}
#[derive(Clone, PartialEq, ::prost::Message)]
@@ -1811,6 +1815,26 @@ pub struct PartitionStats {
#[prost(message, repeated, tag = "4")]
pub column_stats:
::prost::alloc::vec::Vec<super::datafusion_common::ColumnStats>,
}
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct RecursiveQueryNode {
+ #[prost(string, tag = "1")]
+ pub name: ::prost::alloc::string::String,
+ #[prost(message, optional, boxed, tag = "2")]
+ pub static_term:
::core::option::Option<::prost::alloc::boxed::Box<LogicalPlanNode>>,
+ #[prost(message, optional, boxed, tag = "3")]
+ pub recursive_term: ::core::option::Option<
+ ::prost::alloc::boxed::Box<LogicalPlanNode>,
+ >,
+ #[prost(bool, tag = "4")]
+ pub is_distinct: bool,
+}
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct CteWorkTableScanNode {
+ #[prost(string, tag = "1")]
+ pub name: ::prost::alloc::string::String,
+ #[prost(message, optional, tag = "2")]
+ pub schema: ::core::option::Option<super::datafusion_common::Schema>,
+}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord,
::prost::Enumeration)]
#[repr(i32)]
pub enum BuiltInWindowFunction {
diff --git a/datafusion/proto/src/logical_plan/mod.rs
b/datafusion/proto/src/logical_plan/mod.rs
index a55fecec98..50636048eb 100644
--- a/datafusion/proto/src/logical_plan/mod.rs
+++ b/datafusion/proto/src/logical_plan/mod.rs
@@ -21,8 +21,8 @@ use std::sync::Arc;
use crate::protobuf::logical_plan_node::LogicalPlanType::CustomScan;
use crate::protobuf::{
- ColumnUnnestListItem, ColumnUnnestListRecursion, CustomTableScanNode,
- SortExprNodeCollection,
+ ColumnUnnestListItem, ColumnUnnestListRecursion, CteWorkTableScanNode,
+ CustomTableScanNode, SortExprNodeCollection,
};
use crate::{
convert_required, into_required,
@@ -34,6 +34,7 @@ use crate::{
use crate::protobuf::{proto_error, ToProtoError};
use arrow::datatypes::{DataType, Schema, SchemaRef};
+use datafusion::datasource::cte_worktable::CteWorkTable;
#[cfg(feature = "parquet")]
use datafusion::datasource::file_format::parquet::ParquetFormat;
use datafusion::datasource::file_format::{
@@ -68,7 +69,9 @@ use datafusion_expr::{
DistinctOn, DropView, Expr, LogicalPlan, LogicalPlanBuilder, ScalarUDF,
SortExpr,
Statement, WindowUDF,
};
-use datafusion_expr::{AggregateUDF, ColumnUnnestList, FetchType, SkipType,
Unnest};
+use datafusion_expr::{
+ AggregateUDF, ColumnUnnestList, FetchType, RecursiveQuery, SkipType,
Unnest,
+};
use self::to_proto::{serialize_expr, serialize_exprs};
use crate::logical_plan::to_proto::serialize_sorts;
@@ -905,6 +908,41 @@ impl AsLogicalPlan for LogicalPlanNode {
options: into_required!(unnest.options)?,
}))
}
+ LogicalPlanType::RecursiveQuery(recursive_query_node) => {
+ let static_term = recursive_query_node
+ .static_term
+ .as_ref()
+ .ok_or_else(|| DataFusionError::Internal(String::from(
+ "Protobuf deserialization error, RecursiveQueryNode
was missing required field static_term.",
+ )))?
+ .try_into_logical_plan(ctx, extension_codec)?;
+
+ let recursive_term = recursive_query_node
+ .recursive_term
+ .as_ref()
+ .ok_or_else(|| DataFusionError::Internal(String::from(
+ "Protobuf deserialization error, RecursiveQueryNode
was missing required field recursive_term.",
+ )))?
+ .try_into_logical_plan(ctx, extension_codec)?;
+
+ Ok(LogicalPlan::RecursiveQuery(RecursiveQuery {
+ name: recursive_query_node.name.clone(),
+ static_term: Arc::new(static_term),
+ recursive_term: Arc::new(recursive_term),
+ is_distinct: recursive_query_node.is_distinct,
+ }))
+ }
+ LogicalPlanType::CteWorkTableScan(cte_work_table_scan_node) => {
+ let CteWorkTableScanNode { name, schema } =
cte_work_table_scan_node;
+ let schema = convert_required!(*schema)?;
+ let cte_work_table = CteWorkTable::new(name.as_str(),
Arc::new(schema));
+ LogicalPlanBuilder::scan(
+ name.as_str(),
+ provider_as_source(Arc::new(cte_work_table)),
+ None,
+ )?
+ .build()
+ }
}
}
@@ -1061,6 +1099,20 @@ impl AsLogicalPlan for LogicalPlanNode {
},
))),
})
+ } else if let Some(cte_work_table) =
source.downcast_ref::<CteWorkTable>()
+ {
+ let name = cte_work_table.name().to_string();
+ let schema = cte_work_table.schema();
+ let schema: protobuf::Schema = schema.as_ref().try_into()?;
+
+ Ok(LogicalPlanNode {
+ logical_plan_type:
Some(LogicalPlanType::CteWorkTableScan(
+ protobuf::CteWorkTableScanNode {
+ name,
+ schema: Some(schema),
+ },
+ )),
+ })
} else {
let mut bytes = vec![];
extension_codec
@@ -1630,9 +1682,27 @@ impl AsLogicalPlan for LogicalPlanNode {
LogicalPlan::DescribeTable(_) => Err(proto_error(
"LogicalPlan serde is not yet implemented for DescribeTable",
)),
- LogicalPlan::RecursiveQuery(_) => Err(proto_error(
- "LogicalPlan serde is not yet implemented for RecursiveQuery",
- )),
+ LogicalPlan::RecursiveQuery(recursive) => {
+ let static_term = LogicalPlanNode::try_from_logical_plan(
+ recursive.static_term.as_ref(),
+ extension_codec,
+ )?;
+ let recursive_term = LogicalPlanNode::try_from_logical_plan(
+ recursive.recursive_term.as_ref(),
+ extension_codec,
+ )?;
+
+ Ok(LogicalPlanNode {
+ logical_plan_type:
Some(LogicalPlanType::RecursiveQuery(Box::new(
+ protobuf::RecursiveQueryNode {
+ name: recursive.name.clone(),
+ static_term: Some(Box::new(static_term)),
+ recursive_term: Some(Box::new(recursive_term)),
+ is_distinct: recursive.is_distinct,
+ },
+ ))),
+ })
+ }
}
}
}
diff --git a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs
b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs
index ecfbaee235..8445cdc761 100644
--- a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs
+++ b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs
@@ -23,6 +23,7 @@ use arrow::datatypes::{
IntervalUnit, Schema, SchemaRef, TimeUnit, UnionFields, UnionMode,
DECIMAL256_MAX_PRECISION,
};
+use arrow::util::pretty::pretty_format_batches;
use datafusion::datasource::file_format::json::JsonFormatFactory;
use datafusion_common::parsers::CompressionTypeVariant;
use prost::Message;
@@ -2523,3 +2524,30 @@ fn roundtrip_window() {
roundtrip_expr_test(test_expr6, ctx.clone());
roundtrip_expr_test(text_expr7, ctx);
}
+
+#[tokio::test]
+async fn roundtrip_recursive_query() {
+ let query = "WITH RECURSIVE cte AS (
+ SELECT 1 as n
+ UNION ALL
+ SELECT n + 1 FROM cte WHERE n < 5
+ )
+ SELECT * FROM cte;";
+
+ let ctx = SessionContext::new();
+ let dataframe = ctx.sql(query).await.unwrap();
+ let plan = dataframe.logical_plan().clone();
+ let output = dataframe.collect().await.unwrap();
+ let bytes = logical_plan_to_bytes(&plan).unwrap();
+
+ let ctx = SessionContext::new();
+ let logical_round_trip = logical_plan_from_bytes(&bytes, &ctx).unwrap();
+ assert_eq!(format!("{plan:?}"), format!("{logical_round_trip:?}"));
+ let dataframe =
ctx.execute_logical_plan(logical_round_trip).await.unwrap();
+ let output_round_trip = dataframe.collect().await.unwrap();
+
+ assert_eq!(
+ format!("{}", pretty_format_batches(&output).unwrap()),
+ format!("{}", pretty_format_batches(&output_round_trip).unwrap())
+ );
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
