Re: [PR] feat: Python UDFs: per-session inlining toggle and strict refusal setting [datafusion-python]
timsaucer merged PR #1546: URL: https://github.com/apache/datafusion-python/pull/1546 -- 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]
Re: [PR] feat: Python UDFs: per-session inlining toggle and strict refusal setting [datafusion-python]
ntjohnson1 commented on code in PR #1546: URL: https://github.com/apache/datafusion-python/pull/1546#discussion_r3293069846 ## python/datafusion/ipc.py: ## @@ -125,6 +158,67 @@ def get_worker_ctx() -> SessionContext | None: return getattr(_local, "ctx", None) +def set_sender_ctx(ctx: SessionContext) -> None: +"""Install this driver's :class:`SessionContext` for outbound pickles. + +Controls how :func:`pickle.dumps` encodes :class:`Expr` instances on +this thread. The most useful application is propagating a session +configured with +:meth:`SessionContext.with_python_udf_inlining` so the toggle takes +effect through pickle (which otherwise calls +:meth:`Expr.to_bytes` with no context and uses the default codec). + +Idempotent: overwrites any previous value. Stored in a thread-local +slot, so worker threads on the driver may install their own contexts. +Does not affect :meth:`Expr.to_bytes` calls that pass an explicit +``ctx`` — those continue to use the supplied context. + +Examples: +>>> from datafusion import SessionContext +>>> from datafusion.ipc import ( +... set_sender_ctx, get_sender_ctx, clear_sender_ctx, +... ) +>>> driver = SessionContext().with_python_udf_inlining(enabled=False) +>>> set_sender_ctx(driver) +>>> get_sender_ctx() is driver +True +>>> clear_sender_ctx() Review Comment: Does skipping this step cause other doc tests to fail and this is basically cleanup? If not looks unneeded for this doc string -- 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]
Re: [PR] feat: Python UDFs: per-session inlining toggle and strict refusal setting [datafusion-python]
ntjohnson1 commented on code in PR #1546:
URL:
https://github.com/apache/datafusion-python/pull/1546#discussion_r3293043594
##
crates/core/src/codec.rs:
##
@@ -232,16 +232,39 @@ fn strip_wire_header<'a>(
#[derive(Debug)]
pub struct PythonLogicalCodec {
inner: Arc,
+python_udf_inlining: bool,
Review Comment:
NIT: throughout this UDF seems to be used in two contexts. UDF as short hand
for scalar UDF and UDF as short hand for the broader set of all flavors of user
functions udf,udaf,udwf. Everytime I've been reviewing these I kind of forget
that and have to revisit things to make sure I'm doing the right mapping in my
head. Maybe that's expected and standard. No change request more of just an FYI
in case there are thoughts on how to resolve this overloading for potentially
making it easier to bring on future people.
##
python/datafusion/ipc.py:
##
@@ -53,6 +53,36 @@ def init_worker():
payloads are not portable across Python minor versions. See
:meth:`datafusion.Expr.to_bytes` for examples of what travels by
value vs. by reference.
+
+On the driver side, call :func:`set_sender_ctx` to control how
+:func:`pickle.dumps` encodes expressions — for example, to apply
+:meth:`SessionContext.with_python_udf_inlining` to every pickled
+expression on this thread:
+
+.. code-block:: python
Review Comment:
NIT: I don't think doc test hits code-block, I think it needs the >>> prefix
and since this looks standalone executable might be nice to make it in that
format
##
crates/core/src/codec.rs:
##
@@ -301,48 +324,104 @@ impl LogicalExtensionCodec for PythonLogicalCodec {
}
fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec) ->
Result<()> {
-if try_encode_python_scalar_udf(node, buf)? {
+if self.python_udf_inlining && try_encode_python_scalar_udf(node,
buf)? {
return Ok(());
}
self.inner.try_encode_udf(node, buf)
}
fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result>
{
-if let Some(udf) = try_decode_python_scalar_udf(buf)? {
-return Ok(udf);
+if self.python_udf_inlining {
+if let Some(udf) = try_decode_python_scalar_udf(buf)? {
+return Ok(udf);
+}
+} else {
+refuse_if_inline(buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", name)?;
}
self.inner.try_decode_udf(name, buf)
}
fn try_encode_udaf(&self, node: &AggregateUDF, buf: &mut Vec) ->
Result<()> {
-if try_encode_python_udaf(node, buf)? {
+if self.python_udf_inlining && try_encode_python_udaf(node, buf)? {
return Ok(());
}
self.inner.try_encode_udaf(node, buf)
}
fn try_decode_udaf(&self, name: &str, buf: &[u8]) ->
Result> {
-if let Some(udaf) = try_decode_python_udaf(buf)? {
-return Ok(udaf);
+if self.python_udf_inlining {
+if let Some(udaf) = try_decode_python_udaf(buf)? {
+return Ok(udaf);
+}
+} else {
+refuse_if_inline(buf, PY_AGG_UDF_FAMILY, "aggregate UDF", name)?;
}
self.inner.try_decode_udaf(name, buf)
}
fn try_encode_udwf(&self, node: &WindowUDF, buf: &mut Vec) ->
Result<()> {
-if try_encode_python_udwf(node, buf)? {
+if self.python_udf_inlining && try_encode_python_udwf(node, buf)? {
return Ok(());
}
self.inner.try_encode_udwf(node, buf)
}
fn try_decode_udwf(&self, name: &str, buf: &[u8]) ->
Result> {
-if let Some(udwf) = try_decode_python_udwf(buf)? {
-return Ok(udwf);
+if self.python_udf_inlining {
+if let Some(udwf) = try_decode_python_udwf(buf)? {
+return Ok(udwf);
+}
+} else {
+refuse_if_inline(buf, PY_WINDOW_UDF_FAMILY, "window UDF", name)?;
}
self.inner.try_decode_udwf(name, buf)
}
}
+/// Strict-mode gate: if `buf` is a well-framed inline payload for
+/// `family`, return the strict-refusal error; otherwise return
+/// `Ok(())` so the caller can delegate to its `inner` codec.
+///
+/// Routing through [`read_framed_payload`] (rather than a bare
+/// `starts_with` probe) means malformed inline bytes — wrong
+/// wire-format version, mismatched Python version, truncated header —
+/// surface *their* diagnostic instead of the strict-mode message.
+/// The strict message implies sender intent ("inlining is disabled"),
+/// so it should fire only when the bytes really would have decoded.
+///
+/// Fast path: short-circuit on the family-magic prefix before
+/// acquiring the GIL. Plans with many non-Python UDFs would otherwise
+/// pay a GIL acquisition per decode call just to confirm "not a
+/// Python UDF". `read_framed_payload` itself rejects buffers that
+/// don't start with `family`, so this is purely an optimization.
+fn
Re: [PR] feat: Python UDFs: per-session inlining toggle and strict refusal setting [datafusion-python]
Copilot commented on code in PR #1546:
URL:
https://github.com/apache/datafusion-python/pull/1546#discussion_r3290100174
##
crates/core/src/codec.rs:
##
@@ -301,48 +324,102 @@ impl LogicalExtensionCodec for PythonLogicalCodec {
}
fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec) ->
Result<()> {
-if try_encode_python_scalar_udf(node, buf)? {
+if self.python_udf_inlining && try_encode_python_scalar_udf(node,
buf)? {
return Ok(());
}
self.inner.try_encode_udf(node, buf)
}
fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result>
{
-if let Some(udf) = try_decode_python_scalar_udf(buf)? {
-return Ok(udf);
+if self.python_udf_inlining {
+if let Some(udf) = try_decode_python_scalar_udf(buf)? {
+return Ok(udf);
+}
+} else {
+refuse_if_inline(buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", name)?;
}
self.inner.try_decode_udf(name, buf)
}
fn try_encode_udaf(&self, node: &AggregateUDF, buf: &mut Vec) ->
Result<()> {
-if try_encode_python_udaf(node, buf)? {
+if self.python_udf_inlining && try_encode_python_udaf(node, buf)? {
return Ok(());
}
self.inner.try_encode_udaf(node, buf)
}
fn try_decode_udaf(&self, name: &str, buf: &[u8]) ->
Result> {
-if let Some(udaf) = try_decode_python_udaf(buf)? {
-return Ok(udaf);
+if self.python_udf_inlining {
+if let Some(udaf) = try_decode_python_udaf(buf)? {
+return Ok(udaf);
+}
+} else {
+refuse_if_inline(buf, PY_AGG_UDF_FAMILY, "aggregate UDF", name)?;
}
self.inner.try_decode_udaf(name, buf)
}
fn try_encode_udwf(&self, node: &WindowUDF, buf: &mut Vec) ->
Result<()> {
-if try_encode_python_udwf(node, buf)? {
+if self.python_udf_inlining && try_encode_python_udwf(node, buf)? {
return Ok(());
}
self.inner.try_encode_udwf(node, buf)
}
fn try_decode_udwf(&self, name: &str, buf: &[u8]) ->
Result> {
-if let Some(udwf) = try_decode_python_udwf(buf)? {
-return Ok(udwf);
+if self.python_udf_inlining {
+if let Some(udwf) = try_decode_python_udwf(buf)? {
+return Ok(udwf);
+}
+} else {
+refuse_if_inline(buf, PY_WINDOW_UDF_FAMILY, "window UDF", name)?;
}
self.inner.try_decode_udwf(name, buf)
}
}
+/// Strict-mode gate: if `buf` is a well-framed inline payload for
+/// `family`, return the strict-refusal error; otherwise return
+/// `Ok(())` so the caller can delegate to its `inner` codec.
+///
+/// Routing through [`read_framed_payload`] (rather than a bare
+/// `starts_with` probe) means malformed inline bytes — wrong
+/// wire-format version, mismatched Python version, truncated header —
+/// surface *their* diagnostic instead of the strict-mode message.
+/// The strict message implies sender intent ("inlining is disabled"),
+/// so it should fire only when the bytes really would have decoded.
+///
+/// Fast path: short-circuit on the family-magic prefix before
+/// acquiring the GIL. Plans with many non-Python UDFs would otherwise
+/// pay a GIL acquisition per decode call just to confirm "not a
+/// Python UDF". `read_framed_payload` itself rejects buffers that
+/// don't start with `family`, so this is purely an optimization.
+fn refuse_if_inline(buf: &[u8], family: &[u8], kind: &str, name: &str) ->
Result<()> {
+if !buf.starts_with(family) {
+return Ok(());
+}
+Python::attach(|py| match read_framed_payload(py, buf, family, kind)? {
+Some(_) => Err(refuse_inline_payload(kind, name)),
+None => Ok(()),
+})
+}
+
+/// Build the error returned by a strict codec when it receives an
+/// inline Python-UDF payload it has been told not to deserialize.
+fn refuse_inline_payload(kind: &str, name: &str) ->
datafusion::error::DataFusionError {
+// `Execution`, not `Plan`: this is a wire-format decode refusal at
+// codec time, not a planner-stage failure. Downstream error
+// classification keys off the variant — surfacing this as a planner
+// error would mis-route it into "fix your SQL" buckets.
+datafusion::error::DataFusionError::Execution(format!(
+"Refusing to deserialize inline Python {kind} '{name}': Python UDF \
+ inlining is disabled on this session. Ask the sender to re-encode \
+ with inlining disabled (so the UDF travels by name), or register \
+ '{name}' on this receiver's session and enable inlining on both \
+ sides — receivers cannot re-encode bytes they did not produce."
Review Comment:
The strict-mode refusal message is a bit misleading: it suggests registering
the UDF and enabling inlining on both sides, but registering
