@@ -232,16 +232,39 @@ fn strip_wire_header<'a>( |
232 | 232 | #[derive(Debug)] |
233 | 233 | pub struct PythonLogicalCodec { |
234 | 234 | inner: Arc<dyn LogicalExtensionCodec>, |
| 235 | + python_udf_inlining: bool, |
235 | 236 | } |
236 | 237 |
|
237 | 238 | impl PythonLogicalCodec { |
238 | 239 | pub fn new(inner: Arc<dyn LogicalExtensionCodec>) -> Self { |
239 | | - Self { inner } |
| 240 | + Self { |
| 241 | + inner, |
| 242 | + python_udf_inlining: true, |
| 243 | + } |
240 | 244 | } |
241 | 245 |
|
242 | 246 | pub fn inner(&self) -> &Arc<dyn LogicalExtensionCodec> { |
243 | 247 | &self.inner |
244 | 248 | } |
| 249 | + |
| 250 | + /// Toggle inline encoding of Python UDFs. See |
| 251 | + /// `SessionContext.with_python_udf_inlining` (Python) for full |
| 252 | + /// behavior and use cases. |
| 253 | + /// |
| 254 | + /// Security scope: strict mode (`false`) narrows only the codec |
| 255 | + /// layer — it stops `Expr::from_bytes` from invoking |
| 256 | + /// `cloudpickle.loads` on the inline `DFPY*` payload. It does |
| 257 | + /// **not** make `pickle.loads(untrusted_bytes)` safe; treat every |
| 258 | + /// `pickle.loads` on untrusted input as unsafe regardless of this |
| 259 | + /// setting. |
| 260 | + pub fn with_python_udf_inlining(mut self, enabled: bool) -> Self { |
| 261 | + self.python_udf_inlining = enabled; |
| 262 | + self |
| 263 | + } |
| 264 | + |
| 265 | + pub fn python_udf_inlining(&self) -> bool { |
| 266 | + self.python_udf_inlining |
| 267 | + } |
245 | 268 | } |
246 | 269 |
|
247 | 270 | impl Default for PythonLogicalCodec { |
@@ -301,48 +324,104 @@ impl LogicalExtensionCodec for PythonLogicalCodec { |
301 | 324 | } |
302 | 325 |
|
303 | 326 | fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec<u8>) -> Result<()> { |
304 | | - if try_encode_python_scalar_udf(node, buf)? { |
| 327 | + if self.python_udf_inlining && try_encode_python_scalar_udf(node, buf)? { |
305 | 328 | return Ok(()); |
306 | 329 | } |
307 | 330 | self.inner.try_encode_udf(node, buf) |
308 | 331 | } |
309 | 332 |
|
310 | 333 | fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result<Arc<ScalarUDF>> { |
311 | | - if let Some(udf) = try_decode_python_scalar_udf(buf)? { |
312 | | - return Ok(udf); |
| 334 | + if self.python_udf_inlining { |
| 335 | + if let Some(udf) = try_decode_python_scalar_udf(buf)? { |
| 336 | + return Ok(udf); |
| 337 | + } |
| 338 | + } else { |
| 339 | + refuse_if_inline(buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", name)?; |
313 | 340 | } |
314 | 341 | self.inner.try_decode_udf(name, buf) |
315 | 342 | } |
316 | 343 |
|
317 | 344 | fn try_encode_udaf(&self, node: &AggregateUDF, buf: &mut Vec<u8>) -> Result<()> { |
318 | | - if try_encode_python_udaf(node, buf)? { |
| 345 | + if self.python_udf_inlining && try_encode_python_udaf(node, buf)? { |
319 | 346 | return Ok(()); |
320 | 347 | } |
321 | 348 | self.inner.try_encode_udaf(node, buf) |
322 | 349 | } |
323 | 350 |
|
324 | 351 | fn try_decode_udaf(&self, name: &str, buf: &[u8]) -> Result<Arc<AggregateUDF>> { |
325 | | - if let Some(udaf) = try_decode_python_udaf(buf)? { |
326 | | - return Ok(udaf); |
| 352 | + if self.python_udf_inlining { |
| 353 | + if let Some(udaf) = try_decode_python_udaf(buf)? { |
| 354 | + return Ok(udaf); |
| 355 | + } |
| 356 | + } else { |
| 357 | + refuse_if_inline(buf, PY_AGG_UDF_FAMILY, "aggregate UDF", name)?; |
327 | 358 | } |
328 | 359 | self.inner.try_decode_udaf(name, buf) |
329 | 360 | } |
330 | 361 |
|
331 | 362 | fn try_encode_udwf(&self, node: &WindowUDF, buf: &mut Vec<u8>) -> Result<()> { |
332 | | - if try_encode_python_udwf(node, buf)? { |
| 363 | + if self.python_udf_inlining && try_encode_python_udwf(node, buf)? { |
333 | 364 | return Ok(()); |
334 | 365 | } |
335 | 366 | self.inner.try_encode_udwf(node, buf) |
336 | 367 | } |
337 | 368 |
|
338 | 369 | fn try_decode_udwf(&self, name: &str, buf: &[u8]) -> Result<Arc<WindowUDF>> { |
339 | | - if let Some(udwf) = try_decode_python_udwf(buf)? { |
340 | | - return Ok(udwf); |
| 370 | + if self.python_udf_inlining { |
| 371 | + if let Some(udwf) = try_decode_python_udwf(buf)? { |
| 372 | + return Ok(udwf); |
| 373 | + } |
| 374 | + } else { |
| 375 | + refuse_if_inline(buf, PY_WINDOW_UDF_FAMILY, "window UDF", name)?; |
341 | 376 | } |
342 | 377 | self.inner.try_decode_udwf(name, buf) |
343 | 378 | } |
344 | 379 | } |
345 | 380 |
|
| 381 | +/// Strict-mode gate: if `buf` is a well-framed inline payload for |
| 382 | +/// `family`, return the strict-refusal error; otherwise return |
| 383 | +/// `Ok(())` so the caller can delegate to its `inner` codec. |
| 384 | +/// |
| 385 | +/// Routing through [`read_framed_payload`] (rather than a bare |
| 386 | +/// `starts_with` probe) means malformed inline bytes — wrong |
| 387 | +/// wire-format version, mismatched Python version, truncated header — |
| 388 | +/// surface *their* diagnostic instead of the strict-mode message. |
| 389 | +/// The strict message implies sender intent ("inlining is disabled"), |
| 390 | +/// so it should fire only when the bytes really would have decoded. |
| 391 | +/// |
| 392 | +/// Fast path: short-circuit on the family-magic prefix before |
| 393 | +/// acquiring the GIL. Plans with many non-Python UDFs would otherwise |
| 394 | +/// pay a GIL acquisition per decode call just to confirm "not a |
| 395 | +/// Python UDF". `read_framed_payload` itself rejects buffers that |
| 396 | +/// don't start with `family`, so this is purely an optimization. |
| 397 | +fn refuse_if_inline(buf: &[u8], family: &[u8], kind: &str, name: &str) -> Result<()> { |
| 398 | + if !buf.starts_with(family) { |
| 399 | + return Ok(()); |
| 400 | + } |
| 401 | + Python::attach(|py| match read_framed_payload(py, buf, family, kind)? { |
| 402 | + Some(_) => Err(refuse_inline_payload(kind, name)), |
| 403 | + None => Ok(()), |
| 404 | + }) |
| 405 | +} |
| 406 | + |
| 407 | +/// Build the error returned by a strict codec when it receives an |
| 408 | +/// inline Python-UDF payload it has been told not to deserialize. |
| 409 | +fn refuse_inline_payload(kind: &str, name: &str) -> datafusion::error::DataFusionError { |
| 410 | + // `Execution`, not `Plan`: this is a wire-format decode refusal at |
| 411 | + // codec time, not a planner-stage failure. Downstream error |
| 412 | + // classification keys off the variant — surfacing this as a planner |
| 413 | + // error would mis-route it into "fix your SQL" buckets. |
| 414 | + datafusion::error::DataFusionError::Execution(format!( |
| 415 | + "Refusing to deserialize inline Python {kind} '{name}': Python UDF \ |
| 416 | + inlining is disabled on this session. Two remediations: \ |
| 417 | + (1) ask the sender to re-encode with inlining disabled so '{name}' \ |
| 418 | + travels by name, and register '{name}' on this receiver; or \ |
| 419 | + (2) enable inlining on this receiver (accepts the cloudpickle \ |
| 420 | + execution risk on inbound payloads). Receivers cannot re-encode \ |
| 421 | + bytes they did not produce." |
| 422 | + )) |
| 423 | +} |
| 424 | + |
346 | 425 | /// `PhysicalExtensionCodec` mirror of [`PythonLogicalCodec`] parked |
347 | 426 | /// on the same `SessionContext`. Carries the Python-aware encoding |
348 | 427 | /// hooks for physical-layer types (`ExecutionPlan`, `PhysicalExpr`) |
@@ -358,16 +437,33 @@ impl LogicalExtensionCodec for PythonLogicalCodec { |
358 | 437 | #[derive(Debug)] |
359 | 438 | pub struct PythonPhysicalCodec { |
360 | 439 | inner: Arc<dyn PhysicalExtensionCodec>, |
| 440 | + python_udf_inlining: bool, |
361 | 441 | } |
362 | 442 |
|
363 | 443 | impl PythonPhysicalCodec { |
364 | 444 | pub fn new(inner: Arc<dyn PhysicalExtensionCodec>) -> Self { |
365 | | - Self { inner } |
| 445 | + Self { |
| 446 | + inner, |
| 447 | + python_udf_inlining: true, |
| 448 | + } |
366 | 449 | } |
367 | 450 |
|
368 | 451 | pub fn inner(&self) -> &Arc<dyn PhysicalExtensionCodec> { |
369 | 452 | &self.inner |
370 | 453 | } |
| 454 | + |
| 455 | + /// Toggle inline encoding of Python UDFs on this physical codec. |
| 456 | + /// |
| 457 | + /// Mirrors [`PythonLogicalCodec::with_python_udf_inlining`]; see |
| 458 | + /// that method for the full security and portability discussion. |
| 459 | + pub fn with_python_udf_inlining(mut self, enabled: bool) -> Self { |
| 460 | + self.python_udf_inlining = enabled; |
| 461 | + self |
| 462 | + } |
| 463 | + |
| 464 | + pub fn python_udf_inlining(&self) -> bool { |
| 465 | + self.python_udf_inlining |
| 466 | + } |
371 | 467 | } |
372 | 468 |
|
373 | 469 | impl Default for PythonPhysicalCodec { |
@@ -391,15 +487,19 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec { |
391 | 487 | } |
392 | 488 |
|
393 | 489 | fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec<u8>) -> Result<()> { |
394 | | - if try_encode_python_scalar_udf(node, buf)? { |
| 490 | + if self.python_udf_inlining && try_encode_python_scalar_udf(node, buf)? { |
395 | 491 | return Ok(()); |
396 | 492 | } |
397 | 493 | self.inner.try_encode_udf(node, buf) |
398 | 494 | } |
399 | 495 |
|
400 | 496 | fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result<Arc<ScalarUDF>> { |
401 | | - if let Some(udf) = try_decode_python_scalar_udf(buf)? { |
402 | | - return Ok(udf); |
| 497 | + if self.python_udf_inlining { |
| 498 | + if let Some(udf) = try_decode_python_scalar_udf(buf)? { |
| 499 | + return Ok(udf); |
| 500 | + } |
| 501 | + } else { |
| 502 | + refuse_if_inline(buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", name)?; |
403 | 503 | } |
404 | 504 | self.inner.try_decode_udf(name, buf) |
405 | 505 | } |
@@ -417,29 +517,37 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec { |
417 | 517 | } |
418 | 518 |
|
419 | 519 | fn try_encode_udaf(&self, node: &AggregateUDF, buf: &mut Vec<u8>) -> Result<()> { |
420 | | - if try_encode_python_udaf(node, buf)? { |
| 520 | + if self.python_udf_inlining && try_encode_python_udaf(node, buf)? { |
421 | 521 | return Ok(()); |
422 | 522 | } |
423 | 523 | self.inner.try_encode_udaf(node, buf) |
424 | 524 | } |
425 | 525 |
|
426 | 526 | fn try_decode_udaf(&self, name: &str, buf: &[u8]) -> Result<Arc<AggregateUDF>> { |
427 | | - if let Some(udaf) = try_decode_python_udaf(buf)? { |
428 | | - return Ok(udaf); |
| 527 | + if self.python_udf_inlining { |
| 528 | + if let Some(udaf) = try_decode_python_udaf(buf)? { |
| 529 | + return Ok(udaf); |
| 530 | + } |
| 531 | + } else { |
| 532 | + refuse_if_inline(buf, PY_AGG_UDF_FAMILY, "aggregate UDF", name)?; |
429 | 533 | } |
430 | 534 | self.inner.try_decode_udaf(name, buf) |
431 | 535 | } |
432 | 536 |
|
433 | 537 | fn try_encode_udwf(&self, node: &WindowUDF, buf: &mut Vec<u8>) -> Result<()> { |
434 | | - if try_encode_python_udwf(node, buf)? { |
| 538 | + if self.python_udf_inlining && try_encode_python_udwf(node, buf)? { |
435 | 539 | return Ok(()); |
436 | 540 | } |
437 | 541 | self.inner.try_encode_udwf(node, buf) |
438 | 542 | } |
439 | 543 |
|
440 | 544 | fn try_decode_udwf(&self, name: &str, buf: &[u8]) -> Result<Arc<WindowUDF>> { |
441 | | - if let Some(udwf) = try_decode_python_udwf(buf)? { |
442 | | - return Ok(udwf); |
| 545 | + if self.python_udf_inlining { |
| 546 | + if let Some(udwf) = try_decode_python_udwf(buf)? { |
| 547 | + return Ok(udwf); |
| 548 | + } |
| 549 | + } else { |
| 550 | + refuse_if_inline(buf, PY_WINDOW_UDF_FAMILY, "window UDF", name)?; |
443 | 551 | } |
444 | 552 | self.inner.try_decode_udwf(name, buf) |
445 | 553 | } |
@@ -476,6 +584,9 @@ pub(crate) fn try_encode_python_scalar_udf(node: &ScalarUDF, buf: &mut Vec<u8>) |
476 | 584 | /// the caller to delegate to its `inner` codec (and eventually the |
477 | 585 | /// `FunctionRegistry`). |
478 | 586 | pub(crate) fn try_decode_python_scalar_udf(buf: &[u8]) -> Result<Option<Arc<ScalarUDF>>> { |
| 587 | + if !buf.starts_with(PY_SCALAR_UDF_FAMILY) { |
| 588 | + return Ok(None); |
| 589 | + } |
479 | 590 | Python::attach(|py| -> Result<Option<Arc<ScalarUDF>>> { |
480 | 591 | let Some(payload) = read_framed_payload(py, buf, PY_SCALAR_UDF_FAMILY, "scalar UDF")? |
481 | 592 | else { |
@@ -732,6 +843,9 @@ pub(crate) fn try_encode_python_udwf(node: &WindowUDF, buf: &mut Vec<u8>) -> Res |
732 | 843 | } |
733 | 844 |
|
734 | 845 | pub(crate) fn try_decode_python_udwf(buf: &[u8]) -> Result<Option<Arc<WindowUDF>>> { |
| 846 | + if !buf.starts_with(PY_WINDOW_UDF_FAMILY) { |
| 847 | + return Ok(None); |
| 848 | + } |
735 | 849 | Python::attach(|py| -> Result<Option<Arc<WindowUDF>>> { |
736 | 850 | let Some(payload) = read_framed_payload(py, buf, PY_WINDOW_UDF_FAMILY, "window UDF")? |
737 | 851 | else { |
@@ -814,6 +928,9 @@ pub(crate) fn try_encode_python_udaf(node: &AggregateUDF, buf: &mut Vec<u8>) -> |
814 | 928 | } |
815 | 929 |
|
816 | 930 | pub(crate) fn try_decode_python_udaf(buf: &[u8]) -> Result<Option<Arc<AggregateUDF>>> { |
| 931 | + if !buf.starts_with(PY_AGG_UDF_FAMILY) { |
| 932 | + return Ok(None); |
| 933 | + } |
817 | 934 | Python::attach(|py| -> Result<Option<Arc<AggregateUDF>>> { |
818 | 935 | let Some(payload) = read_framed_payload(py, buf, PY_AGG_UDF_FAMILY, "aggregate UDF")? |
819 | 936 | else { |
|