DRAFT implement trait in proto de/serialization to allow finer control

datafusion-examples/examples/composed_extension_codec.rs

@@ -41,6 +41,7 @@ use datafusion::physical_plan::{DisplayAs, ExecutionPlan};
 use datafusion::prelude::SessionContext;
 use datafusion_proto::physical_plan::{
  AsExecutionPlan, ComposedPhysicalExtensionCodec, PhysicalExtensionCodec,
+ TaskContextWithPhysicalCodec,
 };
 use datafusion_proto::protobuf;
 
@@ -52,26 +53,30 @@ async fn main() {
  let exec_plan = Arc::new(ParentExec {
  input: Arc::new(ChildExec {}),
  });
- let ctx = SessionContext::new();
 
+ let task_ctx = SessionContext::new().task_ctx();
  // Position in this list is important as it will be used for decoding.
  // If new codec is added it should go to last position.
- let composed_codec = ComposedPhysicalExtensionCodec::new(vec![
+ let composed_codec = Arc::new(ComposedPhysicalExtensionCodec::new(vec![
  Arc::new(ParentPhysicalExtensionCodec {}),
  Arc::new(ChildPhysicalExtensionCodec {}),
- ]);
+ ]));
+ let mut parser = TaskContextWithPhysicalCodec {
+ task_ctx: task_ctx.as_ref(),
+ codec: composed_codec.as_ref(),
+ };
 
  // serialize execution plan to proto
  let proto: protobuf::PhysicalPlanNode =
  protobuf::PhysicalPlanNode::try_from_physical_plan(
  exec_plan.clone(),
- &composed_codec,
+ &mut parser,
  )
  .expect("to proto");
 
  // deserialize proto back to execution plan
  let result_exec_plan: Arc<dyn ExecutionPlan> = proto
- .try_into_physical_plan(&ctx.task_ctx(), &composed_codec)
+ .try_into_physical_plan(&mut parser)
  .expect("from proto");
 
  // assert that the original and deserialized execution plans are equal
@@ -137,7 +142,6 @@ impl PhysicalExtensionCodec for ParentPhysicalExtensionCodec {
  &self,
  buf: &[u8],
  inputs: &[Arc<dyn ExecutionPlan>],
- _ctx: &TaskContext,
  ) -> Result<Arc<dyn ExecutionPlan>> {
  if buf == "ParentExec".as_bytes() {
  Ok(Arc::new(ParentExec {
@@ -213,7 +217,6 @@ impl PhysicalExtensionCodec for ChildPhysicalExtensionCodec {
  &self,
  buf: &[u8],
  _inputs: &[Arc<dyn ExecutionPlan>],
- _ctx: &TaskContext,
  ) -> Result<Arc<dyn ExecutionPlan>> {
  if buf == "ChildExec".as_bytes() {
  Ok(Arc::new(ChildExec {}))

datafusion/ffi/src/plan_properties.rs

@@ -87,9 +87,11 @@ unsafe extern "C" fn output_partitioning_fn_wrapper(
  let private_data = properties.private_data as *const PlanPropertiesPrivateData;
  let props = &(*private_data).props;
 
- let codec = DefaultPhysicalExtensionCodec {};
- let partitioning_data =
- rresult_return!(serialize_partitioning(props.output_partitioning(), &codec));
+ let mut codec = DefaultPhysicalExtensionCodec {};
+ let partitioning_data = rresult_return!(serialize_partitioning(
+ props.output_partitioning(),
+ &mut codec
+ ));
  let output_partitioning = partitioning_data.encode_to_vec();
 
  ROk(output_partitioning.into())
@@ -117,11 +119,11 @@ unsafe extern "C" fn output_ordering_fn_wrapper(
  let private_data = properties.private_data as *const PlanPropertiesPrivateData;
  let props = &(*private_data).props;
 
- let codec = DefaultPhysicalExtensionCodec {};
+ let mut codec = DefaultPhysicalExtensionCodec {};
  let output_ordering = match props.output_ordering() {
  Some(ordering) => {
  let physical_sort_expr_nodes = rresult_return!(
- serialize_physical_sort_exprs(ordering.to_owned(), &codec)
+ serialize_physical_sort_exprs(ordering.to_owned(), &mut codec)
  );
  let ordering_data = PhysicalSortExprNodeCollection {
  physical_sort_expr_nodes,
@@ -181,8 +183,7 @@ impl TryFrom<FFI_PlanProperties> for PlanProperties {
 
  // TODO Extend FFI to get the registry and codex
  let default_ctx = SessionContext::new();
- let task_context = default_ctx.task_ctx();
- let codex = DefaultPhysicalExtensionCodec {};
+ let mut task_context = default_ctx.task_ctx();
 
  let ffi_orderings = unsafe { (ffi_props.output_ordering)(&ffi_props) };
 
@@ -191,9 +192,8 @@ impl TryFrom<FFI_PlanProperties> for PlanProperties {
  .map_err(|e| DataFusionError::External(Box::new(e)))?;
  let sort_exprs = parse_physical_sort_exprs(
  &proto_output_ordering.physical_sort_expr_nodes,
- &task_context,
+ &mut task_context,
  &schema,
- &codex,
  )?;
 
  let partitioning_vec =
@@ -203,9 +203,8 @@ impl TryFrom<FFI_PlanProperties> for PlanProperties {
  .map_err(|e| DataFusionError::External(Box::new(e)))?;
  let partitioning = parse_protobuf_partitioning(
  Some(&proto_output_partitioning),
- &task_context,
+ &mut task_context,
  &schema,
- &codex,
  )?
  .ok_or(DataFusionError::Plan(
  "Unable to deserialize partitioning protobuf in FFI_PlanProperties"

datafusion/ffi/src/udaf/accumulator_args.rs

@@ -63,11 +63,11 @@ impl TryFrom<AccumulatorArgs<'_>> for FFI_AccumulatorArgs {
  WrappedSchema(FFI_ArrowSchema::try_from(args.return_field.as_ref())?);
  let schema = WrappedSchema(FFI_ArrowSchema::try_from(args.schema)?);
 
- let codec = DefaultPhysicalExtensionCodec {};
+ let mut codec = DefaultPhysicalExtensionCodec {};
  let ordering_req =
- serialize_physical_sort_exprs(args.order_bys.to_owned(), &codec)?;
+ serialize_physical_sort_exprs(args.order_bys.to_owned(), &mut codec)?;
 
- let expr = serialize_physical_exprs(args.exprs, &codec)?;
+ let expr = serialize_physical_exprs(args.exprs, &mut codec)?;
 
  let physical_expr_def = PhysicalAggregateExprNode {
  expr,
@@ -121,17 +121,12 @@ impl TryFrom<FFI_AccumulatorArgs> for ForeignAccumulatorArgs {
  let schema = Schema::try_from(&value.schema.0)?;
 
  let default_ctx = SessionContext::new();
- let task_ctx = default_ctx.task_ctx();
- let codex = DefaultPhysicalExtensionCodec {};
+ let mut task_ctx = default_ctx.task_ctx();
 
- let order_bys = parse_physical_sort_exprs(
- &proto_def.ordering_req,
- &task_ctx,
- &schema,
- &codex,
- )?;
+ let order_bys =
+ parse_physical_sort_exprs(&proto_def.ordering_req, &mut task_ctx, &schema)?;
 
- let exprs = parse_physical_exprs(&proto_def.expr, &task_ctx, &schema, &codex)?;
+ let exprs = parse_physical_exprs(&proto_def.expr, &mut task_ctx, &schema)?;
 
  let expr_fields = exprs
  .iter()

datafusion/ffi/src/udwf/partition_evaluator_args.rs

@@ -96,8 +96,8 @@ impl TryFrom<PartitionEvaluatorArgs<'_>> for FFI_PartitionEvaluatorArgs {
 
  let schema = Arc::new(Schema::new(fields));
 
- let codec = DefaultPhysicalExtensionCodec {};
- let input_exprs = serialize_physical_exprs(args.input_exprs(), &codec)?
+ let mut codec = DefaultPhysicalExtensionCodec {};
+ let input_exprs = serialize_physical_exprs(args.input_exprs(), &mut codec)?
  .into_iter()
  .map(|expr_node| expr_node.encode_to_vec().into())
  .collect();
@@ -136,8 +136,7 @@ impl TryFrom<FFI_PartitionEvaluatorArgs> for ForeignPartitionEvaluatorArgs {
  type Error = DataFusionError;
 
  fn try_from(value: FFI_PartitionEvaluatorArgs) -> Result<Self> {
- let default_ctx = SessionContext::new();
- let codec = DefaultPhysicalExtensionCodec {};
+ let mut task_ctx = SessionContext::new().task_ctx();
 
  let schema: SchemaRef = value.schema.into();
 
@@ -148,9 +147,7 @@ impl TryFrom<FFI_PartitionEvaluatorArgs> for ForeignPartitionEvaluatorArgs {
  .collect::<std::result::Result<Vec<_>, prost::DecodeError>>()
  .map_err(|e| exec_datafusion_err!("Failed to decode PhysicalExprNode: {e}"))?
  .iter()
- .map(|expr_node| {
- parse_physical_expr(expr_node, &default_ctx.task_ctx(), &schema, &codec)
- })
+ .map(|expr_node| parse_physical_expr(expr_node, &mut task_ctx, &schema))
  .collect::<Result<Vec<_>>>()?;
 
  let input_fields = input_exprs

datafusion/proto/src/bytes/mod.rs

@@ -21,7 +21,8 @@ use crate::logical_plan::{
  self, AsLogicalPlan, DefaultLogicalExtensionCodec, LogicalExtensionCodec,
 };
 use crate::physical_plan::{
- AsExecutionPlan, DefaultPhysicalExtensionCodec, PhysicalExtensionCodec,
+ AsExecutionPlan, DefaultPhysicalExtensionCodec, PhysicalDeserializer,
+ PhysicalSerializer, TaskContextWithPhysicalCodec,
 };
 use crate::protobuf;
 use datafusion_common::{plan_datafusion_err, Result};
@@ -275,8 +276,8 @@ pub fn logical_plan_from_json_with_extension_codec(
 
 /// Serialize a PhysicalPlan as bytes
 pub fn physical_plan_to_bytes(plan: Arc<dyn ExecutionPlan>) -> Result<Bytes> {
- let extension_codec = DefaultPhysicalExtensionCodec {};
- physical_plan_to_bytes_with_extension_codec(plan, &extension_codec)
+ let mut extension_codec = DefaultPhysicalExtensionCodec {};
+ physical_plan_to_bytes_with_extension_codec(plan, &mut extension_codec)
 }
 
 /// Serialize a PhysicalPlan as JSON
@@ -291,12 +292,11 @@ pub fn physical_plan_to_json(plan: Arc<dyn ExecutionPlan>) -> Result<String> {
 }
 
 /// Serialize a PhysicalPlan as bytes, using the provided extension codec
-pub fn physical_plan_to_bytes_with_extension_codec(
+pub fn physical_plan_to_bytes_with_extension_codec<S: PhysicalSerializer>(
  plan: Arc<dyn ExecutionPlan>,
-extension_codec: &dyn PhysicalExtensionCodec,
+parser: &mut S,
 ) -> Result<Bytes> {
- let protobuf =
- protobuf::PhysicalPlanNode::try_from_physical_plan(plan, extension_codec)?;
+ let protobuf = protobuf::PhysicalPlanNode::try_from_physical_plan(plan, parser)?;
  let mut buffer = BytesMut::new();
  protobuf
  .encode(&mut buffer)
@@ -322,16 +322,19 @@ pub fn physical_plan_from_bytes(
  ctx: &TaskContext,
 ) -> Result<Arc<dyn ExecutionPlan>> {
  let extension_codec = DefaultPhysicalExtensionCodec {};
- physical_plan_from_bytes_with_extension_codec(bytes, ctx, &extension_codec)
+ let mut parser = TaskContextWithPhysicalCodec {
+ task_ctx: ctx,
+ codec: &extension_codec,
+ };
+ physical_plan_from_bytes_with_extension_codec(bytes, &mut parser)
 }
 
 /// Deserialize a PhysicalPlan from bytes
-pub fn physical_plan_from_bytes_with_extension_codec(
+pub fn physical_plan_from_bytes_with_extension_codec<S: PhysicalDeserializer>(
  bytes: &[u8],
- ctx: &TaskContext,
- extension_codec: &dyn PhysicalExtensionCodec,
+ parser: &mut S,
 ) -> Result<Arc<dyn ExecutionPlan>> {
  let protobuf = protobuf::PhysicalPlanNode::decode(bytes)
  .map_err(|e| plan_datafusion_err!("Error decoding expr as protobuf: {e}"))?;
- protobuf.try_into_physical_plan(ctx, extension_codec)
+ protobuf.try_into_physical_plan(parser)
 }