From ccdbada2afa15fa56ecebc8f74c02f63296365c4 Mon Sep 17 00:00:00 2001
From: Philipp Oppermann <dev@phil-opp.com>
Date: Wed, 5 Apr 2023 17:58:47 +0200
Subject: [PATCH] Make `send_output` work with both `PyBytes` and arrow arrays

---
 apis/python/node/src/lib.rs | 34 ++++++++++++++--------------------
 1 file changed, 14 insertions(+), 20 deletions(-)

diff --git a/apis/python/node/src/lib.rs b/apis/python/node/src/lib.rs
index 32c713be..2d30cb89 100644
--- a/apis/python/node/src/lib.rs
+++ b/apis/python/node/src/lib.rs
@@ -120,31 +120,25 @@ impl Node {
     pub fn send_output(
         &mut self,
         output_id: String,
-        data: &PyBytes,
-        metadata: Option<&PyDict>,
-        _py: Python,
-    ) -> eyre::Result<()> {
-        let data = data.as_bytes();
-        self.send_output_slice(output_id, data.len(), data, metadata)
-    }
-
-    pub fn send_output_arrow(
-        &mut self,
-        output_id: String,
-        data_arrow: PyObject,
+        data: PyObject,
         metadata: Option<&PyDict>,
         py: Python,
     ) -> eyre::Result<()> {
-        let data = arrow::array::ArrayData::from_pyarrow(data_arrow.as_ref(py))
-            .wrap_err("failed to read data as Arrow array")?;
-        if data.buffers().len() != 1 {
-            eyre::bail!("output arrow array must contain a single buffer");
-        }
+        if let Ok(py_bytes) = data.downcast::<PyBytes>(py) {
+            let data = py_bytes.as_bytes();
+            self.send_output_slice(output_id, data.len(), data, metadata)
+        } else if let Ok(arrow_array) = arrow::array::ArrayData::from_pyarrow(data.as_ref(py)) {
+            if arrow_array.buffers().len() != 1 {
+                eyre::bail!("output arrow array must contain a single buffer");
+            }
 
-        let len = data.len();
-        let slice = &data.buffer(0)[..len];
+            let len = arrow_array.len();
+            let slice = &arrow_array.buffer(0)[..len];
 
-        self.send_output_slice(output_id, len, slice, metadata)
+            self.send_output_slice(output_id, len, slice, metadata)
+        } else {
+            eyre::bail!("invalid `data` type, must by `PyBytes` or arrow array")
+        }
     }
 }