Skip to content

Commit 9c12919

Browse files
authored
Update proto to support to/from json with an extension codec (#14561)
#14560
1 parent df0d966 commit 9c12919

File tree

1 file changed

+27
-8
lines changed
  • datafusion/proto/src/bytes

1 file changed

+27
-8
lines changed

datafusion/proto/src/bytes/mod.rs

Lines changed: 27 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -199,11 +199,7 @@ pub fn logical_plan_to_bytes(plan: &LogicalPlan) -> Result<Bytes> {
199199
#[cfg(feature = "json")]
200200
pub fn logical_plan_to_json(plan: &LogicalPlan) -> Result<String> {
201201
let extension_codec = DefaultLogicalExtensionCodec {};
202-
let protobuf =
203-
protobuf::LogicalPlanNode::try_from_logical_plan(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}"))
202+
logical_plan_to_json_with_extension_codec(plan, &extension_codec)
207203
}
208204

209205
/// Serialize a LogicalPlan as bytes, using the provided extension codec
@@ -220,13 +216,24 @@ pub fn logical_plan_to_bytes_with_extension_codec(
220216
Ok(buffer.into())
221217
}
222218

219+
/// Serialize a LogicalPlan as JSON using the provided extension codec
220+
#[cfg(feature = "json")]
221+
pub fn logical_plan_to_json_with_extension_codec(
222+
plan: &LogicalPlan,
223+
extension_codec: &dyn LogicalExtensionCodec,
224+
) -> Result<String> {
225+
let protobuf =
226+
protobuf::LogicalPlanNode::try_from_logical_plan(plan, extension_codec)
227+
.map_err(|e| plan_datafusion_err!("Error serializing plan: {e}"))?;
228+
serde_json::to_string(&protobuf)
229+
.map_err(|e| plan_datafusion_err!("Error serializing plan: {e}"))
230+
}
231+
223232
/// Deserialize a LogicalPlan from JSON
224233
#[cfg(feature = "json")]
225234
pub fn logical_plan_from_json(json: &str, ctx: &SessionContext) -> Result<LogicalPlan> {
226-
let back: protobuf::LogicalPlanNode = serde_json::from_str(json)
227-
.map_err(|e| plan_datafusion_err!("Error serializing plan: {e}"))?;
228235
let extension_codec = DefaultLogicalExtensionCodec {};
229-
back.try_into_logical_plan(ctx, &extension_codec)
236+
logical_plan_from_json_with_extension_codec(json, ctx, &extension_codec)
230237
}
231238

232239
/// Deserialize a LogicalPlan from bytes
@@ -249,6 +256,18 @@ pub fn logical_plan_from_bytes_with_extension_codec(
249256
protobuf.try_into_logical_plan(ctx, extension_codec)
250257
}
251258

259+
/// Deserialize a LogicalPlan from JSON
260+
#[cfg(feature = "json")]
261+
pub fn logical_plan_from_json_with_extension_codec(
262+
json: &str,
263+
ctx: &SessionContext,
264+
extension_codec: &dyn LogicalExtensionCodec,
265+
) -> Result<LogicalPlan> {
266+
let back: protobuf::LogicalPlanNode = serde_json::from_str(json)
267+
.map_err(|e| plan_datafusion_err!("Error deserializing plan: {e}"))?;
268+
back.try_into_logical_plan(ctx, extension_codec)
269+
}
270+
252271
/// Serialize a PhysicalPlan as bytes
253272
pub fn physical_plan_to_bytes(plan: Arc<dyn ExecutionPlan>) -> Result<Bytes> {
254273
let extension_codec = DefaultPhysicalExtensionCodec {};

0 commit comments

Comments
 (0)