jorisvandenbossche commented on a change in pull request #9802:
URL: https://github.com/apache/arrow/pull/9802#discussion_r603276197
##########
File path: python/pyarrow/_dataset.pyx
##########
@@ -421,6 +421,73 @@ cdef class Dataset(_Weakrefable):
return pyarrow_wrap_schema(self.dataset.schema())
+cdef class InMemoryDataset(Dataset):
+ """A Dataset wrapping in-memory data.
+
+ Parameters
+ ----------
+ source
+ The data for this dataset. Can be a RecordBatch, Table, list of
+ RecordBatch/Table, iterable of RecordBatch, or a RecordBatchReader.
+ If an iterable is provided, the schema must also be provided.
+ schema : Schema, optional
+ Only required if passing an iterable as the source.
+ """
+
+ cdef:
+ CInMemoryDataset* in_memory_dataset
+
+ def __init__(self, source, Schema schema=None):
+ cdef:
+ RecordBatchReader reader
+ shared_ptr[CInMemoryDataset] in_memory_dataset
+
+ if isinstance(source, (pa.RecordBatch, pa.Table)):
+ source = [source]
+
+ if isinstance(source, (list, tuple)):
+ batches = []
+ for item in source:
+ if isinstance(item, pa.RecordBatch):
+ batches.append(item)
+ elif isinstance(item, pa.Table):
+ batches.extend(item.to_batches())
+ else:
+ raise TypeError(
+ 'Expected a list of tables or batches. The given list '
+ 'contains a ' + type(item).__name__)
+ if schema is None:
+ schema = item.schema
+ elif not schema.equals(item.schema):
+ raise ArrowTypeError(
+ f'Item has schema {item.schema} which which does not '
Review comment:
```suggestion
f'Item has schema {item.schema} which does not '
```
##########
File path: python/pyarrow/dataset.py
##########
@@ -713,12 +735,15 @@ def write_dataset(data, base_dir, basename_template=None,
format=None,
elif isinstance(data, (pa.Table, pa.RecordBatch)):
schema = schema or data.schema
data = [data]
- elif isinstance(data, list):
+ elif isinstance(data, (list, tuple)):
schema = schema or data[0].schema
Review comment:
Could also pass this to an InMemoryDataset? Then
`_filesystemdataset_write` could potentially be simplified to not have to deal
with a list of batches/tables
##########
File path: python/pyarrow/_dataset.pyx
##########
@@ -421,6 +421,73 @@ cdef class Dataset(_Weakrefable):
return pyarrow_wrap_schema(self.dataset.schema())
+cdef class InMemoryDataset(Dataset):
+ """A Dataset wrapping in-memory data.
+
+ Parameters
+ ----------
+ source
+ The data for this dataset. Can be a RecordBatch, Table, list of
+ RecordBatch/Table, iterable of RecordBatch, or a RecordBatchReader.
+ If an iterable is provided, the schema must also be provided.
+ schema : Schema, optional
+ Only required if passing an iterable as the source.
+ """
+
+ cdef:
+ CInMemoryDataset* in_memory_dataset
+
+ def __init__(self, source, Schema schema=None):
+ cdef:
+ RecordBatchReader reader
+ shared_ptr[CInMemoryDataset] in_memory_dataset
+
+ if isinstance(source, (pa.RecordBatch, pa.Table)):
+ source = [source]
+
+ if isinstance(source, (list, tuple)):
+ batches = []
+ for item in source:
+ if isinstance(item, pa.RecordBatch):
+ batches.append(item)
+ elif isinstance(item, pa.Table):
+ batches.extend(item.to_batches())
+ else:
+ raise TypeError(
+ 'Expected a list of tables or batches. The given list '
+ 'contains a ' + type(item).__name__)
+ if schema is None:
+ schema = item.schema
+ elif not schema.equals(item.schema):
+ raise ArrowTypeError(
+ f'Item has schema {item.schema} which which does not '
Review comment:
And I would maybe add some `\n` around the inserted schema to improve
readability of the message
--
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.
For queries about this service, please contact Infrastructure at:
[email protected]
