datafusion_proto/bytes/
mod.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Serialization / Deserialization to Bytes
19use crate::logical_plan::to_proto::serialize_expr;
20use crate::logical_plan::{
21    self, AsLogicalPlan, DefaultLogicalExtensionCodec, LogicalExtensionCodec,
22};
23use crate::physical_plan::{
24    DefaultPhysicalExtensionCodec, DefaultPhysicalProtoConverter, PhysicalExtensionCodec,
25    PhysicalPlanDecodeContext, PhysicalProtoConverterExtension,
26};
27use crate::protobuf;
28use datafusion_common::{Result, plan_datafusion_err};
29use datafusion_execution::TaskContext;
30use datafusion_expr::{Expr, LogicalPlan};
31use prost::{
32    Message,
33    bytes::{Bytes, BytesMut},
34};
35use std::sync::Arc;
36
37use datafusion_physical_plan::ExecutionPlan;
38
39/// Encodes something (such as [`Expr`]) to/from a stream of
40/// bytes.
41///
42/// ```
43/// use datafusion_expr::{col, lit, Expr};
44/// use datafusion_proto::bytes::Serializeable;
45///
46/// // Create a new `Expr` a < 32
47/// let expr = col("a").lt(lit(5i32));
48///
49/// // Convert it to an opaque form
50/// let bytes = expr.to_bytes().unwrap();
51///
52/// // Decode bytes from somewhere (over network, etc.)
53/// let decoded_expr = Expr::from_bytes(&bytes).unwrap();
54/// assert_eq!(expr, decoded_expr);
55/// ```
56pub trait Serializeable: Sized {
57    /// Convert `self` to an opaque byte stream
58    fn to_bytes(&self) -> Result<Bytes>;
59
60    /// Convert `bytes` (the output of [`to_bytes`]) back into an object. This
61    /// will error if the serialized bytes contain any user defined functions,
62    /// in which case use [`from_bytes_with_ctx`]
63    ///
64    /// [`to_bytes`]: Self::to_bytes
65    /// [`from_bytes_with_ctx`]: Self::from_bytes_with_ctx
66    fn from_bytes(bytes: &[u8]) -> Result<Self> {
67        Self::from_bytes_with_ctx(bytes, &TaskContext::default())
68    }
69
70    /// Convert `bytes` (the output of [`to_bytes`]) back into an object,
71    /// resolving user defined functions with the specified `ctx`
72    ///
73    /// [`to_bytes`]: Self::to_bytes
74    fn from_bytes_with_ctx(bytes: &[u8], ctx: &TaskContext) -> Result<Self>;
75}
76
77impl Serializeable for Expr {
78    fn to_bytes(&self) -> Result<Bytes> {
79        let mut buffer = BytesMut::new();
80        let extension_codec = DefaultLogicalExtensionCodec {};
81        let protobuf: protobuf::LogicalExprNode = serialize_expr(self, &extension_codec)
82            .map_err(|e| plan_datafusion_err!("Error encoding expr as protobuf: {e}"))?;
83
84        protobuf
85            .encode(&mut buffer)
86            .map_err(|e| plan_datafusion_err!("Error encoding protobuf as bytes: {e}"))?;
87
88        let bytes: Bytes = buffer.into();
89
90        // The produced byte stream may lead to "recursion limit" errors, see
91        // https://github.com/apache/datafusion/issues/3968
92        // Until the underlying prost issue ( https://github.com/tokio-rs/prost/issues/736 )
93        // is fixed, verify the bytes can be decoded without hitting that limit.
94        protobuf::LogicalExprNode::decode(bytes.as_ref())
95            .map_err(|e| plan_datafusion_err!("Error decoding expr as protobuf: {e}"))?;
96
97        Ok(bytes)
98    }
99
100    fn from_bytes_with_ctx(bytes: &[u8], ctx: &TaskContext) -> Result<Self> {
101        let protobuf = protobuf::LogicalExprNode::decode(bytes)
102            .map_err(|e| plan_datafusion_err!("Error decoding expr as protobuf: {e}"))?;
103
104        let extension_codec = DefaultLogicalExtensionCodec {};
105        logical_plan::from_proto::parse_expr(&protobuf, ctx, &extension_codec)
106            .map_err(|e| plan_datafusion_err!("Error parsing protobuf into Expr: {e}"))
107    }
108}
109
110/// Serialize a LogicalPlan as bytes
111pub fn logical_plan_to_bytes(plan: &LogicalPlan) -> Result<Bytes> {
112    let extension_codec = DefaultLogicalExtensionCodec {};
113    logical_plan_to_bytes_with_extension_codec(plan, &extension_codec)
114}
115
116/// Serialize a LogicalPlan as JSON
117#[cfg(feature = "json")]
118pub fn logical_plan_to_json(plan: &LogicalPlan) -> Result<String> {
119    let extension_codec = DefaultLogicalExtensionCodec {};
120    logical_plan_to_json_with_extension_codec(plan, &extension_codec)
121}
122
123/// Serialize a LogicalPlan as bytes, using the provided extension codec
124pub fn logical_plan_to_bytes_with_extension_codec(
125    plan: &LogicalPlan,
126    extension_codec: &dyn LogicalExtensionCodec,
127) -> Result<Bytes> {
128    let protobuf =
129        protobuf::LogicalPlanNode::try_from_logical_plan(plan, extension_codec)?;
130    let mut buffer = BytesMut::new();
131    protobuf
132        .encode(&mut buffer)
133        .map_err(|e| plan_datafusion_err!("Error encoding protobuf as bytes: {e}"))?;
134    Ok(buffer.into())
135}
136
137/// Serialize a LogicalPlan as JSON using the provided extension codec
138#[cfg(feature = "json")]
139pub fn logical_plan_to_json_with_extension_codec(
140    plan: &LogicalPlan,
141    extension_codec: &dyn LogicalExtensionCodec,
142) -> Result<String> {
143    let protobuf =
144        protobuf::LogicalPlanNode::try_from_logical_plan(plan, extension_codec)
145            .map_err(|e| plan_datafusion_err!("Error serializing plan: {e}"))?;
146    serde_json::to_string(&protobuf)
147        .map_err(|e| plan_datafusion_err!("Error serializing plan: {e}"))
148}
149
150/// Deserialize a LogicalPlan from JSON
151#[cfg(feature = "json")]
152pub fn logical_plan_from_json(json: &str, ctx: &TaskContext) -> Result<LogicalPlan> {
153    let extension_codec = DefaultLogicalExtensionCodec {};
154    logical_plan_from_json_with_extension_codec(json, ctx, &extension_codec)
155}
156
157/// Deserialize a LogicalPlan from bytes
158pub fn logical_plan_from_bytes(bytes: &[u8], ctx: &TaskContext) -> Result<LogicalPlan> {
159    let extension_codec = DefaultLogicalExtensionCodec {};
160    logical_plan_from_bytes_with_extension_codec(bytes, ctx, &extension_codec)
161}
162
163/// Deserialize a LogicalPlan from bytes
164pub fn logical_plan_from_bytes_with_extension_codec(
165    bytes: &[u8],
166    ctx: &TaskContext,
167    extension_codec: &dyn LogicalExtensionCodec,
168) -> Result<LogicalPlan> {
169    let protobuf = protobuf::LogicalPlanNode::decode(bytes)
170        .map_err(|e| plan_datafusion_err!("Error decoding expr as protobuf: {e}"))?;
171    protobuf.try_into_logical_plan(ctx, extension_codec)
172}
173
174/// Deserialize a LogicalPlan from JSON
175#[cfg(feature = "json")]
176pub fn logical_plan_from_json_with_extension_codec(
177    json: &str,
178    ctx: &TaskContext,
179    extension_codec: &dyn LogicalExtensionCodec,
180) -> Result<LogicalPlan> {
181    let back: protobuf::LogicalPlanNode = serde_json::from_str(json)
182        .map_err(|e| plan_datafusion_err!("Error deserializing plan: {e}"))?;
183    back.try_into_logical_plan(ctx, extension_codec)
184}
185
186/// Serialize a PhysicalPlan as bytes
187pub fn physical_plan_to_bytes(plan: Arc<dyn ExecutionPlan>) -> Result<Bytes> {
188    let extension_codec = DefaultPhysicalExtensionCodec {};
189    let proto_converter = DefaultPhysicalProtoConverter {};
190    physical_plan_to_bytes_with_proto_converter(plan, &extension_codec, &proto_converter)
191}
192
193/// Serialize a PhysicalPlan as JSON
194#[cfg(feature = "json")]
195#[expect(
196    clippy::needless_pass_by_value,
197    reason = "Preserve the existing public API"
198)]
199pub fn physical_plan_to_json(plan: Arc<dyn ExecutionPlan>) -> Result<String> {
200    let extension_codec = DefaultPhysicalExtensionCodec {};
201    let proto_converter = DefaultPhysicalProtoConverter {};
202    let protobuf = proto_converter
203        .execution_plan_to_proto(&plan, &extension_codec)
204        .map_err(|e| plan_datafusion_err!("Error serializing plan: {e}"))?;
205    serde_json::to_string(&protobuf)
206        .map_err(|e| plan_datafusion_err!("Error serializing plan: {e}"))
207}
208
209/// Serialize a PhysicalPlan as bytes, using the provided extension codec
210pub fn physical_plan_to_bytes_with_extension_codec(
211    plan: Arc<dyn ExecutionPlan>,
212    extension_codec: &dyn PhysicalExtensionCodec,
213) -> Result<Bytes> {
214    let proto_converter = DefaultPhysicalProtoConverter {};
215    physical_plan_to_bytes_with_proto_converter(plan, extension_codec, &proto_converter)
216}
217
218/// Serialize a PhysicalPlan as bytes, using the provided extension codec
219/// and protobuf converter.
220#[expect(clippy::needless_pass_by_value)] // Taking the plan by value is part of the public API
221pub fn physical_plan_to_bytes_with_proto_converter(
222    plan: Arc<dyn ExecutionPlan>,
223    extension_codec: &dyn PhysicalExtensionCodec,
224    proto_converter: &dyn PhysicalProtoConverterExtension,
225) -> Result<Bytes> {
226    let protobuf = proto_converter.execution_plan_to_proto(&plan, extension_codec)?;
227    let mut buffer = BytesMut::new();
228    protobuf
229        .encode(&mut buffer)
230        .map_err(|e| plan_datafusion_err!("Error encoding protobuf as bytes: {e}"))?;
231    Ok(buffer.into())
232}
233
234/// Deserialize a PhysicalPlan from JSON
235#[cfg(feature = "json")]
236pub fn physical_plan_from_json(
237    json: &str,
238    ctx: &TaskContext,
239) -> Result<Arc<dyn ExecutionPlan>> {
240    let back: protobuf::PhysicalPlanNode = serde_json::from_str(json)
241        .map_err(|e| plan_datafusion_err!("Error serializing plan: {e}"))?;
242    let extension_codec = DefaultPhysicalExtensionCodec {};
243    let proto_converter = DefaultPhysicalProtoConverter {};
244    let decode_ctx = PhysicalPlanDecodeContext::new(ctx, &extension_codec);
245    proto_converter.proto_to_execution_plan(&back, &decode_ctx)
246}
247
248/// Deserialize a PhysicalPlan from bytes
249pub fn physical_plan_from_bytes(
250    bytes: &[u8],
251    ctx: &TaskContext,
252) -> Result<Arc<dyn ExecutionPlan>> {
253    let extension_codec = DefaultPhysicalExtensionCodec {};
254    let proto_converter = DefaultPhysicalProtoConverter {};
255    physical_plan_from_bytes_with_proto_converter(
256        bytes,
257        ctx,
258        &extension_codec,
259        &proto_converter,
260    )
261}
262
263/// Deserialize a PhysicalPlan from bytes
264pub fn physical_plan_from_bytes_with_extension_codec(
265    bytes: &[u8],
266    ctx: &TaskContext,
267    extension_codec: &dyn PhysicalExtensionCodec,
268) -> Result<Arc<dyn ExecutionPlan>> {
269    let proto_converter = DefaultPhysicalProtoConverter {};
270    physical_plan_from_bytes_with_proto_converter(
271        bytes,
272        ctx,
273        extension_codec,
274        &proto_converter,
275    )
276}
277
278/// Deserialize a PhysicalPlan from bytes
279pub fn physical_plan_from_bytes_with_proto_converter(
280    bytes: &[u8],
281    ctx: &TaskContext,
282    extension_codec: &dyn PhysicalExtensionCodec,
283    proto_converter: &dyn PhysicalProtoConverterExtension,
284) -> Result<Arc<dyn ExecutionPlan>> {
285    let protobuf = protobuf::PhysicalPlanNode::decode(bytes)
286        .map_err(|e| plan_datafusion_err!("Error decoding expr as protobuf: {e}"))?;
287    let decode_ctx = PhysicalPlanDecodeContext::new(ctx, extension_codec);
288    proto_converter.proto_to_execution_plan(&protobuf, &decode_ctx)
289}
datafusion_proto/bytes/mod.rs

datafusion_proto/bytes/
mod.rs
