Skip to content

Commit

Permalink
Add back support for sending PyBytes as output for backwards compat…
Browse files Browse the repository at this point in the history
…ibility

Sending bytes using arrow is possible through new `send_output_arrow` method.
  • Loading branch information
phil-opp committed Apr 4, 2023
1 parent 9a3cc66 commit e3c1b4b
Showing 1 changed file with 27 additions and 4 deletions.
31 changes: 27 additions & 4 deletions apis/python/node/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -117,11 +117,22 @@ impl Node {
pub fn send_output(
&mut self,
output_id: String,
data: PyObject,
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,
metadata: Option<&PyDict>,
py: Python,
) -> eyre::Result<()> {
let data = arrow::array::ArrayData::from_pyarrow(data.as_ref(py))
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");
Expand All @@ -130,12 +141,24 @@ impl Node {
let len = data.len();
let slice = &data.buffer(0)[..len];

self.send_output_slice(output_id, len, slice, metadata)
}
}

impl Node {
fn send_output_slice(
&mut self,
output_id: String,
len: usize,
data: &[u8],
metadata: Option<&PyDict>,
) -> eyre::Result<()> {
let metadata = pydict_to_metadata(metadata)?;
self.node
.send_output(output_id.into(), metadata, len, |out| {
out.copy_from_slice(slice);
out.copy_from_slice(data);
})
.wrap_err("Could not send output")
.wrap_err("failed to send output")
}

pub fn id(&self) -> String {
Expand Down

0 comments on commit e3c1b4b

Please sign in to comment.