adriangb commented on code in PR #24628:
URL: https://github.com/apache/datafusion/pull/24628#discussion_r4011285447
##########
datafusion/physical-expr-common/src/physical_expr.rs:
##########
@@ -781,9 +931,211 @@ pub mod proto_decode {
})
}
+ /// The wire parts of an extension expression, as returned by
+ /// [`PhysicalExprDecodeCtx::decode_extension`].
+ ///
+ /// Destructure with `..` so a part added later does not break the pattern.
+ #[derive(Debug)]
+ #[non_exhaustive]
+ pub struct ExtensionExprParts<'n> {
+ /// The expression's opaque payload, borrowed from the node it was read
+ /// from.
+ pub payload: &'n [u8],
+ /// The expression's children, already decoded.
+ pub children: Vec<Arc<dyn PhysicalExpr>>,
+ }
+
+ /// A third-party [`PhysicalExpr`] that can be decoded by name, without a
+ /// `PhysicalExtensionCodec`.
+ ///
+ /// Extension expressions reach the wire as a `PhysicalExtensionExprNode`:
+ /// an opaque payload plus the encoded children. Historically that node
+ /// carried no type discriminator, so the codec itself had to serve as one:
+ /// `ComposedPhysicalExtensionCodec` writes the *position* of the encoding
+ /// codec into the payload, so both sides must register the same codecs in
+ /// the same order.
+ ///
+ /// Implementing this trait names the expression on the wire instead.
+ /// Register the type in a [`PhysicalExprRegistry`] on the session and the
+ /// name routes straight to [`Self::try_from_proto`].
+ ///
+ /// The encode half is [`PhysicalExpr::try_to_proto`], via
+ /// [`PhysicalExprEncodeCtx::encode_extension`]:
+ ///
+ /// ```ignore
+ /// impl ExtensionPhysicalExpr for MyExpr {
+ /// const EXPR_NAME: &'static str = "my_crate.MyExpr";
+ ///
+ /// fn try_from_proto(
+ /// node: &PhysicalExprNode,
+ /// ctx: &PhysicalExprDecodeCtx<'_>,
+ /// ) -> Result<Arc<dyn PhysicalExpr>> {
+ /// let ExtensionExprParts { payload, children, .. } =
Review Comment:
I see now how we intend to use `ExtensionExprParts`. Still seems weird. Not
clear to me what `MyExprProto` is here. Is this assuming that the expression
serializes itself as proto? It seems to me we could skip a step here and have
`MyExpr::try_from_proto(&PhysicalExprNode, &PhysicalExprDecodeCtx)`?
##########
datafusion/physical-expr-common/src/physical_expr.rs:
##########
@@ -704,12 +784,45 @@ pub mod proto_decode {
self.schema
}
+ /// The session state this decode is running under, so expressions that
+ /// need the function registry or session configuration can reach it —
+ /// the expression-side counterpart of
+ /// `ExecutionPlanDecodeCtx::task_ctx`.
+ ///
+ /// The session type is `datafusion_execution::TaskContext`, which this
Review Comment:
Don't love this. Could we get around it by moving things around crates or
otherwise refactoring so we don't have cycles?
##########
datafusion/physical-expr-common/src/physical_expr.rs:
##########
@@ -781,9 +931,211 @@ pub mod proto_decode {
})
}
+ /// The wire parts of an extension expression, as returned by
+ /// [`PhysicalExprDecodeCtx::decode_extension`].
+ ///
+ /// Destructure with `..` so a part added later does not break the pattern.
+ #[derive(Debug)]
+ #[non_exhaustive]
+ pub struct ExtensionExprParts<'n> {
+ /// The expression's opaque payload, borrowed from the node it was read
+ /// from.
+ pub payload: &'n [u8],
+ /// The expression's children, already decoded.
+ pub children: Vec<Arc<dyn PhysicalExpr>>,
Review Comment:
Why do we return the half decoded payload? It seems to me we should go from
`PhysicalExprNode -> PhysicalExpr` same as decode?
##########
datafusion/proto/tests/cases/plans/expr_registry.rs:
##########
@@ -0,0 +1,778 @@
+// 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.
+
+//! Extension `PhysicalExpr`s decoded by name through a
+//! [`PhysicalExprRegistry`] instead of a [`PhysicalExtensionCodec`].
+
+use std::any::Any;
+use std::fmt::{Display, Formatter};
+use std::sync::Arc;
+
+use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
+use arrow::record_batch::RecordBatch;
+use datafusion::execution::TaskContext;
+use datafusion::physical_plan::ExecutionPlan;
+use datafusion::physical_plan::empty::EmptyExec;
+use datafusion::physical_plan::expressions::{BinaryExpr, col, lit};
+use datafusion::physical_plan::filter::FilterExec;
+use datafusion::physical_plan::proto::{ExtensionExprParts,
PhysicalExprRegistry};
+use datafusion::prelude::{SessionConfig, SessionContext};
+use datafusion_common::{Result, internal_datafusion_err, internal_err};
+use datafusion_expr::ColumnarValue;
+use datafusion_expr::Operator;
+use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
+use datafusion_physical_expr_common::physical_expr::proto_decode::{
+ ExtensionPhysicalExpr, PhysicalExprDecode, PhysicalExprDecodeCtx,
+};
+use
datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx;
+use datafusion_proto::bytes::{
+ physical_plan_from_bytes_with_extension_codec,
+ physical_plan_to_bytes_with_extension_codec,
+};
+use datafusion_proto::physical_plan::{
+ DeduplicatingProtoConverter, DefaultPhysicalProtoConverter,
PhysicalExtensionCodec,
+ PhysicalPlanDecodeContext, PhysicalProtoConverterExtension,
+};
+use datafusion_proto::protobuf::{PhysicalExprNode, physical_expr_node};
+use prost::Message;
+
+/// Generates an extension expression type that decodes through the registry.
+///
+/// Each generated type passes its child through untouched and carries a `tag`
+/// of its own, so a round-trip has both an opaque payload and a child
+/// expression to reconstruct. Several types are generated because the
+/// registry's identity rules — idempotent for one type, an error for two types
+/// claiming one name — can only be exercised with more than one type.
+macro_rules! tag_expr_type {
+ ($name:ident, $wire_name:literal) => {
+ tag_expr_type!(@type $name);
+
+ impl ExtensionPhysicalExpr for $name {
+ const EXPR_NAME: &'static str = $wire_name;
+
+ fn try_from_proto(
Review Comment:
This is a lot of faff to essentially hold a name and delegate to
`$name::try_from_proto`. Should we try to make the name part of the expression
or something instead?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
