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

Test projected partitions

Browse Source
This commit is contained in:
Krisztián Szűcs committed 2019-12-12 15:10:36 +01:00
1 parent 27dbe56ab9
commit 48b5556efa
6 files changed
+40 -43
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

cpp/src/arrow/dataset/dataset.h
-2
View File
@@ -128,8 +128,6 @@ class ARROW_DS_EXPORT SimpleDataSource : public DataSource {
static Result<DataSourcePtr> Make(
std::vector<std::shared_ptr<RecordBatch>> record_batches);
explicit SimpleDataSource();
DataFragmentIterator GetFragmentsImpl(ScanOptionsPtr options) override;
std::string type_name() const override { return "simple"; }
python/pyarrow/_dataset.pyx
+8 -12
View File
@@ -113,6 +113,10 @@ cdef class PartitionScheme:
result = self.scheme.Parse(tobytes(path))
return Expression.wrap(GetResultValue(result))
@property
def schema(self):
return pyarrow_wrap_schema(self.scheme.schema())
cdef class DefaultPartitionScheme(PartitionScheme):
@@ -132,7 +136,7 @@ cdef class DefaultPartitionScheme(PartitionScheme):
cdef class SchemaPartitionScheme(PartitionScheme):
cdef:
CSchemaPartitionScheme* schema_scheme # hmmm...
CSchemaPartitionScheme* schema_scheme
def __init__(self, Schema schema not None):
cdef shared_ptr[CSchemaPartitionScheme] scheme
@@ -145,10 +149,6 @@ cdef class SchemaPartitionScheme(PartitionScheme):
PartitionScheme.init(self, sp)
self.schema_scheme = <CSchemaPartitionScheme*> sp.get()
@property
def schema(self):
return pyarrow_wrap_schema(self.schema_scheme.schema())
cdef class HivePartitionScheme(PartitionScheme):
@@ -166,10 +166,6 @@ cdef class HivePartitionScheme(PartitionScheme):
PartitionScheme.init(self, sp)
self.hive_scheme = <CHivePartitionScheme*> sp.get()
@property
def schema(self):
return pyarrow_wrap_schema(self.hive_scheme.schema())
cdef class FileSystemDiscoveryOptions:
@@ -235,7 +231,7 @@ cdef class DataSourceDiscovery:
self.init(sp)
return self
cdef inline shared_ptr[CDataSourceDiscovery] unwrap(self):
cdef inline shared_ptr[CDataSourceDiscovery] unwrap(self) nogil:
return self.wrapped
@property
@@ -541,7 +537,7 @@ cdef class ScanTask:
self.init(sp)
return self
cdef inline shared_ptr[CScanTask] unwrap(self):
cdef inline shared_ptr[CScanTask] unwrap(self) nogil:
return self.wrapped
def execute(self):
@@ -608,7 +604,7 @@ cdef class ScannerBuilder:
self.init(sp)
return self
cdef inline shared_ptr[CScannerBuilder] unwrap(self):
cdef inline shared_ptr[CScannerBuilder] unwrap(self) nogil:
return self.wrapped
def project(self, columns):
python/pyarrow/_fs.pxd
+3 -3
View File
@@ -41,14 +41,14 @@ cdef class FileStats:
@staticmethod
cdef wrap(CFileStats stats)
cdef inline CFileStats unwrap(self)
cdef inline CFileStats unwrap(self) nogil
cdef class Selector:
cdef:
CSelector selector
cdef inline CSelector unwrap(self)
cdef inline CSelector unwrap(self) nogil
cdef class FileSystem:
@@ -61,7 +61,7 @@ cdef class FileSystem:
@staticmethod
cdef wrap(shared_ptr[CFileSystem]& sp)
cdef inline shared_ptr[CFileSystem] unwrap(self)
cdef inline shared_ptr[CFileSystem] unwrap(self) nogil
cdef class LocalFileSystem(FileSystem):
python/pyarrow/_fs.pyx
+3 -3
View File
@@ -53,7 +53,7 @@ cdef class FileStats:
self.stats = stats
return self
cdef inline CFileStats unwrap(self):
cdef inline CFileStats unwrap(self) nogil:
return self.stats
def __repr__(self):
@@ -150,7 +150,7 @@ cdef class Selector:
self.recursive = recursive
self.allow_non_existent = allow_non_existent
cdef inline CSelector unwrap(self):
cdef inline CSelector unwrap(self) nogil:
return self.selector
@property
@@ -213,7 +213,7 @@ cdef class FileSystem:
self.init(sp)
return self
cdef inline shared_ptr[CFileSystem] unwrap(self):
cdef inline shared_ptr[CFileSystem] unwrap(self) nogil:
return self.wrapped
def get_target_stats(self, paths_or_selector):
python/pyarrow/includes/libarrow_dataset.pxd
+1 -4
View File
@@ -312,6 +312,7 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
cdef cppclass CPartitionScheme "arrow::dataset::PartitionScheme":
c_string type_name() const
CResult[CExpressionPtr] Parse(const c_string& path) const
const shared_ptr[CSchema]& schema()
ctypedef shared_ptr[CPartitionScheme] CPartitionSchemePtr \
"arrow::dataset::PartitionSchemePtr"
@@ -323,14 +324,10 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
cdef cppclass CSchemaPartitionScheme \
"arrow::dataset::SchemaPartitionScheme"(CPartitionScheme):
CSchemaPartitionScheme(shared_ptr[CSchema] schema)
const shared_ptr[CSchema]& schema()
cdef cppclass CHivePartitionScheme \
"arrow::dataset::HivePartitionScheme"(CPartitionScheme):
CHivePartitionScheme(shared_ptr[CSchema] schema)
const shared_ptr[CSchema]& schema()
# vector[CUnconvertedKey] GetUnconvertedKeys(
# const c_string& path) const;
cdef cppclass CFileSystemDiscoveryOptions \
"arrow::dataset::FileSystemDiscoveryOptions":
python/pyarrow/tests/test_dataset.py
+25 -19
View File
@@ -173,24 +173,24 @@ def test_dataset(simple_data_source, tree_data_source, record_batch, schema):
assert result.equals(expected)
def test_scanner_builder(dataset):
def test_scanner_builder(simple_data_source, schema):
dataset = ds.Dataset([simple_data_source], schema)
builder = ds.ScannerBuilder(dataset, memory_pool=pa.default_memory_pool())
scanner = builder.finish()
assert isinstance(scanner, ds.Scanner)
assert len(list(scanner.scan())) == 3
# assert len(list(scanner.scan())) == 3
with pytest.raises(pa.ArrowInvalid):
dataset.new_scan().project(['unknown'])
builder = dataset.new_scan(memory_pool=pa.default_memory_pool())
scanner = builder.project(['i64', 'i64', 'i64']).finish()
scanner = builder.project(['i64']).finish()
assert isinstance(scanner, ds.Scanner)
assert len(list(scanner.scan())) == 3
# assert len(list(scanner.scan())) == 3
for task in scanner.scan():
for batch in task.execute():
assert isinstance(batch, pa.RecordBatch)
# FIXME(kszucs)
# assert batch.num_columns == 1
assert batch.num_columns == 1
def test_abstract_classes():
@@ -307,15 +307,20 @@ def test_expression(schema):
@pytest.mark.parametrize('paths_or_selector', [
fs.Selector('subdir', recursive=True),
[
'subdir',
'subdir/1',
'subdir/1/xxx',
'subdir/1/xxx/file0.parquet',
'subdir/2',
'subdir/2/yyy',
'subdir/2/yyy/file1.parquet',
]
])
def test_file_system_discovery(mockfs, paths_or_selector, record_batch):
format = ds.ParquetFileFormat()
options = ds.FileSystemDiscoveryOptions('/')
assert options.partition_base_dir == '/'
options = ds.FileSystemDiscoveryOptions('subdir')
assert options.partition_base_dir == 'subdir'
assert options.ignore_prefixes == ['.', '_']
assert options.exclude_invalid_files is True
@@ -325,17 +330,14 @@ def test_file_system_discovery(mockfs, paths_or_selector, record_batch):
assert isinstance(discovery.inspect(), pa.Schema)
assert isinstance(discovery.finish(), ds.FileSystemDataSource)
assert isinstance(discovery.partition_scheme, ds.DefaultPartitionScheme)
assert discovery.root_partition.equals(ds.ScalarExpression(True))
scheme = ds.SchemaPartitionScheme(
discovery.partition_scheme = ds.SchemaPartitionScheme(
pa.schema([
pa.field('group', pa.int32()),
pa.field('key', pa.string())
])
)
discovery.partition_scheme = scheme
assert isinstance(discovery.partition_scheme, ds.SchemaPartitionScheme)
assert discovery.root_partition.equals(ds.ScalarExpression(True))
data_source = discovery.finish()
assert isinstance(data_source, ds.DataSource)
@@ -345,13 +347,17 @@ def test_file_system_discovery(mockfs, paths_or_selector, record_batch):
scanner = dataset.new_scan().finish()
assert len(list(scanner.scan())) == 2
for task in scanner.scan():
assert isinstance(task, ds.ScanTask)
for task, group, key in zip(scanner.scan(), [1, 2], ['xxx', 'yyy']):
expected_group_column = pa.array([group] * 5, type=pa.int32())
expected_key_column = pa.array([key] * 5, type=pa.string())
for batch in task.execute():
# FIXME(kszucs)
continue
# assert batch.equals(record_batch)
assert batch.num_columns == 4
assert batch[0].equals(record_batch[0])
assert batch[1].equals(record_batch[1])
assert batch[2].equals(expected_group_column)
assert batch[3].equals(expected_key_column)
table = scanner.to_table()
assert isinstance(table, pa.Table)
assert len(table) == 10
assert table.num_columns == 4