jonded94 commented on code in PR #8790:
URL: https://github.com/apache/arrow-rs/pull/8790#discussion_r2505380122
##########
arrow-pyarrow/src/lib.rs:
##########
@@ -484,6 +487,137 @@ impl IntoPyArrow for ArrowArrayStreamReader {
}
}
+/// This is a convenience wrapper around `Vec<RecordBatch>` that tries to
simplify conversion from
+/// and to `pyarrow.Table`.
+///
+/// This could be used in circumstances where you either want to consume a
`pyarrow.Table` directly
+/// (although technically, since `pyarrow.Table` implements the
ArrayStreamReader PyCapsule
+/// interface, one could also consume a `PyArrowType<ArrowArrayStreamReader>`
instead) or, more
+/// importantly, where one wants to export a `pyarrow.Table` from a
`Vec<RecordBatch>` from the Rust
+/// side.
+///
+/// ```ignore
+/// #[pyfunction]
+/// fn return_table(...) -> PyResult<PyArrowType<Table>> {
+/// let batches: Vec<RecordBatch>;
+/// let schema: SchemaRef;
+/// PyArrowType(Table::try_new(batches, schema).map_err(|err|
err.into_py_err(py))?)
+/// }
+/// ```
+#[derive(Clone)]
+pub struct Table {
+ record_batches: Vec<RecordBatch>,
+ schema: SchemaRef,
+}
+
+impl Table {
+ pub unsafe fn new_unchecked(record_batches: Vec<RecordBatch>, schema:
SchemaRef) -> Self {
+ Self {
+ record_batches,
+ schema,
+ }
+ }
+
+ pub fn try_new(
+ record_batches: Vec<RecordBatch>,
+ schema: Option<SchemaRef>,
Review Comment:
Applied the suggestion
##########
arrow-pyarrow/src/lib.rs:
##########
@@ -484,6 +487,137 @@ impl IntoPyArrow for ArrowArrayStreamReader {
}
}
+/// This is a convenience wrapper around `Vec<RecordBatch>` that tries to
simplify conversion from
+/// and to `pyarrow.Table`.
+///
+/// This could be used in circumstances where you either want to consume a
`pyarrow.Table` directly
+/// (although technically, since `pyarrow.Table` implements the
ArrayStreamReader PyCapsule
+/// interface, one could also consume a `PyArrowType<ArrowArrayStreamReader>`
instead) or, more
+/// importantly, where one wants to export a `pyarrow.Table` from a
`Vec<RecordBatch>` from the Rust
+/// side.
+///
+/// ```ignore
+/// #[pyfunction]
+/// fn return_table(...) -> PyResult<PyArrowType<Table>> {
+/// let batches: Vec<RecordBatch>;
+/// let schema: SchemaRef;
+/// PyArrowType(Table::try_new(batches, schema).map_err(|err|
err.into_py_err(py))?)
+/// }
+/// ```
+#[derive(Clone)]
+pub struct Table {
+ record_batches: Vec<RecordBatch>,
+ schema: SchemaRef,
+}
+
+impl Table {
+ pub unsafe fn new_unchecked(record_batches: Vec<RecordBatch>, schema:
SchemaRef) -> Self {
+ Self {
+ record_batches,
+ schema,
+ }
+ }
+
+ pub fn try_new(
+ record_batches: Vec<RecordBatch>,
+ schema: Option<SchemaRef>,
+ ) -> Result<Self, ArrowError> {
+ let schema = match schema {
+ Some(s) => s,
+ None => {
+ record_batches
+ .get(0)
+ .ok_or_else(|| ArrowError::SchemaError(
+ "If no schema is supplied explicitly, there must be at
least one RecordBatch!".to_owned()
+ ))?
+ .schema()
+ .clone()
+ }
+ };
+ for record_batch in &record_batches {
+ if schema != record_batch.schema() {
+ return Err(ArrowError::SchemaError(
+ "All record batches must have the same schema.".to_owned(),
+ ));
+ }
+ }
+ Ok(Self {
+ record_batches,
+ schema,
+ })
+ }
+
+ pub fn record_batches(&self) -> &[RecordBatch] {
+ &self.record_batches
+ }
+
+ pub fn schema(&self) -> SchemaRef {
+ self.schema.clone()
+ }
+
+ pub fn into_inner(self) -> (Vec<RecordBatch>, SchemaRef) {
+ (self.record_batches, self.schema)
+ }
+}
+
+impl TryFrom<ArrowArrayStreamReader> for Table {
Review Comment:
Applied the suggestion
--
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]
