Skip to content

Commit

Permalink
refactor: better obj creation (#117)
Browse files Browse the repository at this point in the history
  • Loading branch information
tshauck authored Apr 8, 2024
1 parent 5d9a881 commit 9887f5c
Showing 1 changed file with 15 additions and 21 deletions.
36 changes: 15 additions & 21 deletions src/execution_result.rs
Original file line number Diff line number Diff line change
Expand Up @@ -81,27 +81,23 @@ impl PyExecutionResult {

let mut stream = FFI_ArrowArrayStream::new(Box::new(dataframe_record_batch_stream));

Python::with_gil(|py| unsafe {
let stream_reader =
ArrowArrayStreamReader::from_raw(&mut stream).map_err(BioBearError::from)?;
let stream_reader =
unsafe { ArrowArrayStreamReader::from_raw(&mut stream).map_err(BioBearError::from) }?;

stream_reader.into_pyarrow(py)
})
stream_reader.into_pyarrow(py)
}

/// Convert to Arrow Table
fn to_arrow(&self, py: Python) -> PyResult<PyObject> {
let batches = self.collect(py)?.to_object(py);
let schema = self.schema().into_py(py);

Python::with_gil(|py| {
// Instantiate pyarrow Table object and use its from_batches method
let table_class = py.import("pyarrow")?.getattr("Table")?;
// Instantiate pyarrow Table object and use its from_batches method
let table_class = py.import("pyarrow")?.getattr("Table")?;

let args = PyTuple::new(py, &[batches, schema]);
let table: PyObject = table_class.call_method1("from_batches", args)?.into();
Ok(table)
})
let args = PyTuple::new(py, &[batches, schema]);
let table: PyObject = table_class.call_method1("from_batches", args)?.into();
Ok(table)
}

/// Convert to a Polars DataFrame
Expand All @@ -111,16 +107,14 @@ impl PyExecutionResult {

let schema = schema.into_py(py);

Python::with_gil(|py| {
let table_class = py.import("pyarrow")?.getattr("Table")?;
let args = (batches, schema);
let table: PyObject = table_class.call_method1("from_batches", args)?.into();
let table_class = py.import("pyarrow")?.getattr("Table")?;
let args = (batches, schema);
let table: PyObject = table_class.call_method1("from_batches", args)?.into();

let module = py.import("polars")?;
let args = (table,);
let result = module.call_method1("from_arrow", args)?.into();
let module = py.import("polars")?;
let args = (table,);
let result = module.call_method1("from_arrow", args)?.into();

Ok(result)
})
Ok(result)
}
}

0 comments on commit 9887f5c

Please sign in to comment.