mirrors/arrow
0
0 Fork 0
mirror of https://github.com/apache/arrow.git synced 2026-10-29 17:15:39 +08:00
Code Activity

Remove DataFragment and ScanOptions

Browse Source
This commit is contained in:
Krisztián Szűcs committed 2019-12-12 12:27:38 +01:00
1 parent 13eaf46a08
commit f52e735a0d
6 files changed
+27 -423
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

cpp/src/arrow/dataset/dataset.cc
+4
View File
@@ -96,6 +96,10 @@ DataFragmentIterator DataSource::GetFragments(ScanOptionsPtr scan_options) {
return GetFragmentsImpl(std::move(simplified_scan_options));
}
SimpleDataSource::SimpleDataSource(std::vector<std::shared_ptr<RecordBatch>> record_batches) {
SimpleDataSource({std::make_shared<SimpleDataFragment>(record_batches)});
}
DataFragmentIterator SimpleDataSource::GetFragmentsImpl(ScanOptionsPtr scan_options) {
return MakeVectorIterator(fragments_);
}
cpp/src/arrow/dataset/dataset.h
+2
View File
@@ -125,6 +125,8 @@ class ARROW_DS_EXPORT SimpleDataSource : public DataSource {
explicit SimpleDataSource(DataFragmentVector fragments)
: fragments_(std::move(fragments)) {}
explicit SimpleDataSource(std::vector<std::shared_ptr<RecordBatch>> record_batches);
DataFragmentIterator GetFragmentsImpl(ScanOptionsPtr options) override;
std::string type() const override { return "simple_data_source"; }
python/pyarrow/_dataset.pyx
+10 -319
View File
@@ -294,205 +294,6 @@ cdef class FileSystemDataSourceDiscovery(DataSourceDiscovery):
self.filesystem_discovery = <CFileSystemDataSourceDiscovery*> sp.get()
cdef class FileSource:
cdef:
shared_ptr[CFileSource] wrapped
CFileSource* source
cdef readonly:
FileSystem fs
def __init__(self, path, FileSystem filesystem not None, compression=None):
cdef shared_ptr[CFileSource] source
# need to hold a reference for the filesystem
self.fs = filesystem
source.reset(new CFileSource(
tobytes(path),
self.fs.unwrap().get(),
_get_compression_type(compression)
))
self.init(source)
cdef init(self, shared_ptr[CFileSource]& sp):
self.wrapped = sp
self.source = sp.get()
@staticmethod
cdef wrap(shared_ptr[CFileSource]& sp):
cdef FileSource self = FileSource.__new__(FileSource)
self.init(sp)
return self
cdef inline shared_ptr[CFileSource] unwrap(self):
return self.wrapped
def equals(self, FileSource other):
return deref(self.source) == deref(other.unwrap())
def __eq__(self, other):
try:
return self.equals(other)
except TypeError:
return NotImplemented
@property
def compression(self):
return self.wrapped.get().compression()
@property
def path(self):
return frombytes(self.wrapped.get().path())
# def open(self):
# CStatus Open(shared_ptr[CRandomAccessFile]* out)
cdef class DataFragment:
"""A granular piece of a Dataset such as an individual file.
It can be read/scanned separately from other fragments.
A DataFragment yields a collection of RecordBatch, encapsulated in one or
more ScanTasks.
"""
cdef:
shared_ptr[CDataFragment] wrapped
CDataFragment* fragment
def __init__(self):
_forbid_instantiation(self.__class__)
cdef void init(self, const shared_ptr[CDataFragment]& sp):
self.wrapped = sp
self.fragment = sp.get()
@staticmethod
cdef wrap(shared_ptr[CDataFragment]& sp):
cdef DataFragment self
typ = frombytes(sp.get().type())
if typ == 'simple_data_fragment':
self = SimpleDataFragment.__new__(SimpleDataFragment)
elif typ == 'parquet_data_fragment':
self = ParquetDataFragment.__new__(ParquetDataFragment)
else:
raise TypeError(typ)
self.init(sp)
return self
cdef inline shared_ptr[CDataFragment] unwrap(self):
return self.wrapped
def scan(self, ScanContext context=None):
"""Returns an iterator of ScanTasks
Each of ScanTask yields RecordBatches from this DataFragment.
Parameters
----------
context : ScanContext, default None
Returns
-------
scan_tasks: Iterator[ScanTask]
"""
cdef:
CResult[CScanTaskIterator] iterator_result
CScanTaskIterator iterator
CScanTaskPtr task
context = context or ScanContext()
iterator_result = self.fragment.Scan(context.unwrap())
iterator = move(GetResultValue(move(iterator_result)))
while True:
iterator.Next(&task)
if task.get() == nullptr:
raise StopIteration()
else:
yield ScanTask.wrap(task)
@property
def splittable(self):
"""Get whether the DataFragment supports parallel scanning."""
return self.fragment.splittable()
@property
def scan_options(self):
"""Get the scan options used for scanning the data fragment.
Filtering, schema reconciliation, and partition options. None indicates
that no filtering or schema reconciliation will be performed and all
partitions will be scanned.
"""
cdef shared_ptr[CScanOptions] options = self.fragment.scan_options()
if options.get() == nullptr:
return None
else:
return ScanOptions.wrap(options)
cdef class SimpleDataFragment(DataFragment):
"""A fragment that yields ScanTasks out of a fixed set of RecordBatches."""
cdef:
CSimpleDataFragment* simple_fragment
def __init__(self, record_batches):
"""Create a simple data fragment from record batches.
Parameters
----------
record_batches : iterator of RecordBatch
"""
cdef:
RecordBatch batch
vector[shared_ptr[CRecordBatch]] batches
shared_ptr[CSimpleDataFragment] simple_fragment
for batch in record_batches:
batches.push_back(batch.sp_batch)
simple_fragment = make_shared[CSimpleDataFragment](batches)
self.init(<shared_ptr[CDataFragment]> simple_fragment)
cdef void init(self, const shared_ptr[CDataFragment]& sp):
DataFragment.init(self, sp)
self.simple_fragment = <CSimpleDataFragment*> sp.get()
cdef class FileDataFragment(DataFragment):
cdef:
CFileDataFragment* file_fragment
cdef void init(self, const shared_ptr[CDataFragment]& sp):
DataFragment.init(self, sp)
self.file_fragment = <CFileDataFragment*> sp.get()
@property
def source(self):
cdef shared_ptr[CFileSource] source
source = make_shared[CFileSource](self.file_fragment.source())
return FileSource.wrap(source)
cdef class ParquetDataFragment(FileDataFragment):
def __init__(self, FileSource source not None,
ParquetScanOptions scan_options not None):
cdef shared_ptr[CDataFragment] fragment
fragment.reset(new CParquetDataFragment(
deref(source.unwrap()),
scan_options.unwrap()
))
self.init(fragment)
cdef class DataSource:
"""Basic component of a Dataset which yields zero or more DataFragments.
@@ -531,32 +332,6 @@ cdef class DataSource:
cdef shared_ptr[CDataSource] unwrap(self):
return self.wrapped
def fragments(self, ScanOptions scan_options=None):
"""Get the data fragments of this data source.
Parameters
----------
scan_options : ScanOptions, default None
Controls filtering and schema inference.
Returns
-------
fragments : iterator of DataFragments
"""
cdef:
CDataFragmentIterator iterator
CDataFragmentPtr fragment
scan_options = scan_options or ScanOptions()
iterator = self.source.GetFragments(scan_options.unwrap())
while True:
iterator.Next(&fragment)
if fragment.get() == nullptr:
raise StopIteration()
else:
yield DataFragment.wrap(fragment)
@property
def partition_expression(self):
cdef shared_ptr[CExpression] expression
@@ -567,22 +342,23 @@ cdef class DataSource:
return Expression.wrap(expression)
# Rename it to InMemoryDataSource
cdef class SimpleDataSource(DataSource):
"""A DataSource consisting of a flat sequence of DataFragments."""
cdef:
CSimpleDataSource* simple_source
def __init__(self, data_fragments):
def __init__(self, record_batches):
cdef:
DataFragment fragment
CDataFragmentVector fragments
RecordBatch batch
vector[shared_ptr[CRecordBatch]] batches
shared_ptr[CSimpleDataSource] simple_source
for fragment in data_fragments:
fragments.push_back(fragment.unwrap())
for batch in record_batches:
batches.push_back(pyarrow_unwrap_batch(batch))
simple_source = make_shared[CSimpleDataSource](fragments)
simple_source = make_shared[CSimpleDataSource](batches)
self.init(<shared_ptr[CDataSource]> simple_source)
cdef void init(self, const shared_ptr[CDataSource]& sp):
@@ -707,72 +483,6 @@ cdef class Dataset:
return pyarrow_wrap_schema(self.dataset.schema())
cdef class ScanOptions:
cdef:
shared_ptr[CScanOptions] wrapped
CScanOptions* options
def __init__(self, Schema schema=None):
self.init(CScanOptions.Defaults())
if schema is not None:
self.schema = schema
cdef init(self, const shared_ptr[CScanOptions]& sp):
self.wrapped = sp
self.options = sp.get()
@staticmethod
cdef wrap(const shared_ptr[CScanOptions]& sp):
cdef ScanOptions self = ScanOptions.__new__(ScanOptions)
self.init(sp)
return self
cdef inline shared_ptr[CScanOptions] unwrap(self):
return self.wrapped
@property
def schema(self):
"""Schema to which record batches will be reconciled"""
if self.options.schema == nullptr:
return None
else:
return pyarrow_wrap_schema(self.options.schema)
@schema.setter
def schema(self, Schema value):
self.options.schema = pyarrow_unwrap_schema(value)
@property
def use_threads(self):
"""Make use of the thread pool found in the scan context"""
return self.options.use_threads
@use_threads.setter
def use_threads(self, bint value):
self.options.use_threads = value
cdef class FileScanOptions(ScanOptions):
pass
cdef class ParquetScanOptions(FileScanOptions):
pass
cdef class WriteOptions:
pass
cdef class FileWriteOptions(WriteOptions):
pass
cdef class ParquetWriterOptions(FileWriteOptions):
pass
cdef class ScanContext:
cdef:
@@ -983,28 +693,9 @@ cdef class Scanner:
shared_ptr[CScanner] wrapped
CScanner* scanner
def __init__(self, data_sources, ScanOptions scan_options not None,
ScanContext scan_context not None):
"""Create the scanner
Parameters
----------
data_sources : list of DataSource
scan_options : ScanOptions
scan_context : ScanContext
"""
cdef:
DataSource source
CDataSourceVector sources
shared_ptr[CScanner] scanner
for source in data_sources:
sources.push_back(source.unwrap())
scanner = make_shared[CScanner](
sources, scan_options.unwrap(), scan_context.unwrap()
)
self.init(scanner)
def __init__(self):
raise TypeError('Scanner cannot be initialized directly, use '
'ScannerBuilder instead')
cdef void init(self, shared_ptr[CScanner]& sp):
self.wrapped = sp
python/pyarrow/dataset.py
-8
View File
@@ -28,15 +28,11 @@ from pyarrow._dataset import ( # noqa
CastOptions,
CompareOperator,
ComparisonExpression,
DataFragment,
Dataset,
DataSource,
Expression,
FieldExpression,
FileDataFragment,
FileFormat,
FileScanOptions,
FileSource,
FileSystemDataSource,
FileSystemDataSourceDiscovery,
FileSystemDiscoveryOptions,
@@ -45,18 +41,14 @@ from pyarrow._dataset import ( # noqa
IsValidExpression,
NotExpression,
OrExpression,
ParquetDataFragment,
ParquetFileFormat,
ParquetScanOptions,
PartitionScheme,
ScalarExpression,
ScanContext,
Scanner,
ScannerBuilder,
ScanOptions,
ScanTask,
SchemaPartitionScheme,
SimpleDataFragment,
SimpleDataSource,
TreeDataSource,
)
python/pyarrow/includes/libarrow_dataset.pxd
+5 -5
View File
@@ -204,17 +204,17 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
cdef cppclass CSimpleDataSource "arrow::dataset::SimpleDataSource"(
CDataSource):
pass
cdef cppclass CTreeDataSource "arrow::dataset::TreeDataSource"(
CDataSource):
pass
CSimpleDataSource(CRecordBatchIterator record_batches)
ctypedef shared_ptr[CDataSource] CDataSourcePtr \
"arrow::dataset::DataSourcePtr"
ctypedef vector[CDataSourcePtr] CDataSourceVector \
"arrow::dataset::DataSourceVector"
cdef cppclass CTreeDataSource "arrow::dataset::TreeDataSource"(
CDataSource):
CTreeDataSource(CDataSourceVector children)
cdef cppclass CDataset "arrow::dataset::Dataset":
@staticmethod
CResult[shared_ptr[CDataset]] Make(CDataSourceVector sources,
python/pyarrow/tests/test_dataset.py
+6 -91
View File
@@ -23,6 +23,7 @@ import pyarrow.fs as fs
try:
import pyarrow.dataset as ds
except ImportError:
raise
ds = None
# Marks all of the tests in this module
@@ -73,13 +74,8 @@ def table(record_batch):
@pytest.fixture
def simple_data_fragment(record_batch):
return ds.SimpleDataFragment([record_batch] * 5)
@pytest.fixture
def simple_data_source(simple_data_fragment):
return ds.SimpleDataSource([simple_data_fragment] * 4)
def simple_data_source(record_batch):
return ds.SimpleDataSource([record_batch] * 4)
@pytest.fixture
@@ -92,19 +88,6 @@ def dataset(simple_data_source, tree_data_source, schema):
return ds.Dataset([simple_data_source, tree_data_source], schema)
def test_scan_options():
schema = pa.schema([
pa.field('a', pa.string()),
pa.field('b', pa.string())
])
options = ds.ScanOptions(schema)
assert isinstance(options, ds.ScanOptions)
assert options.schema.equals(schema)
options = ds.ScanOptions()
assert options.schema is None
def test_scan_context():
context = ds.ScanContext()
assert isinstance(context.memory_pool, pa.MemoryPool)
@@ -112,45 +95,16 @@ def test_scan_context():
assert isinstance(context.memory_pool, pa.MemoryPool)
def test_simple_data_fragment(record_batch):
batches = [record_batch] * 5
fragment = ds.SimpleDataFragment(batches)
assert isinstance(fragment, ds.SimpleDataFragment)
assert fragment.splittable is False
assert isinstance(fragment.scan_options, ds.ScanOptions)
tasks = list(fragment.scan())
assert len(tasks) == 5
for task, batch in zip(tasks, batches):
assert isinstance(task, ds.ScanTask)
result_batches = list(task.scan())
assert len(result_batches) == 1
assert result_batches[0].equals(batch)
def test_simple_data_source(simple_data_fragment):
source = ds.SimpleDataSource([simple_data_fragment] * 4)
def test_simple_data_source(record_batch):
source = ds.SimpleDataSource([record_batch] * 4)
assert isinstance(source, ds.SimpleDataSource)
assert source.partition_expression is None
result = source.fragments()
fragments = list(result)
assert len(fragments) == 4
for fragment in fragments:
assert isinstance(fragment, ds.DataFragment)
def test_tree_data_source(simple_data_source):
source = ds.TreeDataSource([simple_data_source] * 2)
assert isinstance(source, ds.TreeDataSource)
result = source.fragments()
fragments = list(result)
assert len(fragments) == 8
for fragment in fragments:
assert isinstance(fragment, ds.DataFragment)
def test_filesystem_data_source(mockfs):
file_stats = mockfs.get_target_stats([
@@ -168,7 +122,6 @@ def test_filesystem_data_source(mockfs):
path_partitions=path_partitions,
file_format=file_format
)
assert len(list(source.fragments())) == 2
source_partition = ds.ComparisonExpression(
ds.CompareOperator.Equal,
@@ -194,7 +147,6 @@ def test_filesystem_data_source(mockfs):
path_partitions=path_partitions,
file_format=file_format
)
assert len(list(source.fragments())) == 2
assert source.partition_expression.equals(source_partition)
@@ -225,22 +177,6 @@ def test_dataset(simple_data_source, tree_data_source, schema):
assert isinstance(record_batch, pa.RecordBatch)
def test_scanner(schema, simple_data_source):
sources = [simple_data_source]
# FIXME(kszucs): if schema is not set to options it segfaults
options = ds.ScanOptions(schema=schema)
context = ds.ScanContext()
scanner = ds.Scanner(sources, options, context)
for task in scanner.scan():
assert isinstance(task, ds.ScanTask)
for record_batch in task.scan(): # call it execute?
assert isinstance(record_batch, pa.RecordBatch)
table = scanner.to_table()
assert isinstance(table, pa.Table)
def test_scanner_builder(dataset):
context = ds.ScanContext()
builder = ds.ScannerBuilder(dataset, context)
@@ -255,31 +191,10 @@ def test_scanner_builder(dataset):
assert isinstance(scanner, ds.Scanner)
def test_file_source(mockfs):
source0 = ds.FileSource('/path/to/file.parquet', mockfs,
compression=None)
source1 = ds.FileSource('/path/to/file.parquet', mockfs,
compression=None)
source2 = ds.FileSource('/path/to/file.parquet.gz', mockfs,
compression='gzip')
assert source1.path == '/path/to/file.parquet'
assert source1.fs == mockfs
assert source1.compression == 0 # None
assert source2.path == '/path/to/file.parquet.gz'
assert source2.fs == mockfs
assert source2.compression == 2 # 'gzip'
assert source1 != source2
assert source0 == source1
fragment = ds.ParquetDataFragment(source1, ds.ParquetScanOptions())
assert isinstance(fragment, ds.ParquetDataFragment)
assert fragment.source == source1
def test_abstract_classes():
classes = [
ds.FileFormat,
ds.DataFragment,
ds.Scanner,
ds.DataSource,
ds.Expression,
ds.PartitionScheme,